diff --git a/core/journal/src/prepare_journal.rs b/core/journal/src/prepare_journal.rs index 802b36d819..6acb94726c 100644 --- a/core/journal/src/prepare_journal.rs +++ b/core/journal/src/prepare_journal.rs @@ -350,9 +350,12 @@ impl PrepareJournal { // `PrepareHeader` projection in consensus builds prepares with // `..Default::default()` so the integrity fields are always 0. // Until a producer computes them, verification here would be - // trivially-passing noise. Without it, a body bit-flip that - // leaves the header valid is replayed silently as corrupt - // state. Committed bytes are meant to be byte-identical across + // trivially-passing noise. The per-message receive gate now rejects + // transit body corruption on a follower before apply, but this + // recovery scan still does not re-verify body integrity read back + // from disk, so an at-rest body bit-flip that leaves the header + // valid is replayed silently. Committed bytes are meant to be + // byte-identical across // replicas (deterministic apply, timestamp replicated not // re-projected), so once the producer computes the integrity fields // they should agree on every node and this check can be turned on diff --git a/core/partitions/src/iggy_index_reader.rs b/core/partitions/src/iggy_index_reader.rs index a105a69821..cd618c19fa 100644 --- a/core/partitions/src/iggy_index_reader.rs +++ b/core/partitions/src/iggy_index_reader.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::iggy_index::{IGGY_INDEX_SIZE, IggyIndex}; +use crate::iggy_index::{IGGY_INDEX_SIZE, IggyIndex, IggyIndexCache}; use bytes::Buf; use compio::fs::{File, OpenOptions}; use compio::io::AsyncReadAtExt; @@ -114,4 +114,37 @@ impl IggyIndexReader { .await?, )) } + + /// Load every whole entry into an [`IggyIndexCache`] for offset / timestamp + /// lower-bound lookups. Reads the whole file in one pass (index files are + /// tiny: one sparse entry per flushed chunk). A trailing partial entry + /// (torn write) is ignored (see [`Self::entry_count`]). + /// + /// # Errors + /// + /// Returns an error if the file metadata or bytes cannot be read. + pub async fn load_all(&self) -> Result { + let count = usize::try_from(self.entry_count().await?) + .map_err(|_| IggyError::CannotReadFileMetadata)?; + if count == 0 { + return Ok(IggyIndexCache::empty()); + } + + // `with_capacity` (len 0): `read_exact_at` fills the spare capacity in + // place and advances the length (see `read_entry_at`). + let buffer = Vec::with_capacity(count * IGGY_INDEX_SIZE); + let (result, buffer): (std::io::Result<()>, Vec) = + self.file.read_exact_at(buffer, 0).await.into(); + result.map_err(|_| IggyError::CannotReadFile)?; + + let mut cache = IggyIndexCache::with_capacity(count); + let mut view = buffer.as_slice(); + for _ in 0..count { + let offset = view.get_u64_le(); + let timestamp = view.get_u64_le(); + let position = view.get_u64_le(); + cache.insert(offset, timestamp, position); + } + Ok(cache) + } } diff --git a/core/partitions/src/iggy_partition.rs b/core/partitions/src/iggy_partition.rs index 8d2c3a4bfc..a32c4ba444 100644 --- a/core/partitions/src/iggy_partition.rs +++ b/core/partitions/src/iggy_partition.rs @@ -57,7 +57,8 @@ use server_common::{ MESSAGE_ALIGN, Message, SegmentStorage, iobuf::{Frozen, Owned}, send_messages2::{ - convert_request_message, decode_prepare_slice, stamp_prepare_for_persistence, + ChecksumMode, convert_request_message, decode_prepare_slice, decode_prepare_slice_trusted, + stamp_prepare_for_persistence, verify_received_send_messages, }, sharding::IggyNamespace, }; @@ -723,8 +724,9 @@ where /// can run the disk read + offset persist off the partition borrow. The /// in-memory journal tier is read here directly (mem reads never yield); /// the disk tier is captured as owned descriptors in [`DiskReadPlan`]. + #[allow(clippy::too_many_lines)] pub(crate) fn build_poll_plan( - &self, + &mut self, consumer: PollingConsumer, args: &PollingArgs, ) -> PollPlan { @@ -815,13 +817,22 @@ where } let (start_segment, start_position) = self.disk_poll_start(&query); + // Cap resident sealed read handles: touch this poll's start segment so + // the LRU keeps the hot set and drops the least-recently-used fd + + // index (a no-op for the active segment, whose handle never caches). + self.log.touch_sealed_read_state(start_segment); // Snapshot only the segments the disk walk visits (`start_segment..`), - // so `start_position` applies to the first snapshotted segment. + // so `start_position` applies to the first snapshotted segment. A sealed + // segment carries its shared read-state handle (fd + sparse index) so + // the off-borrow read reuses (or fills) it; the active segment opens + // fresh and resolves from its resident index. let segments = self.log.segments()[start_segment..] .iter() - .map(|segment| DiskSegment { + .zip(self.log.sealed_read_state()[start_segment..].iter()) + .map(|(segment, read_state)| DiskSegment { start_offset: segment.start_offset, persisted: segment.size.as_bytes_u64(), + read_state: segment.sealed.then(|| Rc::clone(read_state)), }) .collect(); let disk = DiskReadPlan { @@ -1119,7 +1130,13 @@ where ); let message = if message.header().operation == Operation::SendMessages { - match convert_request_message(namespace, message) { + // Skip the batch-checksum pass: on the partition ingest path + // nothing reads it before `stamp_prepare_for_persistence` + // recomputes it over the stamped header. An already-canonical + // batch (native v2, or the plane's pre-encrypt convert output) + // returns early above, so Skip only affects the legacy + // transcode, whose output goes straight to project/stamp. + match convert_request_message(namespace, message, ChecksumMode::Skip) { Ok(message) => message, Err(error) => { emit_partition_diag( @@ -1502,6 +1519,32 @@ where ); } } + // First blob-integrity check on the replicated path. The consensus + // layer never validates the body (PrepareHeader integrity fields are + // inert zeros) and the batch checksum is recomputed locally at stamp, + // so a follower must verify each message's stamp-invariant per-message + // checksum before journaling transit bytes. Follower-only: the primary + // (and single-node self-replicate) produced these bytes and already + // checked the client batch at ingest, so they must not pay this pass. + // Fail closed on mismatch - drop without journaling, forwarding, or + // acking; the primary retransmits on prepare-timeout. + if is_backup + && header.operation == Operation::SendMessages + && let Err(error) = verify_received_send_messages(message.as_slice()) + { + emit_partition_diag( + tracing::Level::WARN, + &PartitionDiagEvent::new( + self.diag_ctx(), + "rejecting replicated send_messages: per-message checksum mismatch", + ) + .with_operation(header.operation) + .with_op(header.op) + .with_error(error.to_string()), + ); + return; + } + // Durability-before-ack: clone for chain-replicate, forward only // AFTER apply_replicated_operation persists. Forward-first would // give downstream an op whose WAL entry we never wrote, that violates @@ -1911,11 +1954,15 @@ where } continue; } - // A resident committed SendMessages entry decoded once at append - // (the offset index) with its checksum stamped over these exact - // bytes, so it must decode again here. Guard the invariant for a - // future disk read-back path that could make decode fallible. - let Ok(batch) = decode_prepare_slice(entry.as_slice()) else { + // Resident committed SendMessages entry: this node stamped it + // in `append_messages` (recomputing the batch checksum over these + // exact bytes), so a validating re-decode would only re-hash ~1 + // MiB to confirm our own write. Trust the structural decode; the + // batch-checksum recompute belongs at network ingress (repair + // validation + the follower receive gate), not on locally-stamped + // bytes. Guard the invariant for a future disk read-back path that + // could make decode fallible. + let Ok(batch) = decode_prepare_slice_trusted(entry.as_slice()) else { tracing::error!( target: "iggy.partitions.diag", namespace_raw = self.namespace().inner(), @@ -2315,8 +2362,11 @@ where let Some(entry) = self.log.journal().inner.entry(prepare_header).await else { return Err(IggyError::InvalidCommand); }; - let batch = - decode_prepare_slice(entry.as_slice()).map_err(|_| IggyError::InvalidCommand)?; + // Trusted (no batch-hash): the entry was read back from this replica's + // own journal, where it was stamped/validated at append; only header + // stats are needed, so re-hashing the ~1 MiB blob is redundant. + let batch = decode_prepare_slice_trusted(entry.as_slice()) + .map_err(|_| IggyError::InvalidCommand)?; let message_count = batch.message_count(); if message_count == 0 { return Ok(None); @@ -2788,12 +2838,10 @@ where let mut deleted_messages = 0u64; for _ in 0..removable { // The removable run is always a prefix (oldest first), so the next - // victim is index 0 once the previous one is gone. - let segment = self.log.segments_mut().remove(0); - let mut storage = self.log.storages_mut().remove(0); - self.log.indexes_mut().remove(0); - self.log.messages_writers_mut().remove(0); - self.log.index_writers_mut().remove(0); + // victim is the front once the previous one is gone. + let Some((segment, mut storage)) = self.log.retire_front() else { + break; + }; let (messages_path, index_path) = storage.segment_and_index_paths(); let _ = storage.shutdown(); @@ -2947,11 +2995,9 @@ where // Drain every segment (including the active one) and unlink its files. let segment_count = self.log.segments().len(); for _ in 0..segment_count { - self.log.segments_mut().remove(0); - let mut storage = self.log.storages_mut().remove(0); - self.log.indexes_mut().remove(0); - self.log.messages_writers_mut().remove(0); - self.log.index_writers_mut().remove(0); + let Some((_, mut storage)) = self.log.retire_front() else { + break; + }; let (messages_path, index_path) = storage.segment_and_index_paths(); let _ = storage.shutdown(); @@ -3433,7 +3479,7 @@ fn nth_oldest_sealed_end(segments: &[Segment], count: u32) -> Option { #[cfg(test)] mod tests { use super::*; - use crate::poll_plan::DiskReadOutcome; + use crate::poll_plan::{DiskReadOutcome, SealedSegmentHandle}; use bytes::Bytes; use compio::io::AsyncWriteAtExt; use consensus::LocalPipeline; @@ -3889,10 +3935,12 @@ mod tests { DiskSegment { start_offset: 0, persisted: 512, + read_state: None, }, DiskSegment { start_offset: 5, persisted: later_len, + read_state: None, }, ], start_position: 0, @@ -3972,10 +4020,12 @@ mod tests { DiskSegment { start_offset: 0, persisted: corrupt_len, + read_state: None, }, DiskSegment { start_offset: 5, persisted: later_len, + read_state: None, }, ], start_position: 0, @@ -3998,6 +4048,249 @@ mod tests { let _ = std::fs::remove_dir_all(&dir); } + /// A sealed-segment poll opens the file once and caches the read fd; a later + /// poll of the same segment reuses the cached descriptor. Proven by + /// unlinking the file after the first read: a fresh open-by-path would now + /// fail, so a successful second read can only come from the cached fd (which + /// reads the still-open, unlinked inode). + #[compio::test] + async fn read_disk_caches_and_reuses_sealed_segment_fd() { + let namespace = IggyNamespace::new(1, 1, 0); + + let dir = std::env::temp_dir().join(format!( + "iggy-read-disk-fdcache-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock after epoch") + .as_nanos(), + )); + compio::fs::create_dir_all(&dir) + .await + .expect("create temp partition dir"); + let partition_dir = dir.to_string_lossy().into_owned(); + + let record = build_segment_record(namespace, 0); + let record_len = record.len() as u64; + let path = format!("{partition_dir}/{:0>20}.log", 0u64); + { + let mut file = compio::fs::File::create(&path) + .await + .expect("create segment file"); + let (written, _) = file.write_all_at(record, 0).await.into(); + written.expect("write segment record"); + file.sync_all().await.expect("flush segment file"); + } + + let handle = SealedSegmentHandle::default(); + // The pump touches the poll's start segment before cloning its handle + // into the plan, so a cache-eligible handle is always tracked. + handle.tracked.set(true); + assert!(handle.fd.borrow().is_none(), "fd cache slot starts empty"); + + let plan = DiskReadPlan { + partition_dir: Some(partition_dir.clone()), + segments: vec![DiskSegment { + start_offset: 0, + persisted: record_len, + read_state: Some(Rc::clone(&handle)), + }], + start_position: 0, + namespace_raw: namespace.inner(), + }; + let first = plan + .read_disk(MessageLookup::Offset { + offset: 0, + count: 1, + ceiling: u64::MAX, + }) + .await; + assert!( + matches!(first, DiskReadOutcome::Matched { .. }), + "first sealed poll must match the batch", + ); + assert!( + handle.fd.borrow().is_some(), + "first sealed poll must populate the read-fd cache slot", + ); + + // Unlink the file: a fresh open-by-path would fail now, so the second + // read succeeding proves the cached fd was reused. + std::fs::remove_file(&path).expect("unlink segment file"); + + let plan = DiskReadPlan { + partition_dir: Some(partition_dir.clone()), + segments: vec![DiskSegment { + start_offset: 0, + persisted: record_len, + read_state: Some(Rc::clone(&handle)), + }], + start_position: 0, + namespace_raw: namespace.inner(), + }; + let second = plan + .read_disk(MessageLookup::Offset { + offset: 0, + count: 1, + ceiling: u64::MAX, + }) + .await; + assert!( + matches!(second, DiskReadOutcome::Matched { .. }), + "cached fd must serve the read after the segment path is unlinked", + ); + + let _ = std::fs::remove_dir_all(&dir); + } + + /// An untracked handle (a sealed segment the walk crosses without being the + /// poll's start segment, or a slot evicted mid-poll) opens its file + /// transiently: the read succeeds but no fd is retained, so the sealed LRU + /// cap stays a true bound on resident descriptors. + #[compio::test] + async fn read_disk_does_not_retain_fd_for_untracked_handle() { + let namespace = IggyNamespace::new(1, 1, 0); + + let dir = std::env::temp_dir().join(format!( + "iggy-read-disk-untracked-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock after epoch") + .as_nanos(), + )); + compio::fs::create_dir_all(&dir) + .await + .expect("create temp partition dir"); + let partition_dir = dir.to_string_lossy().into_owned(); + + let record = build_segment_record(namespace, 0); + let record_len = record.len() as u64; + let path = format!("{partition_dir}/{:0>20}.log", 0u64); + { + let mut file = compio::fs::File::create(&path) + .await + .expect("create segment file"); + let (written, _) = file.write_all_at(record, 0).await.into(); + written.expect("write segment record"); + file.sync_all().await.expect("flush segment file"); + } + + let handle = SealedSegmentHandle::default(); + let plan = DiskReadPlan { + partition_dir: Some(partition_dir.clone()), + segments: vec![DiskSegment { + start_offset: 0, + persisted: record_len, + read_state: Some(Rc::clone(&handle)), + }], + start_position: 0, + namespace_raw: namespace.inner(), + }; + let outcome = plan + .read_disk(MessageLookup::Offset { + offset: 0, + count: 1, + ceiling: u64::MAX, + }) + .await; + assert!( + matches!(outcome, DiskReadOutcome::Matched { .. }), + "the transient open must still serve the read", + ); + assert!( + handle.fd.borrow().is_none(), + "an untracked handle must not retain the fd", + ); + + let _ = std::fs::remove_dir_all(&dir); + } + + /// A sealed-segment poll reloads the dropped sparse index from the `.index` + /// file and resolves the start byte from it, skipping the full-segment scan. + /// Proven by prefixing the `.log` with bytes a scan from position 0 would + /// fault on: only an index that jumps straight to the batch reads it. + #[compio::test] + async fn read_disk_reloads_sealed_index_to_skip_scan() { + let namespace = IggyNamespace::new(1, 1, 0); + + let dir = std::env::temp_dir().join(format!( + "iggy-read-disk-idxreload-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock after epoch") + .as_nanos(), + )); + compio::fs::create_dir_all(&dir) + .await + .expect("create temp partition dir"); + let partition_dir = dir.to_string_lossy().into_owned(); + + // `.log`: an undecodable prefix (a scan from byte 0 faults on it) then a + // valid batch at offset 5. `.index`: one sparse entry mapping offset 5 + // to the batch's byte position, so the poll jumps past the prefix. + let prefix = vec![0xABu8; 512]; + let prefix_len = prefix.len() as u64; + let batch = build_segment_record(namespace, 5); + let mut log_bytes = prefix; + log_bytes.extend_from_slice(&batch); + let log_len = log_bytes.len() as u64; + let log_path = format!("{partition_dir}/{:0>20}.log", 0u64); + { + let mut file = compio::fs::File::create(&log_path) + .await + .expect("create segment log"); + let (written, _) = file.write_all_at(log_bytes, 0).await.into(); + written.expect("write segment log"); + file.sync_all().await.expect("flush segment log"); + } + + let index_bytes = crate::iggy_index::IggyIndexCache::serialize( + &crate::iggy_index::IggyIndex::new(5, 0, prefix_len), + ); + let index_path = format!("{partition_dir}/{:0>20}.index", 0u64); + { + let mut file = compio::fs::File::create(&index_path) + .await + .expect("create segment index"); + let (written, _) = file.write_all_at(index_bytes, 0).await.into(); + written.expect("write segment index"); + file.sync_all().await.expect("flush segment index"); + } + + let handle = SealedSegmentHandle::default(); + let plan = DiskReadPlan { + partition_dir: Some(partition_dir.clone()), + segments: vec![DiskSegment { + start_offset: 0, + persisted: log_len, + read_state: Some(Rc::clone(&handle)), + }], + // Byte 0, exactly what disk_poll_start returns for a sealed segment + // whose resident index was dropped. + start_position: 0, + namespace_raw: namespace.inner(), + }; + let outcome = plan + .read_disk(MessageLookup::Offset { + offset: 5, + count: 1, + ceiling: u64::MAX, + }) + .await; + assert!( + matches!(outcome, DiskReadOutcome::Matched { .. }), + "the reloaded sparse index must skip the prefix; a scan from byte 0 would fault", + ); + assert!( + handle.index.borrow().is_some(), + "the sealed poll must cache the reloaded sparse index", + ); + + let _ = std::fs::remove_dir_all(&dir); + } + fn repair_config() -> PartitionsConfig { PartitionsConfig { messages_required_to_save: 1, diff --git a/core/partitions/src/iggy_partitions.rs b/core/partitions/src/iggy_partitions.rs index 14e7f0fe75..b15f30407a 100644 --- a/core/partitions/src/iggy_partitions.rs +++ b/core/partitions/src/iggy_partitions.rs @@ -26,7 +26,7 @@ use iggy_binary_protocol::{ Command2, ConsensusHeader, Operation, PrepareHeader, PrepareOkHeader, RequestHeader, }; use message_bus::MessageBus; -use server_common::send_messages2::{convert_request_message, encrypt_batch_request}; +use server_common::send_messages2::{ChecksumMode, convert_request_message, encrypt_batch_request}; use server_common::sharding::{IggyNamespace, LocalIdx, ShardId}; #[cfg(debug_assertions)] use std::cell::Cell; @@ -391,9 +391,10 @@ where } /// Build an owned [`PollPlan`] for a partition poll synchronously, under a - /// single [`Self::with_partition`] borrow (the in-memory journal tier + the - /// resident-tail straddle snapshot are read here; mem reads never yield). - /// Returns `None` for a missing or tombstoned namespace. + /// single pump-only `&mut` borrow (the in-memory journal tier + the + /// resident-tail straddle snapshot are read here, and the sealed-read-handle + /// LRU is touched; mem reads never yield). Returns `None` for a missing or + /// tombstoned namespace. /// /// Pairs with [`PollPlan::execute`], which runs the disk read + /// offset persist/apply off the borrow on the owned plan. Splitting the @@ -406,9 +407,11 @@ where consumer: PollingConsumer, args: &PollingArgs, ) -> Option { - self.with_partition(namespace, |partition| { - partition.build_poll_plan(consumer, args) - }) + // `build_poll_plan` touches the partition's sealed-read-handle LRU, so it + // needs `&mut`. Sound on the pump: it is fully synchronous (no `.await` + // inside), so no sibling task can realloc the partitions vec under it. + let partition = self.get_mut_by_ns(namespace)?; + Some(partition.build_poll_plan(consumer, args)) } /// Read a consumer's stored offset + the partition commit offset. Fully @@ -509,7 +512,11 @@ where let message = if message.header().operation == Operation::SendMessages && let Some(encryptor) = &self.config().encryptor { - let canonical = convert_request_message(namespace, message) + // Compute the batch checksum: this canonical output is validated by + // `encrypt_batch_request`'s decode before re-encryption, and the + // re-encrypted batch (checksum kept by `encrypt_batch_request`) then + // re-enters `convert` as the canonical-vs-legacy discriminator. + let canonical = convert_request_message(namespace, message, ChecksumMode::Compute) .and_then(|message| encrypt_batch_request(message, encryptor)); match canonical { Ok(message) => message, diff --git a/core/partitions/src/journal.rs b/core/partitions/src/journal.rs index feb78c24b1..be257a60b3 100644 --- a/core/partitions/src/journal.rs +++ b/core/partitions/src/journal.rs @@ -19,7 +19,7 @@ use iggy_binary_protocol::{Operation, PrepareHeader}; use journal::{Journal, Storage}; use server_common::{ iobuf::{Frozen, Owned}, - send_messages2::{COMMAND_HEADER_SIZE, SendMessages2Ref, decode_prepare_slice}, + send_messages2::{COMMAND_HEADER_SIZE, SendMessages2Ref, decode_prepare_slice_trusted}, }; use std::io; use std::{ @@ -519,8 +519,12 @@ impl PartitionJournal { // One decode feeds both the offset/timestamp index (keyed on // `origin_timestamp`) and the surfaced accounting meta (`base_timestamp`, // size, count); the two timestamps are distinct fields, do not conflate. + // Trusted (no batch-hash): every entry reaching append was just stamped + // by `stamp_prepare_for_persistence` (its checksum recomputed over this + // exact blob) or re-appended from an already-validated resident entry, + // so re-hashing the ~1 MiB blob here only to read the header is waste. let (index_offset_timestamp, meta) = if header.operation == Operation::SendMessages { - match decode_prepare_slice(entry.as_slice()) { + match decode_prepare_slice_trusted(entry.as_slice()) { Ok(batch) if batch.message_count() != 0 => { let message_count = batch.message_count(); let meta = RetainedBatchMeta { @@ -854,7 +858,10 @@ fn try_push_resident_entry( if header.operation != Operation::SendMessages { return; } - let Ok(batch) = decode_prepare_slice(prepare.as_slice()) else { + // Resident entries were locally stamped in `append_messages` or validated + // at repair ingress, so a validating re-decode would only re-hash our own + // write. See the invariant note at the committed-prefix flush walk. + let Ok(batch) = decode_prepare_slice_trusted(prepare.as_slice()) else { return; }; let Some(selection) = select_batch_slice(&batch, query, *matched_messages) else { diff --git a/core/partitions/src/log.rs b/core/partitions/src/log.rs index 4fd5673db5..5fbd8126b6 100644 --- a/core/partitions/src/log.rs +++ b/core/partitions/src/log.rs @@ -18,11 +18,13 @@ use crate::iggy_index::{IGGY_INDEX_SIZE, IggyIndexCache}; use crate::iggy_index_writer::IggyIndexWriter; use crate::messages_writer::MessagesWriter; +use crate::poll_plan::SealedSegmentHandle; use crate::segment::Segment; use iggy_common::{IggyByteSize, IggyMessagesBatch}; use journal::{Journal, Storage}; use ringbuffer::AllocRingBuffer; use server_common::{IggyMessagesBatchSetInFlight, SegmentStorage}; +use std::collections::VecDeque; use std::fmt::Debug; use std::rc::Rc; @@ -30,6 +32,14 @@ const SEGMENTS_CAPACITY: usize = 1024; const ACCESS_MAP_CAPACITY: usize = 8; const SIZE_16MB: usize = 16 * 1024 * 1024; +/// Max sealed segments per partition that keep a resident read handle (fd + +/// sparse index). Without a cap every sealed segment a reader ever touched pins +/// one fd for the partition's lifetime; the server-wide budget is this cap times +/// the partition count, so keep it small. 12 covers a lagging consumer's working +/// set (the recent sealed segments it re-reads) with room for a few concurrent +/// readers before an LRU eviction forces a re-open. +const SEALED_READ_STATE_CAP: usize = 12; + /// Tracking metadata for the journal's current state. /// /// Replaces the server journal's `Inner` struct — lives in the `SegmentedLog` @@ -112,6 +122,16 @@ where storage: Vec, messages_writers: Vec>>, index_writers: Vec>>, + // Parallel to `segments`: a shared read-state handle (fd + sparse index) + // per segment, filled lazily on the first sealed-segment poll and cloned + // into the off-borrow poll plan. Maintained in lockstep with `segments` + // (push/remove together). + sealed_read_state: Vec, + // LRU of sealed-segment `start_offset`s (most-recently-used at the front) + // bounding how many `sealed_read_state` handles stay resident, capped at + // `SEALED_READ_STATE_CAP`. Keyed by offset (stable), not slot index (which + // shifts on retire). See `touch_sealed_read_state`. + sealed_lru: VecDeque, in_flight: IggyMessagesBatchSetInFlight, } @@ -131,6 +151,8 @@ where indexes: Vec::with_capacity(SEGMENTS_CAPACITY), messages_writers: Vec::with_capacity(SEGMENTS_CAPACITY), index_writers: Vec::with_capacity(SEGMENTS_CAPACITY), + sealed_read_state: Vec::with_capacity(SEGMENTS_CAPACITY), + sealed_lru: VecDeque::with_capacity(SEALED_READ_STATE_CAP + 1), in_flight: IggyMessagesBatchSetInFlight::default(), } } @@ -149,11 +171,99 @@ where &self.segments } - pub const fn segments_mut(&mut self) -> &mut Vec { + /// Mutable segment views. Length mutation lives in + /// [`Self::add_persisted_segment`] / [`Self::retire_front`] only, so the + /// parallel vecs cannot desync from the outside. + pub fn segments_mut(&mut self) -> &mut [Segment] { &mut self.segments } - pub const fn storages_mut(&mut self) -> &mut Vec { + /// Shared read-state handles, parallel to [`Self::segments`]. Cloned into + /// the poll plan for sealed segments (see [`SealedSegmentHandle`]). + pub fn sealed_read_state(&self) -> &[SealedSegmentHandle] { + &self.sealed_read_state + } + + /// Record a sealed-segment access and enforce [`SEALED_READ_STATE_CAP`] + /// (LRU). `slot` indexes [`Self::segments`]; an out-of-range slot or an + /// unsealed (active) segment is a no-op, so the poll path passes its start + /// segment unconditionally. The LRU is keyed by `start_offset` - stable + /// across retire, unlike the slot index. The touched segment moves to the + /// most-recently-used front and its handle is marked tracked (eligible to + /// cache a read fd, see `SealedSegmentReadState::tracked`); once more than + /// the cap distinct sealed segments are tracked, the least-recently-used + /// one's handle is untracked and dropped (replaced with a fresh empty + /// handle) so its fd + sparse index free. An in-flight poll holding a clone + /// of the dropped handle keeps it alive until it finishes (see + /// [`SealedSegmentHandle`]). + pub fn touch_sealed_read_state(&mut self, slot: usize) { + let Some(touched) = self.segments.get(slot) else { + return; + }; + if !touched.sealed { + return; + } + let start_offset = touched.start_offset; + self.sealed_read_state[slot].tracked.set(true); + if let Some(pos) = self + .sealed_lru + .iter() + .position(|&offset| offset == start_offset) + { + self.sealed_lru.remove(pos); + } + self.sealed_lru.push_front(start_offset); + if self.sealed_lru.len() > SEALED_READ_STATE_CAP { + let Some(evicted) = self.sealed_lru.pop_back() else { + return; + }; + if let Some(evicted_slot) = self + .segments + .iter() + .position(|segment| segment.start_offset == evicted) + { + // Untrack before orphaning so an in-flight poll holding the old + // handle stops caching fds into it. + self.sealed_read_state[evicted_slot].tracked.set(false); + self.sealed_read_state[evicted_slot] = SealedSegmentHandle::default(); + } + } + } + + /// Retire the oldest segment: pop the front of every parallel vec in + /// lockstep and purge the segment's sealed-LRU entry. Returns the pieces + /// the caller still needs (stats + file unlink), or `None` on an empty + /// log. The read-state handle is dropped here; an in-flight poll holding a + /// clone keeps it alive until it finishes (a cached fd reads the unlinked + /// inode). + pub fn retire_front(&mut self) -> Option<(Segment, SegmentStorage)> { + if self.segments.is_empty() { + return None; + } + self.debug_assert_lockstep(); + let segment = self.segments.remove(0); + let storage = self.storage.remove(0); + self.indexes.remove(0); + self.messages_writers.remove(0); + self.index_writers.remove(0); + self.sealed_read_state.remove(0); + self.sealed_lru + .retain(|&offset| offset != segment.start_offset); + Some((segment, storage)) + } + + fn debug_assert_lockstep(&self) { + debug_assert!( + self.segments.len() == self.storage.len() + && self.segments.len() == self.indexes.len() + && self.segments.len() == self.messages_writers.len() + && self.segments.len() == self.index_writers.len() + && self.segments.len() == self.sealed_read_state.len(), + "segment parallel vecs out of lockstep" + ); + } + + pub fn storages_mut(&mut self) -> &mut [SegmentStorage] { &mut self.storage } @@ -165,7 +275,7 @@ where &self.messages_writers } - pub const fn messages_writers_mut(&mut self) -> &mut Vec>> { + pub fn messages_writers_mut(&mut self) -> &mut [Option>] { &mut self.messages_writers } @@ -173,7 +283,7 @@ where &self.index_writers } - pub const fn index_writers_mut(&mut self) -> &mut Vec>> { + pub fn index_writers_mut(&mut self) -> &mut [Option>] { &mut self.index_writers } @@ -205,7 +315,7 @@ where &self.indexes } - pub const fn indexes_mut(&mut self) -> &mut Vec> { + pub fn indexes_mut(&mut self) -> &mut [Option] { &mut self.indexes } @@ -259,6 +369,8 @@ where self.indexes.push(None); self.messages_writers.push(messages_writer); self.index_writers.push(index_writer); + self.sealed_read_state.push(SealedSegmentHandle::default()); + self.debug_assert_lockstep(); } pub fn set_segment_indexes(&mut self, segment_index: usize, indexes: IggyIndexCache) { @@ -297,3 +409,176 @@ where &self.journal } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::journal::{PartitionJournal, PartitionJournalMemStorage}; + + type TestLog = + SegmentedLog, PartitionJournalMemStorage>; + + /// Push a sealed segment with a resident (index-filled) read handle and + /// return a clone of that handle, standing in for an in-flight poll's clone. + /// Goes through `add_persisted_segment` so every parallel vec stays in + /// lockstep (segment start offsets equal their slot indexes in every test + /// that pushes ascending offsets from an empty log). + fn push_resident_sealed(log: &mut TestLog, start_offset: u64) -> SealedSegmentHandle { + log.add_persisted_segment( + Segment { + start_offset, + sealed: true, + ..Segment::default() + }, + SegmentStorage::default(), + None, + None, + ); + let slot = log.sealed_read_state.len() - 1; + *log.sealed_read_state[slot].index.borrow_mut() = Some(IggyIndexCache::with_capacity(1)); + Rc::clone(&log.sealed_read_state[slot]) + } + + #[test] + fn touch_sealed_read_state_evicts_least_recently_used_past_cap() { + let mut log = TestLog::default(); + let handles: Vec<_> = (0..=SEALED_READ_STATE_CAP as u64) + .map(|offset| push_resident_sealed(&mut log, offset)) + .collect(); + // Ascending touch order: slot/offset 0 is the least-recently used. + for slot in 0..=SEALED_READ_STATE_CAP { + log.touch_sealed_read_state(slot); + } + + // Slot 0 dropped: the pump handle was replaced with a fresh empty one. + assert!( + !Rc::ptr_eq(&handles[0], &log.sealed_read_state()[0]), + "least-recently-used handle must be dropped past the cap", + ); + assert!( + log.sealed_read_state()[0].index.borrow().is_none(), + "the dropped slot resets to an empty handle", + ); + // In-flight safety: the dropped handle's clone stays alive and still + // sees its cached index, so a poll holding it finishes without a UAF. + assert_eq!( + Rc::strong_count(&handles[0]), + 1, + "the dropped handle survives for an in-flight poll's clone", + ); + assert!( + handles[0].index.borrow().is_some(), + "the in-flight clone keeps reading the cached index", + ); + assert!( + !handles[0].tracked.get(), + "eviction untracks the orphaned handle so in-flight polls stop caching fds into it", + ); + // Every more-recently-used slot is retained (same Rc). + for (handle, resident) in handles.iter().zip(log.sealed_read_state().iter()).skip(1) { + assert!( + Rc::ptr_eq(handle, resident), + "recently-used handles stay resident", + ); + } + } + + #[test] + fn touch_sealed_read_state_reorders_eviction_on_reaccess() { + let mut log = TestLog::default(); + let mut handles: Vec<_> = (0..SEALED_READ_STATE_CAP as u64) + .map(|offset| push_resident_sealed(&mut log, offset)) + .collect(); + for slot in 0..SEALED_READ_STATE_CAP { + log.touch_sealed_read_state(slot); + } + // Re-access slot 0 -> now most-recently used, so slot 1 becomes the + // least-recently used and next to be evicted. + log.touch_sealed_read_state(0); + + handles.push(push_resident_sealed(&mut log, SEALED_READ_STATE_CAP as u64)); + log.touch_sealed_read_state(SEALED_READ_STATE_CAP); + + assert!( + !Rc::ptr_eq(&handles[1], &log.sealed_read_state()[1]), + "the least-recently-used segment is evicted, not the re-accessed one", + ); + assert!( + Rc::ptr_eq(&handles[0], &log.sealed_read_state()[0]), + "the re-accessed segment stays resident", + ); + } + + #[test] + fn evicted_sealed_slot_is_empty_and_refillable() { + let mut log = TestLog::default(); + for offset in 0..=SEALED_READ_STATE_CAP as u64 { + push_resident_sealed(&mut log, offset); + } + for slot in 0..=SEALED_READ_STATE_CAP { + log.touch_sealed_read_state(slot); + } + + // The evicted slot holds a fresh empty handle, so the next poll re-opens + // instead of reusing a stale descriptor. + let evicted = &log.sealed_read_state()[0]; + assert!(evicted.fd.borrow().is_none()); + assert!(evicted.index.borrow().is_none()); + + // Re-filling it (what the next sealed poll does) works. + *log.sealed_read_state()[0].index.borrow_mut() = Some(IggyIndexCache::with_capacity(1)); + assert!(log.sealed_read_state()[0].index.borrow().is_some()); + } + + #[test] + fn touch_sealed_read_state_ignores_active_and_out_of_range_slots() { + let mut log = TestLog::default(); + log.add_persisted_segment( + Segment { + start_offset: 0, + sealed: false, + ..Segment::default() + }, + SegmentStorage::default(), + None, + None, + ); + + // Active (unsealed) segment: no LRU entry, handle stays untracked, so + // its fd is never retained outside the cap. + log.touch_sealed_read_state(0); + assert!(log.sealed_lru.is_empty()); + assert!(!log.sealed_read_state()[0].tracked.get()); + + // Out-of-range slot (the purge drain window empties the vec across + // awaits): must be a no-op, not a panic. + log.touch_sealed_read_state(1); + assert!(log.sealed_lru.is_empty()); + } + + #[test] + fn retire_front_purges_lru_entry_and_keeps_vecs_lockstep() { + let mut log = TestLog::default(); + push_resident_sealed(&mut log, 0); + push_resident_sealed(&mut log, 5); + log.touch_sealed_read_state(0); + log.touch_sealed_read_state(1); + + let (segment, _storage) = log.retire_front().expect("log has segments"); + assert_eq!(segment.start_offset, 0); + assert!( + !log.sealed_lru.contains(&0), + "retire must purge the segment's LRU entry", + ); + assert!(log.sealed_lru.contains(&5), "the survivor's entry stays"); + assert_eq!(log.segments().len(), 1); + assert_eq!(log.sealed_read_state().len(), 1); + assert_eq!(log.storages().len(), 1); + + assert!(log.retire_front().is_some()); + assert!( + log.retire_front().is_none(), + "an empty log retires nothing instead of panicking", + ); + } +} diff --git a/core/partitions/src/poll_plan.rs b/core/partitions/src/poll_plan.rs index 4792760d37..e1392a6691 100644 --- a/core/partitions/src/poll_plan.rs +++ b/core/partitions/src/poll_plan.rs @@ -24,11 +24,15 @@ //! synchronously under the borrow into the owned types here, drops the borrow, //! then [`PollPlan::execute`] runs the disk read + the in-memory auto-commit //! apply on owned data alone: consumer offsets are already `Arc`, the journal -//! tail is a point-in-time `Frozen` snapshot, and segment files are re-opened -//! by path. No value in this module holds a partition reference, so executing a -//! plan is sound on a detached task concurrently with the pump's own writes. +//! tail is a point-in-time `Frozen` snapshot, and each sealed segment carries a +//! shared [`SealedSegmentReadState`] handle (a plain `Rc`, not a partition +//! reference) whose read fd + sparse index the read reuses or fills on a miss. +//! No value in this module holds a partition reference, so executing a plan is +//! sound on a detached task concurrently with the pump's own writes. use crate::PollFragments; +use crate::iggy_index::IggyIndexCache; +use crate::iggy_index_reader::IggyIndexReader; use crate::journal::{MessageLookup, push_selected_batch_fragments, select_batch_slice}; use compio::io::AsyncReadAtExt; use iggy_common::{ @@ -36,14 +40,46 @@ use iggy_common::{ }; use server_common::iobuf::{Frozen, Owned}; use server_common::send_messages2::{COMMAND_HEADER_SIZE, decode_batch_slice}; +use std::cell::{Cell, RefCell}; use std::hash::Hash; +use std::rc::Rc; use std::sync::Arc; use std::sync::atomic::Ordering; use tracing::warn; -/// Owned, borrow-free inputs for the disk tier of a poll (see module docs). -/// Segment files are re-opened by path because sealed segments drop their -/// writer at rotation. +/// Per-sealed-segment read state, shared as a cheap `Rc` handle between the +/// owning partition and the off-borrow [`DiskReadPlan`] (a plain `Rc`, never a +/// partition reference, so the read runs off the pump). Both slots fill lazily +/// on the first sealed poll and are reused after. The pump drops its handle +/// when the segment is retired, so the state frees once any in-flight poll +/// holding a clone finishes (a cached fd meanwhile reads the unlinked inode, +/// which is fine). The active segment is never cached. +#[derive(Debug, Default)] +pub struct SealedSegmentReadState { + /// Read-only descriptor; compio `File` clones share the kernel fd, so a hit + /// avoids the per-poll `openat` (an `io_uring` op prone to io-wq punts) and + /// preserves kernel readahead. `None` until the first sealed poll opens it. + pub(crate) fd: RefCell>, + /// Sparse offset/timestamp index reloaded from the `.index` file the + /// segment dropped at rotation, so a poll resolves the start byte in + /// O(log n) instead of scanning the whole segment from byte 0 (the stall). + /// `None` until the first sealed poll loads it. + pub(crate) index: RefCell>, + /// Whether the owning partition's sealed LRU currently tracks this handle. + /// Gates the fd store-back in `resolve_segment_file`: a walk crosses every + /// sealed segment from the poll's start onward, but only the start segment + /// is LRU-touched, so an untracked fill would retain a descriptor the + /// `SEALED_READ_STATE_CAP` budget never counts. Set on touch, cleared on + /// evict; plain `Cell`, all access is same-thread (`Rc` handle). + pub(crate) tracked: Cell, +} + +pub type SealedSegmentHandle = Rc; + +/// Owned, borrow-free inputs for the disk tier of a poll (see module docs). A +/// sealed segment reuses its cached [`SealedSegmentReadState`] (read fd + sparse +/// index); the active segment (and any cache miss) opens by path and resolves +/// from its resident index, because sealed segments drop both at rotation. pub struct DiskReadPlan { pub(crate) partition_dir: Option, /// Segments to walk, snapshotted from the poll's starting segment onward @@ -57,6 +93,10 @@ pub struct DiskReadPlan { pub struct DiskSegment { pub(crate) start_offset: u64, pub(crate) persisted: u64, + /// Shared read state, cloned from the owning partition at plan time for a + /// SEALED segment; `None` for the active segment, which always opens fresh + /// and resolves from its resident index. See [`SealedSegmentReadState`]. + pub(crate) read_state: Option, } /// Owned auto-commit input, applied off the partition borrow after a poll (see @@ -406,7 +446,21 @@ impl DiskReadPlan { // `start_position` applies to the first snapshotted segment; each later // segment is walked from byte 0 (reset at the end of every iteration). - let mut position = self.start_position; + // + // A sealed first segment dropped its resident index at rotation, so + // `disk_poll_start` fell back to byte 0. Reload the sparse index (once, + // then cached) and resolve the start byte so the walk skips straight to + // the target instead of scanning the whole segment - the poll stall. A + // miss or load failure keeps `start_position` (the pre-existing + // full-scan fallback). The active segment carries no read state, so its + // resident-index-resolved `start_position` is left untouched. + let mut position = match self.segments.first() { + Some(first) => self + .resolve_sealed_start(first, query, partition_dir) + .await + .unwrap_or(self.start_position), + None => self.start_position, + }; let mut fragments = PollFragments::new(); let mut last_matching_offset = None; let mut matched: u32 = 0; @@ -427,7 +481,7 @@ impl DiskReadPlan { continue; } let path = format!("{partition_dir}/{:0>20}.log", segment.start_offset); - let Some(file) = self.open_segment_with_retry(&path).await else { + let Some(file) = self.resolve_segment_file(segment, &path).await else { // Open exhausted retries: the segment may hold present-but- // unreadable data. Stop here rather than walking past it. faulted = true; @@ -486,6 +540,99 @@ impl DiskReadPlan { } } + /// Resolve the read-only descriptor for `segment`'s file. A sealed segment + /// clones its cached fd on a hit (sharing the kernel fd, no syscall) and, on + /// a miss, opens by path and stores the fd back so later polls skip the + /// `openat`. The active segment (no cache slot) always opens fresh. Returns + /// `None` only when the open exhausts its retries (the caller fails closed). + async fn resolve_segment_file( + &self, + segment: &DiskSegment, + path: &str, + ) -> Option { + let Some(handle) = &segment.read_state else { + return self.open_segment_with_retry(path).await; + }; + // Borrow only to clone the `Option` out, never across the await. + if let Some(cached) = handle.fd.borrow().clone() { + return Some(cached); + } + let file = self.open_segment_with_retry(path).await?; + // Store back only while the pump tracks this handle; an untracked + // fill (walk-through segment, or a slot evicted mid-poll) would pin an + // fd outside the LRU budget, so it opens transiently instead. Benign + // race: a concurrent poll of the same segment may have filled the slot + // while this open was in flight; overwriting with an equivalent fd + // (same inode) is harmless. + if handle.tracked.get() { + *handle.fd.borrow_mut() = Some(file.clone()); + } + Some(file) + } + + /// Resolve the start byte for the poll's target segment from its sparse + /// index, loading the `.index` file on the first sealed poll and caching it + /// on the shared handle. Returns `None` (keep the byte-0 fallback) for the + /// active segment (no handle), a below-range query, or a load failure. + async fn resolve_sealed_start( + &self, + segment: &DiskSegment, + query: MessageLookup, + partition_dir: &str, + ) -> Option { + // TODO: a per-consumer cursor hint (the previous sealed poll's resolved + // position) could seed this so a sequentially advancing consumer skips + // the sparse-index lookup on repeated polls of the same segment. + let handle = segment.read_state.as_ref()?; + // Cache hit: resolve under a short borrow, never across the await. + let cached = handle + .index + .borrow() + .as_ref() + .map(|index| resolve_index_position(index, query)); + if let Some(resolved) = cached { + return resolved; + } + let path = format!("{partition_dir}/{:0>20}.index", segment.start_offset); + let index = self.load_sealed_index(&path).await?; + let resolved = resolve_index_position(&index, query); + *handle.index.borrow_mut() = Some(index); + resolved + } + + /// Load a sealed segment's sparse index from its `.index` file. `None` on a + /// missing/unreadable file so the caller falls back to a byte-0 scan (the + /// pre-existing behavior); the load is retried on the next poll. + async fn load_sealed_index(&self, path: &str) -> Option { + match IggyIndexReader::new(path).await { + Ok(reader) => match reader.load_all().await { + Ok(index) => Some(index), + Err(error) => { + warn!( + target: "iggy.partitions.diag", + plane = "partitions", + namespace_raw = self.namespace_raw, + path, + %error, + "disk poll: failed to read sparse index; scanning from segment start" + ); + None + } + }, + Err(error) => { + warn!( + target: "iggy.partitions.diag", + plane = "partitions", + namespace_raw = self.namespace_raw, + path, + %error, + "disk poll: failed to open sparse index; scanning from segment start" + ); + None + } + } + } + /// Open a segment file for a disk poll, retrying transient IO failures (fd /// pressure under heavy parallel load) so one failed syscall does not /// silently collapse the poll into an empty result. @@ -544,6 +691,18 @@ impl DiskReadPlan { } } +/// Byte position of the sparse-index entry at or below the query's offset / +/// timestamp, or `None` when the query is below the first indexed entry (the +/// caller then scans from the segment start). Mirrors `disk_poll_start`'s +/// resident-index resolution for the sealed, off-pump path. +fn resolve_index_position(index: &IggyIndexCache, query: MessageLookup) -> Option { + match query { + MessageLookup::Offset { offset, .. } => index.offset_lower_bound(offset), + MessageLookup::Timestamp { timestamp, .. } => index.timestamp_lower_bound(timestamp), + } + .map(|entry| entry.position) +} + impl AutoCommitCtx { /// The offset key (kind + numeric id) this auto-commit targets, for the /// replicated `StoreConsumerOffset2` op the serving shard submits. diff --git a/core/server-ng/config.toml b/core/server-ng/config.toml index 7a185a80da..3bae432902 100644 --- a/core/server-ng/config.toml +++ b/core/server-ng/config.toml @@ -668,8 +668,7 @@ ports = { tcp = 8091, quic = 8081, http = 3001, websocket = 8093, tcp_replica = # - numa settings: # + "numa:auto": Use all available numa node, cores # + "numa:nodes=0,1;cores=4;no_ht=true": Use NUMA node 0 and 1, each nodes use 4 cores, and no hyperthreads -# TODO(hubcio): revert to "numa:auto" once multi-shard server-ng is stable. -cpu_allocation = 1 +cpu_allocation = "numa:auto" # Whether shard threads are pinned to dedicated CPU cores (default: true). # Pinned cores are drawn from the process's allowed CPU set (affinity/cpuset diff --git a/core/server-ng/src/dispatch.rs b/core/server-ng/src/dispatch.rs index 89160a63a1..160f380363 100644 --- a/core/server-ng/src/dispatch.rs +++ b/core/server-ng/src/dispatch.rs @@ -185,10 +185,10 @@ where { let shard_handle = Rc::clone(shard_handle); // Runs synchronously on the shard pump (see `process_lifecycle` -> - // `on_partition_read`). `build_poll_snapshot` takes the partition borrow via - // `with_partition` (closure-scoped, debug `BorrowGuard`) and returns an owned - // `PollPlan`; only owned data crosses into `spawn_poll_io`. A fully-resident - // poll replies here without spawning. See the `poll_plan` module docs. + // `on_partition_read`). `build_poll_snapshot` takes a pump-only `&mut` + // partition borrow (synchronous, so no sibling task can realloc under it) and + // returns an owned `PollPlan`; only owned data crosses into `spawn_poll_io`. A + // fully-resident poll replies here without spawning. See the `poll_plan` module docs. Rc::new(move |namespace, read, reply| { let Some(shard) = upgrade_shard_handle(&shard_handle) else { return; diff --git a/core/server-ng/src/responses.rs b/core/server-ng/src/responses.rs index 9cf3378ed6..4d07cef55f 100644 --- a/core/server-ng/src/responses.rs +++ b/core/server-ng/src/responses.rs @@ -1432,7 +1432,7 @@ where /// Size of the in-storage (`IggyMessage2`) per-message header inside a /// `SendMessages2` batch blob: `checksum`(8) + `id`(16) + `offset_delta`(4) /// + `timestamp_delta`(4) + `user_headers_length`(4) + `payload_length`(4) -/// + reserved(8). See `server_common::send_messages2::from_legacy_request`. +/// + reserved(8). See `server_common::send_messages2::SendMessages2Owned::from_messages`. const STORED_MESSAGE_HEADER_SIZE: usize = 48; /// Build the `PolledMessages` reply body from the owning shard's poll diff --git a/core/server_common/src/send_messages2.rs b/core/server_common/src/send_messages2.rs index a64cd06641..6afd33dd86 100644 --- a/core/server_common/src/send_messages2.rs +++ b/core/server_common/src/send_messages2.rs @@ -173,84 +173,12 @@ impl SendMessages2Owned { header[32..36].copy_from_slice(&user_headers_length.to_le_bytes()); header[36..40].copy_from_slice(&payload_length.to_le_bytes()); - let checksum = calculate_checksum_parts(&header[8..], &message.payload, user_headers); - header[0..8].copy_from_slice(&checksum.to_le_bytes()); - + let msg_start = blob.len(); blob.extend_from_slice(&header); blob.extend_from_slice(&message.payload); blob.extend_from_slice(user_headers); - } - - let blob = blob.freeze(); - let mut header = SendMessages2Header::new( - namespace.partition_id() as u64, - origin_timestamp, - u64::try_from(COMMAND_HEADER_SIZE + blob.len()) - .map_err(|_| IggyError::InvalidCommand)?, - message_count, - ); - header.batch_checksum = calculate_batch_checksum(&header, &blob); - - Ok(Self { header, blob }) - } - - pub fn from_legacy_request(namespace: IggyNamespace, body: &[u8]) -> Result { - let (message_count, messages) = legacy_messages_slice(body)?; - let mut parsed = Vec::with_capacity(message_count as usize); - let mut origin_timestamp = u64::MAX; - let mut cursor = 0usize; - - while cursor < messages.len() && parsed.len() < message_count as usize { - let legacy = LegacyMessageRef::decode(&messages[cursor..])?; - origin_timestamp = origin_timestamp.min(legacy.origin_timestamp); - cursor += legacy.total_size; - parsed.push(legacy); - } - - if parsed.len() != message_count as usize || cursor != messages.len() { - return Err(IggyError::InvalidCommand); - } - - if origin_timestamp == u64::MAX { - origin_timestamp = 0; - } - - let mut blob = BytesMut::with_capacity(messages.len()); - for (index, legacy) in parsed.iter().enumerate() { - let id = if legacy.id == 0 { - random_id::get_uuid() - } else { - legacy.id - }; - let offset_delta = u32::try_from(index).map_err(|_| IggyError::InvalidCommand)?; - let timestamp_delta = legacy - .origin_timestamp - .checked_sub(origin_timestamp) - .ok_or(IggyError::InvalidCommand)?; - if timestamp_delta > MAX_TIMESTAMP_DELTA_MICROS { - return Err(IggyError::InvalidMessageTimestampDelta(timestamp_delta)); - } - let timestamp_delta = - u32::try_from(timestamp_delta).map_err(|_| IggyError::InvalidCommand)?; - let user_headers_length = - u32::try_from(legacy.user_headers.len()).map_err(|_| IggyError::InvalidCommand)?; - let payload_length = - u32::try_from(legacy.payload.len()).map_err(|_| IggyError::InvalidCommand)?; - - let mut header = [0u8; MESSAGE_HEADER_SIZE]; - header[8..24].copy_from_slice(&id.to_le_bytes()); - header[24..28].copy_from_slice(&offset_delta.to_le_bytes()); - header[28..32].copy_from_slice(×tamp_delta.to_le_bytes()); - header[32..36].copy_from_slice(&user_headers_length.to_le_bytes()); - header[36..40].copy_from_slice(&payload_length.to_le_bytes()); - - let checksum = - calculate_checksum_parts(&header[8..], legacy.payload, legacy.user_headers); - header[0..8].copy_from_slice(&checksum.to_le_bytes()); - - blob.extend_from_slice(&header); - blob.extend_from_slice(legacy.payload); - blob.extend_from_slice(legacy.user_headers); + let checksum = XxHash3_64::oneshot(&blob[msg_start + 8..]); + blob[msg_start..msg_start + 8].copy_from_slice(&checksum.to_le_bytes()); } let blob = blob.freeze(); @@ -573,12 +501,12 @@ pub fn encrypt_batch_request( header[28..32].copy_from_slice(&view.header.timestamp_delta.to_le_bytes()); header[32..36].copy_from_slice(&user_headers_length.to_le_bytes()); header[36..40].copy_from_slice(&payload_length.to_le_bytes()); - let checksum = calculate_checksum_parts(&header[8..], &encrypted_payload, user_headers); - header[0..8].copy_from_slice(&checksum.to_le_bytes()); - + let msg_start = blob.len(); blob.extend_from_slice(&header); blob.extend_from_slice(&encrypted_payload); blob.extend_from_slice(user_headers); + let checksum = XxHash3_64::oneshot(&blob[msg_start + 8..]); + blob[msg_start..msg_start + 8].copy_from_slice(&checksum.to_le_bytes()); } let blob = blob.freeze(); @@ -590,9 +518,27 @@ pub fn encrypt_batch_request( SendMessages2Owned { header, blob }.encode_request(request_header) } +/// Whether the legacy transcode stamps a batch checksum onto its output. +/// +/// The recompute is an `XxHash3` batch-checksum pass, needed only when a reader +/// validates the transcoded batch before [`stamp_prepare_for_persistence`] +/// recomputes it: the encrypt ingest path re-decodes the canonicalized batch +/// (`encrypt_batch_request`'s validating decode, then the second `convert` its +/// output re-enters as the canonical-vs-legacy discriminator). The partition +/// ingest path has no such reader, so it skips the pass and the checksum stays +/// zero until stamp. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChecksumMode { + /// Compute the batch checksum for the transcoded batch. + Compute, + /// Leave the batch checksum zero; `stamp_prepare_for_persistence` fills it. + Skip, +} + pub fn convert_request_message( namespace: IggyNamespace, message: Message, + checksum: ChecksumMode, ) -> Result, IggyError> { let request_header = *message.header(); let total_size = request_header.size as usize; @@ -600,12 +546,127 @@ pub fn convert_request_message( if decode_batch_slice(body).is_ok() { return Ok(message); } - SendMessages2Owned::from_legacy_request(namespace, body)?.encode_request(request_header) + transcode_legacy_request(namespace, body, request_header, checksum) } -/// Decode one batch slice (`[256B command header][blob]`), validating the -/// batch checksum. The persisted segment-file record and the request slice -/// share this layout, so both decode through here. +/// Transcode a legacy `SendMessages` request body directly into the canonical +/// `[RequestHeader][256B SendMessages2Header][blob]` form, writing each message +/// record straight into the final aligned buffer. +/// +/// Fused replacement for the `from_legacy_request(..).encode_request(..)` +/// two-step: a size walk over the legacy input sizes the single output +/// allocation, then a write walk lays down each canonical record in place. This +/// drops the intermediate blob allocation and the full-blob copy the two-step +/// paid. Output bytes are identical to that path. +/// +/// `checksum` selects whether the output carries a batch checksum (see +/// [`ChecksumMode`]); [`ChecksumMode::Skip`] leaves it zero for the partition +/// ingest path, where stamp recomputes it. +fn transcode_legacy_request( + namespace: IggyNamespace, + body: &[u8], + mut request_header: RequestHeader, + checksum: ChecksumMode, +) -> Result, IggyError> { + let (message_count, messages) = legacy_messages_slice(body)?; + let mut parsed = Vec::with_capacity(message_count as usize); + let mut origin_timestamp = u64::MAX; + let mut cursor = 0usize; + let mut blob_len = 0usize; + + while cursor < messages.len() && parsed.len() < message_count as usize { + let legacy = LegacyMessageRef::decode(&messages[cursor..])?; + origin_timestamp = origin_timestamp.min(legacy.origin_timestamp); + cursor += legacy.total_size; + blob_len = blob_len + .checked_add(MESSAGE_HEADER_SIZE + legacy.payload.len() + legacy.user_headers.len()) + .ok_or(IggyError::InvalidCommand)?; + parsed.push(legacy); + } + + if parsed.len() != message_count as usize || cursor != messages.len() { + return Err(IggyError::InvalidCommand); + } + + if origin_timestamp == u64::MAX { + origin_timestamp = 0; + } + + let header_size = std::mem::size_of::(); + let batch_length = COMMAND_HEADER_SIZE + .checked_add(blob_len) + .ok_or(IggyError::InvalidCommand)?; + let total_size = header_size + .checked_add(batch_length) + .ok_or(IggyError::InvalidCommand)?; + request_header.size = u32::try_from(total_size).map_err(|_| IggyError::InvalidCommand)?; + + let mut buffer = Owned::::zeroed(total_size); + let bytes = buffer.as_mut_slice(); + bytes[0..header_size].copy_from_slice(bytemuck::bytes_of(&request_header)); + + let mut write = PREPARE_SPLIT_POINT; + for (index, legacy) in parsed.iter().enumerate() { + let id = if legacy.id == 0 { + random_id::get_uuid() + } else { + legacy.id + }; + let offset_delta = u32::try_from(index).map_err(|_| IggyError::InvalidCommand)?; + let timestamp_delta = legacy + .origin_timestamp + .checked_sub(origin_timestamp) + .ok_or(IggyError::InvalidCommand)?; + if timestamp_delta > MAX_TIMESTAMP_DELTA_MICROS { + return Err(IggyError::InvalidMessageTimestampDelta(timestamp_delta)); + } + let timestamp_delta = + u32::try_from(timestamp_delta).map_err(|_| IggyError::InvalidCommand)?; + let user_headers_length = + u32::try_from(legacy.user_headers.len()).map_err(|_| IggyError::InvalidCommand)?; + let payload_length = + u32::try_from(legacy.payload.len()).map_err(|_| IggyError::InvalidCommand)?; + + let mut header = [0u8; MESSAGE_HEADER_SIZE]; + header[8..24].copy_from_slice(&id.to_le_bytes()); + header[24..28].copy_from_slice(&offset_delta.to_le_bytes()); + header[28..32].copy_from_slice(×tamp_delta.to_le_bytes()); + header[32..36].copy_from_slice(&user_headers_length.to_le_bytes()); + header[36..40].copy_from_slice(&payload_length.to_le_bytes()); + let msg_start = write; + bytes[write..write + MESSAGE_HEADER_SIZE].copy_from_slice(&header); + write += MESSAGE_HEADER_SIZE; + bytes[write..write + legacy.payload.len()].copy_from_slice(legacy.payload); + write += legacy.payload.len(); + bytes[write..write + legacy.user_headers.len()].copy_from_slice(legacy.user_headers); + write += legacy.user_headers.len(); + // The cover is [msg_start + 8 .. write], including the 8 reserved zero + // header bytes. This relies on the stack header being zero-initialized, + // not on the output buffer being pre-zeroed. + let checksum = XxHash3_64::oneshot(&bytes[msg_start + 8..write]); + bytes[msg_start..msg_start + 8].copy_from_slice(&checksum.to_le_bytes()); + } + + let mut command = SendMessages2Header::new( + namespace.partition_id() as u64, + origin_timestamp, + batch_length as u64, + message_count, + ); + if checksum == ChecksumMode::Compute { + command.batch_checksum = calculate_batch_checksum( + &command, + &bytes[PREPARE_SPLIT_POINT..PREPARE_SPLIT_POINT + blob_len], + ); + } + command.encode_into(&mut bytes[header_size..header_size + COMMAND_HEADER_SIZE]); + + Message::try_from(buffer).map_err(|_| IggyError::InvalidCommand) +} + +/// Decode one batch slice (`[256B command header][blob]`), validating the batch +/// checksum and every per-message checksum. The persisted segment-file record +/// and the request slice share this layout, so both decode through here. pub fn decode_batch_slice(body: &[u8]) -> Result, IggyError> { if body.len() < COMMAND_HEADER_SIZE { return Err(IggyError::InvalidCommand); @@ -618,7 +679,8 @@ pub fn decode_batch_slice(body: &[u8]) -> Result, IggyError } let blob = &body[COMMAND_HEADER_SIZE..COMMAND_HEADER_SIZE + blob_len]; - let expected_checksum = calculate_batch_checksum(&header, blob); + let batch = SendMessages2Ref { header, blob }; + let expected_checksum = verify_and_recompute_batch_checksum(&batch)?; if header.batch_checksum != expected_checksum { return Err(IggyError::InvalidBatchChecksum( header.batch_checksum, @@ -627,10 +689,11 @@ pub fn decode_batch_slice(body: &[u8]) -> Result, IggyError )); } - Ok(SendMessages2Ref { header, blob }) + Ok(batch) } -/// Decode a `Prepare` message from a slice of bytes. +/// Decode a `Prepare` message from a slice of bytes, validating the batch +/// checksum and every per-message checksum. /// /// `bytes` must be 16-byte aligned (`PrepareHeader` has `u128` fields). Source /// from `Frozen` / `Owned` / `Message`. @@ -638,9 +701,42 @@ pub fn decode_batch_slice(body: &[u8]) -> Result, IggyError /// /// # Errors /// -/// `IggyError::InvalidCommand` on: short buffer, bad bit pattern, `size` -/// outside `[header_size, bytes.len()]`, short/checksum-mismatched body. +/// `IggyError::InvalidCommand` on a short buffer, bad bit pattern, `size` +/// outside `[header_size, bytes.len()]`, or frames that do not tile the batch; +/// `InvalidBatchChecksum` / `InvalidMessageChecksum` on an integrity mismatch. pub fn decode_prepare_slice(bytes: &[u8]) -> Result, IggyError> { + decode_prepare_slice_inner(bytes, true) +} + +/// Like [`decode_prepare_slice`] but skips the per-message checksum +/// verification and batch-checksum recompute, extracting only the header meta. +/// Every cheap structural check (length, 16-byte alignment, `size` bounds, blob +/// length) is still enforced. +/// +/// INVARIANT: `bytes` MUST be node-local self-stamped - +/// [`stamp_prepare_for_persistence`] recomputed the batch checksum over the +/// exact blob on THIS node - or already integrity-checked at their network +/// ingress. There is no consensus-layer blob validation: the `PrepareHeader` +/// integrity fields are inert zeros. A replicated `SendMessages` prepare is +/// gated per-message on receipt by [`verify_received_send_messages`], and a +/// repaired prepare is validated via [`decode_prepare_slice`]; both run BEFORE +/// the bytes reach any trusted decode. NEVER call this on unvalidated network +/// bytes - it would let a corrupted blob pass undetected. The full-body +/// per-message checksum pass dominates produce-path CPU, so trusted call sites +/// that only read header meta skip it. +/// +/// # Errors +/// +/// Same structural errors as [`decode_prepare_slice`], minus +/// `InvalidBatchChecksum` and `InvalidMessageChecksum`. +pub fn decode_prepare_slice_trusted(bytes: &[u8]) -> Result, IggyError> { + decode_prepare_slice_inner(bytes, false) +} + +fn decode_prepare_slice_inner( + bytes: &[u8], + validate_checksum: bool, +) -> Result, IggyError> { let header_size = std::mem::size_of::(); if bytes.len() < header_size { return Err(IggyError::InvalidCommand); @@ -677,16 +773,19 @@ pub fn decode_prepare_slice(bytes: &[u8]) -> Result, IggyEr } let blob = &blob[..blob_len]; - let expected_checksum = calculate_batch_checksum(&header, blob); - if header.batch_checksum != expected_checksum { - return Err(IggyError::InvalidBatchChecksum( - header.batch_checksum, - expected_checksum, - header.base_offset, - )); + let batch = SendMessages2Ref { header, blob }; + if validate_checksum { + let expected_checksum = verify_and_recompute_batch_checksum(&batch)?; + if header.batch_checksum != expected_checksum { + return Err(IggyError::InvalidBatchChecksum( + header.batch_checksum, + expected_checksum, + header.base_offset, + )); + } } - Ok(SendMessages2Ref { header, blob }) + Ok(batch) } pub fn stamp_prepare_for_persistence( @@ -711,6 +810,35 @@ pub fn stamp_prepare_for_persistence( Ok((message, command, command.message_count)) } +/// Verify every per-message checksum in a received `SendMessages` prepare. +/// +/// The FIRST blob-integrity check on the replicated path: the `PrepareHeader` +/// integrity fields are inert zeros and the batch checksum is recomputed +/// locally at stamp, so transit corruption of a message body would otherwise +/// reach apply undetected. Backups call this before journaling a replicated +/// prepare; on a mismatch the caller fails closed (drop, no `PrepareOk`) and the +/// primary retransmits on prepare-timeout. +/// +/// The stored `batch_checksum` is not consulted (a received prepare is +/// pre-stamp, `base_offset` / `base_timestamp` zero): integrity rests on the +/// per-message checksums, recomputed over the stamp-invariant cover +/// (`header[8..48] || payload || user_headers`), which excludes the 256B command +/// header, so it holds whether or not this node has stamped yet. Shares the +/// frame walk with the validating decoders via +/// [`verify_and_recompute_batch_checksum`], discarding its recomputed batch +/// value. +/// +/// # Errors +/// +/// [`IggyError::InvalidCommand`] if the records do not tile `message_count` +/// exactly (a length-field corruption desyncs the walk); +/// [`IggyError::InvalidMessageChecksum`] on the first per-message mismatch. +pub fn verify_received_send_messages(bytes: &[u8]) -> Result<(), IggyError> { + let batch = decode_prepare_slice_trusted(bytes)?; + verify_and_recompute_batch_checksum(&batch)?; + Ok(()) +} + fn legacy_messages_slice(body: &[u8]) -> Result<(u32, &[u8]), IggyError> { if body.len() < 4 { return Err(IggyError::InvalidCommand); @@ -778,26 +906,84 @@ impl<'a> LegacyMessageRef<'a> { } } -// Hash in storage order: header tail, payload, user headers (the message -// sections follow the legacy wire layout). -fn calculate_checksum_parts(header_tail: &[u8], payload: &[u8], user_headers: &[u8]) -> u64 { +/// Batch checksum v2: streaming `XxHash3_64` over the six batch header meta +/// fields followed by each message's stored 8-byte checksum field in message +/// order - NOT the message bodies. +/// +/// Bodies are bound only transitively: each per-message checksum already covers +/// `header[8..48] || payload || user_headers`, so hashing the checksum fields +/// binds every body byte IFF a reader also re-verifies the per-message +/// checksums. Stamp (produce) hashes `N * 8` bytes instead of the whole blob; +/// validating decoders pay the one body pass as the per-message verify in +/// [`verify_and_recompute_batch_checksum`], which hashes the checksum-field +/// bytes in the same order so its recompute matches a compute here. +/// +/// Assumes a well-formed blob whose frames tile exactly; every compute site +/// builds the blob and satisfies this. +fn calculate_batch_checksum(header: &SendMessages2Header, blob: &[u8]) -> u64 { let mut hasher = XxHash3_64::new(); - hasher.write(header_tail); - hasher.write(payload); - hasher.write(user_headers); + write_batch_header_fields(&mut hasher, header); + let batch = SendMessages2Ref { + header: *header, + blob, + }; + for framed in batch.iter_with_offsets() { + hasher.write(&blob[framed.start..framed.start + 8]); + } hasher.finish() } -fn calculate_batch_checksum(header: &SendMessages2Header, blob: &[u8]) -> u64 { - let mut hasher = XxHash3_64::new(); +fn write_batch_header_fields(hasher: &mut XxHash3_64, header: &SendMessages2Header) { hasher.write(&header.partition_id.to_le_bytes()); hasher.write(&header.base_offset.to_le_bytes()); hasher.write(&header.base_timestamp.to_le_bytes()); hasher.write(&header.origin_timestamp.to_le_bytes()); hasher.write(&header.batch_length.to_le_bytes()); hasher.write(&header.message_count.to_le_bytes()); - hasher.write(blob); - hasher.finish() +} + +/// Verify every per-message checksum in `batch` and return the recomputed v2 +/// batch checksum (see [`calculate_batch_checksum`]) from a single frame walk. +/// +/// The per-message pass is the equal-integrity half of v2: the batch value +/// binds bodies only through the checksum fields, so a validating decode must +/// re-verify each message here or body corruption that leaves the checksum +/// field intact would pass. This is the one full-body pass a validating decode +/// pays; the caller then compares the returned value against the stored +/// `batch_checksum`. +/// +/// # Errors +/// +/// [`IggyError::InvalidMessageChecksum`] on the first per-message mismatch; +/// [`IggyError::InvalidCommand`] if the frames do not tile `message_count` +/// exactly. +fn verify_and_recompute_batch_checksum(batch: &SendMessages2Ref<'_>) -> Result { + let blob = batch.blob(); + let mut hasher = XxHash3_64::new(); + write_batch_header_fields(&mut hasher, &batch.header); + let mut verified = 0u32; + let mut covered = 0usize; + for framed in batch.iter_with_offsets() { + // Cover (`header[8..48] || payload || user_headers`) hashed raw from the + // blob, byte-exact with the encoder's, so a flipped body byte fails even + // when the stored checksum field is left intact. + let stored = framed.message.header.checksum; + let expected = XxHash3_64::oneshot(&blob[framed.start + 8..framed.end]); + if expected != stored { + return Err(IggyError::InvalidMessageChecksum( + stored, + expected, + batch.header.base_offset + u64::from(framed.message.header.offset_delta), + )); + } + hasher.write(&blob[framed.start..framed.start + 8]); + verified += 1; + covered = framed.end; + } + if verified != batch.message_count() || covered != blob.len() { + return Err(IggyError::InvalidCommand); + } + Ok(hasher.finish()) } fn read_u32(bytes: &[u8], offset: usize) -> Result { @@ -827,7 +1013,8 @@ fn read_u128(bytes: &[u8], offset: usize) -> Result { #[cfg(test)] mod tests { use super::*; - use iggy_binary_protocol::Command2; + use iggy_binary_protocol::{Command2, Operation}; + use iggy_common::Aes256GcmEncryptor; fn aligned_prepare_bytes(size: u32) -> Owned { let mut owned = Owned::::zeroed(std::mem::size_of::()); @@ -839,6 +1026,90 @@ mod tests { owned } + /// Assemble an already-stamped batch into a `Prepare`: + /// `[PrepareHeader][256B batch header][blob]`, copying `owned`'s header and + /// blob verbatim. Shared by every real-batch fixture. + fn prepare_from_owned(owned: &SendMessages2Owned) -> Owned { + let header_size = std::mem::size_of::(); + let total = header_size + owned.header.total_size(); + let mut buffer = Owned::::zeroed(total); + { + let prepare: &mut PrepareHeader = + bytemuck::checked::try_from_bytes_mut(&mut buffer.as_mut_slice()[..header_size]) + .expect("zeroed bytes form a valid PrepareHeader"); + prepare.command = Command2::Prepare; + prepare.size = u32::try_from(total).expect("prepare size fits u32"); + } + let bytes = buffer.as_mut_slice(); + owned + .header + .encode_into(&mut bytes[header_size..header_size + COMMAND_HEADER_SIZE]); + bytes[PREPARE_SPLIT_POINT..PREPARE_SPLIT_POINT + owned.blob.len()] + .copy_from_slice(&owned.blob); + buffer + } + + /// A checksum-consistent STAMPED `Prepare` carrying real per-message records, + /// stamped at a non-zero `base_offset` / `base_timestamp` with a v2 + /// `batch_checksum` over the final header fields + per-message checksum fields. + fn valid_prepare_bytes() -> Owned { + let namespace = IggyNamespace::new(1, 1, 7); + let mut owned = SendMessages2Owned::from_messages(namespace, &sample_messages()) + .expect("build send batch"); + owned.header.base_offset = 10; + owned.header.base_timestamp = 20; + owned.header.batch_checksum = owned.header.checksum_for_blob(&owned.blob); + prepare_from_owned(&owned) + } + + #[test] + fn decode_prepare_slice_trusted_matches_validating_for_valid_batch() { + // The trusted variant must surface byte-identical header meta to the + // validating decode for a checksum-consistent batch; only the + // per-message and batch-checksum passes are skipped. + let owned = valid_prepare_bytes(); + + let validated = decode_prepare_slice(owned.as_slice()).expect("valid batch decodes"); + let trusted = + decode_prepare_slice_trusted(owned.as_slice()).expect("valid batch decodes trusted"); + + assert_eq!(validated.header.base_offset, trusted.header.base_offset); + assert_eq!( + validated.header.base_timestamp, + trusted.header.base_timestamp + ); + assert_eq!( + validated.header.origin_timestamp, + trusted.header.origin_timestamp + ); + assert_eq!(validated.header.batch_length, trusted.header.batch_length); + assert_eq!(validated.message_count(), trusted.message_count()); + assert_eq!(validated.header.total_size(), trusted.header.total_size()); + assert_eq!(validated.blob(), trusted.blob()); + } + + #[test] + fn decode_prepare_slice_trusted_skips_batch_checksum() { + // A stored batch_checksum mutated after stamping fails the validating + // decode but passes the trusted one: exactly why the trusted variant is + // confined to locally-produced bytes (see its doc invariant). + let mut owned = valid_prepare_bytes(); + let corrupt_index = std::mem::size_of::() + BATCH_CHECKSUM_OFFSET; + owned.as_mut_slice()[corrupt_index] ^= 0xFF; + + assert!( + matches!( + decode_prepare_slice(owned.as_slice()), + Err(IggyError::InvalidBatchChecksum(..)) + ), + "validating decode must reject a mutated batch checksum", + ); + assert!( + decode_prepare_slice_trusted(owned.as_slice()).is_ok(), + "trusted decode skips the batch-checksum recomputation", + ); + } + #[test] fn decode_prepare_slice_size_below_header_size_does_not_panic() { // Regression: without the `total_size < header_size` guard, @@ -869,4 +1140,365 @@ mod tests { ); let _ = decode_prepare_slice(misaligned); } + + fn sample_messages() -> IggyMessages2 { + let mut messages = IggyMessages2::with_capacity(2); + messages.push(IggyMessage2 { + header: IggyMessage2Header { + id: 7, + origin_timestamp: 1_000, + ..Default::default() + }, + payload: Bytes::from_static(b"first-payload"), + user_headers: None, + }); + messages.push(IggyMessage2 { + header: IggyMessage2Header { + id: 8, + origin_timestamp: 1_050, + ..Default::default() + }, + payload: Bytes::from_static(b"second-payload"), + user_headers: Some(Bytes::from_static(b"user-header-bytes")), + }); + messages + } + + /// `[PrepareHeader][256B batch header][blob]` carrying real per-message + /// records + checksums from the production encoder, left pre-stamp + /// (`base_offset` / `base_timestamp` zero) as a follower receives it. + fn prepare_with_messages(messages: &IggyMessages2) -> Owned { + let namespace = IggyNamespace::new(1, 1, 7); + let owned = + SendMessages2Owned::from_messages(namespace, messages).expect("build send batch"); + prepare_from_owned(&owned) + } + + #[test] + fn verify_received_send_messages_accepts_clean_batch() { + let owned = prepare_with_messages(&sample_messages()); + verify_received_send_messages(owned.as_slice()) + .expect("a clean batch passes the receive gate"); + } + + #[test] + fn verify_received_send_messages_rejects_flipped_payload_byte() { + let mut owned = prepare_with_messages(&sample_messages()); + // First payload begins right after the first message's 48B header. + let payload_index = PREPARE_SPLIT_POINT + MESSAGE_HEADER_SIZE; + owned.as_mut_slice()[payload_index] ^= 0xFF; + assert!( + matches!( + verify_received_send_messages(owned.as_slice()), + Err(IggyError::InvalidMessageChecksum(..)) + ), + "a flipped payload byte must fail the per-message checksum", + ); + } + + #[test] + fn verify_received_send_messages_rejects_flipped_stored_checksum() { + let mut owned = prepare_with_messages(&sample_messages()); + // The first message's stored checksum is the first 8 bytes of the blob. + owned.as_mut_slice()[PREPARE_SPLIT_POINT] ^= 0xFF; + assert!( + matches!( + verify_received_send_messages(owned.as_slice()), + Err(IggyError::InvalidMessageChecksum(..)) + ), + "a flipped stored checksum must fail the per-message check", + ); + } + + #[test] + fn checksum_oneshot_matches_streaming_reference() { + // Formula pin: the per-message checksum is XxHash3-64 (default seed) + // over `header[8..48] || payload || user_headers` as one byte stream. + // The encoders hash the concatenation in a single oneshot pass; this + // streaming reference feeds the same parts separately. Both must agree + // for every shape, or checksums at rest stop verifying. + fn streaming_reference(header_tail: &[u8], payload: &[u8], user_headers: &[u8]) -> u64 { + let mut hasher = XxHash3_64::new(); + hasher.write(header_tail); + hasher.write(payload); + hasher.write(user_headers); + hasher.finish() + } + + let header_tail: Vec = (0u8..40).collect(); + let kilobyte: Vec = (0..1024u32).map(|index| (index % 251) as u8).collect(); + let cases: &[(&[u8], &[u8])] = &[ + (&[], &[]), + (b"payload-bytes", &[]), + (b"payload-bytes", b"user-header-bytes"), + (&kilobyte, &[]), + (&kilobyte, &kilobyte[..7]), + (&kilobyte[..1023], &kilobyte[..7]), + ]; + for (payload, user_headers) in cases { + let mut concatenated = + Vec::with_capacity(header_tail.len() + payload.len() + user_headers.len()); + concatenated.extend_from_slice(&header_tail); + concatenated.extend_from_slice(payload); + concatenated.extend_from_slice(user_headers); + assert_eq!( + XxHash3_64::oneshot(&concatenated), + streaming_reference(&header_tail, payload, user_headers), + "oneshot must match the streaming reference for payload {} B, user headers {} B", + payload.len(), + user_headers.len(), + ); + } + } + + #[test] + fn batch_checksum_v2_pins_header_fields_then_message_checksum_fields() { + // Formula pin for batch checksum v2: XxHash3-64 (default seed) streaming + // over the six batch header meta fields (LE, in field order) then each + // message's stored 8-byte checksum field in message order - never the + // bodies. This reference walks the blob by the KNOWN input message sizes, + // independent of the production frame decoder, and must equal what the + // encoder stamped, or a stamp will not verify against a read-back + // recompute. + let namespace = IggyNamespace::new(1, 1, 7); + let messages = sample_messages(); + let mut owned = + SendMessages2Owned::from_messages(namespace, &messages).expect("build batch"); + owned.header.base_offset = 100; + owned.header.base_timestamp = 200; + owned.header.batch_checksum = owned.header.checksum_for_blob(&owned.blob); + + let mut hasher = XxHash3_64::new(); + hasher.write(&owned.header.partition_id.to_le_bytes()); + hasher.write(&owned.header.base_offset.to_le_bytes()); + hasher.write(&owned.header.base_timestamp.to_le_bytes()); + hasher.write(&owned.header.origin_timestamp.to_le_bytes()); + hasher.write(&owned.header.batch_length.to_le_bytes()); + hasher.write(&owned.header.message_count.to_le_bytes()); + let mut frame_start = 0usize; + for message in messages.iter() { + hasher.write(&owned.blob[frame_start..frame_start + 8]); + let user_headers = message.user_headers.as_deref().unwrap_or_default(); + frame_start += MESSAGE_HEADER_SIZE + message.payload.len() + user_headers.len(); + } + let reference = hasher.finish(); + + assert_eq!( + frame_start, + owned.blob.len(), + "reference walk must consume the whole blob", + ); + assert_eq!( + owned.header.batch_checksum, reference, + "v2 batch checksum must equal hash(6 header fields || per-message checksum fields)", + ); + } + + #[test] + fn decode_batch_slice_rejects_body_corruption_with_intact_checksum_field() { + // Equal-integrity: v2 binds bodies only through the per-message checksum + // fields, so a flipped body byte that leaves the 8-byte checksum field + // intact keeps the batch value matching. The validating decode must still + // reject it via the per-message verify - the sole at-rest read-back check + // (the poll disk walk) decodes through here. + let namespace = IggyNamespace::new(1, 1, 7); + let owned = + SendMessages2Owned::from_messages(namespace, &sample_messages()).expect("build batch"); + let mut body = vec![0u8; COMMAND_HEADER_SIZE + owned.blob.len()]; + owned.header.encode_into(&mut body[..COMMAND_HEADER_SIZE]); + body[COMMAND_HEADER_SIZE..].copy_from_slice(&owned.blob); + + decode_batch_slice(&body).expect("the clean batch decodes"); + + // First payload byte sits right after the command header and the first + // message's 48B frame header, leaving that frame's checksum field intact. + let payload_index = COMMAND_HEADER_SIZE + MESSAGE_HEADER_SIZE; + body[payload_index] ^= 0xFF; + assert!( + matches!( + decode_batch_slice(&body), + Err(IggyError::InvalidMessageChecksum(..)) + ), + "body corruption with an intact checksum field must fail the per-message verify", + ); + } + + #[test] + fn decode_prepare_slice_rejects_body_corruption_with_intact_checksum_field() { + // The same equal-integrity guarantee at the resident/repair validating + // decode, plus proof that the batch value alone is blind to it. + let mut owned = prepare_with_messages(&sample_messages()); + decode_prepare_slice(owned.as_slice()).expect("the clean prepare decodes"); + + let payload_index = PREPARE_SPLIT_POINT + MESSAGE_HEADER_SIZE; + owned.as_mut_slice()[payload_index] ^= 0xFF; + assert!( + matches!( + decode_prepare_slice(owned.as_slice()), + Err(IggyError::InvalidMessageChecksum(..)) + ), + "body corruption with an intact checksum field must fail the per-message verify", + ); + assert!( + decode_prepare_slice_trusted(owned.as_slice()).is_ok(), + "the intact checksum field leaves the batch value matching, so trusted still passes", + ); + } + + /// Legacy `SendMessages` request body: `[metadata_len=4][message_count]` + /// then `count` skipped index slots, then the 64B-header legacy records. + fn legacy_send_messages_body(messages: &IggyMessages2) -> Vec { + let count = messages.count(); + let mut body = Vec::new(); + body.extend_from_slice(&4u32.to_le_bytes()); + body.extend_from_slice(&count.to_le_bytes()); + body.extend_from_slice(&vec![0u8; count as usize * INDEX_SIZE]); + for message in messages.iter() { + let user_headers = message.user_headers.as_deref().unwrap_or_default(); + let mut header = [0u8; LEGACY_MESSAGE_HEADER_SIZE]; + header[8..24].copy_from_slice(&message.header.id.to_le_bytes()); + header[40..48].copy_from_slice(&message.header.origin_timestamp.to_le_bytes()); + header[48..52].copy_from_slice(&(user_headers.len() as u32).to_le_bytes()); + header[52..56].copy_from_slice(&(message.payload.len() as u32).to_le_bytes()); + body.extend_from_slice(&header); + body.extend_from_slice(&message.payload); + body.extend_from_slice(user_headers); + } + body + } + + fn legacy_request_message(body: &[u8]) -> Message { + let header_size = std::mem::size_of::(); + let total = header_size + body.len(); + let mut buffer = Owned::::zeroed(total); + { + let header: &mut RequestHeader = + bytemuck::checked::try_from_bytes_mut(&mut buffer.as_mut_slice()[..header_size]) + .expect("zeroed bytes form a valid RequestHeader"); + header.command = Command2::Request; + header.operation = Operation::SendMessages; + header.client = 1; + header.session = 1; + header.request = 1; + header.size = u32::try_from(total).expect("size fits u32"); + } + buffer.as_mut_slice()[header_size..].copy_from_slice(body); + Message::try_from(buffer).expect("legacy request message is valid") + } + + #[test] + fn convert_request_message_transcodes_legacy_to_canonical_bytes() { + // Golden: the fused legacy transcode must emit the exact canonical batch + // the native builder (`from_messages`) produces for the same messages - + // command header + blob, byte for byte. Explicit non-zero ids keep it + // deterministic (no `random_id` substitution). + let namespace = IggyNamespace::new(1, 1, 3); + let messages = sample_messages(); + + let owned = + SendMessages2Owned::from_messages(namespace, &messages).expect("build canonical batch"); + let mut expected_body = vec![0u8; COMMAND_HEADER_SIZE + owned.blob.len()]; + owned + .header + .encode_into(&mut expected_body[..COMMAND_HEADER_SIZE]); + expected_body[COMMAND_HEADER_SIZE..].copy_from_slice(&owned.blob); + + let legacy = legacy_request_message(&legacy_send_messages_body(&messages)); + let converted = convert_request_message(namespace, legacy, ChecksumMode::Compute) + .expect("legacy body transcodes"); + let header_size = std::mem::size_of::(); + let actual_body = &converted.as_slice()[header_size..converted.header().size as usize]; + + assert_eq!( + actual_body, expected_body, + "legacy transcode must be byte-identical to the canonical native batch", + ); + + // And the emitted batch is self-consistent: it validates through the + // batch-checksum decode and yields the original messages. + let decoded = decode_batch_slice(actual_body).expect("transcoded batch checksum is valid"); + assert_eq!(decoded.message_count(), messages.count()); + let payloads: Vec<&[u8]> = decoded.iter().map(|view| view.payload).collect(); + assert_eq!( + payloads, + vec![&b"first-payload"[..], &b"second-payload"[..]] + ); + } + + #[test] + fn convert_request_message_skip_leaves_batch_checksum_zero_until_stamp() { + // The partition ingest path passes Skip: the transcoded batch must carry + // a zero checksum (stamp fills it) and be otherwise byte-identical to the + // Compute output - the flag toggles nothing but that one hash. + let namespace = IggyNamespace::new(1, 1, 3); + let messages = sample_messages(); + let body = legacy_send_messages_body(&messages); + let header_size = std::mem::size_of::(); + + let computed = convert_request_message( + namespace, + legacy_request_message(&body), + ChecksumMode::Compute, + ) + .expect("compute transcode"); + let skipped = + convert_request_message(namespace, legacy_request_message(&body), ChecksumMode::Skip) + .expect("skip transcode"); + + let computed_body = &computed.as_slice()[header_size..computed.header().size as usize]; + let skipped_body = &skipped.as_slice()[header_size..skipped.header().size as usize]; + + let skipped_header = SendMessages2Header::decode(&skipped_body[..COMMAND_HEADER_SIZE]) + .expect("decode skipped header"); + assert_eq!( + skipped_header.batch_checksum, 0, + "skip leaves the batch checksum zero until stamp", + ); + + // Patch only the 8-byte batch_checksum field into the skipped body; it + // must then equal the computed body, proving nothing else diverges. + let mut patched = skipped_body.to_vec(); + patched[BATCH_CHECKSUM_OFFSET..BATCH_CHECKSUM_OFFSET + 8] + .copy_from_slice(&computed_body[BATCH_CHECKSUM_OFFSET..BATCH_CHECKSUM_OFFSET + 8]); + assert_eq!( + patched.as_slice(), + computed_body, + "skip and compute differ only in the batch_checksum field", + ); + } + + #[test] + fn encrypt_ingest_path_stays_canonical_through_flag_split() { + // Mirror the plane encrypt ingest sequence: convert(Compute) -> the + // validating decode encrypt performs on its input -> encrypt -> the + // validating decode the second convert performs as its discriminator -> + // convert(Skip) (the partition convert), which sees an already-canonical + // batch and returns it unchanged. Every decode must succeed. + let namespace = IggyNamespace::new(1, 1, 3); + let messages = sample_messages(); + let header_size = std::mem::size_of::(); + + let legacy = legacy_request_message(&legacy_send_messages_body(&messages)); + let canonical = convert_request_message(namespace, legacy, ChecksumMode::Compute) + .expect("pre-encrypt transcode"); + let canonical_body = &canonical.as_slice()[header_size..canonical.header().size as usize]; + decode_batch_slice(canonical_body).expect("encrypt input decode validates the checksum"); + + let encryptor = + EncryptorKind::Aes256Gcm(Aes256GcmEncryptor::new(&[7u8; 32]).expect("valid 32B key")); + let encrypted = encrypt_batch_request(canonical, &encryptor).expect("encrypt batch"); + let encrypted_body: Vec = + encrypted.as_slice()[header_size..encrypted.header().size as usize].to_vec(); + decode_batch_slice(&encrypted_body) + .expect("encrypt output drives the 2nd-convert discriminator"); + + let repassed = convert_request_message(namespace, encrypted, ChecksumMode::Skip) + .expect("second convert passes the canonical batch"); + let repassed_body = &repassed.as_slice()[header_size..repassed.header().size as usize]; + assert_eq!( + repassed_body, + encrypted_body.as_slice(), + "an already-canonical encrypted batch passes the partition convert untouched", + ); + } } diff --git a/core/simulator/src/client.rs b/core/simulator/src/client.rs index f511c10006..028adb38cf 100644 --- a/core/simulator/src/client.rs +++ b/core/simulator/src/client.rs @@ -72,7 +72,7 @@ pub struct SimClient { partition_counter: Cell, /// Deterministic per-message id source for produced messages. The real SDK /// sends `id: 0` and lets the server mint a random UUID - /// (`SendMessages2::from_legacy_request` -> `random_id::get_uuid`); that + /// (`transcode_legacy_request` -> `random_id::get_uuid`); that /// mint is unseeded, so under the deterministic executor a produce's /// replicated body bytes (and their checksums) would differ run to run, /// silently breaking seeded replay. Stamping a deterministic id here keeps @@ -516,7 +516,7 @@ impl SimClient { /// client). VSR clients resolve to an explicit partition before sending, so /// the sim always emits `WirePartitioning::PartitionId`: that is the shape /// the shell's `resolve_partition_request_namespace` decodes, and the raw - /// path converts it to `SendMessages2` via `from_legacy_request`. + /// path converts it to `SendMessages2` via `transcode_legacy_request`. /// /// # Panics /// Panics if a namespace id exceeds `u32` or the request buffer is invalid.