From 51b7d6956e81d5f2b93e8c619c0238cafcd29041 Mon Sep 17 00:00:00 2001 From: forhappy Date: Thu, 15 Oct 2026 12:53:08 -0700 Subject: [PATCH] perf(staging): bound residual decoding --- crab/src/git/push.rs | 2 +- crates/crab-staging/README.md | 11 + crates/crab-staging/src/coalesced_read.rs | 332 +++++++++++++++++ crates/crab-staging/src/lib.rs | 411 ++++++++++++---------- 4 files changed, 560 insertions(+), 196 deletions(-) create mode 100644 crates/crab-staging/src/coalesced_read.rs diff --git a/crab/src/git/push.rs b/crab/src/git/push.rs index 5e46f30aa..2a2381029 100644 --- a/crab/src/git/push.rs +++ b/crab/src/git/push.rs @@ -11738,7 +11738,7 @@ impl PushPipeline { } while let Some(payloads) = read_plan - .read_next(batch_size, XORB_PACK_READ_PAYLOAD_LIMIT) + .read_next_bounded(batch_size, XORB_PACK_READ_PAYLOAD_LIMIT) .await? { check_cancelled(&self.cancel)?; diff --git a/crates/crab-staging/README.md b/crates/crab-staging/README.md index 7b954a0dc..4010dacae 100644 --- a/crates/crab-staging/README.md +++ b/crates/crab-staging/README.md @@ -101,6 +101,17 @@ Drop clears the reader marker and discards an uncommitted snapshot when no sibling reader prevents cleanup. Deferred journals remain for exclusive maintenance; native and uncaptured leases remain owned throughout. +The coalesced residual reader uses `CoalescedChunkReadPlan::read_next_bounded` after all +selected requests are appended. One prepared source is opened and fully verified +once, then retained as a format-bounded serialized body while decoded batches obey +the caller's count and byte limits. A single oversized chunk may exceed the byte +target to make progress; its format limit still applies. Returned chunks own their +allocations. Source descriptors are pinned before worker admission, and cancelling +a pending read preserves its result for the next call. The public `read_next` +method retains the complete-source behavior shipped in v1.2.4; it shares the same +source and chunk verification. Canonical callers should migrate to the bounded +method; removing the released method requires a breaking library release. + Native intent recovery may preserve caller-attested committed recipes under canonical history owners before rolling back an old Git index state. `preserve_publication_history` accepts only verified recipe roots belonging to diff --git a/crates/crab-staging/src/coalesced_read.rs b/crates/crab-staging/src/coalesced_read.rs new file mode 100644 index 000000000..18fdfffd9 --- /dev/null +++ b/crates/crab-staging/src/coalesced_read.rs @@ -0,0 +1,332 @@ +use std::fs::File; +use std::iter::Peekable; +use std::sync::Arc; +use std::vec::IntoIter; + +use super::{ + Bytes, IndexedCoalescedReadGroup, MerkleHash, Path, PreparedChunkLocator, ReaderPool, Result, + StagedChunkLocator, StagingAreaReadOnly, StagingError, StagingMetrics, XorbParser, + decode_prepared_chunk, load_prepared_source, lock_index, pin_segment_read_groups, push_plan, + read_located_staged_chunks_batch, read_prepared_staged_chunks_batch, +}; + +enum CoalescedChunkAuthority { + Segments(Vec<(u64, StagedChunkLocator)>), + Prepared(Vec<(u64, MerkleHash, PreparedChunkLocator)>), +} + +/// One bounded residual-read unit. Prepared units contain exactly one source xorb. +struct CoalescedChunkReadGroup { + authority: CoalescedChunkAuthority, +} + +/// Connection-local residual planner backed by SQLite temporary storage. +/// +/// Appends are bounded by the caller's recipe page. Reads remove one bounded +/// segment batch or one complete prepared source xorb at a time. +pub struct CoalescedChunkReadPlan<'a> { + staging: &'a StagingAreaReadOnly, + next_sequence: u64, + prepared: Option, + pending: Option>>, + buffered: Peekable>, +} + +impl<'a> CoalescedChunkReadPlan<'a> { + pub(super) fn new(staging: &'a StagingAreaReadOnly) -> Result { + lock_index(&staging.index)?.begin_coalesced_read_plan()?; + Ok(Self { + staging, + next_sequence: 0, + prepared: None, + pending: None, + buffered: Vec::new().into_iter().peekable(), + }) + } + + /// Append a bounded request page with an opaque caller context per chunk. + pub fn append(&mut self, requests: &[(MerkleHash, u64, u64)]) -> Result<()> { + let mut indexed = Vec::with_capacity(requests.len()); + for (hash, size, context) in requests { + indexed.push((self.next_sequence, <[u8; 32]>::from(*hash), *size, *context)); + self.next_sequence = self.next_sequence.checked_add(1).ok_or_else(|| { + StagingError::StagingCorrupt("residual request sequence overflow".to_owned()) + })?; + } + lock_index(&self.staging.index)?.append_coalesced_read_requests(&indexed) + } + + /// Read one segment batch or complete prepared source. + /// + /// This preserves the complete-source API shipped in v1.2.4. New callers + /// should use `read_next_bounded`; removing this method requires a breaking release. + pub async fn read_next( + &self, + max_segment_chunks: usize, + max_segment_bytes: u64, + ) -> Result>> { + // The complete-source method is the v1.2.4 public contract. Mixing it + // with a partially consumed bounded source would silently skip chunks. + if self.prepared.is_some() || self.pending.is_some() || self.buffered.len() > 0 { + return Err(StagingError::StagingCorrupt( + "bounded read is not drained".to_owned(), + )); + } + let group = lock_index(&self.staging.index)? + .take_coalesced_read_group(max_segment_chunks, max_segment_bytes)? + .map(|group| match group { + IndexedCoalescedReadGroup::Segments(chunks) => CoalescedChunkReadGroup { + authority: CoalescedChunkAuthority::Segments( + chunks + .into_iter() + .map(|(context, chunk)| (context, StagedChunkLocator::from(chunk))) + .collect(), + ), + }, + IndexedCoalescedReadGroup::Prepared(chunks) => CoalescedChunkReadGroup { + authority: CoalescedChunkAuthority::Prepared( + chunks + .into_iter() + .map(|(context, hash, locator)| { + (context, MerkleHash::from(hash), locator) + }) + .collect(), + ), + }, + }); + match group { + Some(group) => read_coalesced_chunk_group( + &self.staging.root, + &self.staging.readers, + self.staging.metrics.as_deref(), + group, + ) + .await + .map(Some), + None => Ok(None), + } + } + /// Read a bounded decoded batch while retaining at most one serialized source. + /// + /// One oversized chunk may exceed the byte target to ensure progress; the + /// xorb format bounds that chunk. Finish appending requests before reading + /// to coalesce every occurrence of one source into one verified open. + /// Retain selected recipe leases or a push snapshot until reading completes. + /// Cancelling a pending call keeps its worker result for the next call. + pub async fn read_next_bounded( + &mut self, + max_chunks: usize, + max_bytes: u64, + ) -> Result>> { + if self.buffered.len() == 0 { + if self.pending.is_none() { + let task = if let Some(cursor) = self.prepared.take() { + tokio::task::spawn_blocking(move || cursor.batch(max_chunks, max_bytes)) + } else { + let index = lock_index(&self.staging.index)?; + let Some(group) = index.take_coalesced_read_group(max_chunks, max_bytes)? + else { + return Ok(None); + }; + let metrics = self.staging.metrics.clone(); + match group { + IndexedCoalescedReadGroup::Segments(chunks) => { + let locators = chunks + .iter() + .map(|(_, chunk)| StagedChunkLocator::from(*chunk)) + .collect::>(); + let pinned = pin_segment_read_groups(&self.staging.readers, &locators)?; + tokio::spawn(async move { + let payloads = read_located_staged_chunks_batch( + pinned, + metrics.as_deref(), + &locators, + ) + .await?; + let chunks = chunks + .into_iter() + .zip(payloads) + .map(|((context, chunk), (_, data))| { + (context, MerkleHash::from(chunk.chunk_hash), data) + }) + .collect(); + Ok(DecodedBatch { next: None, chunks }) + }) + } + IndexedCoalescedReadGroup::Prepared(chunks) => { + let (_, _, first) = chunks.first().ok_or_else(|| { + StagingError::StagingCorrupt("empty prepared read group".to_owned()) + })?; + let identity = (first.xorb_hash, first.payload_hash, first.xorb_bytes); + // Pin before releasing locator ownership or queuing the + // worker; retirement cannot invalidate this descriptor. + let file = File::open(push_plan::prepared_xorb_path( + &self.staging.root, + &MerkleHash::from(identity.0), + ))?; + tokio::task::spawn_blocking(move || { + let activity = PreparedReadActivity::new(metrics); + let parser = load_prepared_source(file, identity)?; + if let Some(metrics) = activity.0.as_deref() { + metrics.add_staging_bytes_read(identity.2); + metrics.add_prepared_source_xorb_bytes_read(identity.2); + } + PreparedCursor { + parser, + chunks: chunks + .into_iter() + .map(|(context, hash, locator)| { + (context, MerkleHash::from(hash), locator) + }) + .collect::>() + .into_iter() + .peekable(), + } + .batch(max_chunks, max_bytes) + }) + } + } + }; + // Await the stored handle by reference. Cancelling this read keeps + // its owned worker and result available to the next call. + self.pending = Some(task); + } + let pending = self.pending.as_mut().ok_or_else(|| { + StagingError::Internal("bounded read worker disappeared".to_owned()) + })?; + let result = pending.await; + self.pending = None; + let batch = result.map_err(std::io::Error::other)??; + self.prepared = batch.next; + self.buffered = batch.chunks.into_iter().peekable(); + } + // A cancelled call may have queued a larger batch. Bound the returned + // prefix using this call's limits and retain the remaining owned bytes. + let mut out = Vec::new(); + let mut bytes = 0_u64; + while let Some(chunk) = self.buffered.next_if(|(_, _, data)| { + out.is_empty() + || (out.len() < max_chunks.max(1) + && bytes.saturating_add(data.len() as u64) <= max_bytes.max(1)) + }) { + bytes = bytes.saturating_add(chunk.2.len() as u64); + out.push(chunk); + } + Ok(Some(out)) + } +} + +impl Drop for CoalescedChunkReadPlan<'_> { + fn drop(&mut self) { + if let Ok(index) = lock_index(&self.staging.index) { + let _ = index.clear_coalesced_read_plan(); + } + } +} + +struct DecodedBatch { + next: Option, + chunks: Vec<(u64, MerkleHash, Bytes)>, +} + +// The worker owns activity accounting, so dropping an awaiting future cannot +// leave a reader permanently active after its pinned source is released. +struct PreparedReadActivity(Option>); + +impl PreparedReadActivity { + fn new(metrics: Option>) -> Self { + if let Some(metrics) = metrics.as_deref() { + metrics.inc_prepared_source_xorb_open(); + metrics.prepared_source_reader_started(); + } + Self(metrics) + } +} + +impl Drop for PreparedReadActivity { + fn drop(&mut self) { + if let Some(metrics) = self.0.as_deref() { + metrics.prepared_source_reader_finished(); + } + } +} + +struct PreparedCursor { + parser: XorbParser, + chunks: Peekable>, +} + +impl PreparedCursor { + fn batch(mut self, max_chunks: usize, max_bytes: u64) -> Result { + let mut out = Vec::new(); + let mut bytes = 0_u64; + while let Some((context, hash, locator)) = self.chunks.next_if(|(_, _, locator)| { + out.is_empty() + || (out.len() < max_chunks.max(1) + && bytes.saturating_add(u64::from(locator.size)) <= max_bytes.max(1)) + }) { + let data = + decode_prepared_chunk(&self.parser, hash, locator.chunk_index, locator.size)?; + bytes = bytes.saturating_add(u64::from(locator.size)); + out.push((context, hash, data)); + } + let next = self.chunks.peek().is_some().then_some(self); + Ok(DecodedBatch { next, chunks: out }) + } +} + +async fn read_coalesced_chunk_group( + root: &Path, + readers: &ReaderPool, + metrics: Option<&dyn StagingMetrics>, + group: CoalescedChunkReadGroup, +) -> Result> { + match group.authority { + CoalescedChunkAuthority::Segments(chunks) => { + let locators = chunks + .iter() + .map(|(_, locator)| *locator) + .collect::>(); + let payloads = read_located_staged_chunks_batch( + pin_segment_read_groups(readers, &locators)?, + metrics, + &locators, + ) + .await?; + Ok(chunks + .into_iter() + .zip(payloads) + .map(|((position, locator), (_, payload))| (position, locator.hash, payload)) + .collect()) + } + CoalescedChunkAuthority::Prepared(chunks) => { + let contexts = chunks + .iter() + .map(|(context, _, _)| *context) + .collect::>(); + let indexed = chunks + .into_iter() + .enumerate() + .map(|(position, (_, hash, locator))| (position, hash, locator)) + .collect::>(); + Ok(read_prepared_staged_chunks_batch(root, metrics, &indexed) + .await? + .into_iter() + .map(|(position, payload)| { + indexed + .get(position) + .and_then(|(_, hash, _)| { + contexts + .get(position) + .map(|context| (*context, *hash, payload)) + }) + .ok_or_else(|| { + StagingError::Internal( + "prepared read returned an unknown position".to_owned(), + ) + }) + }) + .collect::>>()?) + } + } +} diff --git a/crates/crab-staging/src/lib.rs b/crates/crab-staging/src/lib.rs index 025bbbc30..564ee280d 100644 --- a/crates/crab-staging/src/lib.rs +++ b/crates/crab-staging/src/lib.rs @@ -38,6 +38,7 @@ //! `HashMap` lookups, never across `.await`. mod add_push_plan; +mod coalesced_read; pub mod config; pub mod error; mod index; @@ -67,6 +68,7 @@ mod prop_orphan_sweep_idempotence; #[path = "../tests/unit/prop_torn_tail_recovery.rs"] mod prop_torn_tail_recovery; +pub use coalesced_read::CoalescedChunkReadPlan; pub use config::StagingConfig; pub use error::{Result, StagingError}; pub use metrics::StagingMetrics; @@ -1367,88 +1369,6 @@ pub(crate) struct StagedChunkLocator { pub locator: ChunkLocator, } -enum CoalescedChunkAuthority { - Segments(Vec<(u64, StagedChunkLocator)>), - Prepared(Vec<(u64, MerkleHash, PreparedChunkLocator)>), -} - -/// One bounded residual-read unit. Prepared units contain exactly one source xorb. -struct CoalescedChunkReadGroup { - authority: CoalescedChunkAuthority, -} - -/// Connection-local residual planner backed by SQLite temporary storage. -/// -/// Appends are bounded by the caller's recipe page. Reads remove one bounded -/// segment batch or one complete prepared source xorb at a time. -pub struct CoalescedChunkReadPlan<'a> { - staging: &'a StagingAreaReadOnly, - next_sequence: u64, -} - -impl CoalescedChunkReadPlan<'_> { - /// Append a bounded request page with an opaque caller context per chunk. - pub fn append(&mut self, requests: &[(MerkleHash, u64, u64)]) -> Result<()> { - let mut indexed = Vec::with_capacity(requests.len()); - for (hash, size, context) in requests { - indexed.push((self.next_sequence, <[u8; 32]>::from(*hash), *size, *context)); - self.next_sequence = self.next_sequence.checked_add(1).ok_or_else(|| { - StagingError::StagingCorrupt("residual request sequence overflow".to_owned()) - })?; - } - lock_index(&self.staging.index)?.append_coalesced_read_requests(&indexed) - } - - /// Read and remove the next bounded residual group. - pub async fn read_next( - &self, - max_segment_chunks: usize, - max_segment_bytes: u64, - ) -> Result>> { - let group = lock_index(&self.staging.index)? - .take_coalesced_read_group(max_segment_chunks, max_segment_bytes)? - .map(|group| match group { - IndexedCoalescedReadGroup::Segments(chunks) => CoalescedChunkReadGroup { - authority: CoalescedChunkAuthority::Segments( - chunks - .into_iter() - .map(|(context, chunk)| (context, StagedChunkLocator::from(chunk))) - .collect(), - ), - }, - IndexedCoalescedReadGroup::Prepared(chunks) => CoalescedChunkReadGroup { - authority: CoalescedChunkAuthority::Prepared( - chunks - .into_iter() - .map(|(context, hash, locator)| { - (context, MerkleHash::from(hash), locator) - }) - .collect(), - ), - }, - }); - match group { - Some(group) => read_coalesced_chunk_group( - &self.staging.root, - &self.staging.readers, - self.staging.metrics.as_deref(), - group, - ) - .await - .map(Some), - None => Ok(None), - } - } -} - -impl Drop for CoalescedChunkReadPlan<'_> { - fn drop(&mut self) { - if let Ok(index) = lock_index(&self.staging.index) { - let _ = index.clear_coalesced_read_plan(); - } - } -} - struct IndexedPreparedXorbStatsRef { payload_hash: [u8; 32], bytes: u64, @@ -1539,6 +1459,74 @@ async fn read_staged_chunks_batch( type PreparedXorbIdentity = ([u8; 32], [u8; 32], u64); +fn load_prepared_source(file: File, identity: PreparedXorbIdentity) -> Result { + let (xorb_hash, payload_hash, expected_bytes) = identity; + if expected_bytes > MAX_XORB_SIZE as u64 { + return Err(StagingError::StagingCorrupt(format!( + "prepared xorb {} declares {expected_bytes} bytes above the {MAX_XORB_SIZE}-byte format bound", + MerkleHash::from(xorb_hash).hex() + ))); + } + let file_bytes = file.metadata()?.len(); + if file_bytes != expected_bytes { + return Err(StagingError::StagingCorrupt(format!( + "prepared xorb {} payload changed: expected {expected_bytes} bytes, found {file_bytes} bytes", + MerkleHash::from(xorb_hash).hex() + ))); + } + let capacity = usize::try_from(expected_bytes).map_err(|_| { + StagingError::StagingCorrupt(format!( + "prepared xorb {} size does not fit memory", + MerkleHash::from(xorb_hash).hex() + )) + })?; + let mut payload = Vec::with_capacity(capacity); + file.take(expected_bytes.saturating_add(1)) + .read_to_end(&mut payload)?; + let payload = Bytes::from(payload); + let actual_payload_hash = blake3::hash(&payload); + let size_matches = u64::try_from(payload.len()) == Ok(expected_bytes); + let hash_matches = *actual_payload_hash.as_bytes() == payload_hash; + if !size_matches || !hash_matches { + return Err(StagingError::StagingCorrupt(format!( + "prepared xorb {} payload changed: expected {} bytes/{}, found {} bytes/{} (size_matches={size_matches}, hash_matches={hash_matches})", + MerkleHash::from(xorb_hash).hex(), + expected_bytes, + MerkleHash::from(payload_hash).hex(), + payload.len(), + actual_payload_hash.to_hex() + ))); + } + let parser = XorbParser::parse(payload)?; + if parser.hash() != MerkleHash::from(xorb_hash) { + return Err(StagingError::StagingCorrupt(format!( + "prepared xorb {} metadata hash changed", + MerkleHash::from(xorb_hash).hex() + ))); + } + parser.verify_payload_digest()?; + + Ok(parser) +} + +fn decode_prepared_chunk( + parser: &XorbParser, + expected_hash: MerkleHash, + chunk_index: u32, + expected_size: u32, +) -> Result { + let chunk = parser.get_chunk(chunk_index)?; + if chunk.hash != expected_hash || chunk.data.len() != expected_size as usize { + return Err(StagingError::StagingCorrupt(format!( + "prepared xorb chunk {} does not match its indexed placement", + expected_hash.hex() + ))); + } + // Raw chunks can share the encoded body. Detach each result so a consumed + // batch cannot retain old source bodies while the caller packs another one. + Ok(Bytes::copy_from_slice(&chunk.data)) +} + async fn read_prepared_staged_chunks_batch( root: &Path, metrics: Option<&dyn StagingMetrics>, @@ -1568,59 +1556,17 @@ async fn read_prepared_staged_chunks_batch( MerkleHash::from(xorb_hash).hex() ))); } - let file = File::open(&path)?; - let file_bytes = file.metadata()?.len(); - if file_bytes != expected_bytes { - return Err(StagingError::StagingCorrupt(format!( - "prepared xorb {} payload changed: expected {expected_bytes} bytes, found {file_bytes} bytes", - MerkleHash::from(xorb_hash).hex() - ))); - } - let capacity = usize::try_from(expected_bytes).map_err(|_| { - StagingError::StagingCorrupt(format!( - "prepared xorb {} size does not fit memory", - MerkleHash::from(xorb_hash).hex() - )) - })?; - let mut payload = Vec::with_capacity(capacity); - file.take(expected_bytes.saturating_add(1)) - .read_to_end(&mut payload)?; - let payload = Bytes::from(payload); - let actual_payload_hash = blake3::hash(&payload); - let size_matches = u64::try_from(payload.len()) == Ok(expected_bytes); - let hash_matches = *actual_payload_hash.as_bytes() == payload_hash; - if !size_matches || !hash_matches { - return Err(StagingError::StagingCorrupt(format!( - "prepared xorb {} payload changed: expected {} bytes/{}, found {} bytes/{} (size_matches={size_matches}, hash_matches={hash_matches})", - MerkleHash::from(xorb_hash).hex(), - expected_bytes, - MerkleHash::from(payload_hash).hex(), - payload.len(), - actual_payload_hash.to_hex() - ))); - } - let parser = XorbParser::parse(payload)?; - if parser.hash() != MerkleHash::from(xorb_hash) { - return Err(StagingError::StagingCorrupt(format!( - "prepared xorb {} metadata hash changed", - MerkleHash::from(xorb_hash).hex() - ))); - } - parser.verify_payload_digest()?; + let parser = load_prepared_source( + File::open(&path)?, + (xorb_hash, payload_hash, expected_bytes), + )?; let mut decoded = Vec::with_capacity(wanted.len()); for (position, expected_hash, chunk_index, expected_size) in wanted { - let chunk = parser.get_chunk(chunk_index)?; - if chunk.hash != expected_hash || chunk.data.len() != expected_size as usize { - return Err(StagingError::StagingCorrupt(format!( - "prepared xorb chunk {} does not match its indexed placement", - expected_hash.hex() - ))); - } - // Parser chunks are slices of the full serialized xorb. The - // caller can retain this batch while packing other sources, - // so detach each requested chunk before dropping the parser. - decoded.push((position, Bytes::copy_from_slice(&chunk.data))); + decoded.push(( + position, + decode_prepared_chunk(&parser, expected_hash, chunk_index, expected_size)?, + )); } Ok::<_, StagingError>((expected_bytes, decoded)) }); @@ -1640,62 +1586,6 @@ async fn read_prepared_staged_chunks_batch( Ok(out) } -async fn read_coalesced_chunk_group( - root: &Path, - readers: &ReaderPool, - metrics: Option<&dyn StagingMetrics>, - group: CoalescedChunkReadGroup, -) -> Result> { - match group.authority { - CoalescedChunkAuthority::Segments(chunks) => { - let locators = chunks - .iter() - .map(|(_, locator)| *locator) - .collect::>(); - let payloads = read_located_staged_chunks_batch( - pin_segment_read_groups(readers, &locators)?, - metrics, - &locators, - ) - .await?; - Ok(chunks - .into_iter() - .zip(payloads) - .map(|((position, locator), (_, payload))| (position, locator.hash, payload)) - .collect()) - } - CoalescedChunkAuthority::Prepared(chunks) => { - let contexts = chunks - .iter() - .map(|(context, _, _)| *context) - .collect::>(); - let indexed = chunks - .into_iter() - .enumerate() - .map(|(position, (_, hash, locator))| (position, hash, locator)) - .collect::>(); - Ok(read_prepared_staged_chunks_batch(root, metrics, &indexed) - .await? - .into_iter() - .map(|(position, payload)| { - indexed - .get(position) - .and_then(|(_, hash, _)| { - contexts - .get(position) - .map(|context| (*context, *hash, payload)) - }) - .ok_or_else(|| { - StagingError::Internal( - "prepared read returned an unknown position".to_owned(), - ) - }) - }) - .collect::>>()?) - } - } -} - async fn verify_pending_recipe_payload( root: &Path, index: &Mutex, @@ -3927,11 +3817,7 @@ impl StagingAreaReadOnly { /// Begin one disk-backed residual plan for bounded recipe-page appends. pub fn begin_coalesced_chunk_read_plan(&self) -> Result> { - lock_index(&self.index)?.begin_coalesced_read_plan()?; - Ok(CoalescedChunkReadPlan { - staging: self, - next_sequence: 0, - }) + CoalescedChunkReadPlan::new(self) } /// Check if a chunk exists in staging. @@ -5464,8 +5350,7 @@ mod tests { ); } - #[tokio::test(flavor = "multi_thread")] - async fn prepared_residual_plan_reads_each_unique_source_once() { + async fn prepared_residual_fixture() -> (tempfile::TempDir, Vec<(MerkleHash, Vec)>, Bytes) { let chunks = (0..8) .map(|index| make_chunk(100 + index, 4096)) .collect::>(); @@ -5536,6 +5421,16 @@ mod tests { .expect("remove pending segment authority"); } + (tmp, chunks, xorb_bytes) + } + + #[tokio::test(flavor = "multi_thread")] + async fn prepared_residual_plan_reads_each_unique_source_once() { + let (tmp, chunks, xorb_bytes) = prepared_residual_fixture().await; + let chunk_pairs = chunks + .iter() + .map(|(hash, data)| (*hash, data.len() as u64)) + .collect::>(); let metrics = Arc::new(CountingMetrics::default()); let mut staging = StagingAreaReadOnly::open(tmp.path().to_path_buf()) .await @@ -5574,6 +5469,132 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread")] + async fn bounded_prepared_reads_detach_batches_and_reopen_after_drop() { + let (tmp, chunks, xorb_bytes) = prepared_residual_fixture().await; + let metrics = Arc::new(CountingMetrics::default()); + let mut staging = StagingAreaReadOnly::open(tmp.path().to_owned()) + .await + .unwrap(); + staging.set_metrics(metrics.clone()); + for stop_early in [true, false] { + let mut plan = staging.begin_coalesced_chunk_read_plan().unwrap(); + let requests = chunks + .iter() + .map(|(hash, data)| (*hash, data.len() as u64, 7)) + .collect::>(); + plan.append(&requests).unwrap(); + let mut actual = Vec::new(); + let mut batch_sizes = Vec::new(); + while let Some(batch) = plan.read_next_bounded(2, 8192).await.unwrap() { + batch_sizes.push(batch.len()); + for (context, hash, data) in batch { + assert_eq!(context, 7); + assert!( + data.is_unique(), + "decoded batches cannot retain the source allocation" + ); + actual.push((hash, data.to_vec())); + } + if stop_early { + break; + } + } + if !stop_early { + assert_eq!((batch_sizes, actual), (vec![2; 4], chunks.clone())); + } + } + assert_eq!( + metrics.prepared_source_counts(), + (2, 2 * xorb_bytes.len() as u64, 1) + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn bounded_prepared_reads_reject_changed_payload_before_returning_chunks() { + let (tmp, chunks, _) = prepared_residual_fixture().await; + let staging = StagingAreaReadOnly::open(tmp.path().to_owned()) + .await + .unwrap(); + let mut plan = staging.begin_coalesced_chunk_read_plan().unwrap(); + plan.append(&[(chunks[0].0, chunks[0].1.len() as u64, 0)]) + .unwrap(); + let files = std::fs::read_dir(tmp.path().join("push-plans/payloads")) + .unwrap() + .flat_map(|directory| std::fs::read_dir(directory.unwrap().path()).unwrap()) + .map(|entry| entry.unwrap().path()) + .collect::>(); + assert_eq!(files.len(), 1); + let mut bytes = std::fs::read(&files[0]).unwrap(); + bytes[0] ^= 1; + std::fs::write(&files[0], bytes).unwrap(); + assert!(matches!( + plan.read_next_bounded(2, 8192).await, + Err(StagingError::StagingCorrupt(_)) + )); + } + + #[test] + fn cancelled_bounded_read_resumes_its_queued_source() { + use std::future::{Future, poll_fn}; + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .max_blocking_threads(1) + .enable_all() + .build() + .unwrap(); + let (tmp, chunks, _) = runtime.block_on(prepared_residual_fixture()); + let staging = runtime + .block_on(StagingAreaReadOnly::open(tmp.path().to_owned())) + .unwrap(); + let mut plan = staging.begin_coalesced_chunk_read_plan().unwrap(); + let requests = chunks + .iter() + .map(|(hash, data)| (*hash, data.len() as u64, 0)) + .collect::>(); + plan.append(&requests).unwrap(); + let (release, wait) = std::sync::mpsc::channel(); + let (ready, started) = std::sync::mpsc::channel(); + let blocker = runtime.handle().spawn_blocking(move || { + ready.send(()).unwrap(); + let _ = wait.recv(); + }); + started.recv().unwrap(); + runtime.block_on(async { + // Keep verification queued so cancellation occurs after source + // ownership is captured and before its worker can return a batch. + let mut read = Box::pin(plan.read_next_bounded(2, 8192)); + poll_fn(|cx| { + assert!(read.as_mut().poll(cx).is_pending()); + std::task::Poll::Ready(()) + }) + .await; + drop(read); + #[cfg(unix)] + for directory in std::fs::read_dir(tmp.path().join("push-plans/payloads")).unwrap() { + for entry in std::fs::read_dir(directory.unwrap().path()).unwrap() { + std::fs::remove_file(entry.unwrap().path()).unwrap(); + } + } + release.send(()).unwrap(); + let mut actual = Vec::new(); + while let Some(batch) = plan.read_next_bounded(1, 4096).await.unwrap() { + assert_eq!( + batch.len(), + 1, + "resuming with a smaller budget must bound the queued result" + ); + actual.extend( + batch + .into_iter() + .map(|(_, hash, data)| (hash, data.to_vec())), + ); + } + assert_eq!(actual, chunks); + blocker.await.unwrap(); + }); + } + fn insert_extra_prepared_xorb_row( root: &std::path::Path, _file_hash: MerkleHash,