spate_json/framing.rs
1//! Streaming NDJSON record framing.
2//!
3//! [`NdjsonFramer`] is the JSON connector's
4//! [`RecordFramer`](spate_core::framing::RecordFramer): it cuts a byte *stream*
5//! (an object streamed by `spate-s3`, an HTTP body, a file tail) into
6//! newline-delimited records, so a streaming source can hand the JSON
7//! deserializer one document per payload. Wire it into a streaming source at
8//! assembly time:
9//!
10//! ```no_run
11//! # use spate_json::NdjsonFramer;
12//! # struct S3Source;
13//! # impl S3Source {
14//! # fn with_framer<F: Fn() -> Box<dyn spate_core::framing::RecordFramer> + Send + Sync>(self, _: F) -> Self { self }
15//! # }
16//! # let source = S3Source;
17//! let max_record_bytes = 64 << 20;
18//! let source = source.with_framer(move || Box::new(NdjsonFramer::new(max_record_bytes)));
19//! ```
20//!
21//! This is the source-side counterpart to the deserializer's in-memory
22//! [`ndjson`](crate::JsonFraming::Ndjson) framing, which splits a whole payload
23//! already resident in RAM (a Kafka message). `NdjsonFramer` is bounded and
24//! chunk-fed instead, for byte streams that never fit in memory. A source that
25//! frames one record per payload pairs with the deserializer's `single` mode
26//! (the [`PerRecord`](spate_core::framing::FramingContract::PerRecord) contract).
27
28use spate_core::framing::RecordFramer;
29use std::collections::VecDeque;
30use std::io;
31
32/// Newline-delimited record framing (NDJSON / JSON Lines) as a streaming
33/// [`RecordFramer`].
34///
35/// The framer knows nothing about JSON beyond the newline convention NDJSON
36/// guarantees: a JSON document never contains an unescaped newline, so every
37/// `\n` is an unambiguous record boundary.
38///
39/// Framing rules (pinned; changing any of them changes record indexes, which
40/// is a resume-compatibility break for sources that checkpoint by index):
41///
42/// - Records are split on `\n`. Exactly one trailing `\r` is stripped (CRLF
43/// input); nothing else is trimmed.
44/// - Lines that are empty or all-ASCII-whitespace are skipped and do not
45/// consume a record index.
46/// - An unterminated final line (no closing `\n`) is a record.
47#[derive(Debug)]
48pub struct NdjsonFramer {
49 /// Bytes of the current, not-yet-terminated line.
50 partial: Vec<u8>,
51 /// Completed records, in order.
52 ready: VecDeque<Vec<u8>>,
53 /// Decoded bytes seen (metrics).
54 decoded_bytes: u64,
55 /// Upper bound on one record line (decoded bytes). Without it, a stream
56 /// holding no newline at all would buffer unboundedly and abort the
57 /// process instead of failing the pipeline with a policy error.
58 max_record_bytes: usize,
59}
60
61impl NdjsonFramer {
62 /// A line framer bounding each record at `max_record_bytes` decoded bytes.
63 #[must_use]
64 pub fn new(max_record_bytes: usize) -> NdjsonFramer {
65 NdjsonFramer {
66 partial: Vec::new(),
67 ready: VecDeque::new(),
68 decoded_bytes: 0,
69 max_record_bytes,
70 }
71 }
72
73 /// Extend the current line, enforcing the record-size cap before the bytes
74 /// are buffered.
75 fn push_partial(&mut self, bytes: &[u8]) -> io::Result<()> {
76 if self.partial.len() + bytes.len() > self.max_record_bytes {
77 return Err(io::Error::new(
78 io::ErrorKind::InvalidData,
79 format!(
80 "record line exceeds the configured max_record_bytes ({}); \
81 is the stream really newline-delimited?",
82 self.max_record_bytes
83 ),
84 ));
85 }
86 self.partial.extend_from_slice(bytes);
87 Ok(())
88 }
89
90 /// Complete the current line: strip one `\r`, skip whitespace-only.
91 fn complete_line(&mut self) {
92 if self.partial.last() == Some(&b'\r') {
93 self.partial.pop();
94 }
95 if self.partial.iter().all(u8::is_ascii_whitespace) {
96 self.partial.clear();
97 return;
98 }
99 self.ready.push_back(std::mem::take(&mut self.partial));
100 }
101}
102
103impl RecordFramer for NdjsonFramer {
104 fn push(&mut self, bytes: &[u8]) -> io::Result<()> {
105 self.decoded_bytes += bytes.len() as u64;
106 let mut rest = bytes;
107 while let Some(nl) = rest.iter().position(|&b| b == b'\n') {
108 self.push_partial(&rest[..nl])?;
109 self.complete_line();
110 rest = &rest[nl + 1..];
111 }
112 self.push_partial(rest)
113 }
114
115 fn finish(&mut self) -> io::Result<()> {
116 if !self.partial.is_empty() {
117 self.complete_line();
118 }
119 Ok(())
120 }
121
122 fn pop(&mut self) -> Option<Vec<u8>> {
123 self.ready.pop_front()
124 }
125
126 fn decoded_bytes(&self) -> u64 {
127 self.decoded_bytes
128 }
129
130 fn resync_delimiter(&self) -> Option<u8> {
131 Some(b'\n')
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use proptest::prelude::*;
139
140 /// A cap far above anything the tests feed, so only the dedicated cap test
141 /// exercises it.
142 const TEST_CAP: usize = 1 << 20;
143
144 /// Reference implementation of the line-framing rules on a whole byte
145 /// stream, used as the proptest oracle.
146 fn reference_frames(decoded: &[u8]) -> Vec<Vec<u8>> {
147 decoded
148 .split(|&b| b == b'\n')
149 .map(|line| line.strip_suffix(b"\r").unwrap_or(line))
150 .filter(|line| !line.iter().all(u8::is_ascii_whitespace))
151 .map(<[u8]>::to_vec)
152 .collect()
153 }
154
155 /// Feed `chunks` through a fresh `NdjsonFramer` and collect every record.
156 fn frame_all(max: usize, chunks: &[&[u8]]) -> io::Result<Vec<Vec<u8>>> {
157 let mut f = NdjsonFramer::new(max);
158 for chunk in chunks {
159 f.push(chunk)?;
160 }
161 f.finish()?;
162 let mut out = Vec::new();
163 while let Some(r) = f.pop() {
164 out.push(r);
165 }
166 Ok(out)
167 }
168
169 #[test]
170 fn splits_strips_cr_and_skips_blank_lines() {
171 let stream = b"{\"a\":1}\r\n\n \n{\"b\":2}\n\t\r\n{\"c\":3}";
172 let records = frame_all(TEST_CAP, &[stream]).unwrap();
173 assert_eq!(
174 records,
175 vec![
176 b"{\"a\":1}".to_vec(),
177 b"{\"b\":2}".to_vec(),
178 b"{\"c\":3}".to_vec()
179 ],
180 "one CR stripped, whitespace-only lines skipped, unterminated final line kept"
181 );
182 }
183
184 #[test]
185 fn empty_stream_frames_no_records() {
186 assert!(frame_all(TEST_CAP, &[b""]).unwrap().is_empty());
187 }
188
189 #[test]
190 fn decoded_bytes_counts_every_fed_byte() {
191 let mut f = NdjsonFramer::new(TEST_CAP);
192 f.push(b"one\n").unwrap();
193 f.push(b"two").unwrap();
194 f.finish().unwrap();
195 assert_eq!(f.decoded_bytes(), 7);
196 }
197
198 #[test]
199 fn a_line_over_the_record_cap_is_an_error_not_an_allocation() {
200 // The third push carries the (never-terminated) line past the cap; the
201 // error surfaces at the offending push, before the bytes are buffered.
202 let mut f = NdjsonFramer::new(8);
203 f.push(b"1234").unwrap();
204 f.push(b"5678").unwrap();
205 let err = f.push(b"9").unwrap_err();
206 assert!(err.to_string().contains("max_record_bytes"), "{err}");
207
208 // A line exactly at the cap is fine.
209 let records = frame_all(8, &[b"12345678\n"]).unwrap();
210 assert_eq!(records, vec![b"12345678".to_vec()]);
211 }
212
213 #[test]
214 fn record_framer_is_object_safe() {
215 // Compiles only if the trait is dyn-compatible (the seam's contract).
216 let mut framer: Box<dyn RecordFramer> = Box::new(NdjsonFramer::new(TEST_CAP));
217 framer.push(b"x\n").unwrap();
218 framer.finish().unwrap();
219 assert_eq!(framer.pop(), Some(b"x".to_vec()));
220 }
221
222 /// Arbitrary line content: no `\n` (the separator), but everything else
223 /// including `\r` and whitespace runs.
224 fn arb_line() -> impl Strategy<Value = Vec<u8>> {
225 proptest::collection::vec(
226 prop_oneof![
227 any::<u8>().prop_filter("no newline", |b| *b != b'\n'),
228 Just(b' '),
229 Just(b'\t'),
230 Just(b'\r'),
231 ],
232 0..80,
233 )
234 }
235
236 fn arb_stream() -> impl Strategy<Value = Vec<u8>> {
237 (
238 proptest::collection::vec(arb_line(), 0..40),
239 any::<bool>(), // terminated by a final newline or not
240 )
241 .prop_map(|(lines, terminated)| {
242 let mut stream = lines.join(&b'\n');
243 if terminated && !stream.is_empty() {
244 stream.push(b'\n');
245 }
246 stream
247 })
248 }
249
250 proptest! {
251 /// The framer is a pure function of the byte stream: any two chunkings
252 /// of the same bytes yield the reference framing. This determinism is
253 /// what makes resume-by-record-index safe for every streaming source
254 /// that reuses `NdjsonFramer`.
255 #[test]
256 fn line_framing_is_chunking_independent(stream in arb_stream()) {
257 let expected = reference_frames(&stream);
258 for split in 0..=stream.len() {
259 let framed = frame_all(TEST_CAP, &[&stream[..split], &stream[split..]]).unwrap();
260 prop_assert_eq!(&framed, &expected, "split at {}", split);
261 }
262 }
263
264 /// Framing the stream whole yields the records of the bytes through
265 /// any delimiter followed by the records of the bytes after it, each
266 /// framed on its own, under a record cap small enough to fail some
267 /// runs.
268 #[test]
269 fn framing_composes_across_a_cut_after_the_delimiter(
270 stream in arb_stream(),
271 cap in 1..96usize,
272 ) {
273 let delim = NdjsonFramer::new(cap)
274 .resync_delimiter()
275 .expect("NDJSON declares a resync delimiter");
276 let whole = frame_all(cap, &[&stream]).ok();
277 for d in (0..stream.len()).filter(|&d| stream[d] == delim) {
278 let head = frame_all(cap, &[&stream[..=d]]).ok();
279 let tail = frame_all(cap, &[&stream[d + 1..]]).ok();
280 let halves = head.zip(tail).map(|(mut head, tail)| {
281 head.extend(tail);
282 head
283 });
284 prop_assert_eq!(&halves, &whole, "cut after byte {}", d);
285 }
286 }
287 }
288}