1use 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}