From 6bed8deb54cc59ae2826cc1af60eef15a0bbadaa Mon Sep 17 00:00:00 2001 From: "Xinyao (Morry) Niu" Date: Wed, 22 Jul 2026 18:06:09 +0800 Subject: [PATCH 1/2] fix(storage): validate checkpoint source offsets Require every decoded physical line to begin exactly where the previous line ended so tail candidates cannot publish raw records with gaps or overlaps. Cover deterministic chunk boundaries, large declared counts, and multi-record checkpoints that complete across repeated EOF observations. --- src/storage/message_log/decoder.rs | 10 ++ src/storage/message_log/decoder/tests.rs | 198 +++++++++++++++++++++++ 2 files changed, 208 insertions(+) diff --git a/src/storage/message_log/decoder.rs b/src/storage/message_log/decoder.rs index fea757a..78ef3b1 100644 --- a/src/storage/message_log/decoder.rs +++ b/src/storage/message_log/decoder.rs @@ -45,6 +45,7 @@ struct PendingCheckpoint { pub(crate) struct CheckpointDecoder { next_seq: u64, committed_end: u64, + next_line_offset: u64, pending: Option, } @@ -59,6 +60,7 @@ impl CheckpointDecoder { Self { next_seq, committed_end, + next_line_offset: committed_end, pending: None, } } @@ -82,6 +84,11 @@ impl CheckpointDecoder { let source_offset = line_end .checked_sub(line_len) .context("message record offset precedes the start of the file")?; + ensure!( + source_offset == self.next_line_offset, + "message record starts at byte {source_offset}, expected {}", + self.next_line_offset + ); let stored = parse_stored_line(path, &line_with_newline)?; let checkpoint = stored.local.checkpoint.clone(); let expected_index = self @@ -137,6 +144,7 @@ impl CheckpointDecoder { }; if let Some(pending) = self.pending.as_mut() { + self.next_line_offset = line_end; pending.records.push(record); if expected_index + 1 < pending.metadata.count { return Ok(DecodeResult::NeedMore); @@ -152,12 +160,14 @@ impl CheckpointDecoder { .checked_add(1) .context("message sequence overflow after singleton checkpoint")?; self.committed_end = line_end; + self.next_line_offset = line_end; return Ok(DecodeResult::Checkpoint(CommittedCheckpoint { records: vec![record], committed_end: line_end, })); }; let count = checkpoint.count; + self.next_line_offset = line_end; self.pending = Some(PendingCheckpoint { metadata: checkpoint, records: vec![record], diff --git a/src/storage/message_log/decoder/tests.rs b/src/storage/message_log/decoder/tests.rs index 1719166..fc3c7cb 100644 --- a/src/storage/message_log/decoder/tests.rs +++ b/src/storage/message_log/decoder/tests.rs @@ -121,6 +121,26 @@ fn decoder_buffers_only_the_current_checkpoint_and_preserves_raw_line() { assert_eq!(checkpoint.committed_end, end); } +#[test] +fn decoder_does_not_reserve_the_declared_checkpoint_count() { + let first = line("m1", 0, 1_000_000, "first"); + let mut decoder = CheckpointDecoder::default(); + + assert!(matches!( + decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64, + ) + .unwrap(), + DecodeResult::NeedMore + )); + let pending = decoder.pending.as_ref().unwrap(); + assert_eq!(pending.records.len(), 1); + assert!(pending.records.capacity() < 16); +} + #[test] fn decoder_validates_a_candidate_group_from_a_tail_offset() { let first = line("m900", 0, 2, "tail call"); @@ -243,6 +263,65 @@ struct CountingReader { reads: Arc, } +struct DeterministicChunkReader { + bytes: Vec, + offset: usize, + step: usize, +} + +impl AsyncRead for DeterministicChunkReader { + fn poll_read( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + buffer: &mut ReadBuf<'_>, + ) -> Poll> { + if self.offset == self.bytes.len() { + return Poll::Ready(Ok(())); + } + let chunk = (self.step.wrapping_mul(17).wrapping_add(5) % 23) + 1; + self.step = self.step.wrapping_add(1); + let end = self + .offset + .saturating_add(chunk.min(buffer.remaining())) + .min(self.bytes.len()); + buffer.put_slice(&self.bytes[self.offset..end]); + self.offset = end; + Poll::Ready(Ok(())) + } +} + +#[tokio::test] +async fn deterministic_chunk_boundaries_preserve_checkpoint_and_raw_order() { + let expected = [ + line("m1", 0, 1, "prefix"), + line("m2", 0, 3, "assistant"), + line("m3", 1, 3, "first result"), + line("m4", 2, 3, "second result"), + line("m5", 0, 1, "suffix"), + ] + .concat(); + let source = DeterministicChunkReader { + bytes: expected.clone(), + offset: 0, + step: 0, + }; + let mut reader = CommittedCheckpointReader::new( + BufReader::with_capacity(29, source), + PathBuf::from("messages.jsonl"), + ); + let mut refs = Vec::new(); + let mut raw = Vec::new(); + while let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await.unwrap() { + for record in checkpoint.records { + refs.push(record.trajectory.message_ref); + raw.extend(record.raw); + } + } + + assert_eq!(refs, ["m1", "m2", "m3", "m4", "m5"]); + assert_eq!(raw, expected); +} + impl AsyncRead for CountingReader { fn poll_read( mut self: Pin<&mut Self>, @@ -311,3 +390,122 @@ async fn partial_line_can_complete_after_a_later_read() { }; assert_eq!(checkpoint.records[0].raw, complete); } + +#[tokio::test] +async fn partial_multi_record_checkpoint_commits_once_after_multiple_eofs() { + let prefix = line("m1", 0, 1, "prefix"); + let first = line("m2", 0, 3, "assistant"); + let second = line("m3", 1, 3, "first result"); + let third = line("m4", 2, 3, "second result"); + let second_split = second.len() / 2; + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("messages.jsonl"); + let initial = [prefix.as_slice(), first.as_slice(), &second[..second_split]].concat(); + tokio::fs::write(&path, &initial).await.unwrap(); + let file = tokio::fs::File::open(&path).await.unwrap(); + let mut reader = CommittedCheckpointReader::new(BufReader::new(file), path.clone()); + + let DecodeResult::Checkpoint(committed_prefix) = reader.next_checkpoint().await.unwrap() else { + panic!("the singleton prefix should be committed"); + }; + assert_eq!(committed_prefix.records[0].trajectory.message_ref, "m1"); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + use tokio::io::AsyncWriteExt; + let mut writer = tokio::fs::OpenOptions::new() + .append(true) + .open(&path) + .await + .unwrap(); + writer.write_all(&second[second_split..]).await.unwrap(); + writer.flush().await.unwrap(); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); + assert_eq!(reader.committed_end(), prefix.len() as u64); + + writer.write_all(&third).await.unwrap(); + writer.flush().await.unwrap(); + let DecodeResult::Checkpoint(checkpoint) = reader.next_checkpoint().await.unwrap() else { + panic!("the completed group should commit after the final append"); + }; + assert_eq!( + checkpoint + .records + .iter() + .map(|record| record.trajectory.message_ref.as_str()) + .collect::>(), + ["m2", "m3", "m4"] + ); + assert_eq!(checkpoint.records[0].source_offset, prefix.len() as u64); + assert_eq!( + checkpoint.committed_end, + (prefix.len() + first.len() + second.len() + third.len()) as u64 + ); + assert!(matches!( + reader.next_checkpoint().await.unwrap(), + DecodeResult::NeedMore + )); +} + +#[test] +fn decoder_rejects_non_contiguous_and_underflowing_line_offsets() { + let first = line("m1", 0, 1, "first"); + let mut decoder = CheckpointDecoder::default(); + let error = decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64 + 7, + ) + .unwrap_err(); + assert!(error.to_string().contains("starts at byte 7, expected 0")); + + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), first.clone(), 1) + .unwrap_err(); + assert!( + error + .to_string() + .contains("record offset precedes the start of the file") + ); + + let DecodeResult::Checkpoint(_) = decoder + .push_complete_line( + Path::new("messages.jsonl"), + first.clone(), + first.len() as u64, + ) + .unwrap() + else { + panic!("the contiguous first line should commit"); + }; + let second = line("m2", 0, 1, "second"); + let overlap_end = first.len() as u64 + second.len() as u64 - 1; + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), second, overlap_end) + .unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "message record starts at byte {}, expected {}", + first.len() - 1, + first.len() + ) + ); +} From c9a00e09108894a4e08bd35685691a9a3a17676c Mon Sep 17 00:00:00 2001 From: "Xinyao (Morry) Niu" Date: Wed, 22 Jul 2026 18:12:21 +0800 Subject: [PATCH 2/2] fix(storage): reject checkpoint sequence overflow early Validate the sequence after a complete checkpoint before retaining its first record. This prevents an unrepresentable next sequence from surfacing only after the checkpoint has reached its commit boundary. Keep commit_pending transactional as a defensive invariant and cover the boundary with a retry test that proves pending and byte offsets remain untouched. --- src/storage/message_log/decoder.rs | 13 +++++++--- src/storage/message_log/decoder/tests.rs | 33 ++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/src/storage/message_log/decoder.rs b/src/storage/message_log/decoder.rs index 78ef3b1..435d863 100644 --- a/src/storage/message_log/decoder.rs +++ b/src/storage/message_log/decoder.rs @@ -129,8 +129,8 @@ impl CheckpointDecoder { usize::try_from(checkpoint.count) .context("message checkpoint count does not fit in memory")?; self.next_seq - .checked_add(checkpoint.count - 1) - .context("message sequence overflow inside checkpoint")?; + .checked_add(checkpoint.count) + .context("message sequence overflow after checkpoint")?; } let expected_seq = self .next_seq @@ -182,16 +182,21 @@ impl CheckpointDecoder { fn commit_pending(&mut self, committed_end: u64) -> Result { let pending = self .pending - .take() + .as_ref() .context("checkpoint decoder has no pending checkpoint to commit")?; ensure!( pending.records.len() as u64 == pending.metadata.count, "checkpoint decoder committed an incomplete checkpoint" ); - self.next_seq = self + let next_seq = self .next_seq .checked_add(pending.metadata.count) .context("message sequence overflow after checkpoint")?; + let pending = self + .pending + .take() + .context("checkpoint decoder lost its pending checkpoint")?; + self.next_seq = next_seq; self.committed_end = committed_end; Ok(DecodeResult::Checkpoint(CommittedCheckpoint { records: pending.records, diff --git a/src/storage/message_log/decoder/tests.rs b/src/storage/message_log/decoder/tests.rs index fc3c7cb..5c43b6c 100644 --- a/src/storage/message_log/decoder/tests.rs +++ b/src/storage/message_log/decoder/tests.rs @@ -141,6 +141,39 @@ fn decoder_does_not_reserve_the_declared_checkpoint_count() { assert!(pending.records.capacity() < 16); } +#[test] +fn sequence_overflow_rejects_checkpoint_before_advancing_decoder_state() { + let initial_seq = u64::MAX - 1; + let initial_offset = 97_u64; + let overflowing = line(&format!("m{initial_seq}"), 0, 2, "first"); + let overflowing_end = initial_offset + overflowing.len() as u64; + let mut decoder = CheckpointDecoder::new(initial_seq, initial_offset); + + let error = decoder + .push_complete_line(Path::new("messages.jsonl"), overflowing, overflowing_end) + .unwrap_err(); + assert!( + error + .to_string() + .contains("message sequence overflow after checkpoint") + ); + assert_eq!(decoder.next_seq, initial_seq); + assert_eq!(decoder.committed_end, initial_offset); + assert_eq!(decoder.next_line_offset, initial_offset); + assert!(decoder.pending.is_none()); + + let valid = line(&format!("m{initial_seq}"), 0, 1, "retry"); + let valid_end = initial_offset + valid.len() as u64; + let DecodeResult::Checkpoint(checkpoint) = decoder + .push_complete_line(Path::new("messages.jsonl"), valid, valid_end) + .unwrap() + else { + panic!("decoder should remain reusable after rejecting the overflow"); + }; + assert_eq!(checkpoint.committed_end, valid_end); + assert_eq!(decoder.next_seq, u64::MAX); +} + #[test] fn decoder_validates_a_candidate_group_from_a_tail_offset() { let first = line("m900", 0, 2, "tail call");