Skip to main content

ursula_stream/
record_index.rs

1//! Exact retained record-ordinal to canonical-offset boundaries for JSON streams.
2
3use serde::Deserialize;
4use serde::Serialize;
5
6#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
7pub struct StreamRecordIndex {
8    first_record: u64,
9    record_offsets: Vec<u64>,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13pub struct StreamRecordRange {
14    pub first_record: u64,
15    pub next_record: u64,
16}
17
18#[derive(Debug)]
19pub(crate) struct PreparedRecordAppend {
20    range: StreamRecordRange,
21    record_offsets: Vec<u64>,
22}
23
24impl PreparedRecordAppend {
25    pub(crate) fn range(&self) -> StreamRecordRange {
26        self.range
27    }
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum RecordIndexError {
32    InvalidBoundaries,
33    ArithmeticOverflow,
34    RecordGone { first_record: u64, next_record: u64 },
35    RecordBeyondTail { next_record: u64 },
36    OffsetNotRecordBoundary,
37}
38
39pub fn is_json_record_content_type(content_type: &str) -> bool {
40    content_type
41        .split(';')
42        .next()
43        .is_some_and(|value| value.trim().eq_ignore_ascii_case("application/json"))
44}
45
46pub fn canonical_json_record_ends(
47    content_type: &str,
48    payload: &[u8],
49) -> Result<Vec<u64>, RecordIndexError> {
50    if !is_json_record_content_type(content_type) {
51        return Ok(Vec::new());
52    }
53    if payload.is_empty() {
54        return Ok(Vec::new());
55    }
56    if payload.last() != Some(&b'\n') {
57        return Err(RecordIndexError::InvalidBoundaries);
58    }
59    payload
60        .iter()
61        .enumerate()
62        .filter_map(|(index, byte)| (*byte == b'\n').then_some(index + 1))
63        .map(|end| u64::try_from(end).map_err(|_| RecordIndexError::ArithmeticOverflow))
64        .collect()
65}
66
67impl StreamRecordIndex {
68    pub fn new() -> Self {
69        Self::default()
70    }
71
72    pub fn restore(
73        first_record: u64,
74        record_offsets: Vec<u64>,
75        retained_offset: u64,
76        tail_offset: u64,
77    ) -> Result<Self, RecordIndexError> {
78        let index = Self {
79            first_record,
80            record_offsets,
81        };
82        index.validate(retained_offset, tail_offset)?;
83        Ok(index)
84    }
85
86    pub fn range(&self) -> Result<StreamRecordRange, RecordIndexError> {
87        let retained = u64::try_from(self.record_offsets.len())
88            .map_err(|_| RecordIndexError::ArithmeticOverflow)?;
89        let next_record = self
90            .first_record
91            .checked_add(retained)
92            .ok_or(RecordIndexError::ArithmeticOverflow)?;
93        Ok(StreamRecordRange {
94            first_record: self.first_record,
95            next_record,
96        })
97    }
98
99    pub fn record_offsets(&self) -> &[u64] {
100        &self.record_offsets
101    }
102
103    pub fn append_relative_ends(
104        &mut self,
105        base_offset: u64,
106        payload_len: u64,
107        relative_ends: &[u64],
108    ) -> Result<StreamRecordRange, RecordIndexError> {
109        let prepared = self.prepare_append(base_offset, payload_len, relative_ends)?;
110        Ok(self.commit_append(prepared))
111    }
112
113    pub(crate) fn prepare_append(
114        &self,
115        base_offset: u64,
116        payload_len: u64,
117        relative_ends: &[u64],
118    ) -> Result<PreparedRecordAppend, RecordIndexError> {
119        validate_relative_ends(payload_len, relative_ends)?;
120        let record_start = self.range()?.next_record;
121        let appended =
122            u64::try_from(relative_ends.len()).map_err(|_| RecordIndexError::ArithmeticOverflow)?;
123        let record_next = record_start
124            .checked_add(appended)
125            .ok_or(RecordIndexError::ArithmeticOverflow)?;
126        if self
127            .record_offsets
128            .last()
129            .is_some_and(|last| *last >= base_offset)
130        {
131            return Err(RecordIndexError::InvalidBoundaries);
132        }
133        let mut starts = Vec::with_capacity(relative_ends.len());
134        let mut previous_end = 0;
135        for end in relative_ends {
136            let start_offset = base_offset
137                .checked_add(previous_end)
138                .ok_or(RecordIndexError::ArithmeticOverflow)?;
139            starts.push(start_offset);
140            previous_end = *end;
141        }
142        Ok(PreparedRecordAppend {
143            range: StreamRecordRange {
144                first_record: record_start,
145                next_record: record_next,
146            },
147            record_offsets: starts,
148        })
149    }
150
151    pub(crate) fn commit_append(&mut self, prepared: PreparedRecordAppend) -> StreamRecordRange {
152        self.record_offsets.extend(prepared.record_offsets);
153        prepared.range
154    }
155
156    pub fn offset_for(&self, record: u64, tail_offset: u64) -> Result<u64, RecordIndexError> {
157        let range = self.range()?;
158        if record < range.first_record {
159            return Err(RecordIndexError::RecordGone {
160                first_record: range.first_record,
161                next_record: range.next_record,
162            });
163        }
164        if record > range.next_record {
165            return Err(RecordIndexError::RecordBeyondTail {
166                next_record: range.next_record,
167            });
168        }
169        if record == range.next_record {
170            return Ok(tail_offset);
171        }
172        let relative = record
173            .checked_sub(range.first_record)
174            .ok_or(RecordIndexError::ArithmeticOverflow)?;
175        let index = usize::try_from(relative).map_err(|_| RecordIndexError::ArithmeticOverflow)?;
176        self.record_offsets
177            .get(index)
178            .copied()
179            .ok_or(RecordIndexError::InvalidBoundaries)
180    }
181
182    pub fn record_for_offset(
183        &self,
184        offset: u64,
185        tail_offset: u64,
186    ) -> Result<u64, RecordIndexError> {
187        let range = self.range()?;
188        if offset == tail_offset {
189            return Ok(range.next_record);
190        }
191        let relative = self
192            .record_offsets
193            .binary_search(&offset)
194            .map_err(|_| RecordIndexError::OffsetNotRecordBoundary)?;
195        range
196            .first_record
197            .checked_add(u64::try_from(relative).map_err(|_| RecordIndexError::ArithmeticOverflow)?)
198            .ok_or(RecordIndexError::ArithmeticOverflow)
199    }
200
201    pub fn retain_from_offset(
202        &mut self,
203        retained_offset: u64,
204        tail_offset: u64,
205    ) -> Result<u64, RecordIndexError> {
206        let removed = if retained_offset == tail_offset {
207            self.record_offsets.len()
208        } else {
209            self.record_offsets
210                .binary_search(&retained_offset)
211                .map_err(|_| RecordIndexError::OffsetNotRecordBoundary)?
212        };
213        let removed_u64 =
214            u64::try_from(removed).map_err(|_| RecordIndexError::ArithmeticOverflow)?;
215        self.first_record = self
216            .first_record
217            .checked_add(removed_u64)
218            .ok_or(RecordIndexError::ArithmeticOverflow)?;
219        self.record_offsets.drain(..removed);
220        Ok(self.first_record)
221    }
222
223    pub fn validate(&self, retained_offset: u64, tail_offset: u64) -> Result<(), RecordIndexError> {
224        let _ = self.range()?;
225        if retained_offset > tail_offset {
226            return Err(RecordIndexError::InvalidBoundaries);
227        }
228        if self.record_offsets.is_empty() {
229            return (retained_offset == tail_offset)
230                .then_some(())
231                .ok_or(RecordIndexError::InvalidBoundaries);
232        }
233        if self.record_offsets.first().copied() != Some(retained_offset)
234            || self
235                .record_offsets
236                .iter()
237                .any(|offset| *offset >= tail_offset)
238            || self.record_offsets.windows(2).any(|pair| {
239                let [left, right] = pair else {
240                    return true;
241                };
242                left >= right
243            })
244        {
245            return Err(RecordIndexError::InvalidBoundaries);
246        }
247        Ok(())
248    }
249}
250
251fn validate_relative_ends(payload_len: u64, relative_ends: &[u64]) -> Result<(), RecordIndexError> {
252    if payload_len == 0 {
253        return relative_ends
254            .is_empty()
255            .then_some(())
256            .ok_or(RecordIndexError::InvalidBoundaries);
257    }
258    if relative_ends.last().copied() != Some(payload_len)
259        || relative_ends.first().copied() == Some(0)
260        || relative_ends.windows(2).any(|pair| {
261            let [left, right] = pair else {
262                return true;
263            };
264            left >= right
265        })
266    {
267        return Err(RecordIndexError::InvalidBoundaries);
268    }
269    Ok(())
270}
271
272#[cfg(test)]
273mod tests {
274    use super::RecordIndexError;
275    use super::StreamRecordIndex;
276    use super::StreamRecordRange;
277    use super::canonical_json_record_ends;
278
279    #[test]
280    fn append_maps_contiguous_ordinals_to_exact_offsets() {
281        let mut index = StreamRecordIndex::new();
282        assert_eq!(
283            index.append_relative_ends(0, 18, &[9, 18]),
284            Ok(StreamRecordRange {
285                first_record: 0,
286                next_record: 2,
287            })
288        );
289        assert_eq!(index.record_offsets(), &[0, 9]);
290        assert_eq!(index.offset_for(0, 18), Ok(0));
291        assert_eq!(index.offset_for(1, 18), Ok(9));
292        assert_eq!(index.offset_for(2, 18), Ok(18));
293
294        assert_eq!(
295            index.append_relative_ends(18, 9, &[9]),
296            Ok(StreamRecordRange {
297                first_record: 2,
298                next_record: 3,
299            })
300        );
301        assert_eq!(index.record_offsets(), &[0, 9, 18]);
302        assert_eq!(index.offset_for(3, 27), Ok(27));
303    }
304
305    #[test]
306    fn retention_drops_offsets_without_renumbering() {
307        let mut index = StreamRecordIndex::new();
308        index
309            .append_relative_ends(0, 27, &[9, 18, 27])
310            .expect("append boundaries");
311        assert_eq!(index.retain_from_offset(18, 27), Ok(2));
312        assert_eq!(index.record_offsets(), &[18]);
313        assert_eq!(
314            index.offset_for(1, 27),
315            Err(RecordIndexError::RecordGone {
316                first_record: 2,
317                next_record: 3,
318            })
319        );
320        assert_eq!(index.offset_for(2, 27), Ok(18));
321    }
322
323    #[test]
324    fn restore_rejects_misaligned_or_non_monotonic_offsets() {
325        assert_eq!(
326            StreamRecordIndex::restore(0, vec![1, 9], 0, 18),
327            Err(RecordIndexError::InvalidBoundaries)
328        );
329        assert_eq!(
330            StreamRecordIndex::restore(0, vec![0, 0], 0, 18),
331            Err(RecordIndexError::InvalidBoundaries)
332        );
333        assert!(StreamRecordIndex::restore(2, vec![18], 18, 27).is_ok());
334    }
335
336    #[test]
337    fn relative_ends_must_cover_the_payload_exactly() {
338        let mut index = StreamRecordIndex::new();
339        assert_eq!(
340            index.append_relative_ends(0, 18, &[9]),
341            Err(RecordIndexError::InvalidBoundaries)
342        );
343        assert_eq!(
344            index.append_relative_ends(0, 18, &[9, 9, 18]),
345            Err(RecordIndexError::InvalidBoundaries)
346        );
347        assert_eq!(
348            index.append_relative_ends(0, 0, &[0]),
349            Err(RecordIndexError::InvalidBoundaries)
350        );
351    }
352
353    #[test]
354    fn canonical_json_payload_exposes_each_ndjson_boundary() {
355        assert_eq!(
356            canonical_json_record_ends(
357                "application/json; charset=utf-8",
358                b"{\"a\":1}\n{\"b\":2}\n"
359            ),
360            Ok(vec![8, 16])
361        );
362        assert_eq!(
363            canonical_json_record_ends("application/octet-stream", b"x"),
364            Ok(vec![])
365        );
366        assert_eq!(
367            canonical_json_record_ends("application/json", b"{}"),
368            Err(RecordIndexError::InvalidBoundaries)
369        );
370    }
371}