watermelon_proto/proto/encoder/
stream.rs

1#[cfg(feature = "std")]
2use std::io;
3
4use bytes::{Buf, BufMut, Bytes, BytesMut};
5
6use crate::util::BufList;
7
8use super::{ClientOp, FrameEncoder};
9
10const WRITE_FLATTEN_THRESHOLD: usize = 4096;
11
12#[derive(Debug)]
13pub struct StreamEncoder {
14    write_buf: BufList<Bytes>,
15    flattened_writes: BytesMut,
16}
17
18impl StreamEncoder {
19    #[must_use]
20    pub fn new() -> Self {
21        Self {
22            write_buf: BufList::new(),
23            flattened_writes: BytesMut::new(),
24        }
25    }
26
27    pub fn enqueue_write_op(&mut self, item: &ClientOp) {
28        super::encode(self, item);
29    }
30
31    #[cfg(test)]
32    fn all_bytes(&mut self) -> alloc::vec::Vec<u8> {
33        self.copy_to_bytes(self.remaining()).to_vec()
34    }
35}
36
37impl Buf for StreamEncoder {
38    fn remaining(&self) -> usize {
39        self.write_buf.remaining() + self.flattened_writes.remaining()
40    }
41
42    fn has_remaining(&self) -> bool {
43        self.write_buf.has_remaining() || self.flattened_writes.has_remaining()
44    }
45
46    fn chunk(&self) -> &[u8] {
47        let chunk = self.write_buf.chunk();
48        if chunk.is_empty() {
49            &self.flattened_writes
50        } else {
51            chunk
52        }
53    }
54
55    #[cfg(feature = "std")]
56    fn chunks_vectored<'a>(&'a self, dst: &mut [io::IoSlice<'a>]) -> usize {
57        let mut n = self.write_buf.chunks_vectored(dst);
58        n += self.flattened_writes.chunks_vectored(&mut dst[n..]);
59        n
60    }
61
62    fn advance(&mut self, cnt: usize) {
63        assert!(cnt <= self.remaining());
64
65        let mid = self.write_buf.remaining().min(cnt);
66        self.write_buf.advance(mid);
67
68        let rem = cnt - mid;
69        if rem == self.flattened_writes.len() {
70            // https://github.com/tokio-rs/bytes/pull/698
71            self.flattened_writes.clear();
72        } else {
73            self.flattened_writes.advance(rem);
74        }
75    }
76
77    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
78        assert!(
79            len <= self.remaining(),
80            "copy_to_bytes out of range ({} <= {})",
81            len,
82            self.remaining()
83        );
84
85        if self.write_buf.remaining() >= len {
86            self.write_buf.copy_to_bytes(len)
87        } else if !self.write_buf.has_remaining() {
88            self.flattened_writes.copy_to_bytes(len)
89        } else {
90            let rem = len - self.write_buf.remaining();
91
92            let mut bufs = BytesMut::with_capacity(len);
93            bufs.put(&mut self.write_buf);
94            bufs.put_slice(&self.flattened_writes[..rem]);
95
96            if self.flattened_writes.remaining() == rem {
97                // https://github.com/tokio-rs/bytes/pull/698
98                self.flattened_writes.clear();
99            } else {
100                self.flattened_writes.advance(rem);
101            }
102
103            bufs.freeze()
104        }
105    }
106}
107
108impl FrameEncoder for StreamEncoder {
109    fn small_write(&mut self, buf: &[u8]) {
110        self.flattened_writes.extend_from_slice(buf);
111    }
112
113    fn write<B>(&mut self, buf: B)
114    where
115        B: Into<Bytes> + AsRef<[u8]>,
116    {
117        let b = buf.as_ref();
118
119        let len = b.len();
120        if len == 0 {
121            return;
122        }
123
124        if len < WRITE_FLATTEN_THRESHOLD {
125            self.flattened_writes.extend_from_slice(b);
126        } else {
127            if !self.flattened_writes.is_empty() {
128                let buf = self.flattened_writes.split().freeze();
129                self.write_buf.push(buf);
130            }
131
132            self.write_buf.push(buf.into());
133        }
134    }
135}
136
137impl Default for StreamEncoder {
138    fn default() -> Self {
139        Self::new()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use core::num::NonZero;
146    #[cfg(feature = "std")]
147    use std::io::IoSlice;
148
149    use bytes::{Buf as _, Bytes};
150
151    use super::StreamEncoder;
152    #[cfg(feature = "std")]
153    use crate::proto::encoder::FrameEncoder;
154    use crate::{
155        MessageBase, QueueGroup, Subject,
156        headers::{HeaderMap, HeaderName, HeaderValue},
157        proto::ClientOp,
158    };
159
160    #[test]
161    fn base() {
162        let encoder = StreamEncoder::new();
163        assert_eq!(0, encoder.remaining());
164        assert!(!encoder.has_remaining());
165    }
166
167    #[cfg(feature = "std")]
168    #[test]
169    fn vectored() {
170        let mut encoder = StreamEncoder::new();
171        let mut bufs = [IoSlice::new(&[]); 64];
172        assert_eq!(0, encoder.chunks_vectored(&mut bufs));
173
174        encoder.small_write(b"1234");
175        let mut bufs = [IoSlice::new(&[]); 64];
176        assert_eq!(1, encoder.chunks_vectored(&mut bufs));
177
178        encoder.small_write(b"5678");
179        let mut bufs = [IoSlice::new(&[]); 64];
180        assert_eq!(1, encoder.chunks_vectored(&mut bufs));
181
182        encoder.write("9");
183        let mut bufs = [IoSlice::new(&[]); 64];
184        assert_eq!(1, encoder.chunks_vectored(&mut bufs));
185
186        encoder.write(vec![b'_'; 8196]);
187        let mut bufs = [IoSlice::new(&[]); 64];
188        assert_eq!(2, encoder.chunks_vectored(&mut bufs));
189
190        encoder.write(vec![b':'; 8196]);
191        let mut bufs = [IoSlice::new(&[]); 64];
192        assert_eq!(3, encoder.chunks_vectored(&mut bufs));
193
194        encoder.write("0");
195        let mut bufs = [IoSlice::new(&[]); 64];
196        assert_eq!(4, encoder.chunks_vectored(&mut bufs));
197    }
198
199    #[test]
200    fn encode_ping() {
201        let mut encoder = StreamEncoder::new();
202        encoder.enqueue_write_op(&ClientOp::Ping);
203        assert_eq!(6, encoder.remaining());
204        assert!(encoder.has_remaining());
205        assert_eq!("PING\r\n".as_bytes(), encoder.all_bytes());
206    }
207
208    #[test]
209    fn encode_pong() {
210        let mut encoder = StreamEncoder::new();
211        encoder.enqueue_write_op(&ClientOp::Pong);
212        assert_eq!(6, encoder.remaining());
213        assert!(encoder.has_remaining());
214        assert_eq!("PONG\r\n".as_bytes(), encoder.all_bytes());
215    }
216
217    #[test]
218    fn encode_subscribe() {
219        let mut encoder = StreamEncoder::new();
220        encoder.enqueue_write_op(&ClientOp::Subscribe {
221            id: 1.into(),
222            subject: Subject::from_static("hello.world"),
223            queue_group: None,
224        });
225        assert_eq!(19, encoder.remaining());
226        assert!(encoder.has_remaining());
227        assert_eq!("SUB hello.world 1\r\n".as_bytes(), encoder.all_bytes());
228    }
229
230    #[test]
231    fn encode_subscribe_with_queue_group() {
232        let mut encoder = StreamEncoder::new();
233        encoder.enqueue_write_op(&ClientOp::Subscribe {
234            id: 1.into(),
235            subject: Subject::from_static("hello.world"),
236            queue_group: Some(QueueGroup::from_static("stuff")),
237        });
238        assert_eq!(25, encoder.remaining());
239        assert!(encoder.has_remaining());
240        assert_eq!(
241            "SUB hello.world stuff 1\r\n".as_bytes(),
242            encoder.all_bytes()
243        );
244    }
245
246    #[test]
247    fn encode_unsubscribe() {
248        let mut encoder = StreamEncoder::new();
249        encoder.enqueue_write_op(&ClientOp::Unsubscribe {
250            id: 1.into(),
251            max_messages: None,
252        });
253        assert_eq!(9, encoder.remaining());
254        assert!(encoder.has_remaining());
255        assert_eq!("UNSUB 1\r\n".as_bytes(), encoder.all_bytes());
256    }
257
258    #[test]
259    fn encode_unsubscribe_with_max_messages() {
260        let mut encoder = StreamEncoder::new();
261        encoder.enqueue_write_op(&ClientOp::Unsubscribe {
262            id: 1.into(),
263            max_messages: Some(NonZero::new(5).unwrap()),
264        });
265        assert_eq!(11, encoder.remaining());
266        assert!(encoder.has_remaining());
267        assert_eq!("UNSUB 1 5\r\n".as_bytes(), encoder.all_bytes());
268    }
269
270    #[test]
271    fn encode_publish() {
272        let mut encoder = StreamEncoder::new();
273        encoder.enqueue_write_op(&ClientOp::Publish {
274            message: MessageBase {
275                subject: Subject::from_static("hello.world"),
276                reply_subject: None,
277                headers: HeaderMap::new(),
278                payload: Bytes::from_static(b"Hello World!"),
279            },
280        });
281        assert_eq!(34, encoder.remaining());
282        assert!(encoder.has_remaining());
283        assert_eq!(
284            "PUB hello.world 12\r\nHello World!\r\n".as_bytes(),
285            encoder.all_bytes()
286        );
287    }
288
289    #[test]
290    fn encode_publish_with_reply_subject() {
291        let mut encoder = StreamEncoder::new();
292        encoder.enqueue_write_op(&ClientOp::Publish {
293            message: MessageBase {
294                subject: Subject::from_static("hello.world"),
295                reply_subject: Some(Subject::from_static("_INBOX.1234")),
296                headers: HeaderMap::new(),
297                payload: Bytes::from_static(b"Hello World!"),
298            },
299        });
300        assert_eq!(46, encoder.remaining());
301        assert!(encoder.has_remaining());
302        assert_eq!(
303            "PUB hello.world _INBOX.1234 12\r\nHello World!\r\n".as_bytes(),
304            encoder.all_bytes()
305        );
306    }
307
308    #[test]
309    fn encode_publish_with_headers() {
310        let mut encoder = StreamEncoder::new();
311        encoder.enqueue_write_op(&ClientOp::Publish {
312            message: MessageBase {
313                subject: Subject::from_static("hello.world"),
314                reply_subject: None,
315                headers: [
316                    (
317                        HeaderName::from_static("Nats-Msg-Id"),
318                        HeaderValue::from_static("abcd"),
319                    ),
320                    (
321                        HeaderName::from_static("Nats-Sequence"),
322                        HeaderValue::from_static("1"),
323                    ),
324                ]
325                .into_iter()
326                .collect(),
327                payload: Bytes::from_static(b"Hello World!"),
328            },
329        });
330        assert_eq!(87, encoder.remaining());
331        assert!(encoder.has_remaining());
332        assert_eq!(
333            "HPUB hello.world 49 61\r\nNATS/1.0\r\nNats-Msg-Id: abcd\r\nNats-Sequence: 1\r\n\r\nHello World!\r\n".as_bytes(),
334            encoder.all_bytes()
335        );
336    }
337}
watermelon_proto/proto/encoder/stream.rs

watermelon_proto/proto/encoder/
stream.rs
