Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cap MessageStream buffer size to 1MB

Browse source 
We temporarily let it grow when the message size exceed the limit,
but restore the buffer's capacity shortly after. This ensures that,
for each connection in its entire lifetime, we only ever use 1MB.
This commit is contained in:
Antonio Scandurra 2022-05-31 11:16:32 +02:00
parent da46d78ea5
commit 339069b1d3
17 changed files with 62 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

49
crates/rpc/src/proto.rs
View file

@ -254,6 +254,8 @@ entity_messages!(
entity_messages!(channel_id, ChannelMessageSent);
const MAX_BUFFER_LEN: usize = 1 * 1024 * 1024;
/// A stream of protobuf messages.
pub struct MessageStream<S> {
stream: S,
@ -293,14 +295,16 @@ where
match message {
Message::Envelope(message) => {
self.encoding_buffer.resize(message.encoded_len(), 0);
self.encoding_buffer.clear();
self.encoding_buffer.reserve(message.encoded_len());
message
.encode(&mut self.encoding_buffer)
.map_err(|err| io::Error::from(err))?;
let buffer =
zstd::stream::encode_all(self.encoding_buffer.as_slice(), COMPRESSION_LEVEL)
.unwrap();
self.encoding_buffer.clear();
self.encoding_buffer.shrink_to(MAX_BUFFER_LEN);
self.stream.send(WebSocketMessage::Binary(buffer)).await?;
}
Message::Ping => {
@ -327,10 +331,12 @@ where
while let Some(bytes) = self.stream.next().await {
match bytes? {
WebSocketMessage::Binary(bytes) => {
self.encoding_buffer.clear();
zstd::stream::copy_decode(bytes.as_slice(), &mut self.encoding_buffer).unwrap();
let envelope = Envelope::decode(self.encoding_buffer.as_slice())
.map_err(io::Error::from)?;
self.encoding_buffer.clear();
self.encoding_buffer.shrink_to(MAX_BUFFER_LEN);
return Ok(Message::Envelope(envelope));
}
WebSocketMessage::Ping(_) => return Ok(Message::Ping),
@ -379,3 +385,40 @@ impl From<Nonce> for u128 {
upper_half | lower_half
}
}
#[cfg(test)]
mod tests {
use super::*;
#[gpui::test]
async fn test_buffer_size() {
let (tx, rx) = futures::channel::mpsc::unbounded();
let mut sink = MessageStream::new(tx.sink_map_err(|_| anyhow!("")));
sink.write(Message::Envelope(Envelope {
payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree {
root_name: "abcdefg".repeat(10),
..Default::default()
})),
..Default::default()
}))
.await
.unwrap();
assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
sink.write(Message::Envelope(Envelope {
payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree {
root_name: "abcdefg".repeat(1000000),
..Default::default()
})),
..Default::default()
}))
.await
.unwrap();
assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
let mut stream = MessageStream::new(rx.map(|msg| anyhow::Ok(msg)));
stream.read().await.unwrap();
assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
stream.read().await.unwrap();
assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
}
}