slog_cloud/
batch.rs

1use std::marker::PhantomData;
2
3use bytes::{Bytes, BytesMut};
4
5/// Trait for batch builders.
6pub trait BatchBuilder {
7    type Item;
8    type Batch;
9
10    /// Push a given item into the builder.
11    ///
12    /// If the item exceeds the maximum batch size, the current batch is
13    /// returned and the item will start a new batch.
14    fn push(&mut self, item: Self::Item) -> Option<Self::Batch>;
15
16    /// Flush the builder and return the accumulated batch, if any.
17    fn flush(&mut self) -> Option<Self::Batch>;
18}
19
20/// Dummy batch builder.
21pub struct DummyBatchBuilder<M> {
22    _pd: PhantomData<M>,
23}
24
25impl<M> DummyBatchBuilder<M> {
26    /// Create a new dummy batch builder.
27    #[inline]
28    pub const fn new() -> Self {
29        Self { _pd: PhantomData }
30    }
31}
32
33impl<M> Default for DummyBatchBuilder<M> {
34    #[inline]
35    fn default() -> Self {
36        Self::new()
37    }
38}
39
40impl<M> BatchBuilder for DummyBatchBuilder<M> {
41    type Item = M;
42    type Batch = M;
43
44    #[inline]
45    fn push(&mut self, item: Self::Item) -> Option<Self::Batch> {
46        Some(item)
47    }
48
49    #[inline]
50    fn flush(&mut self) -> Option<Self::Batch> {
51        None
52    }
53}
54
55/// NDJSON batch builder.
56///
57/// The builder accumulates serialized JSON values into an NDJSON batch of a
58/// given maximum size.
59pub struct NDJSONBatchBuilder {
60    buffer: BytesMut,
61    max_size: usize,
62}
63
64impl NDJSONBatchBuilder {
65    /// Create a new NDJSON batch builder with the given maximum batch size.
66    #[inline]
67    pub fn new(max_size: usize) -> Self {
68        Self {
69            buffer: BytesMut::new(),
70            max_size,
71        }
72    }
73}
74
75impl BatchBuilder for NDJSONBatchBuilder {
76    type Item = Bytes;
77    type Batch = Bytes;
78
79    fn push(&mut self, item: Self::Item) -> Option<Self::Batch> {
80        let buffer_len = self.buffer.len();
81
82        let item_len = item.len();
83
84        let new_len = buffer_len + item_len + 1;
85
86        let res = if buffer_len > 0 && new_len > self.max_size {
87            let batch = self.buffer.split();
88
89            Some(batch.freeze())
90        } else {
91            None
92        };
93
94        self.buffer.reserve(item_len + 1);
95        self.buffer.extend_from_slice(&item);
96        self.buffer.extend_from_slice(b"\n");
97
98        res
99    }
100
101    fn flush(&mut self) -> Option<Self::Batch> {
102        let batch = self.buffer.split();
103
104        if batch.is_empty() {
105            None
106        } else {
107            Some(batch.freeze())
108        }
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use bytes::Bytes;
115
116    use super::{BatchBuilder, NDJSONBatchBuilder};
117
118    fn batch(s: &'static str) -> Option<Bytes> {
119        Some(Bytes::from_static(s.as_bytes()))
120    }
121
122    #[test]
123    fn ndjson_batch_is_emitted_when_next_item_does_not_fit() {
124        let mut builder = NDJSONBatchBuilder::new(8);
125
126        assert_eq!(builder.push(Bytes::from_static(b"aaa")), None);
127        // the batch is exactly 8 bytes now
128        assert_eq!(builder.push(Bytes::from_static(b"bbb")), None);
129        assert_eq!(builder.push(Bytes::from_static(b"c")), batch("aaa\nbbb\n"));
130        assert_eq!(builder.flush(), batch("c\n"));
131        assert_eq!(builder.flush(), None);
132    }
133
134    #[test]
135    fn oversized_ndjson_item_forms_its_own_batch() {
136        let mut builder = NDJSONBatchBuilder::new(4);
137
138        assert_eq!(builder.push(Bytes::from_static(b"a")), None);
139        assert_eq!(builder.push(Bytes::from_static(b"bbbbbb")), batch("a\n"));
140        assert_eq!(builder.push(Bytes::from_static(b"c")), batch("bbbbbb\n"));
141        assert_eq!(builder.flush(), batch("c\n"));
142    }
143}
slog_cloud/batch.rs

slog_cloud/
batch.rs
