rama_core/stream/json/
codec.rs

1use bytes::{Buf, BufMut};
2use core::marker::PhantomData;
3use rama_error::{BoxError, ErrorContext as _};
4use serde::{Serialize, de::DeserializeOwned};
5
6use super::engine::NdjsonEngine;
7use crate::{bytes::BytesMut, stream::json::ParseConfig};
8
9/// NDJson encoder.
10pub struct JsonEncoder<T> {
11    written: bool,
12    _phantom: PhantomData<fn() -> T>,
13}
14
15impl<T> JsonEncoder<T> {
16    #[must_use]
17    pub fn new() -> Self {
18        Self {
19            written: false,
20            _phantom: PhantomData,
21        }
22    }
23}
24
25impl<T> JsonEncoder<T> {
26    #[must_use]
27    /// Use [`JsonEncoder::new`] for new streams.
28    /// This constructor can be used if you wish to start
29    /// the strema with a newline because you are continuing
30    /// to write to a stream where you left of.
31    pub fn new_continued() -> Self {
32        Self {
33            written: true,
34            _phantom: PhantomData,
35        }
36    }
37}
38
39impl<T> Clone for JsonEncoder<T> {
40    fn clone(&self) -> Self {
41        *self
42    }
43}
44
45impl<T> Copy for JsonEncoder<T> {}
46
47impl<T> Default for JsonEncoder<T> {
48    #[inline]
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl<T: Serialize> crate::stream::codec::Encoder<T> for JsonEncoder<T> {
55    type Error = BoxError;
56
57    fn encode(&mut self, data: T, buf: &mut BytesMut) -> Result<(), Self::Error> {
58        if self.written {
59            buf.put_u8(b'\n');
60        }
61        let result = serde_json::to_writer(buf.writer(), &data)
62            .context("serde-json write data to buffer")
63            .into_box_error();
64        self.written = true;
65        result
66    }
67}
68
69/// NDJson decoder decoding ndjson stream of bytes
70/// into json objects.
71pub struct JsonDecoder<T> {
72    engine: NdjsonEngine<T>,
73}
74
75impl<T> JsonDecoder<T> {
76    /// Creates a new fallible NDJSON decoder with default [`ParseConfig`].
77    #[must_use]
78    pub fn new() -> Self {
79        Self {
80            engine: NdjsonEngine::new(),
81        }
82    }
83
84    /// Creates a new fallible NDJSON decoder with the given
85    /// [`ParseConfig`] to control its behavior.
86    ///
87    /// See [`ParseConfig`] for more details.
88    #[must_use]
89    pub fn new_with_config(config: ParseConfig) -> Self {
90        Self {
91            engine: NdjsonEngine::with_config(config),
92        }
93    }
94}
95
96impl<T> Default for JsonDecoder<T> {
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102impl<T: DeserializeOwned> crate::stream::codec::Decoder for JsonDecoder<T> {
103    type Item = T;
104    type Error = BoxError;
105
106    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
107        // If we already buffered parsed values, return them first.
108        if let Some(result) = self.engine.pop() {
109            return Ok(Some(result.context("json-deserialize next value")?));
110        }
111
112        // DO NOT finalize here on empty src; just ask FramedRead to read more.
113        if !src.is_empty() {
114            self.engine.input(&src);
115            src.advance(src.len());
116        }
117
118        match self.engine.pop() {
119            Some(result) => Ok(Some(result.context("json-deserialize next value")?)),
120            None => Ok(None),
121        }
122    }
123
124    // If your trait has a dedicated EOF hook, implement it like this:
125    fn decode_eof(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
126        if !src.is_empty() {
127            self.engine.input(&src);
128            src.advance(src.len());
129        }
130        self.engine.finalize();
131
132        if let Some(result) = self.engine.pop() {
133            return Ok(Some(result.context("json-deserialize next value")?));
134        }
135        Ok(None)
136    }
137}
138#[cfg(test)]
139mod tests {
140    use ahash::{HashSet, HashSetExt as _};
141
142    use super::*;
143    use crate::stream::codec::{Decoder as _, Encoder as _};
144    use serde::{Deserialize, Serialize};
145
146    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
147    struct Item {
148        id: u32,
149        name: String,
150    }
151
152    #[test]
153    fn encode_single_value_no_leading_newline() -> Result<(), BoxError> {
154        let mut enc: JsonEncoder<u32> = JsonEncoder::new();
155        let mut buf = BytesMut::new();
156
157        enc.encode(42, &mut buf)?;
158
159        let s = core::str::from_utf8(&buf)?;
160        assert_eq!(s, "42"); // no newline
161        Ok(())
162    }
163
164    #[test]
165    fn encode_multiple_values_separated_by_newline_without_trailing_newline() -> Result<(), BoxError>
166    {
167        let mut enc: JsonEncoder<u32> = JsonEncoder::new();
168        let mut buf = BytesMut::new();
169
170        enc.encode(1, &mut buf)?;
171        enc.encode(2, &mut buf)?;
172        enc.encode(3, &mut buf)?;
173
174        let bytes = buf.as_ref();
175        // must be "1\n2\n3" with no trailing newline
176        assert_eq!(bytes, b"1\n2\n3");
177        assert_ne!(bytes.last().copied(), Some(b'\n'));
178        Ok(())
179    }
180
181    #[test]
182    fn roundtrip_structs_encode_then_decode_all() -> Result<(), BoxError> {
183        let mut enc: JsonEncoder<Item> = JsonEncoder::new();
184        let mut buf = BytesMut::new();
185
186        let input = vec![
187            Item {
188                id: 1,
189                name: "alice".to_owned(),
190            },
191            Item {
192                id: 2,
193                name: "bob".to_owned(),
194            },
195            Item {
196                id: 3,
197                name: "carol".to_owned(),
198            },
199        ];
200
201        for it in &input {
202            enc.encode(it.clone(), &mut buf)?;
203        }
204
205        let mut dec: JsonDecoder<Item> = JsonDecoder::new();
206        let mut out = Vec::new();
207
208        // First call will feed the entire buffer into the engine and return the first item
209        if let Some(first) = dec.decode(&mut buf)? {
210            out.push(first);
211        }
212
213        // Subsequent calls will keep popping from the engine even if src is now empty
214        while let Some(next) = dec.decode(&mut buf)? {
215            out.push(next);
216        }
217
218        if let Some(next) = dec.decode_eof(&mut buf)? {
219            out.push(next);
220        }
221
222        assert_eq!(out, input);
223        Ok(())
224    }
225
226    #[test]
227    fn decode_incremental_streaming_chunks() -> Result<(), BoxError> {
228        // Prepare an NDJSON stream
229        let mut enc: JsonEncoder<Item> = JsonEncoder::new();
230        let mut full = BytesMut::new();
231
232        let items = vec![
233            Item {
234                id: 10,
235                name: "ten".into(),
236            },
237            Item {
238                id: 20,
239                name: "twenty".into(),
240            },
241            Item {
242                id: 30,
243                name: "thirty".into(),
244            },
245        ];
246        for it in &items {
247            enc.encode(it.clone(), &mut full)?;
248        }
249
250        // Split into irregular chunks to mimic a real stream
251        let all_bytes = full.freeze();
252        let split_points = [1, 7, 13, all_bytes.len()]; // arbitrary cut points
253        let mut chunks = Vec::new();
254        let mut start = 0;
255        for &end in &split_points {
256            chunks.push(all_bytes.slice(start..end));
257            start = end;
258        }
259
260        // Feed chunks one by one
261        let mut dec: JsonDecoder<Item> = JsonDecoder::new();
262        let mut collected = Vec::new();
263        let mut staging = BytesMut::new();
264
265        for chunk in chunks {
266            staging.extend_from_slice(&chunk);
267            // Try to drain as many items as available after each chunk
268            while let Some(item) = dec.decode(&mut staging)? {
269                collected.push(item);
270            }
271        }
272
273        if let Ok(Some(next)) = dec.decode_eof(&mut staging) {
274            collected.push(next);
275        }
276
277        assert_eq!(collected, items);
278        Ok(())
279    }
280
281    #[test]
282    fn decode_reports_error_for_malformed_json_line() {
283        let mut dec: JsonDecoder<serde_json::Value> = JsonDecoder::new();
284        let mut buf = BytesMut::from(&b"{not valid json}\n{\"ok\": true}\n"[..]);
285
286        // First decode should error because the first line is invalid
287        let err = dec
288            .decode(&mut buf)
289            .expect_err("expected an error for malformed json");
290        let msg = format!("{err}");
291        // The error is wrapped with context in the decoder
292        assert!(msg.contains("json-deserialize next value"));
293
294        // After an error, you can still try to keep decoding
295        // The engine has consumed the input already, so call decode again to pop any next valid item
296        let next = dec
297            .decode(&mut buf)
298            .expect("second decode should not error");
299        // Depending on engine behavior, it may or may not yield the second value after the failed one
300        // We only assert that it is either Some(valid) or None, but it must not be an error
301        if let Some(val) = next {
302            assert_eq!(val, serde_json::json!({"ok": true}));
303        }
304    }
305
306    #[derive(Debug, Clone, Deserialize)]
307    #[expect(dead_code)]
308    struct OrderEvent {
309        item: String,
310        quantity: u32,
311        prepaid: bool,
312    }
313
314    #[test]
315    fn decode_order_events() {
316        let inputs = [
317            r#"{"item":"Apple Watch Series 9","quantity":2,"prepaid":true}"#,
318            concat!("\n", r#"{"item":"extra item","quantity":0,"prepaid":true}"#),
319            concat!(
320                "\n",
321                r#"{"item":"Gaming Mousepad XL","quantity":1,"prepaid":false}"#
322            ),
323            concat!(
324                "\n",
325                r#"{"item":"Noise Cancelling Headphones","quantity":3,"prepaid":true}"#
326            ),
327            concat!(
328                "\n",
329                r#"{"item":"Ergonomic Chair","quantity":1,"prepaid":true}"#
330            ),
331            concat!(
332                "\n",
333                r#"{"item":"extra item","quantity":6,"prepaid":false}"#
334            ),
335            concat!(
336                "\n",
337                r#"{"item":"LED Monitor 27\"","quantity":4,"prepaid":false}"#
338            ),
339            concat!(
340                "\n",
341                r#"{"item":"Smartphone Stand","quantity":6,"prepaid":false}"#
342            ),
343            concat!(
344                "\n",
345                r#"{"item":"Mechanical Keyboard","quantity":2,"prepaid":true}"#
346            ),
347            concat!(
348                "\n",
349                r#"{"item":"extra item","quantity":12,"prepaid":true}"#
350            ),
351            concat!(
352                "\n",
353                r#"{"item":"Laptop Sleeve 15.6\"","quantity":3,"prepaid":false}"#
354            ),
355            concat!(
356                "\n",
357                r#"{"item":"USB-C Docking Station","quantity":1,"prepaid":true}"#
358            ),
359            concat!(
360                "\n",
361                r#"{"item":"Wireless Presenter","quantity":1,"prepaid":false}"#
362            ),
363            concat!(
364                "\n",
365                r#"{"item":"extra item","quantity":18,"prepaid":false}"#
366            ),
367            concat!(
368                "\n",
369                r#"{"item":"Foldable Desk Lamp","quantity":5,"prepaid":true}"#
370            ),
371            concat!(
372                "\n",
373                r#"{"item":"Portable SSD 1TB","quantity":2,"prepaid":true}"#
374            ),
375            concat!(
376                "\n",
377                r#"{"item":"Webcam Cover Slide","quantity":10,"prepaid":false}"#
378            ),
379            concat!(
380                "\n",
381                r#"{"item":"extra item","quantity":24,"prepaid":true}"#
382            ),
383            concat!(
384                "\n",
385                r#"{"item":"Bluetooth Speaker","quantity":2,"prepaid":false}"#
386            ),
387            concat!(
388                "\n",
389                r#"{"item":"Fitness Tracker Band","quantity":4,"prepaid":true}"#
390            ),
391            concat!(
392                "\n",
393                r#"{"item":"Laser Pointer","quantity":1,"prepaid":false}"#
394            ),
395            concat!(
396                "\n",
397                r#"{"item":"extra item","quantity":30,"prepaid":false}"#
398            ),
399            concat!(
400                "\n",
401                r#"{"item":"Conference Mic","quantity":2,"prepaid":true}"#
402            ),
403            concat!(
404                "\n",
405                r#"{"item":"Noise-Absorbing Panels","quantity":12,"prepaid":false}"#
406            ),
407            concat!(
408                "\n",
409                r#"{"item":"Desk Organizer Set","quantity":1,"prepaid":true}"#
410            ),
411            concat!(
412                "\n",
413                r#"{"item":"extra item","quantity":36,"prepaid":true}"#
414            ),
415            concat!(
416                "\n",
417                r#"{"item":"Whiteboard Eraser Pack","quantity":6,"prepaid":false}"#
418            ),
419            concat!(
420                "\n",
421                r#"{"item":"Travel Power Adapter","quantity":2,"prepaid":true}"#
422            ),
423        ];
424
425        let mut event_count = 0;
426        let mut unique_events = HashSet::new();
427        let mut dec: JsonDecoder<OrderEvent> = JsonDecoder::new();
428
429        for input in inputs {
430            let mut buf = BytesMut::from(input);
431            while !buf.is_empty() {
432                if let Some(event) = dec.decode(&mut buf).unwrap() {
433                    unique_events.insert(event.item);
434                    event_count += 1;
435                }
436            }
437        }
438        while let Some(event) = dec.decode_eof(&mut Default::default()).unwrap() {
439            unique_events.insert(event.item);
440            event_count += 1;
441        }
442        assert_eq!(28, event_count);
443        assert_eq!(22, unique_events.len());
444    }
445
446    #[test]
447    fn decode_order_events_random_chunks() {
448        let raw_input = [
449            r##"{"item":"Apple Watch Series 9","quantity":2,"prepaid":true}"##,
450            r##"{"item":"extra item","quantity":0,"prepaid":true}"##,
451            r##"{"item":"Gaming Mousepad XL","quantity":1,"prepaid":false}"##,
452            r##"{"item":"Noise Cancelling Headphones","quantity":3,"prepaid":true}"##,
453            r##"{"item":"Ergonomic Chair","quantity":1,"prepaid":true}"##,
454            r##"{"item":"extra item","quantity":6,"prepaid":false}"##,
455            r##"{"item":"LED Monitor 27\"","quantity":4,"prepaid":false}"##,
456            r##"{"item":"Smartphone Stand","quantity":6,"prepaid":false}"##,
457            r##"{"item":"Mechanical Keyboard","quantity":2,"prepaid":true}"##,
458            r##"{"item":"extra item","quantity":12,"prepaid":true}"##,
459            r##"{"item":"Laptop Sleeve 15.6\"","quantity":3,"prepaid":false}"##,
460            r##"{"item":"USB-C Docking Station","quantity":1,"prepaid":true}"##,
461            r##"{"item":"Wireless Presenter","quantity":1,"prepaid":false}"##,
462            r##"{"item":"extra item","quantity":18,"prepaid":false}"##,
463            r##"{"item":"Foldable Desk Lamp","quantity":5,"prepaid":true}"##,
464            r##"{"item":"Portable SSD 1TB","quantity":2,"prepaid":true}"##,
465            r##"{"item":"Webcam Cover Slide","quantity":10,"prepaid":false}"##,
466            r##"{"item":"extra item","quantity":24,"prepaid":true}"##,
467            r##"{"item":"Bluetooth Speaker","quantity":2,"prepaid":false}"##,
468            r##"{"item":"Fitness Tracker Band","quantity":4,"prepaid":true}"##,
469            r##"{"item":"Laser Pointer","quantity":1,"prepaid":false}"##,
470            r##"{"item":"extra item","quantity":30,"prepaid":false}"##,
471            r##"{"item":"Conference Mic","quantity":2,"prepaid":true}"##,
472            r##"{"item":"Noise-Absorbing Panels","quantity":12,"prepaid":false}"##,
473            r##"{"item":"Desk Organizer Set","quantity":1,"prepaid":true}"##,
474            r##"{"item":"extra item","quantity":36,"prepaid":true}"##,
475            r##"{"item":"Whiteboard Eraser Pack","quantity":6,"prepaid":false}"##,
476            r##"{"item":"Travel Power Adapter","quantity":2,"prepaid":true}"##,
477        ]
478        .join("\n");
479
480        // try it 32 times...
481        for _ in 0..32 {
482            let max = raw_input.len();
483            let mut begin = 0;
484
485            let mut event_count = 0;
486            let mut unique_events = HashSet::new();
487            let mut dec: JsonDecoder<OrderEvent> = JsonDecoder::new();
488
489            while begin < max {
490                let end = rand::random_range(begin..=max);
491                let mut buf = BytesMut::from(&raw_input[begin..end]);
492                while !buf.is_empty() {
493                    if let Some(event) = dec.decode(&mut buf).unwrap() {
494                        unique_events.insert(event.item);
495                        event_count += 1;
496                    }
497                }
498                begin = end;
499            }
500
501            while let Some(event) = dec.decode_eof(&mut Default::default()).unwrap() {
502                unique_events.insert(event.item);
503                event_count += 1;
504            }
505
506            assert_eq!(28, event_count);
507            assert_eq!(22, unique_events.len());
508        }
509    }
510}
rama_core/stream/json/codec.rs

rama_core/stream/json/
codec.rs
