From 3a944be2876aedca4150e17c6a260dc4fd8bd528 Mon Sep 17 00:00:00 2001 From: Hubert Gruszecki Date: Mon, 27 Jul 2026 10:55:56 +0200 Subject: [PATCH] perf(server-ng): reach single-node latency parity with legacy server server-ng trailed the legacy server on single-node benchmarks: the default cpu allocation collapsed all work onto one shard, produce hashed each batch three times (convert, stamp, flush revalidation), and a lone consumer crossing a sealed segment degraded to full-segment scans through unbounded read handles. Shard allocation now defaults to numa:auto. Produce computes the batch checksum once and marks locally originated batches trusted; replicated blobs keep a per-message receive gate on followers, so transit integrity still holds end to end. Sealed segment read handles are capped by a per-partition LRU. The batch checksum is redefined to cover the six header meta fields plus each message's stored checksum field instead of the whole blob, binding bodies transitively through the per-message checksums that every validating decode re-verifies. Stamping a batch now hashes N*8 bytes instead of the full blob, removing the last full-blob pass from the produce path. Message data written by earlier server-ng builds fails checksum validation after this change; wipe local_data when upgrading. Per-message checksums are computed in a single oneshot pass over the encoded record, pinned by a streaming-vs-oneshot reference test. --- core/journal/src/prepare_journal.rs | 9 +- core/partitions/src/iggy_index_reader.rs | 35 +- core/partitions/src/iggy_partition.rs | 341 ++++++++- core/partitions/src/iggy_partitions.rs | 23 +- core/partitions/src/journal.rs | 13 +- core/partitions/src/log.rs | 295 +++++++- core/partitions/src/poll_plan.rs | 175 ++++- core/server-ng/config.toml | 3 +- core/server-ng/src/dispatch.rs | 8 +- core/server-ng/src/responses.rs | 2 +- core/server_common/src/send_messages2.rs | 844 ++++++++++++++++++++--- core/simulator/src/client.rs | 4 +- 12 files changed, 1585 insertions(+), 167 deletions(-) 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.