1#![allow(non_snake_case)]
  2
  3pub use ::proto::*;
  4
  5use anyhow::anyhow;
  6use async_tungstenite::tungstenite::Message as WebSocketMessage;
  7use futures::{SinkExt as _, StreamExt as _};
  8use proto::Message as _;
  9use std::time::Instant;
 10use std::{fmt::Debug, io};
 11use zstd::zstd_safe::WriteBuf;
 12
 13const KIB: usize = 1024;
 14const MIB: usize = KIB * 1024;
 15const MAX_BUFFER_LEN: usize = MIB;
 16
 17/// A stream of protobuf messages.
 18pub struct MessageStream<S> {
 19    stream: S,
 20    encoding_buffer: Vec<u8>,
 21}
 22
 23#[derive(Debug)]
 24pub enum Message {
 25    Envelope(Envelope),
 26    Ping,
 27    Pong,
 28}
 29
 30impl<S> MessageStream<S> {
 31    pub fn new(stream: S) -> Self {
 32        Self {
 33            stream,
 34            encoding_buffer: Vec::new(),
 35        }
 36    }
 37}
 38
 39impl<S> MessageStream<S>
 40where
 41    S: futures::Sink<WebSocketMessage, Error = anyhow::Error> + Unpin,
 42{
 43    pub async fn write(&mut self, message: Message) -> Result<(), anyhow::Error> {
 44        #[cfg(any(test, feature = "test-support"))]
 45        const COMPRESSION_LEVEL: i32 = -7;
 46
 47        #[cfg(not(any(test, feature = "test-support")))]
 48        const COMPRESSION_LEVEL: i32 = 4;
 49
 50        match message {
 51            Message::Envelope(message) => {
 52                self.encoding_buffer.reserve(message.encoded_len());
 53                message
 54                    .encode(&mut self.encoding_buffer)
 55                    .map_err(io::Error::from)?;
 56                let buffer =
 57                    zstd::stream::encode_all(self.encoding_buffer.as_slice(), COMPRESSION_LEVEL)
 58                        .unwrap();
 59
 60                self.encoding_buffer.clear();
 61                self.encoding_buffer.shrink_to(MAX_BUFFER_LEN);
 62                self.stream
 63                    .send(WebSocketMessage::Binary(buffer.into()))
 64                    .await?;
 65            }
 66            Message::Ping => {
 67                self.stream
 68                    .send(WebSocketMessage::Ping(Default::default()))
 69                    .await?;
 70            }
 71            Message::Pong => {
 72                self.stream
 73                    .send(WebSocketMessage::Pong(Default::default()))
 74                    .await?;
 75            }
 76        }
 77
 78        Ok(())
 79    }
 80}
 81
 82impl<S> MessageStream<S>
 83where
 84    S: futures::Stream<Item = Result<WebSocketMessage, anyhow::Error>> + Unpin,
 85{
 86    pub async fn read(&mut self) -> Result<(Message, Instant), anyhow::Error> {
 87        while let Some(bytes) = self.stream.next().await {
 88            let received_at = Instant::now();
 89            match bytes? {
 90                WebSocketMessage::Binary(bytes) => {
 91                    zstd::stream::copy_decode(bytes.as_slice(), &mut self.encoding_buffer)?;
 92                    let envelope = Envelope::decode(self.encoding_buffer.as_slice())
 93                        .map_err(io::Error::from)?;
 94
 95                    self.encoding_buffer.clear();
 96                    self.encoding_buffer.shrink_to(MAX_BUFFER_LEN);
 97                    return Ok((Message::Envelope(envelope), received_at));
 98                }
 99                WebSocketMessage::Ping(_) => return Ok((Message::Ping, received_at)),
100                WebSocketMessage::Pong(_) => return Ok((Message::Pong, received_at)),
101                WebSocketMessage::Close(_) => break,
102                _ => {}
103            }
104        }
105        Err(anyhow!("connection closed"))
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[gpui::test]
114    async fn test_buffer_size() {
115        let (tx, rx) = futures::channel::mpsc::unbounded();
116        let mut sink = MessageStream::new(tx.sink_map_err(|_| anyhow!("")));
117        sink.write(Message::Envelope(Envelope {
118            payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree {
119                root_name: "abcdefg".repeat(10),
120                ..Default::default()
121            })),
122            ..Default::default()
123        }))
124        .await
125        .unwrap();
126        assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
127        sink.write(Message::Envelope(Envelope {
128            payload: Some(envelope::Payload::UpdateWorktree(UpdateWorktree {
129                root_name: "abcdefg".repeat(1000000),
130                ..Default::default()
131            })),
132            ..Default::default()
133        }))
134        .await
135        .unwrap();
136        assert!(sink.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
137
138        let mut stream = MessageStream::new(rx.map(anyhow::Ok));
139        stream.read().await.unwrap();
140        assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
141        stream.read().await.unwrap();
142        assert!(stream.encoding_buffer.capacity() <= MAX_BUFFER_LEN);
143    }
144}