diff --git a/bindings/uniffi/src/config.rs b/bindings/uniffi/src/config.rs index ecebff4e5..fcca3fee0 100644 --- a/bindings/uniffi/src/config.rs +++ b/bindings/uniffi/src/config.rs @@ -188,6 +188,10 @@ impl TryFrom for slatedb::DbReaderMode { pub struct ReaderOptions { /// How often the reader polls for new manifests and WAL data, in milliseconds. pub manifest_poll_interval_ms: u64, + + /// How frequently an open reader probes the exact next WAL ID. + #[uniffi(default = 1000)] + pub wal_poll_interval_ms: u64, /// Lifetime of an internally managed checkpoint, in milliseconds. pub checkpoint_lifetime_ms: u64, /// Maximum size of one in-memory table used while replaying WAL data. @@ -206,6 +210,7 @@ impl Default for ReaderOptions { fn default() -> Self { Self { manifest_poll_interval_ms: 10_000, + wal_poll_interval_ms: 1_000, checkpoint_lifetime_ms: 600_000, max_memtable_bytes: 64 * 1024 * 1024, skip_wal_replay: false, @@ -218,6 +223,7 @@ impl From for slatedb::config::DbReaderOptions { fn from(value: ReaderOptions) -> Self { slatedb::config::DbReaderOptions { manifest_poll_interval: Duration::from_millis(value.manifest_poll_interval_ms), + wal_poll_interval: Duration::from_millis(value.wal_poll_interval_ms), checkpoint_lifetime: Duration::from_millis(value.checkpoint_lifetime_ms), max_memtable_bytes: value.max_memtable_bytes, skip_wal_replay: value.skip_wal_replay, @@ -528,6 +534,7 @@ impl From for slatedb::config::GarbageCollectorOptions #[cfg(test)] mod tests { use super::{GarbageCollectorOptions, ReaderOptions}; + use std::time::Duration; #[test] fn boundary_files_are_enabled_by_default() { @@ -572,17 +579,20 @@ mod tests { let reader: slatedb::config::DbReaderOptions = ReaderOptions::default().into(); assert_eq!(reader.object_store_max_retries, None); + assert_eq!(reader.wal_poll_interval, Duration::from_secs(1)); } #[test] fn reader_object_store_max_retries_threads_through() { let reader: slatedb::config::DbReaderOptions = ReaderOptions { object_store_max_retries: Some(5), + wal_poll_interval_ms: 250, ..ReaderOptions::default() } .into(); assert_eq!(reader.object_store_max_retries, Some(5)); + assert_eq!(reader.wal_poll_interval, Duration::from_millis(250)); } } diff --git a/slatedb/benches/db_reader_memory_scaling.rs b/slatedb/benches/db_reader_memory_scaling.rs index baa8ecaa5..53b6af674 100644 --- a/slatedb/benches/db_reader_memory_scaling.rs +++ b/slatedb/benches/db_reader_memory_scaling.rs @@ -1,3 +1,7 @@ +// Benchmarks intentionally use unique object-store namespaces and stdout result records. Those +// operations are forbidden in library code but are the benchmark contract. +#![allow(clippy::disallowed_macros, clippy::disallowed_methods)] + //! Heap-memory scaling benchmarks for DbReader snapshots and incremental WAL replay. //! //! This is separate from `db_reader_scaling` because its global allocator performs diff --git a/slatedb/benches/db_reader_scaling.rs b/slatedb/benches/db_reader_scaling.rs index c55c23c4e..cdde9bc28 100644 --- a/slatedb/benches/db_reader_scaling.rs +++ b/slatedb/benches/db_reader_scaling.rs @@ -1,3 +1,11 @@ +// Benchmarks intentionally use wall-clock timing, unique object-store namespaces, and stdout +// result records. Those operations are forbidden in library code but are the benchmark contract. +#![allow( + clippy::disallowed_macros, + clippy::disallowed_methods, + clippy::disallowed_types +)] + //! End-to-end scaling benchmarks for reader-backed snapshots and incremental WAL replay. //! //! This target deliberately uses fixed repetitions instead of Criterion's adaptive @@ -144,6 +152,7 @@ fn writer_settings() -> Settings { fn quiet_reader_options(max_memtable_bytes: u64) -> DbReaderOptions { DbReaderOptions { manifest_poll_interval: Duration::from_secs(60 * 60), + wal_poll_interval: Duration::from_secs(60 * 60), checkpoint_lifetime: Duration::from_secs(3 * 60 * 60), max_memtable_bytes, ..DbReaderOptions::default() @@ -156,6 +165,7 @@ fn polling_reader_options( ) -> DbReaderOptions { DbReaderOptions { manifest_poll_interval: POLL_INTERVAL, + wal_poll_interval: POLL_INTERVAL, checkpoint_lifetime, max_memtable_bytes, ..DbReaderOptions::default() diff --git a/slatedb/src/blob.rs b/slatedb/src/blob.rs index 572f2f934..4986e356c 100644 --- a/slatedb/src/blob.rs +++ b/slatedb/src/blob.rs @@ -12,3 +12,48 @@ pub(crate) trait ReadOnlyBlob { #[allow(dead_code)] async fn read(&self) -> Result; } + +/// An immutable object held entirely in memory. +pub(crate) struct BytesBlob { + bytes: Bytes, +} + +impl BytesBlob { + pub(crate) fn new(bytes: Bytes) -> Self { + Self { bytes } + } +} + +impl ReadOnlyBlob for BytesBlob { + async fn len(&self) -> Result { + u64::try_from(self.bytes.len()).map_err(|err| { + SlateDBError::WalDataError(std::sync::Arc::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + err, + ))) + }) + } + + async fn read_range(&self, range: Range) -> Result { + let start = usize::try_from(range.start).ok(); + let end = usize::try_from(range.end).ok(); + let Some((start, end)) = start.zip(end) else { + return Err(invalid_range(range, self.bytes.len())); + }; + if start > end || end > self.bytes.len() { + return Err(invalid_range(range, self.bytes.len())); + } + Ok(self.bytes.slice(start..end)) + } + + async fn read(&self) -> Result { + Ok(self.bytes.clone()) + } +} + +fn invalid_range(range: Range, len: usize) -> SlateDBError { + SlateDBError::WalDataError(std::sync::Arc::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid in-memory WAL range {range:?} for object length {len}"), + ))) +} diff --git a/slatedb/src/block_iterator.rs b/slatedb/src/block_iterator.rs index 3a1f5455c..082d82620 100644 --- a/slatedb/src/block_iterator.rs +++ b/slatedb/src/block_iterator.rs @@ -105,7 +105,7 @@ pub(crate) struct BlockIterator { off_off: usize, // first key in the block, because slateDB does not support multi version of keys // so we use `Bytes` temporarily - first_key: Bytes, + first_key: Option, ordering: IterationOrder, } @@ -195,7 +195,7 @@ impl RowEntryIterator for BlockIterator { impl BlockIterator { pub(crate) fn new(block: B, ordering: IterationOrder) -> Self { BlockIterator { - first_key: BlockIterator::decode_first_key(&block), + first_key: None, block, off_off: 0, ordering, @@ -215,23 +215,41 @@ impl BlockIterator { self.off_off >= self.block.offsets().len() } - fn load_at_current_off(&self) -> Result, SlateDBError> { + fn load_at_current_off(&mut self) -> Result, SlateDBError> { if self.is_empty() { return Ok(None); } + self.ensure_first_key()?; let off_off = match self.ordering { Ascending => self.off_off, Descending => self.block.offsets().len() - 1 - self.off_off, }; - let off = self.block.offsets()[off_off]; - let off_usz = off as usize; - // TODO: bounds checks to avoid panics? (paulgb) - let mut cursor = self.block.data().slice(off_usz..); + let off_usz = usize::from(self.block.offsets()[off_off]); + let row_end = self + .block + .offsets() + .get(off_off + 1) + .map(|offset| usize::from(*offset)) + .unwrap_or_else(|| self.block.data().len()); + if off_usz >= row_end || row_end > self.block.data().len() { + return Err(corrupt_block("invalid V1 row boundary")); + } + let mut cursor = self.block.data().slice(off_usz..row_end); let codec = SstRowCodecV0::new(); let sst_row = codec.decode(&mut cursor)?; + if cursor.has_remaining() { + return Err(corrupt_block("V1 row contains trailing bytes")); + } + let first_key = self + .first_key + .as_ref() + .ok_or_else(|| corrupt_block("V1 block first key was not initialized"))?; + if sst_row.key_prefix_len > first_key.len() { + return Err(corrupt_block("V1 row key prefix exceeds first key")); + } Ok(Some(RowEntry::new( - sst_row.restore_full_key(&self.first_key), + sst_row.restore_full_key(first_key), sst_row.value, sst_row.seq, sst_row.create_ts, @@ -239,33 +257,77 @@ impl BlockIterator { ))) } - fn decode_first_key(block: &B) -> Bytes { + fn ensure_first_key(&mut self) -> Result<(), SlateDBError> { + if self.first_key.is_some() { + return Ok(()); + } + self.first_key = Some(Self::decode_first_key(&self.block)?); + Ok(()) + } + + fn decode_first_key(block: &B) -> Result { + if block.offsets().first().copied() != Some(0) { + return Err(corrupt_block("first V1 row offset is not zero")); + } + if block.data().len() < 4 { + return Err(corrupt_block("truncated V1 first key")); + } let mut buf = block.data().slice(..); let overlap_len = buf.get_u16() as usize; - assert_eq!(overlap_len, 0, "first key overlap should be 0"); + if overlap_len != 0 { + return Err(corrupt_block("first V1 key prefix is not zero")); + } let key_len = buf.get_u16() as usize; + if buf.remaining() < key_len { + return Err(corrupt_block("truncated V1 first key suffix")); + } let first_key = &buf[..key_len]; - Bytes::copy_from_slice(first_key) + Ok(Bytes::copy_from_slice(first_key)) } /// Decodes just the key at the given offset index without parsing the full row. /// This is more efficient for binary search where we only need to compare keys. - fn decode_key_at_index(&self, index: usize) -> Result { - let off = self.block.offsets()[index] as usize; + fn decode_key_at_index(&mut self, index: usize) -> Result { + self.ensure_first_key()?; + let Some(&offset) = self.block.offsets().get(index) else { + return Err(corrupt_block("V1 key offset index is out of range")); + }; + let off = usize::from(offset); + if off >= self.block.data().len() { + return Err(corrupt_block("V1 key offset is outside block data")); + } let mut cursor = self.block.data().slice(off..); + if cursor.remaining() < 4 { + return Err(corrupt_block("truncated V1 key lengths")); + } + let key_prefix_len = cursor.get_u16() as usize; let key_suffix_len = cursor.get_u16() as usize; + let first_key = self + .first_key + .as_ref() + .ok_or_else(|| corrupt_block("V1 block first key was not initialized"))?; + if key_prefix_len > first_key.len() || cursor.remaining() < key_suffix_len { + return Err(corrupt_block("invalid V1 key prefix or suffix length")); + } let key_suffix = &cursor[..key_suffix_len]; // Reconstruct the full key from first_key prefix + suffix - let mut full_key = BytesMut::with_capacity(key_prefix_len + key_suffix_len); - full_key.extend_from_slice(&self.first_key[..key_prefix_len]); + let key_len = key_prefix_len + .checked_add(key_suffix_len) + .ok_or_else(|| corrupt_block("V1 key length overflow"))?; + let mut full_key = BytesMut::with_capacity(key_len); + full_key.extend_from_slice(&first_key[..key_prefix_len]); full_key.extend_from_slice(key_suffix); Ok(full_key.freeze()) } } +fn corrupt_block(reason: &'static str) -> SlateDBError { + SlateDBError::CorruptSst { reason, path: None } +} + #[cfg(test)] mod tests { use crate::block_iterator::BlockIterator; @@ -622,7 +684,7 @@ mod tests { assert!(block_builder.add_value(b"prefix_bbb", b"2", None, None)); assert!(block_builder.add_value(b"prefix_ccc", b"3", None, None)); let block = block_builder.build().unwrap(); - let iter = BlockIterator::new_ascending(&block); + let mut iter = BlockIterator::new_ascending(&block); // when: decoding keys at each index // then: full keys are correctly reconstructed diff --git a/slatedb/src/block_iterator_v2.rs b/slatedb/src/block_iterator_v2.rs index e88e0ae85..e2d482097 100644 --- a/slatedb/src/block_iterator_v2.rs +++ b/slatedb/src/block_iterator_v2.rs @@ -1,5 +1,5 @@ use async_trait::async_trait; -use bytes::{Bytes, BytesMut}; +use bytes::{Buf, Bytes, BytesMut}; use crate::block_iterator::BlockLike; use crate::error::SlateDBError; @@ -32,24 +32,16 @@ pub(crate) struct BlockIteratorV2 { impl BlockIteratorV2 { pub(crate) fn new(block: B, ordering: IterationOrder) -> Self { match ordering { - IterationOrder::Ascending => { - let initial_key = if block.offsets().is_empty() { - Bytes::new() - } else { - Self::decode_first_key_at_restart(&block, 0) - }; - - BlockIteratorV2 { - inner: BlockIteratorInner::Ascending(AscendingState { - block, - current_restart_idx: 0, - offset_in_block: 0, - entries_since_restart: 0, - current_key: initial_key, - exhausted: false, - }), - } - } + IterationOrder::Ascending => BlockIteratorV2 { + inner: BlockIteratorInner::Ascending(AscendingState { + block, + current_restart_idx: 0, + offset_in_block: 0, + entries_since_restart: 0, + current_key: Bytes::new(), + exhausted: false, + }), + }, IterationOrder::Descending => BlockIteratorV2 { inner: BlockIteratorInner::Descending(DescendingBlockIteratorV2::new(block)), }, @@ -68,36 +60,66 @@ impl BlockIteratorV2 { } } - fn decode_first_key_at_restart(block: &B, restart_idx: usize) -> Bytes { - let restart_offset = block.offsets()[restart_idx] as usize; - let mut data = block.data().slice(restart_offset..); + fn decode_first_key_at_restart(block: &B, restart_idx: usize) -> Result { + let Some(&restart_offset) = block.offsets().get(restart_idx) else { + return Err(corrupt_block("V2 restart index is out of range")); + }; + let restart_offset = usize::from(restart_offset); + let region_end = block + .offsets() + .get(restart_idx + 1) + .map(|offset| usize::from(*offset)) + .unwrap_or_else(|| block.data().len()); + if restart_offset >= region_end || region_end > block.data().len() { + return Err(corrupt_block("invalid V2 restart region")); + } + let mut data = block.data().slice(restart_offset..region_end); let codec = SstRowCodecV2::new(); - let (shared_bytes, key_suffix) = codec.decode_key_only(&mut data); - assert_eq!(shared_bytes, 0, "restart point should have shared_bytes=0"); - key_suffix + let (shared_bytes, key_suffix) = codec.decode_key_only(&mut data)?; + if shared_bytes != 0 { + return Err(corrupt_block("V2 restart key has a shared prefix")); + } + Ok(key_suffix) } } impl AscendingState { - fn seek_to_restart(&mut self, restart_idx: usize) { + fn seek_to_restart(&mut self, restart_idx: usize) -> Result<(), SlateDBError> { if restart_idx >= self.block.offsets().len() { self.exhausted = true; - return; + return Ok(()); } self.current_restart_idx = restart_idx; self.offset_in_block = self.block.offsets()[restart_idx] as usize; self.entries_since_restart = 0; - self.current_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, restart_idx); + self.current_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, restart_idx)?; self.exhausted = false; + Ok(()) } fn decode_entry_at_current_offset(&self) -> Result<(RowEntry, usize), SlateDBError> { - let mut data = self.block.data().slice(self.offset_in_block..); + let region_end = self.restart_region_end(self.current_restart_idx); + if self.offset_in_block >= region_end || region_end > self.block.data().len() { + return Err(corrupt_block("invalid V2 row boundary")); + } + let mut data = self.block.data().slice(self.offset_in_block..region_end); + let initial_len = data.len(); let codec = SstRowCodecV2::new(); let entry = codec.decode(&mut data)?; - let bytes_consumed = self.block.data().len() - self.offset_in_block - data.len(); - let new_offset = self.offset_in_block + bytes_consumed; + let bytes_consumed = initial_len + .checked_sub(data.remaining()) + .ok_or_else(|| corrupt_block("V2 row cursor moved backwards"))?; + if bytes_consumed == 0 { + return Err(corrupt_block("V2 row consumed no bytes")); + } + let new_offset = self + .offset_in_block + .checked_add(bytes_consumed) + .ok_or_else(|| corrupt_block("V2 row offset overflow"))?; + if usize::try_from(entry.shared_bytes).unwrap_or(usize::MAX) > self.current_key.len() { + return Err(corrupt_block("V2 row shared prefix exceeds previous key")); + } let full_key = entry.restore_full_key(&self.current_key); Ok(( @@ -112,16 +134,30 @@ impl AscendingState { )) } - fn decode_key_at_offset(&self, offset: usize, prev_key: &[u8]) -> Bytes { - let mut data = self.block.data().slice(offset..); + fn decode_key_at_offset( + &self, + offset: usize, + region_end: usize, + prev_key: &[u8], + ) -> Result { + if offset >= region_end || region_end > self.block.data().len() { + return Err(corrupt_block("invalid V2 key boundary")); + } + let mut data = self.block.data().slice(offset..region_end); let codec = SstRowCodecV2::new(); - let (shared_bytes, key_suffix) = codec.decode_key_only(&mut data); + let (shared_bytes, key_suffix) = codec.decode_key_only(&mut data)?; let shared = shared_bytes as usize; + if shared > prev_key.len() { + return Err(corrupt_block("V2 key shared prefix exceeds previous key")); + } - let mut full_key = BytesMut::with_capacity(shared + key_suffix.len()); + let key_len = shared + .checked_add(key_suffix.len()) + .ok_or_else(|| corrupt_block("V2 key length overflow"))?; + let mut full_key = BytesMut::with_capacity(key_len); full_key.extend_from_slice(&prev_key[..shared]); full_key.extend_from_slice(&key_suffix); - full_key.freeze() + Ok(full_key.freeze()) } fn is_empty(&self) -> bool { @@ -135,14 +171,14 @@ impl AscendingState { /// exist before the found restart point. /// Binary search for the first restart index where key >= target. /// Returns `restarts.len()` if no such restart exists. - fn binary_search_restarts(&self, target: &[u8]) -> usize { + fn binary_search_restarts(&self, target: &[u8]) -> Result { let restarts = self.block.offsets(); let mut low = 0; let mut high = restarts.len(); while low < high { let mid = low + (high - low) / 2; - let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, mid); + let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, mid)?; if restart_key.as_ref() < target { low = mid + 1; @@ -151,60 +187,60 @@ impl AscendingState { } } - low + Ok(low) } /// Find the restart region to begin an ascending scan for `target`. /// Backs up one position when target exactly matches a restart point's first key, /// so that duplicate keys straddling restart boundaries aren't missed. - fn find_restart_for_key_ascending(&self, target: &[u8]) -> usize { + fn find_restart_for_key_ascending(&self, target: &[u8]) -> Result { let restarts = self.block.offsets(); if restarts.is_empty() { - return 0; + return Ok(0); } - let low = self.binary_search_restarts(target); + let low = self.binary_search_restarts(target)?; if low < restarts.len() { - let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, low); + let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, low)?; if restart_key.as_ref() == target { - return low.saturating_sub(1); + return Ok(low.saturating_sub(1)); } } - low.saturating_sub(1) + Ok(low.saturating_sub(1)) } /// Find the restart region to begin a descending scan for `target`. /// Returns the last restart whose first key <= target, so that for duplicate keys /// spanning multiple restart regions we start from the last one. - fn find_restart_for_key_descending(&self, target: &[u8]) -> usize { + fn find_restart_for_key_descending(&self, target: &[u8]) -> Result { let restarts = self.block.offsets(); if restarts.is_empty() { - return 0; + return Ok(0); } // binary_search_restarts finds the first restart with key >= target. - let low = self.binary_search_restarts(target); + let low = self.binary_search_restarts(target)?; if low < restarts.len() { - let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, low); + let restart_key = BlockIteratorV2::decode_first_key_at_restart(&self.block, low)?; if restart_key.as_ref() == target { // Scan forward to find the last restart with the same first key. let mut last = low; while last + 1 < restarts.len() { let next_key = - BlockIteratorV2::decode_first_key_at_restart(&self.block, last + 1); + BlockIteratorV2::decode_first_key_at_restart(&self.block, last + 1)?; if next_key.as_ref() != target { break; } last += 1; } - return last; + return Ok(last); } } - low.saturating_sub(1) + Ok(low.saturating_sub(1)) } fn restart_region_end(&self, restart_idx: usize) -> usize { @@ -216,11 +252,24 @@ impl AscendingState { } fn advance_past_current_entry(&mut self) -> Result<(), SlateDBError> { - let mut data = self.block.data().slice(self.offset_in_block..); + let region_end = self.restart_region_end(self.current_restart_idx); + if self.offset_in_block >= region_end || region_end > self.block.data().len() { + return Err(corrupt_block("invalid V2 row boundary")); + } + let mut data = self.block.data().slice(self.offset_in_block..region_end); + let initial_len = data.len(); let codec = SstRowCodecV2::new(); codec.decode(&mut data)?; - let bytes_consumed = self.block.data().len() - self.offset_in_block - data.len(); - self.offset_in_block += bytes_consumed; + let bytes_consumed = initial_len + .checked_sub(data.remaining()) + .ok_or_else(|| corrupt_block("V2 row cursor moved backwards"))?; + if bytes_consumed == 0 { + return Err(corrupt_block("V2 row consumed no bytes")); + } + self.offset_in_block = self + .offset_in_block + .checked_add(bytes_consumed) + .ok_or_else(|| corrupt_block("V2 row offset overflow"))?; self.entries_since_restart += 1; Ok(()) } @@ -274,11 +323,11 @@ impl RowEntryIterator for BlockIteratorV2 { return Ok(()); } - let start_restart_idx = state.find_restart_for_key_ascending(next_key); + let start_restart_idx = state.find_restart_for_key_ascending(next_key)?; // Iterate through restart regions starting from binary search result. for restart_idx in start_restart_idx..state.block.offsets().len() { - state.seek_to_restart(restart_idx); + state.seek_to_restart(restart_idx)?; if state.exhausted || state.current_key.as_ref() >= next_key { return Ok(()); @@ -291,8 +340,11 @@ impl RowEntryIterator for BlockIteratorV2 { while state.offset_in_block < region_end && state.offset_in_block < state.block.data().len() { - let current_key = - state.decode_key_at_offset(state.offset_in_block, &prev_key); + let current_key = state.decode_key_at_offset( + state.offset_in_block, + region_end, + &prev_key, + )?; if current_key.as_ref() >= next_key { state.current_key = prev_key; @@ -331,19 +383,13 @@ impl DescendingBlockIteratorV2 { pub(crate) fn new(block: B) -> Self { let num_restarts = block.offsets().len(); - let initial_key = if num_restarts == 0 { - Bytes::new() - } else { - BlockIteratorV2::decode_first_key_at_restart(&block, 0) - }; - DescendingBlockIteratorV2 { ascending: AscendingState { block, current_restart_idx: 0, offset_in_block: 0, entries_since_restart: 0, - current_key: initial_key, + current_key: Bytes::new(), exhausted: false, }, current_restart_idx: num_restarts as isize - 1, @@ -362,7 +408,7 @@ impl DescendingBlockIteratorV2 { return Ok(()); } - self.ascending.seek_to_restart(restart_idx); + self.ascending.seek_to_restart(restart_idx)?; let region_end = self.ascending.restart_region_end(restart_idx); while self.ascending.offset_in_block < region_end && !self.ascending.is_empty() { @@ -436,7 +482,7 @@ impl RowEntryIterator for DescendingBlockIteratorV2 { } // Find the last restart region whose first key <= next_key, then scan backwards. - let start_restart_idx = self.ascending.find_restart_for_key_descending(next_key); + let start_restart_idx = self.ascending.find_restart_for_key_descending(next_key)?; for restart_idx in (0..=start_restart_idx).rev() { self.current_restart_idx = restart_idx as isize; @@ -468,6 +514,10 @@ impl RowEntryIterator for DescendingBlockIteratorV2 { } } +fn corrupt_block(reason: &'static str) -> SlateDBError { + SlateDBError::CorruptSst { reason, path: None } +} + #[cfg(feature = "bench-internal")] pub mod benches { use super::BlockIteratorV2; diff --git a/slatedb/src/config.rs b/slatedb/src/config.rs index cbc72df91..2ed475f12 100644 --- a/slatedb/src/config.rs +++ b/slatedb/src/config.rs @@ -713,6 +713,10 @@ pub struct Settings { /// memtable will be frozen even if it has not reached `l0_sst_size_bytes`. pub max_wal_flushes_before_l0_flush: u64, + /// Controls concurrent full-object WAL prefetch during recovery. + #[serde(default)] + pub wal_replay: WalReplaySettings, + /// Defines the max total number of SSTs in L0 across the entire key space. Memtables /// will not be flushed if the total L0 count (including in-flight uploads) would exceed /// this value, until compaction can compact the ssts into compacted. @@ -828,6 +832,7 @@ impl std::fmt::Debug for Settings { "max_wal_flushes_before_l0_flush", &self.max_wal_flushes_before_l0_flush, ) + .field("wal_replay", &self.wal_replay) .field("l0_max_ssts", &self.l0_max_ssts) .field("l0_max_ssts_per_key", &self.l0_max_ssts_per_key) .field("l0_flush_parallelism", &self.l0_flush_parallelism) @@ -866,6 +871,7 @@ impl Settings { /// Returns an [`crate::Error`] with [`crate::ErrorKind::Invalid`] describing /// the first invalid setting or combination encountered. pub fn validate(&self) -> Result<(), crate::Error> { + self.wal_replay.validate()?; if self.l0_flush_parallelism == 0 { return Err(SlateDBError::InvalidConfiguration( "l0_flush_parallelism must be at least 1".into(), @@ -1082,6 +1088,7 @@ impl Default for Settings { max_unflushed_bytes: 1_073_741_824, l0_sst_size_bytes: 64 * 1024 * 1024, max_wal_flushes_before_l0_flush: 4096, + wal_replay: WalReplaySettings::default(), l0_max_ssts: 8, l0_max_ssts_per_key: 8, l0_flush_parallelism: 4, @@ -1101,14 +1108,109 @@ impl Default for Settings { } } +/// Limits full-object WAL prefetch during writer and reader recovery. +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(default)] +pub struct WalReplaySettings { + /// Maximum number of WAL objects fetched concurrently. + pub max_concurrent_objects: usize, + + /// Maximum WAL payload memory used by full-object replay. + /// + /// Replay reserves `min(max_inflight_bytes / 4, 64 MiB)` for metadata, the + /// index, and one lazily decoded block. Half of that working partition is + /// reserved for metadata/index and half for the current decoded block. It + /// uses the remainder for retained normal-sized encoded objects. Oversized + /// objects use bounded range replay instead of bypassing the limit. Writers + /// reject a WAL before durable acknowledgement when one of its sections + /// cannot be recovered within these partitions. + pub max_inflight_bytes: usize, +} + +impl WalReplaySettings { + const MAX_DECODE_WORKING_BYTES: usize = 64 * 1024 * 1024; + + pub(crate) fn validate(&self) -> Result<(), SlateDBError> { + if self.max_concurrent_objects == 0 { + return Err(SlateDBError::InvalidConfiguration( + "wal_replay.max_concurrent_objects must be at least 1".into(), + )); + } + if self.max_inflight_bytes == 0 { + return Err(SlateDBError::InvalidConfiguration( + "wal_replay.max_inflight_bytes must be at least 1".into(), + )); + } + if u32::try_from(self.max_inflight_bytes).is_err() { + return Err(SlateDBError::InvalidConfiguration( + "wal_replay.max_inflight_bytes must fit in a 32-bit semaphore permit count".into(), + )); + } + Ok(()) + } + + pub(crate) fn working_memory_limit(&self) -> usize { + (self.max_inflight_bytes / 4).min(Self::MAX_DECODE_WORKING_BYTES) + } + + pub(crate) fn encoded_byte_limit(&self) -> Result { + self.max_inflight_bytes + .checked_sub(self.working_memory_limit()) + .ok_or(SlateDBError::InvalidConfiguration( + "wal replay memory budget cannot reserve decode working memory".into(), + )) + } + + pub(crate) fn block_working_memory_limit(&self) -> usize { + self.working_memory_limit() / 2 + } + + pub(crate) fn metadata_working_memory_limit(&self) -> usize { + self.working_memory_limit() / 2 + } + + pub(crate) fn validate_wal_block_size(&self, block_size: usize) -> Result<(), SlateDBError> { + // In the worst uncompressed case range replay holds the encoded block, + // its decoded bytes, and a materialized `u16` offset vector together. + let required_bytes = block_size + .checked_mul(3) + .and_then(|size| size.checked_add(10)) + .ok_or(SlateDBError::InvalidConfiguration( + "SST block size overflows WAL replay memory accounting".into(), + ))?; + let limit_bytes = self.block_working_memory_limit(); + if required_bytes > limit_bytes { + return Err(SlateDBError::InvalidConfiguration(format!( + "SST block size requires {required_bytes} bytes of WAL replay working memory, limit is {limit_bytes} bytes" + ))); + } + Ok(()) + } +} + +impl Default for WalReplaySettings { + fn default() -> Self { + Self { + max_concurrent_objects: 64, + max_inflight_bytes: 256 * 1024 * 1024, + } + } +} + #[derive(Clone, Debug, Deserialize, Serialize)] pub struct DbReaderOptions { - /// How frequently to poll for new manifest files and WAL data. Refreshing the manifest - /// file allows readers to detect newly compacted data. The reader will also look for - /// new writes to the WAL at this poll interval. Readers using - /// [`crate::DbReaderMode::Checkpoint`] do not poll the manifest or WAL. + /// How frequently to poll for new manifest files. Refreshing the manifest + /// allows readers to detect newly compacted data. Runtime WAL polling uses + /// [`Self::wal_poll_interval`] independently. Readers using + /// [`crate::DbReaderMode::Checkpoint`] do not poll either source. pub manifest_poll_interval: Duration, + /// How frequently an open reader probes the exact next WAL ID after + /// consuming the manifest-known range. Must be nonzero. Runtime WAL + /// polling performs bounded range reads and never discovers WALs via LIST. + #[serde(default = "default_reader_wal_poll_interval")] + pub wal_poll_interval: Duration, + /// For readers using [`crate::DbReaderMode::ManagedCheckpoint`], the client maintains a /// checkpoint against the latest database state. The checkpoint's expire time is set to the /// current time plus this value. This lifetime must always be greater than @@ -1119,6 +1221,10 @@ pub struct DbReaderOptions { /// Defaults to 64MB pub max_memtable_bytes: u64, + /// Controls concurrent full-object WAL prefetch during recovery. + #[serde(default)] + pub wal_replay: WalReplaySettings, + /// Options for the local disk cache. If `root_folder` is set, the reader /// will wrap its object store in a `CachedObjectStore` backed by the /// local filesystem, mirroring the behaviour of `Db`. @@ -1151,8 +1257,10 @@ impl Default for DbReaderOptions { fn default() -> Self { Self { manifest_poll_interval: Duration::from_secs(10), + wal_poll_interval: default_reader_wal_poll_interval(), checkpoint_lifetime: Duration::from_secs(10 * 60), max_memtable_bytes: 64 * 1024 * 1024, + wal_replay: WalReplaySettings::default(), object_store_cache_options: ObjectStoreCacheOptions::default(), skip_wal_replay: false, metric_level: None, @@ -1161,6 +1269,10 @@ impl Default for DbReaderOptions { } } +fn default_reader_wal_poll_interval() -> Duration { + Duration::from_secs(1) +} + /// The compression algorithm to use for SSTables. #[non_exhaustive] #[derive(Clone, Copy, Deserialize, PartialEq, Debug, Serialize)] @@ -2101,6 +2213,82 @@ object_store_cache_options: #[test] fn test_validate_accepts_default_settings() { assert!(Settings::default().validate().is_ok()); + assert_eq!( + Settings::default().wal_replay, + WalReplaySettings { + max_concurrent_objects: 64, + max_inflight_bytes: 256 * 1024 * 1024, + } + ); + } + + #[test] + fn test_validate_rejects_invalid_wal_replay_settings() { + for wal_replay in [ + WalReplaySettings { + max_concurrent_objects: 0, + ..WalReplaySettings::default() + }, + WalReplaySettings { + max_inflight_bytes: 0, + ..WalReplaySettings::default() + }, + ] { + let settings = Settings { + wal_replay, + ..Settings::default() + }; + assert!(settings.validate().is_err()); + } + } + + #[test] + fn wal_replay_rejects_a_block_target_larger_than_its_working_partition() { + let settings = WalReplaySettings { + max_concurrent_objects: 1, + max_inflight_bytes: 64 * 1024, + }; + + assert!(settings.validate_wal_block_size(4096).is_err()); + assert!(settings.validate_wal_block_size(1024).is_ok()); + } + + #[cfg(target_pointer_width = "64")] + #[test] + fn test_validate_rejects_wal_replay_byte_semaphore_overflow() { + let settings = Settings { + wal_replay: WalReplaySettings { + max_inflight_bytes: u32::MAX as usize + 1, + ..WalReplaySettings::default() + }, + ..Settings::default() + }; + assert!(settings.validate().is_err()); + } + + #[test] + fn test_wal_replay_settings_default_when_omitted() { + let mut settings = serde_json::to_value(Settings::default()).unwrap(); + settings.as_object_mut().unwrap().remove("wal_replay"); + let decoded: Settings = serde_json::from_value(settings).unwrap(); + assert_eq!(decoded.wal_replay, WalReplaySettings::default()); + + let mut reader = serde_json::to_value(DbReaderOptions::default()).unwrap(); + reader.as_object_mut().unwrap().remove("wal_replay"); + reader.as_object_mut().unwrap().remove("wal_poll_interval"); + let decoded: DbReaderOptions = serde_json::from_value(reader).unwrap(); + assert_eq!(decoded.wal_replay, WalReplaySettings::default()); + assert_eq!(decoded.wal_poll_interval, Duration::from_secs(1)); + + let partial: WalReplaySettings = serde_json::from_value(serde_json::json!({ + "max_concurrent_objects": 8 + })) + .unwrap(); + assert_eq!(partial.max_concurrent_objects, 8); + assert_eq!( + partial.max_inflight_bytes, + WalReplaySettings::default().max_inflight_bytes + ); } #[test] diff --git a/slatedb/src/db.rs b/slatedb/src/db.rs index 199f60f38..883b18e2b 100644 --- a/slatedb/src/db.rs +++ b/slatedb/src/db.rs @@ -57,7 +57,6 @@ use crate::db_snapshot::DbSnapshot; use crate::db_state::{collect_touched_segments, DbState, SsTableId}; use crate::db_stats::DbStats; use crate::error::SlateDBError; -use crate::iter::IterationOrder; use crate::manifest::{Manifest, VersionedManifest}; use crate::mem_table::KVTableMetadata; use crate::memtable_flusher::{FlushResult, FlushTarget, MemtableFlusher}; @@ -67,12 +66,11 @@ use crate::paths::PathResolver; use crate::prefix_extractor::PrefixExtractor; use crate::reader::{Reader, ScanContext}; use crate::snapshot_manager::SnapshotManager; -use crate::sst_iter::SstIteratorOptions; use crate::tablestore::TableStore; use crate::transaction_manager::TransactionManager; use crate::types::KeyValue; use crate::utils::{format_bytes_si, SafeSender, WatchableOnceCellReader}; -use crate::wal_replay::{WalReplayIterator, WalReplayOptions}; +use crate::wal_replay::{ExactWalReplayIterator, ExactWalReplaySource, WalReplayOptions}; use crate::{DbCacheManagerOps, DbMetadataOps, DbReadOps, DbWriteOps}; use slatedb_common::clock::SystemClock; use slatedb_common::metrics::MetricsRecorderHelper; @@ -510,6 +508,12 @@ impl DbInner { } async fn replay_wal(&self, wal_id_range: Range) -> Result<(), SlateDBError> { + let replay_started = self.system_clock.now(); + let replay_range_start = wal_id_range.start; + let replay_range_end = wal_id_range.end; + let replay_wal_count = replay_range_end.saturating_sub(replay_range_start); + let mut replayed_entries = 0_u64; + let mut replayed_bytes = 0_u64; let mut current_memtable_wal_id = self .state .read() @@ -526,26 +530,16 @@ impl DbInner { |_| -> Result<(), SlateDBError> { Ok(()) } ); - let sst_iter_options = SstIteratorOptions { - max_fetch_tasks: 1, - blocks_to_fetch: 256, - cache_blocks: false, - cache_metadata: false, - eager_spawn: true, - order: IterationOrder::Ascending, - prefix: None, - filter_context: None, - }; - let replay_options = WalReplayOptions { - sst_batch_size: 4, + prefetch: self.settings.wal_replay, max_memtable_bytes: self.settings.l0_sst_size_bytes, - sst_iter_options, min_seq: None, + source: ExactWalReplaySource::WriterOpen, + task_scope: None, }; let db_state = self.state.read().state().core().clone(); - let mut replay_iter = WalReplayIterator::range( + let mut replay_iter = ExactWalReplayIterator::range( wal_id_range, &db_state, replay_options, @@ -590,6 +584,11 @@ impl DbInner { .table .record_touched_segments(touched_segments); } + let metadata = replayed_table.table.metadata(); + replayed_entries = replayed_entries + .saturating_add(u64::try_from(metadata.entry_num).unwrap_or(u64::MAX)); + replayed_bytes = replayed_bytes + .saturating_add(u64::try_from(metadata.entries_size_in_bytes).unwrap_or(u64::MAX)); // Replayed rows come from WAL SSTs in remote storage, so they are already // durable. Update `last_remote_persisted_seq` before replaying to avoid a race with // the memtable flusher. The flusher calls flush_wals() to guarantee all data in the @@ -610,6 +609,20 @@ impl DbInner { let guard = self.state.read(); self.status_manager .report_memtable_segments(collect_touched_segments(&guard.view())); + info!( + "SlateDB WAL replay completed [writer_epoch={}, replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}, last_replayed_wal_id={}, replayed_entries={}, replayed_bytes={}, elapsed_ms={}]", + writer_epoch, + replay_range_start, + replay_range_end, + replay_wal_count, + current_memtable_wal_id, + replayed_entries, + replayed_bytes, + self.system_clock + .now() + .signed_duration_since(replay_started) + .num_milliseconds() + ); Ok(()) } @@ -2260,7 +2273,7 @@ mod tests { REQUEST_COUNT as OBJECT_STORE_REQUEST_COUNT, REQUEST_DURATION_SECONDS as OBJECT_STORE_REQUEST_DURATION_SECONDS, }; - use crate::iter::RowEntryIterator; + use crate::iter::{IterationOrder, RowEntryIterator}; use crate::manifest::store::{ManifestStore, StoredManifest}; use crate::manifest::{ManifestCore, VersionedManifest}; use crate::merge_operator::{ @@ -6708,12 +6721,12 @@ mod tests { ) .await; - let head_arrivals_before = gated_store.head_gate.arrivals(); - gated_store.head_gate.close(); + let get_arrivals_before = gated_store.get_opts_gate.arrivals(); + gated_store.get_opts_gate.close(); fail_parallel::cfg(fp_registry.clone(), "replay-wal-pause", "off").unwrap(); gated_store - .head_gate - .wait_for_arrivals(head_arrivals_before + 1) + .get_opts_gate + .wait_for_arrivals(get_arrivals_before + 1) .await; let db2 = Db::builder(path, base_store.clone()) @@ -6731,7 +6744,7 @@ mod tests { .delete_sst(&SsTableId::Wal(1)) .await .unwrap(); - gated_store.head_gate.release(); + gated_store.get_opts_gate.release(); let err = match w1_handle.await.unwrap() { Ok(_) => panic!("expected W1 open to fail"), @@ -6787,19 +6800,19 @@ mod tests { ) .await; - let head_arrivals_before = gated_store.head_gate.arrivals(); - gated_store.head_gate.close(); + let get_arrivals_before = gated_store.get_opts_gate.arrivals(); + gated_store.get_opts_gate.close(); fail_parallel::cfg(fp_registry.clone(), "replay-wal-pause", "off").unwrap(); gated_store - .head_gate - .wait_for_arrivals(head_arrivals_before + 1) + .get_opts_gate + .wait_for_arrivals(get_arrivals_before + 1) .await; probe_table_store .delete_sst(&SsTableId::Wal(1)) .await .unwrap(); - gated_store.head_gate.release(); + gated_store.get_opts_gate.release(); let err = match w1_handle.await.unwrap() { Ok(_) => panic!("expected W1 open to fail"), @@ -7333,6 +7346,7 @@ mod tests { min_filter_keys, l0_sst_size_bytes, max_wal_flushes_before_l0_flush: 4096, + wal_replay: crate::config::WalReplaySettings::default(), compactor_options, compression_codec: None, object_store_cache_options: ObjectStoreCacheOptions::default(), @@ -10361,6 +10375,63 @@ mod tests { recovered.close().await.unwrap(); } + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn acknowledged_wal_larger_than_full_object_limit_recovers_on_reopen() { + let object_store: Arc = Arc::new(InMemory::new()); + let path = "/tmp/acknowledged_oversized_wal_range_replay"; + let mut settings = test_db_options(0, 1024 * 1024, None); + settings.flush_interval = Some(Duration::from_millis(10)); + settings.wal_replay = crate::config::WalReplaySettings { + max_concurrent_objects: 4, + max_inflight_bytes: 128 * 1024, + }; + let source = Db::builder(path, object_store.clone()) + .with_settings(settings.clone()) + .build() + .await + .unwrap(); + let mut batch = WriteBatch::new(); + for id in 0..64 { + batch.put(format!("oversized-{id:03}"), vec![id as u8; 2048]); + } + + source + .write(batch) + .await + .expect("the large multi-block WAL was not acknowledged"); + let last_flushed_wal_id = source + .inner + .wal_observer + .status() + .unwrap() + .last_flushed_wal_id; + assert!(last_flushed_wal_id > 0); + let listed = source + .inner + .table_store + .list_wal_ssts_for_replay(1..last_flushed_wal_id + 1) + .await + .unwrap(); + assert!(listed.iter().any(|wal| wal.metadata.size > 96 * 1024)); + + // Keep the source open to model a crash. The replacement writer must + // recover the acknowledged WAL rather than relying on clean shutdown. + let recovered = Db::builder(path, object_store) + .with_settings(settings) + .build() + .await + .expect("replacement writer could not range-replay the oversized WAL"); + for id in 0..64 { + assert_eq!( + recovered.get(format!("oversized-{id:03}")).await.unwrap(), + Some(Bytes::from(vec![id as u8; 2048])) + ); + } + + recovered.close().await.unwrap(); + drop(source); + } + /// RFC-0024: WAL replay through a conforming extractor preserves the /// keys and lets segment-aware writes resume after the next open. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/slatedb/src/db/builder.rs b/slatedb/src/db/builder.rs index 2e1dc32d8..aadc18419 100644 --- a/slatedb/src/db/builder.rs +++ b/slatedb/src/db/builder.rs @@ -102,6 +102,7 @@ //! ``` //! use std::collections::{BTreeSet, HashMap}; +use std::num::NonZeroU64; use std::ops::{Bound, RangeBounds}; use std::sync::Arc; @@ -194,6 +195,7 @@ pub struct DbBuilder> { filter_policies: Vec>, metrics_recorder: Arc, segment_extractor: Option>, + writer_epoch: Option, } impl> DbBuilder

{ @@ -223,9 +225,22 @@ impl> DbBuilder

{ filter_policies: default_filter_policies(), metrics_recorder: Arc::new(NoopMetricsRecorder::new()), segment_extractor: None, + writer_epoch: None, } } + /// Requires this writer open to claim exactly `writer_epoch`. + /// + /// SlateDB rejects the open before publishing its WAL fence when the + /// stored writer epoch is greater than or equal to this value. Managed + /// control planes can therefore persist a monotonic epoch before opening + /// a writer and safely reject delayed stale opens. Builders that do not + /// set this value retain the embedded last-open-wins behavior. + pub fn with_writer_epoch(mut self, writer_epoch: NonZeroU64) -> Self { + self.writer_epoch = Some(writer_epoch); + self + } + /// Set the segment extractor (RFC-0024). When configured, every /// write is routed through the extractor and the database tracks /// per-segment LSM state. The extractor must be configured at @@ -435,6 +450,7 @@ impl> DbBuilder

{ let system_clock = self .system_clock .unwrap_or_else(|| Arc::new(DefaultSystemClock::new())); + let open_started = system_clock.now(); let metrics_recorder = self.metrics_recorder.clone(); let recorder = @@ -464,6 +480,7 @@ impl> DbBuilder

{ // under the retry and instrumentation layers, so the same cache // instance can be shared with the compactor and GC below while each // component keeps its own layers. + let stage_started = system_clock.now(); let cached_object_store = CachedObjectStore::from_config( self.main_object_store.clone(), &self.settings.object_store_cache_options, @@ -472,6 +489,17 @@ impl> DbBuilder

{ rand.clone(), ) .await?; + info!( + "SlateDB writer open stage completed [stage=cache_initialization, elapsed_ms={}, total_elapsed_ms={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds() + ); let maybe_cached_main_object_store: Arc = match &cached_object_store { Some(cached_store) => cached_store.clone(), None => self.main_object_store.clone(), @@ -519,6 +547,9 @@ impl> DbBuilder

{ block_format, ..SsTableFormat::default() }; + self.settings + .wal_replay + .validate_wal_block_size(sst_format.block_size)?; // Setup the manifest store and load latest manifest let manifest_store = Arc::new(ManifestStore::new( @@ -529,8 +560,20 @@ impl> DbBuilder

{ &path, retrying_main_object_store.clone(), )); + let stage_started = system_clock.now(); let latest_manifest = StoredManifest::try_load(manifest_store.clone(), system_clock.clone()).await?; + info!( + "SlateDB writer open stage completed [stage=manifest_load, elapsed_ms={}, total_elapsed_ms={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds() + ); if let Some(latest_manifest) = &latest_manifest { latest_manifest @@ -615,11 +658,29 @@ impl> DbBuilder

{ system_clock.clone(), task_executor.clone(), ); + let stage_started = system_clock.now(); let WriterFenceResult { manifest, replay_range, mut wal_writer, - } = fencer.fence(stored_manifest).await?; + } = fencer.fence(stored_manifest, self.writer_epoch).await?; + let replay_start_wal_id = replay_range.start; + let replay_end_wal_id = replay_range.end; + let replay_wal_count = replay_end_wal_id.saturating_sub(replay_start_wal_id); + info!( + "SlateDB writer open stage completed [stage=writer_fence, elapsed_ms={}, total_elapsed_ms={}, replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds(), + replay_start_wal_id, + replay_end_wal_id, + replay_wal_count + ); let (wal_writer, wal_observer) = if DbInner::wal_enabled_in_options(&self.settings) { let wal_observer = wal_writer.observer(); (Some(wal_writer), wal_observer) @@ -646,6 +707,7 @@ impl> DbBuilder

{ let (write_tx, write_rx) = SafeSender::unbounded_channel(reader); // Create the database inner state + let stage_started = system_clock.now(); let memtable_flusher = Arc::new(MemtableFlusher::new(status_manager.as_ref())); let inner = Arc::new( DbInner::new( @@ -816,18 +878,67 @@ impl> DbBuilder

{ // Monitor background tasks task_executor.monitor_on(&tokio_handle)?; + info!( + "SlateDB writer open stage completed [stage=runtime_initialization, elapsed_ms={}, total_elapsed_ms={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds() + ); // Replay WAL + let stage_started = system_clock.now(); + info!( + "SlateDB WAL replay started [replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}]", + replay_start_wal_id, replay_end_wal_id, replay_wal_count + ); inner.replay_wal(replay_range).await?; + info!( + "SlateDB writer open stage completed [stage=wal_replay, elapsed_ms={}, total_elapsed_ms={}, replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds(), + replay_start_wal_id, + replay_end_wal_id, + replay_wal_count + ); // Preload cache if enabled + let stage_started = system_clock.now(); if let Some(cached_obj_store) = &cached_object_store { inner .preload_cache(cached_obj_store, &path_resolver) .await?; } + info!( + "SlateDB writer open stage completed [stage=cache_preload, elapsed_ms={}, total_elapsed_ms={}]", + system_clock + .now() + .signed_duration_since(stage_started) + .num_milliseconds(), + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds() + ); // Create and return the Db instance + info!( + "SlateDB writer open completed [elapsed_ms={}]", + system_clock + .now() + .signed_duration_since(open_started) + .num_milliseconds() + ); Ok(Db { inner, task_executor, @@ -1853,6 +1964,9 @@ impl> DbReaderBuilder

{ block_transformer: self.block_transformer, ..SsTableFormat::default() }; + self.options + .wal_replay + .validate_wal_block_size(sst_format.block_size)?; let path_resolver = PathResolver::new_with_external_ssts(path.clone(), external_ssts); let table_store = Arc::new(TableStore::new_with_fp_registry( ObjectStores::new(retrying_object_store, retrying_wal_object_store), diff --git a/slatedb/src/db_cache/serde.rs b/slatedb/src/db_cache/serde.rs index fbe4412d8..bae181285 100644 --- a/slatedb/src/db_cache/serde.rs +++ b/slatedb/src/db_cache/serde.rs @@ -102,7 +102,7 @@ impl SerializedCachedEntryV1 { fn into_cached_entry(self) -> Result { let item = match self { SerializedCachedEntryV1::Block(encoded) => { - let block = Block::decode(encoded); + let block = Block::decode(encoded)?; CachedItem::Block(Arc::new(block)) } SerializedCachedEntryV1::SsTableIndex(encoded) => { diff --git a/slatedb/src/db_reader.rs b/slatedb/src/db_reader.rs index d76d8c617..f44928a2a 100644 --- a/slatedb/src/db_reader.rs +++ b/slatedb/src/db_reader.rs @@ -9,9 +9,8 @@ use crate::db_iter::DbIteratorGuard; use crate::db_state::{collect_touched_segments, SsTableId}; use crate::db_stats::DbStats; use crate::db_status::{ClosedResultWriter, DbStatus, DbStatusManager}; -use crate::dispatcher::{MessageHandler, MessageHandlerExecutor, MessageTickerDef}; +use crate::dispatcher::{MessageHandler, MessageHandlerExecutor, MessageTickerDef, Notifier}; use crate::error::SlateDBError; -use crate::iter::IterationOrder; use crate::manifest::store::{ManifestStore, StoredManifest}; use crate::manifest::{Manifest, ManifestCore, VersionedManifest}; use crate::mem_table::{ImmutableMemtable, KVTable, WritableKVTable}; @@ -20,17 +19,20 @@ use crate::oracle::DbReaderOracle; use crate::paths::PathResolver; use crate::prefix_extractor::PrefixExtractor; use crate::reader::{DbStateReader, Reader, ScanContext}; -use crate::sst_iter::SstIteratorOptions; +use crate::replay_task_scope::ReplayTaskScope; use crate::tablestore::TableStore; use crate::types::KeyValue; -use crate::utils::IdGenerator; -use crate::wal_replay::{WalReplayIterator, WalReplayOptions}; +use crate::utils::{panic_string, IdGenerator}; +use crate::wal_replay::{ + ExactWalReplayIterator, ExactWalReplaySource, ReplayedMemtable, RuntimeWalReplayError, + RuntimeWalReplayIterator, RuntimeWalReplayOptions, RuntimeWalReplaySource, WalReplayOptions, +}; use crate::{Checkpoint, DbIterator, DbSnapshot}; use crate::{DbCacheManagerOps, DbMetadataOps, DbReadOps}; use async_trait::async_trait; use bytes::Bytes; -use futures::stream::BoxStream; -use log::{info, warn}; +use futures::{stream::BoxStream, FutureExt}; +use log::{error, info, warn}; use object_store::path::Path; use object_store::ObjectStore; use parking_lot::RwLock; @@ -38,6 +40,7 @@ use slatedb_common::clock::SystemClock; use slatedb_common::DbRand; use std::collections::{BTreeSet, HashMap}; use std::ops::Sub; +use std::panic::AssertUnwindSafe; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::LazyLock; use std::sync::{Arc, Weak}; @@ -46,6 +49,8 @@ use tokio::sync::Notify; use uuid::Uuid; pub(crate) const DB_READER_TASK_NAME: &str = "manifest_poller"; +const MAX_RUNTIME_WALS_PER_TURN: u64 = 64; +const MAX_RUNTIME_REPLAY_TURN_TIME: std::time::Duration = std::time::Duration::from_secs(1); /// Determines how a [`DbReader`] chooses and refreshes the database state it reads. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] @@ -103,6 +108,8 @@ pub(crate) struct DbReaderInner { status_manager: DbStatusManager, segment_extractor: Option>, rand: Arc, + /// Root ownership scope for every replay task in this reader tenure. + replay_tasks: ReplayTaskScope, /// Kept alive so the underlying `MetricsRecorder` is not dropped while /// metric handles in `DbStats` (and other stats structs) are still in use. /// See: https://github.com/slatedb/slatedb/issues/1469 @@ -110,9 +117,58 @@ pub(crate) struct DbReaderInner { recorder: slatedb_common::metrics::MetricsRecorderHelper, } -#[derive(Debug)] enum DbReaderMessage { PollManifest, + PollWalTail, + ContinueWalTail, + WalTailComplete(Result), + ManifestRefreshComplete(Result, SlateDBError>), +} + +impl std::fmt::Debug for DbReaderMessage { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::PollManifest => "PollManifest", + Self::PollWalTail => "PollWalTail", + Self::ContinueWalTail => "ContinueWalTail", + Self::WalTailComplete(Ok(_)) => "WalTailComplete(Ok)", + Self::WalTailComplete(Err(_)) => "WalTailComplete(Err)", + Self::ManifestRefreshComplete(Ok(Some(_))) => "ManifestRefreshComplete(Ok(Some))", + Self::ManifestRefreshComplete(Ok(None)) => "ManifestRefreshComplete(Ok(None))", + Self::ManifestRefreshComplete(Err(_)) => "ManifestRefreshComplete(Err)", + }) + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum ReaderReplayPath { + ExactOpen { + include_unmanifested: bool, + source: ExactWalReplaySource, + }, + RuntimeManifest, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RuntimeMissingPolicy { + Error, + ExactNextIsCaughtUp(u64), +} + +struct RuntimeReplayTurn { + manifest_id: u64, + base_last_wal_id: u64, + imm_memtable: ReplayMemtables, + last_wal_id: u64, + last_committed_seq: u64, + caught_up: bool, + reached_turn_limit: bool, +} + +struct ManifestRefreshCandidate { + base_manifest_id: u64, + base_last_wal_id: u64, + state: ReaderState, } #[derive(Clone)] @@ -366,21 +422,27 @@ impl DbReaderInner { } else { (manifest.id(), manifest.manifest().clone()) }; - let replay_new_wals = + let include_unmanifested = !matches!(mode, DbReaderMode::Checkpoint(_)) && !options.skip_wal_replay; let db_stats = DbStats::new(&recorder); + let replay_tasks = ReplayTaskScope::new(); let initial_state = Arc::new( Self::build_reader_state( checkpoint, manifest_id, initial_manifest, ReplayMemtables::default(), - replay_new_wals, + ReaderReplayPath::ExactOpen { + include_unmanifested, + source: ExactWalReplaySource::ReaderOpen, + }, Arc::clone(&table_store), &options, segment_extractor.as_ref(), None, &db_stats, + Some(replay_tasks.clone()), + &system_clock, ) .await?, ); @@ -427,6 +489,7 @@ impl DbReaderInner { status_manager, segment_extractor, rand, + replay_tasks, recorder, }; Ok(inner) @@ -641,57 +704,22 @@ impl DbReaderInner { .report_manifest_and_memtable_segments(versioned_manifest, touched_segments); } - async fn maybe_replay_new_wals(&self) -> Result<(), SlateDBError> { - if self.options.skip_wal_replay { - return Ok(()); - } - let current_state = Arc::clone(&self.state.read()); - let mut imm_memtable = current_state.imm_memtable.clone(); - let generation = Arc::clone(¤t_state.generation); - let mut publish = - |imm_memtable: &ReplayMemtables, last_wal_id: u64, last_committed_seq: u64| { - self.oracle.advance_durable_seq(last_committed_seq); - self.db_stats - .reader_replay_memtables - .set(imm_memtable.len() as i64); - let mut write_guard = self.state.write(); - *write_guard = Arc::new(ReaderState { - generation: Arc::clone(&generation), - imm_memtable: imm_memtable.clone(), - last_wal_id, - last_remote_persisted_seq: last_committed_seq, - }); - drop(write_guard); - self.status_manager - .report_memtable_segments(collect_touched_segments(self.state.read().as_ref())); - }; - - Self::replay_wal_into( - Arc::clone(&self.table_store), - &self.options, - current_state.core(), - &mut imm_memtable, - Some(( - current_state.last_wal_id, - current_state.last_remote_persisted_seq, - )), - true, - self.segment_extractor.as_ref(), - Some(&mut publish), - Some(&self.db_stats), - ) - .await?; - Ok(()) - } - async fn rebuild_checkpoint_state( &self, new_checkpoint: Checkpoint, ) -> Result { let manifest_id = new_checkpoint.manifest_id; let manifest = self.manifest_store.read_manifest(manifest_id).await?; - self.rebuild_state(Some(new_checkpoint), manifest_id, manifest) - .await + self.rebuild_state( + Some(new_checkpoint), + manifest_id, + manifest, + ReaderReplayPath::ExactOpen { + include_unmanifested: !self.options.skip_wal_replay, + source: ExactWalReplaySource::CheckpointRecovery, + }, + ) + .await } async fn rebuild_state( @@ -699,8 +727,29 @@ impl DbReaderInner { checkpoint: Option, manifest_id: u64, manifest: Manifest, + replay_path: ReaderReplayPath, ) -> Result { let prior = self.state.read().clone(); + self.rebuild_state_from( + prior, + checkpoint, + manifest_id, + manifest, + replay_path, + self.replay_tasks.clone(), + ) + .await + } + + async fn rebuild_state_from( + &self, + prior: Arc, + checkpoint: Option, + manifest_id: u64, + manifest: Manifest, + replay_path: ReaderReplayPath, + replay_tasks: ReplayTaskScope, + ) -> Result { let replay_cursor = Some(( prior.last_wal_id.max(manifest.core.replay_after_wal_id), prior @@ -739,40 +788,96 @@ impl DbReaderInner { manifest_id, manifest, imm_memtable, - !self.options.skip_wal_replay, + replay_path, Arc::clone(&self.table_store), &self.options, self.segment_extractor.as_ref(), replay_cursor, &self.db_stats, + Some(replay_tasks), + &self.system_clock, ) .await } + async fn build_latest_manifest_candidate( + &self, + base: Arc, + replay_tasks: ReplayTaskScope, + ) -> Result, SlateDBError> { + let latest_manifest = self.manifest_store.read_latest_manifest().await?; + if latest_manifest.id <= base.generation.manifest_id { + return Ok(None); + } + + let base_manifest_id = base.generation.manifest_id; + let base_last_wal_id = base.last_wal_id; + let state = self + .rebuild_state_from( + base, + None, + latest_manifest.id, + latest_manifest.manifest, + ReaderReplayPath::RuntimeManifest, + replay_tasks, + ) + .await?; + Ok(Some(ManifestRefreshCandidate { + base_manifest_id, + base_last_wal_id, + state, + })) + } + async fn build_reader_state( checkpoint: Option, manifest_id: u64, manifest: Manifest, mut imm_memtable: ReplayMemtables, - replay_new_wals: bool, + replay_path: ReaderReplayPath, table_store: Arc, options: &DbReaderOptions, segment_extractor: Option<&Arc>, replay_cursor: Option<(u64, u64)>, db_stats: &DbStats, + replay_tasks: Option, + system_clock: &Arc, ) -> Result { - let (last_wal_id, last_committed_seq) = Self::replay_wal_into( - Arc::clone(&table_store), - options, - &manifest.core, - &mut imm_memtable, - replay_cursor, - replay_new_wals, - segment_extractor, - None, - Some(db_stats), - ) - .await?; + let (last_wal_id, last_committed_seq) = match replay_path { + ReaderReplayPath::ExactOpen { + include_unmanifested, + source, + } => { + Self::replay_wal_into_exact_with_source( + Arc::clone(&table_store), + options, + &manifest.core, + &mut imm_memtable, + replay_cursor, + include_unmanifested, + source, + segment_extractor, + None, + Some(db_stats), + replay_tasks.clone(), + ) + .await? + } + ReaderReplayPath::RuntimeManifest => { + Self::replay_manifest_range_into( + Arc::clone(&table_store), + options, + &manifest.core, + &mut imm_memtable, + replay_cursor, + segment_extractor, + Some(db_stats), + replay_tasks, + system_clock, + ) + .await? + } + }; db_stats .reader_replay_memtables @@ -786,22 +891,23 @@ impl DbReaderInner { }) } - async fn refresh_latest_manifest(&self) -> Result<(), SlateDBError> { - let latest_manifest = self.manifest_store.read_latest_manifest().await?; - self.apply_latest_manifest(latest_manifest).await - } - + #[cfg(test)] async fn apply_latest_manifest( &self, latest_manifest: VersionedManifest, ) -> Result<(), SlateDBError> { let manifest_id = latest_manifest.id; if manifest_id <= self.state.read().generation.manifest_id { - return self.maybe_replay_new_wals().await; + return Ok(()); } let new_state = self - .rebuild_state(None, manifest_id, latest_manifest.manifest) + .rebuild_state( + None, + manifest_id, + latest_manifest.manifest, + ReaderReplayPath::RuntimeManifest, + ) .await?; self.install_state(new_state); info!("refreshed reader to latest manifest [manifest_id={manifest_id}]"); @@ -868,7 +974,8 @@ impl DbReaderInner { self: &Arc, task_executor: &MessageHandlerExecutor, ) -> Result<(), SlateDBError> { - let poller = ManifestPoller::new(Arc::clone(self)); + let (tail_tx, tail_rx) = async_channel::unbounded(); + let poller = ManifestPoller::new(Arc::clone(self), tail_tx, tail_rx); let (_tx, rx) = async_channel::unbounded(); let result = task_executor.add_handler( DB_READER_TASK_NAME.to_string(), @@ -880,7 +987,8 @@ impl DbReaderInner { result } - async fn replay_wal_into( + #[cfg(test)] + async fn replay_wal_into_exact( table_store: Arc, reader_options: &DbReaderOptions, core: &ManifestCore, @@ -888,20 +996,38 @@ impl DbReaderInner { replay_cursor: Option<(u64, u64)>, replay_new_wals: bool, segment_extractor: Option<&Arc>, - mut publish: Option>, + publish: Option>, db_stats: Option<&DbStats>, ) -> Result<(u64, u64), SlateDBError> { - let sst_iter_options = SstIteratorOptions { - max_fetch_tasks: 1, - blocks_to_fetch: 256, - cache_blocks: true, - cache_metadata: false, - eager_spawn: true, - order: IterationOrder::Ascending, - prefix: None, - filter_context: None, - }; + Self::replay_wal_into_exact_with_source( + table_store, + reader_options, + core, + into_tables, + replay_cursor, + replay_new_wals, + ExactWalReplaySource::ReaderOpen, + segment_extractor, + publish, + db_stats, + None, + ) + .await + } + async fn replay_wal_into_exact_with_source( + table_store: Arc, + reader_options: &DbReaderOptions, + core: &ManifestCore, + into_tables: &mut ReplayMemtables, + replay_cursor: Option<(u64, u64)>, + replay_new_wals: bool, + source: ExactWalReplaySource, + segment_extractor: Option<&Arc>, + mut publish: Option>, + db_stats: Option<&DbStats>, + replay_tasks: Option, + ) -> Result<(u64, u64), SlateDBError> { let (mut replay_after_wal_id, mut last_committed_seq) = replay_cursor.unwrap_or_else(|| { if let Some(latest_replayed_table) = into_tables.front() { @@ -913,21 +1039,39 @@ impl DbReaderInner { (core.replay_after_wal_id, core.last_l0_seq) } }); + let replay_list_counter = db_stats.map(|stats| match source { + ExactWalReplaySource::ReaderOpen => { + Arc::clone(&stats.reader_wal_replay_list_reader_open) + } + ExactWalReplaySource::CheckpointRecovery => { + Arc::clone(&stats.reader_wal_replay_list_checkpoint_recovery) + } + ExactWalReplaySource::WriterOpen => { + unreachable!("writer replay does not use DbReader metrics") + } + }); let wal_id_end = if replay_new_wals { + if let Some(counter) = replay_list_counter.as_ref() { + counter.increment(1); + } table_store.last_seen_wal_id(replay_after_wal_id).await? + 1 } else { core.next_wal_sst_id }; let replay_options = WalReplayOptions { - sst_batch_size: 4, + prefetch: reader_options.wal_replay, max_memtable_bytes: reader_options.max_memtable_bytes as usize, - sst_iter_options, // Skip entries that we already have in `imm_memtable` (that might be above last_l0_seq). min_seq: Some(last_committed_seq), + source, + task_scope: replay_tasks, }; - let mut replay_iter = WalReplayIterator::range( + if let Some(counter) = replay_list_counter { + counter.increment(1); + } + let mut replay_iter = ExactWalReplayIterator::range( (replay_after_wal_id + 1)..wal_id_end, core, replay_options, @@ -938,7 +1082,6 @@ impl DbReaderInner { while let Some(replayed_table) = match replay_iter.next().await { Ok(Some(replayed_table)) => Some(replayed_table), Ok(None) => None, - Err(err) if has_not_found_object_store_error(&err) => None, Err(err) => return Err(err), } { assert!(replayed_table.last_wal_id > replay_after_wal_id); @@ -980,6 +1123,277 @@ impl DbReaderInner { Ok((replay_after_wal_id, last_committed_seq)) } + async fn replay_manifest_range_into( + table_store: Arc, + reader_options: &DbReaderOptions, + core: &ManifestCore, + into_tables: &mut ReplayMemtables, + replay_cursor: Option<(u64, u64)>, + segment_extractor: Option<&Arc>, + db_stats: Option<&DbStats>, + replay_tasks: Option, + system_clock: &Arc, + ) -> Result<(u64, u64), SlateDBError> { + let mut replay_cursor = + replay_cursor.unwrap_or((core.replay_after_wal_id, core.last_l0_seq)); + while replay_cursor.0.saturating_add(1) < core.next_wal_sst_id { + let start = replay_cursor.0.saturating_add(1); + let deadline = system_clock.now() + MAX_RUNTIME_REPLAY_TURN_TIME; + let completed_before = replay_cursor.0; + let mut replay = RuntimeWalReplayIterator::range( + start..core.next_wal_sst_id, + core, + RuntimeWalReplayOptions { + max_memtable_bytes: reader_options.max_memtable_bytes as usize, + min_seq: Some(replay_cursor.1), + max_wals_per_batch: std::num::NonZeroUsize::new( + usize::try_from(MAX_RUNTIME_WALS_PER_TURN) + .expect("runtime WAL turn limit must fit usize"), + ), + source: RuntimeWalReplaySource::Manifest, + task_scope: replay_tasks.clone(), + ..RuntimeWalReplayOptions::default() + }, + Arc::clone(&table_store), + )?; + + loop { + let completed_wals = replay_cursor.0.saturating_sub(completed_before); + if completed_wals >= MAX_RUNTIME_WALS_PER_TURN { + break; + } + let remaining = deadline + .signed_duration_since(system_clock.now()) + .to_std() + .unwrap_or(std::time::Duration::ZERO); + let replayed = tokio::select! { + biased; + _ = system_clock.sleep(remaining) => break, + result = replay.next() => result, + }; + if system_clock.now() >= deadline { + break; + } + let replayed_table = match replayed { + Ok(Some(table)) => table, + Ok(None) => break, + Err(RuntimeWalReplayError::MissingInitialObject { source, .. }) + | Err(RuntimeWalReplayError::Replay(source)) => return Err(source), + }; + Self::apply_runtime_replayed_table( + replayed_table, + into_tables, + &mut replay_cursor, + segment_extractor, + db_stats, + )?; + } + drop(replay); + + if replay_cursor.0.saturating_add(1) < core.next_wal_sst_id { + tokio::task::yield_now().await; + } + } + Ok(replay_cursor) + } + + #[cfg(test)] + async fn replay_runtime_range_into( + table_store: Arc, + reader_options: &DbReaderOptions, + core: &ManifestCore, + into_tables: &mut ReplayMemtables, + replay_cursor: &mut (u64, u64), + wal_range: std::ops::Range, + missing_policy: RuntimeMissingPolicy, + segment_extractor: Option<&Arc>, + db_stats: Option<&DbStats>, + max_wals_per_batch: Option, + replay_tasks: Option, + ) -> Result { + let mut replay = RuntimeWalReplayIterator::range( + wal_range, + core, + RuntimeWalReplayOptions { + max_memtable_bytes: reader_options.max_memtable_bytes as usize, + min_seq: Some(replay_cursor.1), + max_wals_per_batch, + source: max_wals_per_batch.map_or(RuntimeWalReplaySource::Manifest, |_| { + RuntimeWalReplaySource::Tail + }), + task_scope: replay_tasks, + ..RuntimeWalReplayOptions::default() + }, + table_store, + )?; + + loop { + let replayed_table = match replay.next().await { + Ok(Some(table)) => table, + Ok(None) => break, + Err(RuntimeWalReplayError::MissingInitialObject { wal_id, source }) => { + if missing_policy == RuntimeMissingPolicy::ExactNextIsCaughtUp(wal_id) { + return Ok(true); + } + return Err(source); + } + Err(RuntimeWalReplayError::Replay(source)) => return Err(source), + }; + Self::apply_runtime_replayed_table( + replayed_table, + into_tables, + replay_cursor, + segment_extractor, + db_stats, + )?; + } + + Ok(false) + } + + fn apply_runtime_replayed_table( + replayed_table: ReplayedMemtable, + into_tables: &mut ReplayMemtables, + replay_cursor: &mut (u64, u64), + segment_extractor: Option<&Arc>, + db_stats: Option<&DbStats>, + ) -> Result { + assert!(replayed_table.last_wal_id > replay_cursor.0); + let replayed_ssts = replayed_table.last_wal_id - replay_cursor.0; + replay_cursor.0 = replayed_table.last_wal_id; + if let Some(db_stats) = db_stats { + let metadata = replayed_table.table.metadata(); + db_stats.reader_wal_replay_ssts.increment(replayed_ssts); + db_stats + .reader_wal_replay_bytes + .increment(metadata.entries_size_in_bytes as u64); + db_stats.reader_wal_replay_batches.increment(1); + } + if !replayed_table.table.is_empty() && replayed_table.last_seq > replay_cursor.1 { + let first_seq = replayed_table + .table + .table() + .first_seq() + .expect("expected first_seq on non-empty table"); + assert!(first_seq > replay_cursor.1); + replay_cursor.1 = replayed_table.last_seq; + if let Some(extractor) = segment_extractor { + Self::record_replayed_touched_segments(extractor.as_ref(), &replayed_table.table)?; + } + into_tables.prepend(Arc::new(ImmutableMemtable::new( + replayed_table.table, + replayed_table.last_wal_id, + ))); + } + Ok(replayed_ssts) + } + + async fn replay_runtime_turn( + inner: Arc, + base: Arc, + replay_tasks: ReplayTaskScope, + ) -> Result { + let manifest_id = base.generation.manifest_id; + let base_last_wal_id = base.last_wal_id; + let core = base.core(); + let known_end = core.next_wal_sst_id; + let deadline = inner.system_clock.now() + MAX_RUNTIME_REPLAY_TURN_TIME; + let mut imm_memtable = base.imm_memtable.clone(); + let mut cursor = (base.last_wal_id, base.last_remote_persisted_seq); + let mut completed_wals = 0_u64; + let mut caught_up = false; + + 'turn: while completed_wals < MAX_RUNTIME_WALS_PER_TURN + && inner.system_clock.now() < deadline + { + let next_wal_id = cursor + .0 + .checked_add(1) + .ok_or(SlateDBError::InvalidDBState)?; + let remaining = MAX_RUNTIME_WALS_PER_TURN - completed_wals; + let (range_end, missing_policy) = if next_wal_id < known_end { + ( + next_wal_id + .checked_add(remaining) + .ok_or(SlateDBError::InvalidDBState)? + .min(known_end), + RuntimeMissingPolicy::Error, + ) + } else { + ( + next_wal_id + .checked_add(1) + .ok_or(SlateDBError::InvalidDBState)?, + RuntimeMissingPolicy::ExactNextIsCaughtUp(next_wal_id), + ) + }; + let mut replay = RuntimeWalReplayIterator::range( + next_wal_id..range_end, + core, + RuntimeWalReplayOptions { + max_memtable_bytes: inner.options.max_memtable_bytes as usize, + min_seq: Some(cursor.1), + max_wals_per_batch: std::num::NonZeroUsize::new(1), + source: RuntimeWalReplaySource::Tail, + task_scope: Some(replay_tasks.clone()), + ..RuntimeWalReplayOptions::default() + }, + Arc::clone(&inner.table_store), + )?; + + loop { + let remaining = deadline + .signed_duration_since(inner.system_clock.now()) + .to_std() + .unwrap_or(std::time::Duration::ZERO); + let replayed = tokio::select! { + biased; + _ = inner.system_clock.sleep(remaining) => break 'turn, + result = replay.next() => result, + }; + // A result completing at the exact deadline is deliberately + // discarded even if the runtime happened to poll it first. + if inner.system_clock.now() >= deadline { + break 'turn; + } + let replayed_table = match replayed { + Ok(Some(table)) => table, + Ok(None) => break, + Err(RuntimeWalReplayError::MissingInitialObject { wal_id, source }) => { + if missing_policy == RuntimeMissingPolicy::ExactNextIsCaughtUp(wal_id) { + caught_up = true; + break 'turn; + } + return Err(source); + } + Err(RuntimeWalReplayError::Replay(source)) => return Err(source), + }; + completed_wals = completed_wals.saturating_add(Self::apply_runtime_replayed_table( + replayed_table, + &mut imm_memtable, + &mut cursor, + inner.segment_extractor.as_ref(), + Some(&inner.db_stats), + )?); + if completed_wals >= MAX_RUNTIME_WALS_PER_TURN { + break 'turn; + } + } + } + + Ok(RuntimeReplayTurn { + manifest_id, + base_last_wal_id, + imm_memtable, + last_wal_id: cursor.0, + last_committed_seq: cursor.1, + caught_up, + reached_turn_limit: !caught_up + && (completed_wals >= MAX_RUNTIME_WALS_PER_TURN + || inner.system_clock.now() >= deadline), + }) + } + /// Re-derive each replayed entry's segment prefix (RFC-0024) and record the /// table's touched-segment set, mirroring the writer's replay path. Durable /// WAL entries were validated when accepted, so the antichain check is not @@ -1029,10 +1443,23 @@ impl DbReaderInner { struct ManifestPoller { inner: Arc, generations: HashMap>, + tail_tx: async_channel::Sender, + tail_rx: Option>, + tail_in_flight: bool, + continuation_queued: bool, + tail_task: Option>, + tail_scope: Option, + manifest_refresh_in_flight: bool, + manifest_refresh_task: Option>, + manifest_refresh_scope: Option, } impl ManifestPoller { - fn new(inner: Arc) -> Self { + fn new( + inner: Arc, + tail_tx: async_channel::Sender, + tail_rx: async_channel::Receiver, + ) -> Self { let mut generations = HashMap::new(); if inner.mode == DbReaderMode::ManagedCheckpoint { let generation = Arc::clone(&inner.state.read().generation); @@ -1042,11 +1469,211 @@ impl ManifestPoller { .id; generations.insert(checkpoint_id, Arc::downgrade(&generation)); } - let poller = Self { inner, generations }; + let poller = Self { + inner, + generations, + tail_rx: Some(tail_rx), + tail_tx, + tail_in_flight: false, + continuation_queued: false, + tail_task: None, + tail_scope: None, + manifest_refresh_in_flight: false, + manifest_refresh_task: None, + manifest_refresh_scope: None, + }; poller.report_active_checkpoints(); poller } + fn runtime_tailing_enabled(&self) -> bool { + !self.inner.options.skip_wal_replay + && !matches!(self.inner.mode, DbReaderMode::Checkpoint(_)) + } + + fn start_runtime_tail(&mut self) { + if self.tail_in_flight || self.manifest_refresh_in_flight || !self.runtime_tailing_enabled() + { + return; + } + self.tail_in_flight = true; + self.continuation_queued = false; + let inner = Arc::clone(&self.inner); + let base = Arc::clone(&inner.state.read()); + let sender = self.tail_tx.clone(); + let scope = inner.replay_tasks.child(); + let replay_scope = scope.clone(); + self.tail_task = Some(scope.spawn(async move { + let replay = tokio::select! { + biased; + _ = replay_scope.cancelled() => return, + replay = AssertUnwindSafe(DbReaderInner::replay_runtime_turn( + inner, + base, + replay_scope.clone(), + )).catch_unwind() => replay, + }; + let result = match replay { + Ok(result) => result, + Err(panic) => { + let task_name = "reader_runtime_wal_replay".to_string(); + error!( + "runtime WAL replay task panicked [task_name={}, panic={}]", + task_name, + panic_string(&panic), + ); + Err(SlateDBError::BackgroundTaskPanic(task_name)) + } + }; + let _ = sender.send(DbReaderMessage::WalTailComplete(result)).await; + })); + self.tail_scope = Some(scope); + } + + fn start_manifest_refresh(&mut self) { + if self.manifest_refresh_in_flight { + return; + } + self.manifest_refresh_in_flight = true; + let inner = Arc::clone(&self.inner); + let base = Arc::clone(&inner.state.read()); + let sender = self.tail_tx.clone(); + let scope = inner.replay_tasks.child(); + let replay_scope = scope.clone(); + self.manifest_refresh_task = Some(scope.spawn(async move { + let result = tokio::select! { + biased; + _ = replay_scope.cancelled() => return, + result = AssertUnwindSafe(inner.build_latest_manifest_candidate( + base, + replay_scope.clone(), + )).catch_unwind() => match result { + Ok(result) => result, + Err(panic) => { + let task_name = "reader_manifest_wal_replay".to_string(); + error!( + "manifest WAL replay task panicked [task_name={}, panic={}]", + task_name, + panic_string(&panic), + ); + Err(SlateDBError::BackgroundTaskPanic(task_name)) + } + }, + }; + let _ = sender + .send(DbReaderMessage::ManifestRefreshComplete(result)) + .await; + })); + self.manifest_refresh_scope = Some(scope); + } + + fn schedule_tail_continuation(&mut self) { + if self.continuation_queued || !self.runtime_tailing_enabled() { + return; + } + self.continuation_queued = true; + let _ = self.tail_tx.try_send(DbReaderMessage::ContinueWalTail); + } + + async fn finish_tail_task(&mut self) { + if let Some(task) = self.tail_task.take() { + let _ = task.await; + } + if let Some(scope) = self.tail_scope.take() { + scope.shutdown().await; + } + } + + async fn cancel_tail_task(&mut self) { + if let Some(scope) = self.tail_scope.as_ref() { + scope.cancel(); + } + if let Some(task) = self.tail_task.as_ref() { + task.abort(); + } + self.finish_tail_task().await; + self.tail_in_flight = false; + } + + async fn finish_manifest_refresh_task(&mut self) { + if let Some(task) = self.manifest_refresh_task.take() { + let _ = task.await; + } + if let Some(scope) = self.manifest_refresh_scope.take() { + scope.shutdown().await; + } + } + + async fn cancel_manifest_refresh_task(&mut self) { + if let Some(scope) = self.manifest_refresh_scope.as_ref() { + scope.cancel(); + } + if let Some(task) = self.manifest_refresh_task.as_ref() { + task.abort(); + } + self.finish_manifest_refresh_task().await; + self.manifest_refresh_in_flight = false; + } + + async fn apply_manifest_refresh(&mut self, candidate: Option) { + self.manifest_refresh_in_flight = false; + self.finish_manifest_refresh_task().await; + let Some(candidate) = candidate else { + self.inner.db_stats.reader_manifest_polls.increment(1); + self.schedule_tail_continuation(); + return; + }; + let current = Arc::clone(&self.inner.state.read()); + if current.generation.manifest_id != candidate.base_manifest_id + || current.last_wal_id != candidate.base_last_wal_id + { + self.start_manifest_refresh(); + return; + } + + let manifest_id = candidate.state.generation.manifest_id; + self.inner.install_state(candidate.state); + self.inner.db_stats.reader_manifest_polls.increment(1); + info!("refreshed reader to latest manifest [manifest_id={manifest_id}]"); + self.schedule_tail_continuation(); + } + + async fn apply_runtime_tail(&mut self, turn: RuntimeReplayTurn) { + self.tail_in_flight = false; + self.finish_tail_task().await; + let current = Arc::clone(&self.inner.state.read()); + if current.generation.manifest_id != turn.manifest_id + || current.last_wal_id != turn.base_last_wal_id + { + self.schedule_tail_continuation(); + return; + } + + let generation = Arc::clone(¤t.generation); + self.inner + .oracle + .advance_durable_seq(turn.last_committed_seq); + self.inner + .db_stats + .reader_replay_memtables + .set(turn.imm_memtable.len() as i64); + let new_state = Arc::new(ReaderState { + generation, + imm_memtable: turn.imm_memtable, + last_wal_id: turn.last_wal_id, + last_remote_persisted_seq: turn.last_committed_seq, + }); + let touched_segments = collect_touched_segments(new_state.as_ref()); + *self.inner.state.write() = new_state; + self.inner + .status_manager + .report_memtable_segments(touched_segments); + + if turn.reached_turn_limit || !turn.caught_up { + self.schedule_tail_continuation(); + } + } + fn report_active_checkpoints(&self) { self.inner .db_stats @@ -1162,8 +1789,34 @@ impl ManifestPoller { } } +struct TailReplayNotifier { + receiver: async_channel::Receiver, +} + +#[async_trait] +impl Notifier for TailReplayNotifier { + async fn notify(&mut self) -> DbReaderMessage { + match self.receiver.recv().await { + Ok(message) => message, + Err(_) => std::future::pending().await, + } + } +} + impl Drop for ManifestPoller { fn drop(&mut self) { + if let Some(scope) = self.tail_scope.take() { + scope.cancel(); + } + if let Some(task) = self.tail_task.take() { + task.abort(); + } + if let Some(scope) = self.manifest_refresh_scope.take() { + scope.cancel(); + } + if let Some(task) = self.manifest_refresh_task.take() { + task.abort(); + } // The gauge tracks only GC checkpoints actively managed by this // poller. Reset it even if startup, cleanup, or the poller task fails. self.inner.db_stats.reader_active_checkpoints.set(0); @@ -1173,14 +1826,61 @@ impl Drop for ManifestPoller { #[async_trait] impl MessageHandler for ManifestPoller { fn tickers(&mut self) -> Vec> { - vec![MessageTickerDef::new( + let mut tickers = vec![MessageTickerDef::new( self.inner.options.manifest_poll_interval, Box::new(|| DbReaderMessage::PollManifest), - )] + )]; + if self.runtime_tailing_enabled() { + tickers.push(MessageTickerDef::new( + self.inner.options.wal_poll_interval, + Box::new(|| DbReaderMessage::PollWalTail), + )); + } + tickers + } + + fn notifiers(&mut self) -> Vec>> { + self.tail_rx.take().map_or_else(Vec::new, |receiver| { + vec![Box::new(TailReplayNotifier { receiver }) as Box>] + }) } async fn handle(&mut self, message: DbReaderMessage) -> Result<(), SlateDBError> { - assert!(matches!(message, DbReaderMessage::PollManifest)); + let replay_tasks = self.inner.replay_tasks.clone(); + tokio::select! { + biased; + _ = replay_tasks.cancelled() => { + self.cancel_tail_task().await; + self.cancel_manifest_refresh_task().await; + Ok(()) + } + result = async { + match message { + DbReaderMessage::PollWalTail | DbReaderMessage::ContinueWalTail => { + self.continuation_queued = false; + self.start_runtime_tail(); + return Ok(()); + } + DbReaderMessage::WalTailComplete(result) => { + let turn = result?; + self.apply_runtime_tail(turn).await; + return Ok(()); + } + DbReaderMessage::ManifestRefreshComplete(result) => { + match result { + Ok(candidate) => self.apply_manifest_refresh(candidate).await, + Err(error) => { + self.manifest_refresh_in_flight = false; + self.finish_manifest_refresh_task().await; + warn!( + "failed to refresh reader to latest manifest [error={error:?}]" + ); + } + } + return Ok(()); + } + DbReaderMessage::PollManifest => {} + } match self.inner.mode { DbReaderMode::ManagedCheckpoint => { let mut manifest = StoredManifest::load( @@ -1198,8 +1898,7 @@ impl MessageHandler for ManifestPoller { let checkpoint = self.inner.create_checkpoint(&mut manifest).await?; self.inner.reestablish_checkpoint(checkpoint).await?; self.register_current_generation(); - } else { - self.inner.maybe_replay_new_wals().await?; + self.schedule_tail_continuation(); } self.refresh_live_checkpoints(&mut manifest).await?; @@ -1207,17 +1906,15 @@ impl MessageHandler for ManifestPoller { Ok(()) } DbReaderMode::FollowLatest => { - let result = self.inner.refresh_latest_manifest().await; - if let Err(error) = result { - warn!("failed to refresh reader to latest manifest [error={error:?}]"); - } else { - self.inner.db_stats.reader_manifest_polls.increment(1); - } + self.cancel_tail_task().await; + self.start_manifest_refresh(); Ok(()) } // No polling is needed for a pinned checkpoint, so we just return Ok(()). DbReaderMode::Checkpoint(_) => Ok(()), } + } => result, + } } async fn cleanup( @@ -1225,16 +1922,9 @@ impl MessageHandler for ManifestPoller { _messages: BoxStream<'async_trait, DbReaderMessage>, _result: Result<(), SlateDBError>, ) -> Result<(), SlateDBError> { - if self.inner.mode != DbReaderMode::ManagedCheckpoint { - return Ok(()); - } - let mut manifest = StoredManifest::load( - Arc::clone(&self.inner.manifest_store), - self.inner.system_clock.clone(), - ) - .await?; - let checkpoint_ids = self.generations.keys().copied().collect::>(); - if !checkpoint_ids.is_empty() { + self.cancel_tail_task().await; + self.cancel_manifest_refresh_task().await; + if self.inner.mode == DbReaderMode::ManagedCheckpoint { let live_generations = self .generations .values() @@ -1249,11 +1939,10 @@ impl MessageHandler for ManifestPoller { for generation in live_generations { generation.drain().await; } - info!( - "deleting reader established checkpoints for shutdown [checkpoint_ids={:?}]", - checkpoint_ids - ); - manifest.delete_checkpoints(&checkpoint_ids).await?; + // Managed checkpoints have a finite lease and can safely expire. + // Remote cleanup is intentionally not part of reader close: an + // unavailable object store must not prevent retirement proof. + self.generations.clear(); } self.inner.db_stats.reader_active_checkpoints.set(0); Ok(()) @@ -1262,6 +1951,12 @@ impl MessageHandler for ManifestPoller { impl DbReader { fn validate_options(mode: DbReaderMode, options: &DbReaderOptions) -> Result<(), SlateDBError> { + options.wal_replay.validate()?; + if options.wal_poll_interval.is_zero() { + return Err(SlateDBError::InvalidWalPollInterval( + options.wal_poll_interval, + )); + } if mode != DbReaderMode::ManagedCheckpoint { return Ok(()); } @@ -1836,10 +2531,10 @@ impl DbReader { // both reader modes before shutting down any managed task. self.inner.status_manager.write_result(Ok(())); - self.task_executor - .shutdown_task(DB_READER_TASK_NAME) - .await - .map_err(Into::::into)?; + self.inner.replay_tasks.cancel(); + let shutdown_result = self.task_executor.shutdown_task(DB_READER_TASK_NAME).await; + self.inner.replay_tasks.shutdown().await; + shutdown_result.map_err(Into::::into)?; if let Err(e) = self.inner.table_store.close_cache().await { warn!("failed to close block cache [error={:?}]", e); @@ -1953,27 +2648,12 @@ impl DbCacheManagerOps for DbReader { } } -/// Checks if the error or any of its sources is an `object_store::Error::NotFound` error. -fn has_not_found_object_store_error(err: &(dyn std::error::Error + 'static)) -> bool { - let mut current = Some(err); - while let Some(current_err) = current { - if current_err - .downcast_ref::() - .is_some_and(|err| matches!(err, object_store::Error::NotFound { .. })) - || current_err - .downcast_ref::>() - .is_some_and(|err| matches!(err.as_ref(), object_store::Error::NotFound { .. })) - { - return true; - } - current = current_err.source(); - } - false -} - #[cfg(test)] mod tests { - use super::{DbReaderMessage, ManifestPoller, ReaderGeneration, ReaderState, ReplayMemtables}; + use super::{ + DbReaderMessage, ManifestPoller, ReaderGeneration, ReaderState, ReplayMemtables, + RuntimeMissingPolicy, MAX_RUNTIME_REPLAY_TURN_TIME, + }; use crate::block_cache_policy::BlockCachePolicy; use crate::clock::MonotonicClock; use crate::config::{ @@ -1988,7 +2668,7 @@ mod tests { use crate::db_status::DbStatusManager; use crate::dispatcher::MessageHandler; use crate::format::sst::SsTableFormat; - use crate::iter::IterationOrder; + use crate::iter::{IterationOrder, RowEntryIterator}; use crate::manifest::store::{ManifestStore, StoredManifest}; use crate::manifest::{Manifest, ManifestCore, VersionedManifest}; use crate::mem_table::{ImmutableMemtable, WritableKVTable}; @@ -1999,6 +2679,7 @@ mod tests { use crate::proptest_util::rng::new_test_rng; use crate::proptest_util::sample; use crate::reader::Reader; + use crate::replay_task_scope::ReplayTaskScope; use crate::tablestore::{TableStore, TableStoreKind}; use crate::types::RowEntry; use crate::{error::SlateDBError, test_utils, CloseReason, Db}; @@ -2049,6 +2730,24 @@ mod tests { .expect("reader did not install a new manifest generation"); } + async fn finish_manual_manifest_refresh(poller: &mut ManifestPoller) { + let receiver = poller + .tail_rx + .as_ref() + .expect("manual poller must retain its completion receiver") + .clone(); + loop { + let message = tokio::time::timeout(Duration::from_secs(5), receiver.recv()) + .await + .expect("manifest refresh did not complete") + .expect("manifest completion channel closed"); + poller.handle(message).await.unwrap(); + if !poller.manifest_refresh_in_flight { + break; + } + } + } + #[tokio::test] async fn should_get_value_from_db() { let object_store: Arc = Arc::new(InMemory::new()); @@ -2472,8 +3171,10 @@ mod tests { .contains("snapshots are unsupported in FollowLatest mode")); assert!(recording_store.write_kinds().is_empty()); - let mut poller = ManifestPoller::new(Arc::clone(&reader.inner)); + let (tail_tx, tail_rx) = async_channel::unbounded(); + let mut poller = ManifestPoller::new(Arc::clone(&reader.inner), tail_tx, tail_rx); poller.handle(DbReaderMessage::PollManifest).await.unwrap(); + finish_manual_manifest_refresh(&mut poller).await; assert!(reader.manifest().id() >= latest_manifest.id); assert_eq!( @@ -2542,8 +3243,10 @@ mod tests { saved_manifests.push((location, bytes)); } - let mut poller = ManifestPoller::new(Arc::clone(&reader.inner)); + let (tail_tx, tail_rx) = async_channel::unbounded(); + let mut poller = ManifestPoller::new(Arc::clone(&reader.inner), tail_tx, tail_rx); poller.handle(DbReaderMessage::PollManifest).await.unwrap(); + finish_manual_manifest_refresh(&mut poller).await; assert_eq!(reader.manifest().id(), manifest_id); assert_eq!( @@ -2564,6 +3267,7 @@ mod tests { db.close().await.unwrap(); poller.handle(DbReaderMessage::PollManifest).await.unwrap(); + finish_manual_manifest_refresh(&mut poller).await; assert!(reader.manifest().id() > manifest_id); assert_eq!( reader.get(b"key").await.unwrap(), @@ -2572,6 +3276,95 @@ mod tests { reader.close().await.unwrap(); } + #[tokio::test] + async fn close_cancels_a_blocked_manifest_get() { + let inner = Arc::new(InMemory::new()); + let writer_store: Arc = inner.clone(); + let gated = Arc::new(test_utils::GatedObjectStore::new(inner)); + let object_store: Arc = gated.clone(); + let path = Path::from("/tmp/test_reader_close_blocked_manifest_list"); + let provider = TestProvider::new(path.clone(), writer_store); + let db = provider.new_db(Settings::default()).await.unwrap(); + db.put(b"key", b"value").await.unwrap(); + db.flush().await.unwrap(); + + let reader = DbReader::open( + path, + object_store, + DbReaderMode::FollowLatest, + DbReaderOptions { + manifest_poll_interval: Duration::from_secs(60), + wal_poll_interval: Duration::from_secs(60), + checkpoint_lifetime: Duration::ZERO, + ..DbReaderOptions::default() + }, + ) + .await + .unwrap(); + + db.put(b"key", b"updated").await.unwrap(); + db.flush().await.unwrap(); + let (tx, rx) = async_channel::unbounded(); + let mut poller = ManifestPoller::new(Arc::clone(&reader.inner), tx, rx); + gated.get_opts_gate.close(); + poller.handle(DbReaderMessage::PollManifest).await.unwrap(); + tokio::time::timeout( + Duration::from_secs(1), + gated.get_opts_gate.wait_for_arrivals(1), + ) + .await + .expect("manifest refresh did not reach the blocked GET"); + + tokio::time::timeout(Duration::from_secs(1), reader.close()) + .await + .expect("reader close waited for a blocked manifest GET") + .unwrap(); + poller.cancel_manifest_refresh_task().await; + } + + #[tokio::test] + async fn close_cancels_a_blocked_checkpoint_manifest_list() { + let inner = Arc::new(InMemory::new()); + let writer_store: Arc = inner.clone(); + let gated = Arc::new(test_utils::GatedObjectStore::new(inner)); + let reader_store: Arc = gated.clone(); + let path = Path::from("/tmp/test_reader_close_blocked_checkpoint_list"); + let provider = TestProvider::new(path.clone(), writer_store); + let db = provider.new_db(Settings::default()).await.unwrap(); + db.put(b"key", b"value").await.unwrap(); + db.flush().await.unwrap(); + + let reader = DbReader::open( + path, + reader_store, + DbReaderMode::ManagedCheckpoint, + DbReaderOptions { + manifest_poll_interval: Duration::from_secs(60), + wal_poll_interval: Duration::from_secs(60), + ..DbReaderOptions::default() + }, + ) + .await + .unwrap(); + let (tx, rx) = async_channel::unbounded(); + let mut poller = ManifestPoller::new(Arc::clone(&reader.inner), tx, rx); + gated.list_gate.close(); + let poll = tokio::spawn(async move { poller.handle(DbReaderMessage::PollManifest).await }); + tokio::time::timeout(Duration::from_secs(1), gated.list_gate.wait_for_arrivals(1)) + .await + .expect("checkpoint refresh did not reach the blocked LIST"); + + tokio::time::timeout(Duration::from_secs(1), reader.close()) + .await + .expect("reader close waited for a blocked manifest LIST") + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), poll) + .await + .expect("cancelled checkpoint poll remained blocked") + .unwrap() + .unwrap(); + } + #[tokio::test(start_paused = true)] async fn should_reestablish_reader_checkpoint() { let object_store: Arc = Arc::new(InMemory::new()); @@ -2585,6 +3378,7 @@ mod tests { let db = test_provider.new_db(db_options).await.unwrap(); let reader_options = DbReaderOptions { manifest_poll_interval: Duration::from_millis(10), + wal_poll_interval: Duration::from_millis(10), ..DbReaderOptions::default() }; let reader = test_provider @@ -2627,6 +3421,7 @@ mod tests { let _db = test_provider.new_db(Settings::default()).await; let reader_options = DbReaderOptions { manifest_poll_interval: Duration::from_millis(500), + wal_poll_interval: Duration::from_millis(10), checkpoint_lifetime: Duration::from_millis(1000), ..DbReaderOptions::default() }; @@ -2664,10 +3459,12 @@ mod tests { > initial_reader_checkpoint.expire_time.unwrap() ); - // The checkpoint is removed on shutdown + // Shutdown is a bounded local operation. It must not wait for a remote + // manifest write merely to delete the checkpoint; its finite lease is + // reclaimed by normal checkpoint GC after expiry. reader.close().await.unwrap(); let updated_manifest = manifest_store.read_latest_manifest().await.unwrap(); - assert_eq!(0, updated_manifest.manifest.core.checkpoints.len()); + assert_eq!(1, updated_manifest.manifest.core.checkpoints.len()); } // Regression test for https://github.com/slatedb/slatedb/issues/1750. @@ -2855,6 +3652,7 @@ mod tests { let reader_options = DbReaderOptions { manifest_poll_interval: Duration::from_millis(500), + wal_poll_interval: Duration::from_millis(10), checkpoint_lifetime: Duration::from_millis(1000), ..DbReaderOptions::default() }; @@ -2868,7 +3666,7 @@ mod tests { db.put(key, value).await.unwrap(); db.flush().await.unwrap(); - tokio::time::sleep(Duration::from_millis(500)).await; + tokio::time::sleep(Duration::from_millis(20)).await; assert_eq!( reader.get(key).await.unwrap(), Some(Bytes::from_static(value)) @@ -3489,7 +4287,7 @@ mod tests { let mut core = ManifestCore::new(); core.next_wal_sst_id = 5; - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3567,7 +4365,7 @@ mod tests { publications.push((wal_id, seq, tables.len())); }; - let result = DbReaderInner::replay_wal_into( + let result = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &options, &core, @@ -3585,28 +4383,305 @@ mod tests { assert_eq!(vec![(1, 1, 1), (2, 2, 2), (3, 3, 3)], publications); } - #[test] - fn has_not_found_object_store_error_should_walk_nested_error_sources() { - let err = crate::Error::from(SlateDBError::from(object_store::Error::NotFound { - path: "missing-wal".to_string(), - source: Box::new(std::io::Error::other("missing")), - })); + #[tokio::test] + async fn runtime_known_missing_wal_is_error_and_exact_next_probe_reports_caught_up() { + let object_store: Arc = Arc::new(InMemory::new()); + let path = Path::from("/tmp/test_db_reader_runtime_missing_contract"); + let table_store = TestProvider::new(path, object_store).table_store(); + let options = DbReaderOptions::default(); + let mut core = ManifestCore::new(); + core.next_wal_sst_id = 2; + let mut known_cursor = (0, 0); - assert!(super::has_not_found_object_store_error(&err)); + let known_missing = DbReaderInner::replay_runtime_range_into( + Arc::clone(&table_store), + &options, + &core, + &mut ReplayMemtables::default(), + &mut known_cursor, + 1..2, + RuntimeMissingPolicy::Error, + None, + None, + None, + None, + ) + .await; + assert!(known_missing.unwrap_err().has_object_store_not_found()); + + let mut tail_cursor = (0, 0); + let caught_up = DbReaderInner::replay_runtime_range_into( + table_store, + &options, + &core, + &mut ReplayMemtables::default(), + &mut tail_cursor, + 1..2, + RuntimeMissingPolicy::ExactNextIsCaughtUp(1), + None, + None, + None, + None, + ) + .await + .unwrap(); + assert!(caught_up); + assert_eq!(tail_cursor, (0, 0)); } - #[test] - fn has_not_found_object_store_error_should_ignore_non_not_found_errors() { - let err = SlateDBError::from(object_store::Error::NotImplemented { - operation: "test".to_string(), - implementer: "test".to_string(), + #[tokio::test] + async fn manifest_replay_yields_at_the_sixty_four_wal_boundary() { + let object_store: Arc = Arc::new(InMemory::new()); + let path = Path::from("/tmp/test_db_reader_manifest_turn_limit"); + let table_store = TestProvider::new(path, object_store).table_store(); + for wal_id in 1..=65 { + write_wal_sst( + Arc::clone(&table_store), + wal_id, + vec![RowEntry::new_value( + format!("key-{wal_id:03}").as_bytes(), + b"value", + wal_id, + )], + ) + .await + .unwrap(); + } + let mut core = ManifestCore::new(); + core.next_wal_sst_id = 66; + let mut replayed = ReplayMemtables::default(); + let clock: Arc = Arc::new(DefaultSystemClock::new()); + let scope = ReplayTaskScope::new(); + + let cursor = DbReaderInner::replay_manifest_range_into( + table_store, + &DbReaderOptions::default(), + &core, + &mut replayed, + None, + None, + None, + Some(scope.clone()), + &clock, + ) + .await + .unwrap(); + + assert_eq!(cursor, (65, 65)); + assert_eq!(replayed.len(), 2); + scope.shutdown().await; + } + + #[tokio::test] + async fn runtime_tail_yields_after_sixty_four_wals_and_continues_without_listing() { + let recording = Arc::new(test_utils::RecordingObjectStore::new(Arc::new( + InMemory::new(), + ))); + let object_store: Arc = recording.clone(); + let path = Path::from("/tmp/test_db_reader_runtime_turn_limit"); + let provider = TestProvider::new(path, object_store); + let db = provider + .new_db(Settings { + flush_interval: None, + compactor_options: None, + garbage_collector_options: None, + ..Settings::default() + }) + .await + .unwrap(); + let checkpoint = db + .create_checkpoint(CheckpointScope::All, &CheckpointOptions::default()) + .await + .unwrap(); + db.close().await.unwrap(); + let reader = provider + .new_db_reader(DbReaderOptions::default(), Some(checkpoint.id), None) + .await + .unwrap(); + let table_store = provider.table_store(); + let initial_wal_id = reader.inner.state.read().last_wal_id; + for wal_id in initial_wal_id + 1..=initial_wal_id + 65 { + write_wal_sst( + Arc::clone(&table_store), + wal_id, + vec![RowEntry::new_value( + format!("key-{wal_id:03}").as_bytes(), + b"value", + wal_id, + )], + ) + .await + .unwrap(); + } + recording.clear(); + + let first_base = Arc::clone(&reader.inner.state.read()); + let first = DbReaderInner::replay_runtime_turn( + Arc::clone(&reader.inner), + first_base, + reader.inner.replay_tasks.child(), + ) + .await + .unwrap(); + assert_eq!(first.last_wal_id, initial_wal_id + 64); + assert!(first.reached_turn_limit); + assert!(!first.caught_up); + + let second_base = Arc::new(ReaderState { + generation: Arc::clone(&reader.inner.state.read().generation), + imm_memtable: first.imm_memtable, + last_wal_id: first.last_wal_id, + last_remote_persisted_seq: first.last_committed_seq, }); + let second = DbReaderInner::replay_runtime_turn( + Arc::clone(&reader.inner), + second_base, + reader.inner.replay_tasks.child(), + ) + .await + .unwrap(); + assert_eq!(second.last_wal_id, initial_wal_id + 65); + assert!(second.caught_up); + assert!(!second.reached_turn_limit); + assert_eq!(recording.list_calls(), 0); + + reader.close().await.unwrap(); + } + + #[tokio::test(start_paused = true)] + async fn runtime_tail_deadline_wins_without_publishing_a_partial_wal() { + let inner_store: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(test_utils::GatedObjectStore::new(inner_store)); + let object_store: Arc = gated.clone(); + let path = Path::from("/tmp/test_db_reader_runtime_deadline"); + let provider = TestProvider::new(path, object_store); + let db = provider + .new_db(Settings { + flush_interval: None, + compactor_options: None, + garbage_collector_options: None, + ..Settings::default() + }) + .await + .unwrap(); + let reader = provider + .new_db_reader( + DbReaderOptions { + // Keep the test-owned replay turn isolated from the background poller. + skip_wal_replay: true, + ..DbReaderOptions::default() + }, + None, + None, + ) + .await + .unwrap(); + let base = Arc::clone(&reader.inner.state.read()); + let base_memtable_count = base.imm_memtable.len(); + db.put_with_options( + b"deadline-key", + b"deadline-value", + &PutOptions::default(), + &WriteOptions { + await_durable: false, + ..WriteOptions::default() + }, + ) + .await + .unwrap(); + db.flush_with_options(FlushOptions { + flush_type: FlushType::Wal, + }) + .await + .unwrap(); + + let prior_head_arrivals = gated.head_gate.arrivals(); + gated.head_gate.close(); + let blocked = tokio::spawn(DbReaderInner::replay_runtime_turn( + Arc::clone(&reader.inner), + Arc::clone(&base), + reader.inner.replay_tasks.child(), + )); + gated + .head_gate + .wait_for_arrivals(prior_head_arrivals + 1) + .await; + tokio::time::advance(MAX_RUNTIME_REPLAY_TURN_TIME).await; + // Make the object read ready at the exact deadline. The deadline branch + // is biased and must still win. + gated.head_gate.release(); + let expired = blocked.await.unwrap().unwrap(); + assert_eq!(expired.last_wal_id, base.last_wal_id); + assert_eq!(expired.last_committed_seq, base.last_remote_persisted_seq); + assert_eq!(expired.imm_memtable.len(), base_memtable_count); + assert!(expired.reached_turn_limit); + assert!(!expired.caught_up); + assert_eq!(reader.inner.state.read().last_wal_id, base.last_wal_id); + + let completed = DbReaderInner::replay_runtime_turn( + Arc::clone(&reader.inner), + base, + reader.inner.replay_tasks.child(), + ) + .await + .unwrap(); + assert!(completed.last_wal_id > expired.last_wal_id); + assert!(completed.caught_up); + let mut rows = completed + .imm_memtable + .front() + .expect("completed WAL must be present") + .table() + .iter(); + let row = rows.next().await.unwrap().unwrap(); + assert_eq!(row.key, Bytes::from_static(b"deadline-key")); + assert_eq!( + row.value, + crate::types::ValueDeletable::Value(Bytes::from_static(b"deadline-value")) + ); + + reader.close().await.unwrap(); + db.close().await.unwrap(); + } + + #[test] + fn reader_validation_rejects_invalid_wal_replay_settings_in_every_mode() { + for mode in [ + DbReaderMode::Checkpoint(Uuid::nil()), + DbReaderMode::ManagedCheckpoint, + DbReaderMode::FollowLatest, + ] { + let options = DbReaderOptions { + wal_replay: crate::config::WalReplaySettings { + max_concurrent_objects: 0, + ..Default::default() + }, + ..Default::default() + }; + assert!(DbReader::validate_options(mode, &options).is_err()); + } + } - assert!(!super::has_not_found_object_store_error(&err)); + #[test] + fn reader_validation_rejects_zero_wal_poll_interval_in_every_mode() { + for mode in [ + DbReaderMode::Checkpoint(Uuid::nil()), + DbReaderMode::ManagedCheckpoint, + DbReaderMode::FollowLatest, + ] { + let options = DbReaderOptions { + wal_poll_interval: Duration::ZERO, + ..DbReaderOptions::default() + }; + assert!(matches!( + DbReader::validate_options(mode, &options), + Err(SlateDBError::InvalidWalPollInterval(interval)) if interval.is_zero() + )); + } } #[tokio::test] - async fn replay_wal_into_should_treat_missing_wal_sst_as_end_of_iteration() { + async fn exact_replay_should_fail_closed_on_a_missing_wal() { let object_store: Arc = Arc::new(InMemory::new()); let path = Path::from("/tmp/test_db_reader_missing_wal"); let test_provider = TestProvider::new(path, Arc::clone(&object_store)); @@ -3624,7 +4699,7 @@ mod tests { let mut core = ManifestCore::new(); core.next_wal_sst_id = 3; - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let error = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3636,15 +4711,14 @@ mod tests { None, ) .await - .unwrap(); + .unwrap_err(); - assert_eq!(last_wal_id, 0); - assert_eq!(last_committed_seq, 0); + assert!(error.has_object_store_not_found()); assert!(into_tables.is_empty()); } #[tokio::test] - async fn replay_wal_into_should_keep_previously_replayed_tables_before_missing_wal_sst() { + async fn exact_replay_should_surface_a_later_missing_wal() { let object_store: Arc = Arc::new(InMemory::new()); let path = Path::from("/tmp/test_db_reader_missing_wal_after_replay"); let test_provider = TestProvider::new(path, Arc::clone(&object_store)); @@ -3679,7 +4753,7 @@ mod tests { ..DbReaderOptions::default() }; - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let error = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &reader_options, &core, @@ -3691,10 +4765,9 @@ mod tests { None, ) .await - .unwrap(); + .unwrap_err(); - assert_eq!(last_wal_id, 2); - assert_eq!(last_committed_seq, 3); + assert!(error.has_object_store_not_found()); assert_eq!(into_tables.len(), 1); let replayed = into_tables.front().unwrap(); @@ -3718,7 +4791,7 @@ mod tests { let mut into_tables = ReplayMemtables::default(); let core = ManifestCore::new(); - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3752,7 +4825,7 @@ mod tests { let mut into_tables = ReplayMemtables::default(); let core = ManifestCore::new(); - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3801,7 +4874,7 @@ mod tests { core.last_l0_seq = 8; core.next_wal_sst_id = 5; - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3819,7 +4892,7 @@ mod tests { assert_eq!(last_committed_seq, 10); let head_after_first_replay = Arc::clone(into_tables.head.as_ref().unwrap()); - let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into( + let (last_wal_id, last_committed_seq) = DbReaderInner::replay_wal_into_exact( Arc::clone(&table_store), &DbReaderOptions::default(), &core, @@ -3854,6 +4927,7 @@ mod tests { let reader_options = DbReaderOptions { manifest_poll_interval: Duration::from_millis(500), + wal_poll_interval: Duration::from_millis(10), ..DbReaderOptions::default() }; let metrics_recorder = Arc::new(DefaultMetricsRecorder::new()); @@ -3880,13 +4954,16 @@ mod tests { fail_parallel::cfg( Arc::clone(&test_provider.fp_registry), - "probe-wal-ssts", + "runtime-wal-replay", "return", ) .unwrap(); - tokio::time::sleep(Duration::from_millis(20)).await; + tokio::time::sleep(Duration::from_millis(30)).await; let result = reader.get(b"key").await.unwrap_err(); - assert_eq!(result.to_string(), "Unavailable error: io error (oops)"); + assert_eq!( + result.to_string(), + "Unavailable error: io error (runtime WAL replay failpoint)" + ); assert_eq!( Some(0), lookup_metric( @@ -3935,6 +5012,12 @@ mod tests { .new_db_reader(reader_options.clone(), None, None) .await .unwrap(); + let (tail_tx, tail_rx) = async_channel::unbounded(); + let mut skipped_poller = ManifestPoller::new(Arc::clone(&reader.inner), tail_tx, tail_rx); + assert!(!skipped_poller.runtime_tailing_enabled()); + assert_eq!(skipped_poller.tickers().len(), 1); + skipped_poller.start_runtime_tail(); + assert!(!skipped_poller.tail_in_flight); // Should see the L0 flushed data assert_eq!( @@ -4337,6 +5420,7 @@ mod tests { status_manager: DbStatusManager::new(0), segment_extractor: None, rand: test_provider.rand.clone(), + replay_tasks: ReplayTaskScope::new(), recorder, }; @@ -4429,6 +5513,7 @@ mod tests { status_manager: DbStatusManager::new(0), segment_extractor: None, rand: test_provider.rand.clone(), + replay_tasks: ReplayTaskScope::new(), recorder, } } @@ -4662,7 +5747,9 @@ mod tests { #[tokio::test] async fn should_record_incremental_wal_replay_metrics() { - use slatedb_common::metrics::{lookup_metric, DefaultMetricsRecorder}; + use slatedb_common::metrics::{ + lookup_metric, lookup_metric_with_labels, DefaultMetricsRecorder, + }; let object_store: Arc = Arc::new(InMemory::new()); let path = Path::from("/tmp/test_db_reader_wal_replay_metrics"); @@ -4717,6 +5804,42 @@ mod tests { Some(1), lookup_metric(&metrics_recorder, crate::db_stats::READER_REPLAY_MEMTABLES) ); + let list_metric = crate::db_stats::READER_WAL_REPLAY_OBJECT_STORE_CALLS; + assert!(lookup_metric_with_labels( + &metrics_recorder, + list_metric, + &[ + ( + crate::db_stats::REPLAY_SOURCE_LABEL, + crate::db_stats::REPLAY_SOURCE_READER_OPEN, + ), + ( + crate::db_stats::REPLAY_OPERATION_LABEL, + crate::db_stats::REPLAY_OPERATION_LIST, + ), + ], + ) + .is_some_and(|value| value > 0)); + for source in [ + crate::db_stats::REPLAY_SOURCE_RUNTIME_MANIFEST, + crate::db_stats::REPLAY_SOURCE_RUNTIME_TAIL, + ] { + assert_eq!( + lookup_metric_with_labels( + &metrics_recorder, + list_metric, + &[ + (crate::db_stats::REPLAY_SOURCE_LABEL, source), + ( + crate::db_stats::REPLAY_OPERATION_LABEL, + crate::db_stats::REPLAY_OPERATION_LIST, + ), + ], + ), + Some(0), + "runtime replay source unexpectedly performed LIST: {source}", + ); + } assert_eq!( Some(1), lookup_metric( @@ -4761,6 +5884,12 @@ mod tests { .build() .await .unwrap(); + let (tail_tx, tail_rx) = async_channel::unbounded(); + let mut pinned_poller = ManifestPoller::new(Arc::clone(&reader.inner), tail_tx, tail_rx); + assert!(!pinned_poller.runtime_tailing_enabled()); + assert_eq!(pinned_poller.tickers().len(), 1); + pinned_poller.start_runtime_tail(); + assert!(!pinned_poller.tail_in_flight); assert_eq!( Some(0), diff --git a/slatedb/src/db_stats.rs b/slatedb/src/db_stats.rs index f34b0e46e..cb8c8fe5c 100644 --- a/slatedb/src/db_stats.rs +++ b/slatedb/src/db_stats.rs @@ -52,6 +52,15 @@ pub const READER_WAL_REPLAY_BATCHES: &str = db_stat_name!("reader_wal_replay_bat pub const READER_REPLAY_MEMTABLES: &str = db_stat_name!("reader_replay_memtables"); pub const READER_ACTIVE_CHECKPOINTS: &str = db_stat_name!("reader_active_checkpoints"); pub const READER_MANIFEST_POLLS: &str = db_stat_name!("reader_manifest_polls"); +pub const READER_WAL_REPLAY_OBJECT_STORE_CALLS: &str = + db_stat_name!("reader_wal_replay_object_store_calls"); +pub const REPLAY_SOURCE_LABEL: &str = "source"; +pub const REPLAY_OPERATION_LABEL: &str = "operation"; +pub const REPLAY_OPERATION_LIST: &str = "list"; +pub const REPLAY_SOURCE_READER_OPEN: &str = "reader_open"; +pub const REPLAY_SOURCE_CHECKPOINT_RECOVERY: &str = "checkpoint_recovery"; +pub const REPLAY_SOURCE_RUNTIME_MANIFEST: &str = "runtime_manifest"; +pub const REPLAY_SOURCE_RUNTIME_TAIL: &str = "runtime_tail"; /// Label key distinguishing filter metrics for point lookups from those for /// prefix scans. Value is one of [`FILTER_KIND_POINT`] or @@ -96,6 +105,8 @@ pub(crate) struct DbStatsInner { pub(crate) reader_replay_memtables: Arc, pub(crate) reader_active_checkpoints: Arc, pub(crate) reader_manifest_polls: Arc, + pub(crate) reader_wal_replay_list_reader_open: Arc, + pub(crate) reader_wal_replay_list_checkpoint_recovery: Arc, } #[derive(Clone)] @@ -205,7 +216,38 @@ impl DbStats { .counter(READER_MANIFEST_POLLS) .description("Number of successful DbReader manifest polls completed") .register(), + reader_wal_replay_list_reader_open: recorder + .counter(READER_WAL_REPLAY_OBJECT_STORE_CALLS) + .labels(&[ + (REPLAY_SOURCE_LABEL, REPLAY_SOURCE_READER_OPEN), + (REPLAY_OPERATION_LABEL, REPLAY_OPERATION_LIST), + ]) + .register(), + reader_wal_replay_list_checkpoint_recovery: recorder + .counter(READER_WAL_REPLAY_OBJECT_STORE_CALLS) + .labels(&[ + (REPLAY_SOURCE_LABEL, REPLAY_SOURCE_CHECKPOINT_RECOVERY), + (REPLAY_OPERATION_LABEL, REPLAY_OPERATION_LIST), + ]) + .register(), }; + // Register zero-valued runtime series so live monitoring can prove + // runtime replay did not LIST without conflating fencing, GC, or + // maintenance operations from the same object store. + recorder + .counter(READER_WAL_REPLAY_OBJECT_STORE_CALLS) + .labels(&[ + (REPLAY_SOURCE_LABEL, REPLAY_SOURCE_RUNTIME_MANIFEST), + (REPLAY_OPERATION_LABEL, REPLAY_OPERATION_LIST), + ]) + .register(); + recorder + .counter(READER_WAL_REPLAY_OBJECT_STORE_CALLS) + .labels(&[ + (REPLAY_SOURCE_LABEL, REPLAY_SOURCE_RUNTIME_TAIL), + (REPLAY_OPERATION_LABEL, REPLAY_OPERATION_LIST), + ]) + .register(); DbStats { inner: Arc::new(inner), } diff --git a/slatedb/src/db_transaction.rs b/slatedb/src/db_transaction.rs index 78a7c4d2e..d1e832745 100644 --- a/slatedb/src/db_transaction.rs +++ b/slatedb/src/db_transaction.rs @@ -3775,6 +3775,7 @@ mod tests { min_filter_keys, l0_sst_size_bytes, max_wal_flushes_before_l0_flush: 4096, + wal_replay: crate::config::WalReplaySettings::default(), compactor_options, compression_codec: None, object_store_cache_options: crate::config::ObjectStoreCacheOptions::default(), diff --git a/slatedb/src/error.rs b/slatedb/src/error.rs index 651401ba8..b10cdfa5e 100644 --- a/slatedb/src/error.rs +++ b/slatedb/src/error.rs @@ -34,6 +34,24 @@ pub(crate) enum SlateDBError { #[error("empty block")] EmptyBlock, + #[error( + "corrupt SST data: {reason}{}", + .path.as_ref().map(|path| format!(" in {path}")).unwrap_or_default() + )] + CorruptSst { + reason: &'static str, + path: Option, + }, + + #[error( + "WAL replay memory limit exceeded: {kind} requires {required_bytes} bytes, limit is {limit_bytes} bytes" + )] + WalReplayMemoryLimitExceeded { + kind: &'static str, + required_bytes: usize, + limit_bytes: usize, + }, + #[error("empty RowEntry key")] EmptyKey, @@ -251,15 +269,15 @@ pub(crate) enum SlateDBError { #[error("invalid manifest poll interval. interval=`{0:?}`")] InvalidManifestPollInterval(Duration), + #[error("invalid WAL poll interval. interval=`{0:?}`")] + InvalidWalPollInterval(Duration), + #[error("checkpoint lifetime must be at least double the manifest poll interval. lifetime=`{lifetime:?}`, interval=`{interval:?}`")] CheckpointLifetimeTooShort { lifetime: Duration, interval: Duration, }, - #[error("invalid sst batch size. size=`{0}`")] - InvalidSSTBatchSize(usize), - #[error("invalid configuration: {0}")] InvalidConfiguration(String), @@ -336,6 +354,10 @@ impl SlateDBError { SlateDBError::ChecksumMismatch { path: None } => SlateDBError::ChecksumMismatch { path: Some(path.clone()), }, + SlateDBError::CorruptSst { reason, path: None } => SlateDBError::CorruptSst { + reason, + path: Some(path.clone()), + }, other => other, } } @@ -384,6 +406,7 @@ impl SlateDBError { SlateDBError::InvalidFlatbuffer(_) | SlateDBError::EmptyBlock | SlateDBError::EmptyBlockMeta + | SlateDBError::CorruptSst { .. } | SlateDBError::InvalidFilterBlock | SlateDBError::BlockTransformError => Some(RetryReason::BlockDecodeError), _ => None, @@ -716,10 +739,10 @@ impl From for Error { } SlateDBError::InvalidObjectStorePath(_) => Error::invalid(msg), SlateDBError::UnknownConfigurationFormat(_) => Error::invalid(msg), - SlateDBError::InvalidSSTBatchSize(_) => Error::invalid(msg), SlateDBError::InvalidConfiguration(_) => Error::invalid(msg), SlateDBError::InvalidCheckpointLifetime(_) => Error::invalid(msg), - SlateDBError::InvalidManifestPollInterval(_) => Error::invalid(msg), + SlateDBError::InvalidManifestPollInterval(_) + | SlateDBError::InvalidWalPollInterval(_) => Error::invalid(msg), SlateDBError::CheckpointLifetimeTooShort { .. } => Error::invalid(msg), SlateDBError::DbReaderSnapshotUnsupportedInFollowLatest => Error::invalid(msg), SlateDBError::SeekKeyOutOfRange { .. } => Error::invalid(msg), @@ -776,6 +799,8 @@ impl From for Error { SlateDBError::InvalidTransactionalObjectState => Error::data(msg), SlateDBError::EmptyManifest => Error::data(msg), SlateDBError::EmptyBlock => Error::data(msg), + SlateDBError::CorruptSst { .. } => Error::data(msg), + SlateDBError::WalReplayMemoryLimitExceeded { .. } => Error::data(msg), SlateDBError::EmptyKey => Error::data(msg), SlateDBError::EmptyBlockMeta => Error::data(msg), SlateDBError::InvalidFilterBlock => Error::data(msg), diff --git a/slatedb/src/fence.rs b/slatedb/src/fence.rs index ead82336e..734c02e79 100644 --- a/slatedb/src/fence.rs +++ b/slatedb/src/fence.rs @@ -10,6 +10,7 @@ use fail_parallel::{fail_point_send, FailPointTx}; use log::error; use slatedb_common::metrics::MetricsRecorderHelper; use slatedb_common::SystemClock; +use std::num::NonZeroU64; use std::ops::Range; use std::sync::Arc; use std::time::Duration; @@ -85,6 +86,7 @@ impl WriterFencer { pub(crate) async fn fence( self, stored_manifest: StoredManifest, + writer_epoch: Option, ) -> Result { let wal_writer_init = WalWriterInit::load( self.closed_result_reader.clone(), @@ -97,12 +99,25 @@ impl WriterFencer { ) .await?; - let manifest = FenceableManifest::init_writer( - stored_manifest, - self.manifest_update_timeout, - self.system_clock.clone(), - ) - .await?; + let manifest = match writer_epoch { + Some(writer_epoch) => { + FenceableManifest::init_writer_with_epoch( + stored_manifest, + self.manifest_update_timeout, + self.system_clock.clone(), + writer_epoch.get(), + ) + .await? + } + None => { + FenceableManifest::init_writer( + stored_manifest, + self.manifest_update_timeout, + self.system_clock.clone(), + ) + .await? + } + }; self.fail_point_send("FenceManifest"); let mut manifest = manifest.into(); @@ -400,7 +415,7 @@ mod tests { h.put(&db, 1, false).await; let fencer = h.fencer.take().unwrap(); - let result = fencer.fence(h.stored_manifest.take().unwrap()).await; + let result = fencer.fence(h.stored_manifest.take().unwrap(), None).await; assert!(result.is_ok()); h.put(&db, 2, true).await; @@ -447,7 +462,7 @@ mod tests { let fencer = h.fencer.take().unwrap(); let stored_manifest = h.stored_manifest.take().unwrap(); - let jh = tokio::task::spawn(async move { fencer.fence(stored_manifest).await }); + let jh = tokio::task::spawn(async move { fencer.fence(stored_manifest, None).await }); // wait for LoadEmptyWalId pause assert_eq!(h.event_rx.recv().await.unwrap(), "LoadEmptyWalId"); @@ -549,7 +564,7 @@ mod tests { // spawn WriterFencer on another task let fencer = h.fencer.take().unwrap(); let stored_manifest = h.stored_manifest.take().unwrap(); - let jh = tokio::task::spawn(async move { fencer.fence(stored_manifest).await }); + let jh = tokio::task::spawn(async move { fencer.fence(stored_manifest, None).await }); // wait for fencer to load empty wal id and pause h.event_rx.recv().await.unwrap(); diff --git a/slatedb/src/flatbuffer_types.rs b/slatedb/src/flatbuffer_types.rs index 7863e2970..b7f350ab1 100644 --- a/slatedb/src/flatbuffer_types.rs +++ b/slatedb/src/flatbuffer_types.rs @@ -127,7 +127,7 @@ impl SsTableInfoCodec for FlatBufferSsTableInfoCodec { fn decode(&self, bytes: &Bytes) -> Result { let info = flatbuffers::root_with_opts::(&verifier_options(), bytes)?; - Ok(Self::sst_info(&info)) + Self::sst_info_checked(&info) } fn clone_box(&self) -> Box { @@ -137,6 +137,10 @@ impl SsTableInfoCodec for FlatBufferSsTableInfoCodec { impl FlatBufferSsTableInfoCodec { pub(crate) fn sst_info(info: &FbSsTableInfo) -> SsTableInfo { + Self::sst_info_checked(info).expect("validated manifest contains invalid SST info enum") + } + + fn sst_info_checked(info: &FbSsTableInfo) -> Result { let first_entry: Option = info .first_entry() .map(|entry| Bytes::copy_from_slice(entry.bytes())); @@ -144,19 +148,42 @@ impl FlatBufferSsTableInfoCodec { .last_entry() .map(|entry| Bytes::copy_from_slice(entry.bytes())); - SsTableInfo { + let compression_codec = match info.compression_format() { + CompressionFormat::None => None, + #[cfg(feature = "snappy")] + CompressionFormat::Snappy => Some(CompressionCodec::Snappy), + #[cfg(feature = "lz4")] + CompressionFormat::Lz4 => Some(CompressionCodec::Lz4), + #[cfg(feature = "zlib")] + CompressionFormat::Zlib => Some(CompressionCodec::Zlib), + #[cfg(feature = "zstd")] + CompressionFormat::Zstd => Some(CompressionCodec::Zstd), + _ => return Err(corrupt_sst_info("unknown or unavailable compression codec")), + }; + let sst_type = match info.sst_type() { + FbSstType::Compacted => SstType::Compacted, + FbSstType::Wal => SstType::Wal, + _ => return Err(corrupt_sst_info("unknown SST type")), + }; + let filter_format = match info.filter_format() { + FbFilterFormat::Legacy => FilterFormat::Legacy, + FbFilterFormat::Composite => FilterFormat::Composite, + _ => return Err(corrupt_sst_info("unknown filter format")), + }; + + Ok(SsTableInfo { first_entry, last_entry, index_offset: info.index_offset(), index_len: info.index_len(), filter_offset: info.filter_offset(), filter_len: info.filter_len(), - compression_codec: info.compression_format().into(), - sst_type: info.sst_type().into(), + compression_codec, + sst_type, stats_offset: info.stats_offset(), stats_len: info.stats_len(), - filter_format: info.filter_format().into(), - } + filter_format, + }) } pub(crate) fn create_from_sst_info(info: &SsTableInfo) -> Bytes { @@ -1475,16 +1502,6 @@ impl From for FbSstType { } } -impl From for SstType { - fn from(value: FbSstType) -> Self { - match value { - FbSstType::Compacted => SstType::Compacted, - FbSstType::Wal => SstType::Wal, - _ => unreachable!("unknown SstType value: {:?}", value), - } - } -} - impl From for FbFilterFormat { fn from(value: FilterFormat) -> Self { match value { @@ -1494,30 +1511,8 @@ impl From for FbFilterFormat { } } -impl From for FilterFormat { - fn from(value: FbFilterFormat) -> Self { - match value { - FbFilterFormat::Legacy => FilterFormat::Legacy, - FbFilterFormat::Composite => FilterFormat::Composite, - _ => unreachable!("unknown FilterFormat value: {:?}", value), - } - } -} - -impl From for Option { - fn from(value: CompressionFormat) -> Self { - match value { - #[cfg(feature = "snappy")] - CompressionFormat::Snappy => Some(CompressionCodec::Snappy), - #[cfg(feature = "lz4")] - CompressionFormat::Lz4 => Some(CompressionCodec::Lz4), - #[cfg(feature = "zlib")] - CompressionFormat::Zlib => Some(CompressionCodec::Zlib), - #[cfg(feature = "zstd")] - CompressionFormat::Zstd => Some(CompressionCodec::Zstd), - _ => None, - } - } +fn corrupt_sst_info(reason: &'static str) -> SlateDBError { + SlateDBError::CorruptSst { reason, path: None } } impl From for FbCompactionStatus { diff --git a/slatedb/src/format/block.rs b/slatedb/src/format/block.rs index 02b1c79f1..9647f39dc 100644 --- a/slatedb/src/format/block.rs +++ b/slatedb/src/format/block.rs @@ -3,7 +3,7 @@ use crate::format::row::{SstRowCodecV0, SstRowEntry}; use crate::format::sst::{CHECKSUM_SIZE, OFFSET_SIZE}; use crate::types::RowEntry; use crate::utils::clamp_allocated_size_bytes; -use bytes::{Buf, BufMut, Bytes, BytesMut}; +use bytes::{BufMut, Bytes, BytesMut}; pub(crate) const SIZEOF_U16: usize = std::mem::size_of::(); @@ -26,23 +26,74 @@ impl Block { } #[rustfmt::skip] - pub(crate) fn decode(bytes: Bytes) -> Self { - // Get number of elements in the block + pub(crate) fn decode(bytes: Bytes) -> Result { + Self::decode_with_memory_limit(bytes, usize::MAX) + } + + pub(crate) fn decode_with_memory_limit( + bytes: Bytes, + memory_limit: usize, + ) -> Result { let data = bytes.as_ref(); - let entry_offsets_len = (&data[data.len() - SIZEOF_U16..]).get_u16() as usize; - let data_end = data.len() - - SIZEOF_U16 // Entry u16 length - - entry_offsets_len * SIZEOF_U16; // Offset byte array length - let offsets_raw = &data[data_end..data.len() - SIZEOF_U16]; // Entry u16 - let offsets = offsets_raw - .chunks(SIZEOF_U16) - .map(|mut x| x.get_u16()) - .collect(); - let bytes = bytes.slice(0..data_end); - Self { - data: bytes, - offsets, + if data.len() < SIZEOF_U16 { + return Err(corrupt_block("block is missing its offset count")); + } + let entry_offsets_len = u16::from_be_bytes( + data[data.len() - SIZEOF_U16..] + .try_into() + .map_err(|_| corrupt_block("invalid block offset count"))?, + ) as usize; + if entry_offsets_len == 0 { + return Err(SlateDBError::EmptyBlock); + } + let offsets_size = entry_offsets_len + .checked_mul(SIZEOF_U16) + .ok_or_else(|| corrupt_block("block offset table size overflow"))?; + let required_memory = data + .len() + .checked_add(offsets_size) + .ok_or_else(|| corrupt_block("decoded block memory size overflow"))?; + if required_memory > memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "decoded WAL block", + required_bytes: required_memory, + limit_bytes: memory_limit, + }); } + let trailer_size = offsets_size + .checked_add(SIZEOF_U16) + .ok_or_else(|| corrupt_block("block trailer size overflow"))?; + let data_end = data + .len() + .checked_sub(trailer_size) + .ok_or_else(|| corrupt_block("block offset table exceeds block length"))?; + if data_end == 0 { + return Err(corrupt_block("block contains no row data")); + } + let offsets_raw = &data[data_end..data.len() - SIZEOF_U16]; + let mut offsets = Vec::with_capacity(entry_offsets_len); + for raw in offsets_raw.chunks_exact(SIZEOF_U16) { + let offset = u16::from_be_bytes( + raw.try_into() + .map_err(|_| corrupt_block("invalid block row offset"))?, + ); + if usize::from(offset) >= data_end { + return Err(corrupt_block("block row offset is outside row data")); + } + if offsets.last().is_some_and(|previous| *previous >= offset) { + return Err(corrupt_block( + "block row offsets are not strictly increasing", + )); + } + offsets.push(offset); + } + if offsets.first().copied() != Some(0) { + return Err(corrupt_block("first block row offset is not zero")); + } + Ok(Self { + data: bytes.slice(0..data_end), + offsets, + }) } pub(crate) fn clamp_allocated_size(&self) -> Self { @@ -73,6 +124,10 @@ impl Block { } } +fn corrupt_block(reason: &'static str) -> SlateDBError { + SlateDBError::CorruptSst { reason, path: None } +} + pub(super) struct BlockBuilderV1 { offsets: Vec, data: Vec, @@ -329,7 +384,7 @@ mod tests { fn test_block(#[case] test_case: BlockTestCase) { let block = build_block(&test_case); let encoded = block.encode(); - let decoded = Block::decode(encoded); + let decoded = Block::decode(encoded).unwrap(); let block_data = &block.data; let block_offsets = &block.offsets; // Decode the block data using offsets and validate each decoded entry @@ -389,7 +444,7 @@ mod tests { extended_data.put(encoded.as_ref()); extended_data.put_bytes(0u8, case.extra_bytes); let extended_data = extended_data.freeze(); - let block_extended = Block::decode(extended_data.slice(..encoded.len())); + let block_extended = Block::decode(extended_data.slice(..encoded.len())).unwrap(); let block_clamped = block_extended.clamp_allocated_size(); diff --git a/slatedb/src/format/block_v2.rs b/slatedb/src/format/block_v2.rs index 1d3c6851a..06de4d238 100644 --- a/slatedb/src/format/block_v2.rs +++ b/slatedb/src/format/block_v2.rs @@ -328,7 +328,7 @@ mod tests { // when: encoding and decoding let encoded = block.encode(); - let decoded = Block::decode(encoded); + let decoded = Block::decode(encoded).unwrap(); // then: blocks are equal assert_eq!(block.data, decoded.data); @@ -567,7 +567,7 @@ mod tests { // Verify the block can be encoded and decoded let encoded = block.encode(); - let decoded = Block::decode(encoded); + let decoded = Block::decode(encoded).unwrap(); assert_eq!(block.data, decoded.data); assert_eq!(block.offsets, decoded.offsets); diff --git a/slatedb/src/format/row.rs b/slatedb/src/format/row.rs index 10e690571..8bcd0a863 100644 --- a/slatedb/src/format/row.rs +++ b/slatedb/src/format/row.rs @@ -198,16 +198,22 @@ impl SstRowCodecV0 { } pub(crate) fn decode(&self, data: &mut Bytes) -> Result { + require_remaining(data, 4, "truncated V1 row key lengths")?; let key_prefix_len = data.get_u16() as usize; let key_suffix_len = data.get_u16() as usize; + require_remaining(data, key_suffix_len, "truncated V1 row key suffix")?; let key_suffix = data.slice(..key_suffix_len); data.advance(key_suffix_len); // decode seq & flags + require_remaining(data, 9, "truncated V1 row sequence or flags")?; let seq = data.get_u64(); let flags = self.decode_flags(data.get_u8())?; // decode expire_ts & create_ts + let timestamp_bytes = usize::from(flags.contains(RowFlags::HAS_EXPIRE_TS)) * 8 + + usize::from(flags.contains(RowFlags::HAS_CREATE_TS)) * 8; + require_remaining(data, timestamp_bytes, "truncated V1 row timestamp")?; let (expire_ts, create_ts) = if flags.contains(RowFlags::HAS_EXPIRE_TS | RowFlags::HAS_CREATE_TS) { (Some(data.get_i64()), Some(data.get_i64())) @@ -221,19 +227,25 @@ impl SstRowCodecV0 { // skip decoding value for tombstone. if flags.contains(RowFlags::TOMBSTONE) { - return Ok(SstRowEntry::new( + return Ok(SstRowEntry { key_prefix_len, key_suffix, seq, - ValueDeletable::Tombstone, create_ts, - None, - )); + // Tombstones historically discard expiry even if an old + // encoder wrote the flag and field. Preserve that V1 + // decoding contract while still consuming and validating the + // encoded timestamp above. + expire_ts: None, + value: ValueDeletable::Tombstone, + }); } // decode value + require_remaining(data, 4, "truncated V1 row value length")?; let value_len = data.get_u32() as usize; - let value = data.slice(..value_len); + require_remaining(data, value_len, "truncated V1 row value")?; + let value = data.split_to(value_len); Ok(SstRowEntry { key_prefix_len, key_suffix, @@ -266,6 +278,17 @@ impl SstRowCodecV0 { } } +fn require_remaining( + data: &Bytes, + required: usize, + reason: &'static str, +) -> Result<(), SlateDBError> { + if data.len() < required { + return Err(SlateDBError::CorruptSst { reason, path: None }); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/slatedb/src/format/row_codec_v2.rs b/slatedb/src/format/row_codec_v2.rs index f837f5a52..906039d51 100644 --- a/slatedb/src/format/row_codec_v2.rs +++ b/slatedb/src/format/row_codec_v2.rs @@ -2,7 +2,7 @@ use crate::error::SlateDBError; use crate::format::row::RowFlags; use crate::types::ValueDeletable; -use crate::utils::{decode_varint, encode_varint, varint_len}; +use crate::utils::{decode_varint_checked, encode_varint, varint_len}; use bytes::{Buf, BufMut, Bytes, BytesMut}; /// Intermediate representation for V2 row encoding. @@ -170,14 +170,16 @@ impl SstRowCodecV2 { /// Decode a V2 row entry from the data buffer. pub(crate) fn decode(&self, data: &mut impl Buf) -> Result { - let shared_bytes = decode_varint(data); - let unshared_bytes = decode_varint(data) as usize; - let value_len = decode_varint(data) as usize; + let shared_bytes = decode_varint_checked(data)?; + let unshared_bytes = decode_varint_checked(data)? as usize; + let value_len = decode_varint_checked(data)? as usize; // Read key_delta + require_remaining(data, unshared_bytes, "truncated V2 row key suffix")?; let key_suffix = data.copy_to_bytes(unshared_bytes); // Read value + require_remaining(data, value_len, "truncated V2 row value")?; let value_bytes = if value_len > 0 { Some(data.copy_to_bytes(value_len)) } else { @@ -185,10 +187,18 @@ impl SstRowCodecV2 { }; // Read seq & flags + require_remaining(data, 9, "truncated V2 row sequence or flags")?; let seq = data.get_u64(); let flags = self.decode_flags(data.get_u8())?; + if flags.contains(RowFlags::TOMBSTONE) && value_len != 0 { + return Err(corrupt_row("V2 tombstone contains value bytes")); + } + // Read timestamps + let timestamp_bytes = usize::from(flags.contains(RowFlags::HAS_EXPIRE_TS)) * 8 + + usize::from(flags.contains(RowFlags::HAS_CREATE_TS)) * 8; + require_remaining(data, timestamp_bytes, "truncated V2 row timestamp")?; let (expire_ts, create_ts) = if flags.contains(RowFlags::HAS_EXPIRE_TS | RowFlags::HAS_CREATE_TS) { (Some(data.get_i64()), Some(data.get_i64())) @@ -221,14 +231,18 @@ impl SstRowCodecV2 { /// Decode only the key portion for seek optimization. /// Returns (shared_bytes, key_suffix). - pub(crate) fn decode_key_only(&self, data: &mut impl Buf) -> (u32, Bytes) { - let shared_bytes = decode_varint(data); - let unshared_bytes = decode_varint(data) as usize; - let _value_len = decode_varint(data); - + pub(crate) fn decode_key_only( + &self, + data: &mut impl Buf, + ) -> Result<(u32, Bytes), SlateDBError> { + let shared_bytes = decode_varint_checked(data)?; + let unshared_bytes = decode_varint_checked(data)? as usize; + let _value_len = decode_varint_checked(data)?; + + require_remaining(data, unshared_bytes, "truncated V2 row key suffix")?; let key_suffix = data.copy_to_bytes(unshared_bytes); - (shared_bytes, key_suffix) + Ok((shared_bytes, key_suffix)) } fn decode_flags(&self, flags: u8) -> Result { @@ -249,6 +263,21 @@ impl SstRowCodecV2 { } } +fn require_remaining( + data: &impl Buf, + required: usize, + reason: &'static str, +) -> Result<(), SlateDBError> { + if data.remaining() < required { + return Err(corrupt_row(reason)); + } + Ok(()) +} + +fn corrupt_row(reason: &'static str) -> SlateDBError { + SlateDBError::CorruptSst { reason, path: None } +} + #[cfg(test)] mod tests { use super::*; @@ -343,7 +372,7 @@ mod tests { // when: decoding key only let mut slice = buf.as_slice(); - let (decoded_shared, decoded_suffix) = codec.decode_key_only(&mut slice); + let (decoded_shared, decoded_suffix) = codec.decode_key_only(&mut slice).unwrap(); // then: key info is correct assert_eq!(decoded_shared, shared_bytes); diff --git a/slatedb/src/format/sst.rs b/slatedb/src/format/sst.rs index 069f34dfc..0589f3580 100644 --- a/slatedb/src/format/sst.rs +++ b/slatedb/src/format/sst.rs @@ -16,7 +16,7 @@ use flatbuffers::DefaultAllocator; use futures::future::try_join_all; use log::warn; use std::collections::VecDeque; -#[cfg(feature = "zlib")] +#[cfg(any(feature = "zlib", feature = "zstd"))] use std::io::Read; #[cfg(feature = "zlib")] use std::io::Write; @@ -35,6 +35,78 @@ fn is_supported_version(version: u16) -> bool { matches!(version, SST_FORMAT_VERSION | SST_FORMAT_VERSION_V2) } +fn validate_decoded_size(required_bytes: usize, limit_bytes: usize) -> Result<(), SlateDBError> { + if required_bytes > limit_bytes { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "decoded WAL section", + required_bytes, + limit_bytes, + }); + } + Ok(()) +} + +#[cfg(any(feature = "zlib", feature = "zstd"))] +fn measure_bounded_decompressed( + reader: &mut impl Read, + max_output_bytes: usize, +) -> Result { + debug_assert_ne!(max_output_bytes, usize::MAX); + let mut decoded_len = 0usize; + let mut chunk = [0u8; 8 * 1024]; + while decoded_len < max_output_bytes { + let remaining = max_output_bytes - decoded_len; + let chunk_len = remaining.min(chunk.len()); + let read = reader + .read(&mut chunk[..chunk_len]) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + if read == 0 { + return Ok(decoded_len); + } + decoded_len = decoded_len + .checked_add(read) + .ok_or(SlateDBError::BlockDecompressionError)?; + } + + let mut overflow = [0u8; 1]; + if reader + .read(&mut overflow) + .map_err(|_| SlateDBError::BlockDecompressionError)? + != 0 + { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "decoded WAL section", + required_bytes: max_output_bytes.saturating_add(1), + limit_bytes: max_output_bytes, + }); + } + Ok(decoded_len) +} + +#[cfg(any(feature = "zlib", feature = "zstd"))] +fn read_exact_decompressed( + reader: &mut impl Read, + decoded_len: usize, +) -> Result { + // The first pass measured the exact output size using only a fixed stack + // buffer. Allocate exactly that much for the retained decoded block; this + // avoids either reserving the entire replay working budget for a tiny + // block or allowing Vec growth to exceed the limit. + let mut decompressed = vec![0u8; decoded_len]; + reader + .read_exact(&mut decompressed) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + let mut overflow = [0u8; 1]; + if reader + .read(&mut overflow) + .map_err(|_| SlateDBError::BlockDecompressionError)? + != 0 + { + return Err(SlateDBError::BlockDecompressionError); + } + Ok(Bytes::from(decompressed)) +} + #[allow(private_interfaces)] pub(crate) enum BlockBuilder { V1(BlockBuilderV1), @@ -178,6 +250,10 @@ pub(crate) const VERSION_SIZE: usize = SIZEOF_U16; /// /// #[async_trait] /// impl BlockTransformer for XorTransformer { +/// fn max_decoded_len(&self, encoded_len: usize) -> Option { +/// Some(encoded_len) +/// } +/// /// async fn encode(&self, data: Bytes) -> Result { /// let transformed: Vec = data.iter().map(|b| b ^ self.key).collect(); /// Ok(Bytes::from(transformed)) @@ -190,6 +266,15 @@ pub(crate) const VERSION_SIZE: usize = SIZEOF_U16; /// ``` #[async_trait] pub trait BlockTransformer: Send + Sync { + /// Returns a strict upper bound for the decoded allocation produced from an + /// encoded buffer of `encoded_len` bytes. + /// + /// Full-object WAL replay checks this bound before calling [`Self::decode`]. + /// Return `None` when the transform cannot provide a safe bound; WAL replay + /// will then fail closed instead of invoking it. Implementations must never + /// allocate or return more decoded bytes than this bound. + fn max_decoded_len(&self, encoded_len: usize) -> Option; + /// Encode (transform) block data before storage. async fn encode(&self, data: Bytes) -> Result; @@ -613,6 +698,11 @@ pub(crate) type LengthOffsetAndVersion = (u64, u64, u16); pub(crate) type TableInfoAndVersion = (SsTableInfo, u16); +pub(crate) enum StagedSstInfoError { + Initial(SlateDBError), + Later(SlateDBError), +} + #[derive(Clone)] pub(crate) struct SsTableFormat { pub(crate) block_size: usize, @@ -650,10 +740,21 @@ impl SsTableFormat { let header = obj .read_range((obj_len - NUM_FOOTER_BYTES_LONG)..obj_len) .await?; - assert_eq!(header.len(), NUM_FOOTER_BYTES); + if header.len() != NUM_FOOTER_BYTES { + return Err(SlateDBError::CorruptSst { + reason: "invalid SST footer length", + path: None, + }); + } let version = header.slice(8..NUM_FOOTER_BYTES).get_u16(); let sst_metadata_offset = header.slice(0..8).get_u64(); + if sst_metadata_offset > obj_len - NUM_FOOTER_BYTES_LONG { + return Err(SlateDBError::CorruptSst { + reason: "SST metadata offset exceeds footer start", + path: None, + }); + } Ok((obj_len, sst_metadata_offset, version)) } @@ -682,6 +783,100 @@ impl SsTableFormat { SsTableInfo::decode(sst_metadata_bytes, &*self.sst_codec).map(|info| (info, version)) } + pub(crate) async fn read_info_and_version_staged( + &self, + obj: &impl ReadOnlyBlob, + ) -> Result { + let (obj_len, sst_metadata_offset, version) = self + .read_length_and_metadata_offset_and_version(obj) + .await + .map_err(StagedSstInfoError::Initial)?; + self.validate_version(version) + .map_err(StagedSstInfoError::Later)?; + let sst_metadata_bytes = obj + .read_range(sst_metadata_offset..obj_len - NUM_FOOTER_BYTES_LONG) + .await + .map_err(StagedSstInfoError::Later)?; + SsTableInfo::decode(sst_metadata_bytes, &*self.sst_codec) + .map(|info| (info, version)) + .map_err(StagedSstInfoError::Later) + } + + pub(crate) async fn read_wal_info_and_version_bounded( + &self, + obj: &impl ReadOnlyBlob, + object_len: u64, + metadata_offset: u64, + version: u16, + working_memory_limit: usize, + ) -> Result { + self.validate_version(version)?; + let metadata_end = object_len + .checked_sub(NUM_FOOTER_BYTES_LONG) + .ok_or(SlateDBError::EmptySSTable)?; + let encoded_len = usize::try_from(metadata_end.checked_sub(metadata_offset).ok_or( + SlateDBError::CorruptSst { + reason: "WAL metadata range is outside the object", + path: None, + }, + )?) + .map_err(|_| SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + })?; + // The FlatBuffer decoder copies first/last entries. Their combined + // bytes cannot exceed the encoded metadata that contains them. + let conservative_bytes = + encoded_len + .checked_mul(2) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + })?; + if conservative_bytes > working_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: conservative_bytes, + limit_bytes: working_memory_limit, + }); + } + let metadata = obj.read_range(metadata_offset..metadata_end).await?; + if metadata.len() != encoded_len { + return Err(SlateDBError::CorruptSst { + reason: "WAL metadata range returned an unexpected length", + path: None, + }); + } + let info = SsTableInfo::decode(metadata, &*self.sst_codec)?; + let retained_bytes = info + .first_entry + .as_ref() + .map_or(0, Bytes::len) + .checked_add(info.last_entry.as_ref().map_or(0, Bytes::len)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + })?; + let actual_bytes = encoded_len.checked_add(retained_bytes).ok_or( + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + }, + )?; + if actual_bytes > working_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: actual_bytes, + limit_bytes: working_memory_limit, + }); + } + Ok((info, version)) + } + pub(crate) async fn read_filters( &self, info: &SsTableInfo, @@ -804,6 +999,62 @@ impl SsTableFormat { self.decode_index(index_bytes, compression_codec).await } + pub(crate) async fn read_wal_index_bounded( + &self, + info: &SsTableInfo, + obj: &impl ReadOnlyBlob, + working_memory_limit: usize, + ) -> Result { + let index_end = + info.index_offset + .checked_add(info.index_len) + .ok_or(SlateDBError::CorruptSst { + reason: "WAL index range overflow", + path: None, + })?; + let index_bytes = obj.read_range(info.index_offset..index_end).await?; + self.decode_index_bounded(index_bytes, info.compression_codec, working_memory_limit) + .await + } + + pub(crate) async fn read_ranged_wal_index_bounded( + &self, + info: &SsTableInfo, + obj: &impl ReadOnlyBlob, + working_memory_limit: usize, + ) -> Result { + let index_end = + info.index_offset + .checked_add(info.index_len) + .ok_or(SlateDBError::CorruptSst { + reason: "WAL index range overflow", + path: None, + })?; + let encoded_len = usize::try_from(info.index_len).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + } + })?; + let decode_memory_limit = working_memory_limit.checked_sub(encoded_len).ok_or( + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + required_bytes: encoded_len, + limit_bytes: working_memory_limit, + }, + )?; + let index_bytes = obj.read_range(info.index_offset..index_end).await?; + if index_bytes.len() != encoded_len { + return Err(SlateDBError::CorruptSst { + reason: "WAL index range returned an unexpected length", + path: None, + }); + } + self.decode_index_bounded(index_bytes, info.compression_codec, decode_memory_limit) + .await + } + #[cfg(test)] pub(crate) async fn read_index_raw( &self, @@ -839,6 +1090,50 @@ impl SsTableFormat { Ok(SsTableIndexOwned::new(decompressed_bytes)?) } + async fn decode_index_bounded( + &self, + bytes: Bytes, + compression_codec: Option, + working_memory_limit: usize, + ) -> Result { + let index_bytes = self.validate_checksum(bytes)?; + let untransformed_bytes = match &self.block_transformer { + Some(transformer) => { + let decoded_bound = transformer + .max_decoded_len(index_bytes.len()) + .ok_or(SlateDBError::BlockTransformError)?; + validate_decoded_size(decoded_bound, working_memory_limit)?; + let transformed = transformer + .decode(index_bytes) + .await + .map_err(|_| SlateDBError::BlockTransformError)?; + if transformed.len() > decoded_bound { + return Err(SlateDBError::BlockTransformError); + } + transformed + } + None => index_bytes, + }; + let transformed_allocation = if self.block_transformer.is_some() { + untransformed_bytes.len() + } else { + 0 + }; + let remaining = working_memory_limit + .checked_sub(transformed_allocation) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "transformed WAL index", + required_bytes: transformed_allocation, + limit_bytes: working_memory_limit, + })?; + let decompressed = match compression_codec { + Some(codec) => Self::decompress_bounded(untransformed_bytes, codec, remaining)?, + None => untransformed_bytes, + }; + validate_decoded_size(decompressed.len(), remaining)?; + Ok(SsTableIndexOwned::new(decompressed)?) + } + pub(crate) async fn read_stats( &self, info: &SsTableInfo, @@ -880,34 +1175,74 @@ impl SsTableFormat { fn decompress( #[allow(unused_variables)] compressed_data: Bytes, compression_option: CompressionCodec, + ) -> Result { + Self::decompress_bounded(compressed_data, compression_option, usize::MAX) + } + + fn decompress_bounded( + #[allow(unused_variables)] compressed_data: Bytes, + compression_option: CompressionCodec, + _max_output_bytes: usize, ) -> Result { match compression_option { #[cfg(feature = "snappy")] - CompressionCodec::Snappy => Ok(Bytes::from( - snap::raw::Decoder::new() - .decompress_vec(&compressed_data) - .map_err(|_| SlateDBError::BlockDecompressionError)?, - )), + CompressionCodec::Snappy => { + let decoded_len = snap::raw::decompress_len(&compressed_data) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + validate_decoded_size(decoded_len, _max_output_bytes)?; + Ok(Bytes::from( + snap::raw::Decoder::new() + .decompress_vec(&compressed_data) + .map_err(|_| SlateDBError::BlockDecompressionError)?, + )) + } #[cfg(feature = "zlib")] CompressionCodec::Zlib => { + if _max_output_bytes == usize::MAX { + let mut decoder = flate2::read::ZlibDecoder::new(&compressed_data[..]); + let mut decompressed = Vec::new(); + decoder + .read_to_end(&mut decompressed) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + return Ok(Bytes::from(decompressed)); + } let mut decoder = flate2::read::ZlibDecoder::new(&compressed_data[..]); - let mut decompressed = Vec::new(); - decoder - .read_to_end(&mut decompressed) - .map_err(|_| SlateDBError::BlockDecompressionError)?; - Ok(Bytes::from(decompressed)) + let decoded_len = measure_bounded_decompressed(&mut decoder, _max_output_bytes)?; + let mut decoder = flate2::read::ZlibDecoder::new(&compressed_data[..]); + read_exact_decompressed(&mut decoder, decoded_len) } #[cfg(feature = "lz4")] CompressionCodec::Lz4 => { + if compressed_data.len() < size_of::() { + return Err(SlateDBError::BlockDecompressionError); + } + let decoded_len = u32::from_le_bytes( + compressed_data[..size_of::()] + .try_into() + .map_err(|_| SlateDBError::BlockDecompressionError)?, + ) as usize; + validate_decoded_size(decoded_len, _max_output_bytes)?; let decompressed = lz4_flex::block::decompress_size_prepended(&compressed_data) .map_err(|_| SlateDBError::BlockDecompressionError)?; Ok(Bytes::from(decompressed)) } #[cfg(feature = "zstd")] CompressionCodec::Zstd => { - let decompressed = zstd::stream::decode_all(&compressed_data[..]) + if _max_output_bytes == usize::MAX { + let mut decoder = zstd::stream::read::Decoder::new(&compressed_data[..]) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + let mut decompressed = Vec::new(); + decoder + .read_to_end(&mut decompressed) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + return Ok(Bytes::from(decompressed)); + } + let mut decoder = zstd::stream::read::Decoder::new(&compressed_data[..]) .map_err(|_| SlateDBError::BlockDecompressionError)?; - Ok(Bytes::from(decompressed)) + let decoded_len = measure_bounded_decompressed(&mut decoder, _max_output_bytes)?; + let mut decoder = zstd::stream::read::Decoder::new(&compressed_data[..]) + .map_err(|_| SlateDBError::BlockDecompressionError)?; + read_exact_decompressed(&mut decoder, decoded_len) } } } @@ -987,7 +1322,93 @@ impl SsTableFormat { None => untransformed_bytes, }; - Ok(Block::decode(decompressed_bytes)) + Block::decode(decompressed_bytes) + } + + pub(crate) async fn decode_wal_block_bounded( + &self, + bytes: Bytes, + compression_codec: Option, + working_memory_limit: usize, + ) -> Result { + let block_bytes = self.validate_checksum(bytes)?; + let untransformed_bytes = match &self.block_transformer { + Some(transformer) => { + let decoded_bound = transformer + .max_decoded_len(block_bytes.len()) + .ok_or(SlateDBError::BlockTransformError)?; + validate_decoded_size(decoded_bound, working_memory_limit)?; + let transformed = transformer + .decode(block_bytes) + .await + .map_err(|_| SlateDBError::BlockTransformError)?; + if transformed.len() > decoded_bound { + return Err(SlateDBError::BlockTransformError); + } + transformed + } + None => block_bytes, + }; + let transformed_allocation = if self.block_transformer.is_some() { + untransformed_bytes.len() + } else { + 0 + }; + let remaining = working_memory_limit + .checked_sub(transformed_allocation) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "transformed WAL block", + required_bytes: transformed_allocation, + limit_bytes: working_memory_limit, + })?; + let decoded = match compression_codec { + Some(codec) => Self::decompress_bounded(untransformed_bytes, codec, remaining)?, + None => untransformed_bytes, + }; + Block::decode_with_memory_limit(decoded, remaining) + } + + pub(crate) async fn decode_wal_block_from_object( + &self, + info: &SsTableInfo, + index_owned: &SsTableIndexOwned, + block_index: usize, + object_bytes: &Bytes, + working_memory_limit: usize, + ) -> Result { + let index = index_owned.borrow(); + if block_index >= index.block_meta().len() { + return Err(SlateDBError::CorruptSst { + reason: "WAL block index is out of range", + path: None, + }); + } + let start = index.block_meta().get(block_index).offset(); + let end = if block_index + 1 < index.block_meta().len() { + index.block_meta().get(block_index + 1).offset() + } else { + info.filter_offset + }; + let start = usize::try_from(start).map_err(|_| SlateDBError::CorruptSst { + reason: "WAL block start does not fit in memory", + path: None, + })?; + let end = usize::try_from(end).map_err(|_| SlateDBError::CorruptSst { + reason: "WAL block end does not fit in memory", + path: None, + })?; + if start >= end || end > object_bytes.len() { + return Err(SlateDBError::CorruptSst { + reason: "WAL block range is outside the object", + path: None, + }); + } + self.decode_wal_block_bounded( + object_bytes.slice(start..end), + info.compression_codec, + working_memory_limit, + ) + .await } pub(crate) async fn read_block( @@ -1019,6 +1440,9 @@ impl SsTableFormat { /// validate checksum and return the actual data bytes pub(crate) fn validate_checksum(&self, bytes: Bytes) -> Result { + if bytes.len() < CHECKSUM_SIZE { + return Err(SlateDBError::ChecksumMismatch { path: None }); + } let data_bytes = bytes.slice(..bytes.len() - CHECKSUM_SIZE); let mut checksum_bytes = bytes.slice(bytes.len() - CHECKSUM_SIZE..); let checksum = crc32fast::hash(&data_bytes); diff --git a/slatedb/src/lib.rs b/slatedb/src/lib.rs index 38cae17c2..a7767bba1 100644 --- a/slatedb/src/lib.rs +++ b/slatedb/src/lib.rs @@ -160,6 +160,7 @@ mod peeking_iterator; #[cfg(test)] mod proptest_util; mod reader; +mod replay_task_scope; mod retention_iterator; mod retrying_object_store; mod segment_iterator; diff --git a/slatedb/src/manifest/store.rs b/slatedb/src/manifest/store.rs index 016b97f37..2128df67c 100644 --- a/slatedb/src/manifest/store.rs +++ b/slatedb/src/manifest/store.rs @@ -50,6 +50,25 @@ impl FenceableManifest { Ok(Self { inner: fr, clock }) } + pub(crate) async fn init_writer_with_epoch( + stored_manifest: StoredManifest, + manifest_update_timeout: Duration, + system_clock: Arc, + writer_epoch: u64, + ) -> Result { + let clock = system_clock.clone(); + let fr = FenceableTransactionalObject::init_with_epoch( + stored_manifest.inner, + manifest_update_timeout, + system_clock, + writer_epoch, + |m: &Manifest| m.writer_epoch, + |m: &mut Manifest, epoch: u64| m.writer_epoch = epoch, + ) + .await?; + Ok(Self { inner: fr, clock }) + } + pub(crate) async fn init_compactor( stored_manifest: StoredManifest, manifest_update_timeout: Duration, @@ -814,6 +833,48 @@ mod tests { } } + #[tokio::test] + async fn exact_writer_epoch_rejects_equal_and_stale_claims_without_advancing() { + let store = new_memory_manifest_store(); + StoredManifest::create_new_db( + store.clone(), + ManifestCore::new(), + Arc::new(DefaultSystemClock::new()), + ) + .await + .unwrap(); + let timeout = Duration::from_secs(300); + + let stored = StoredManifest::load(store.clone(), Arc::new(DefaultSystemClock::new())) + .await + .unwrap(); + let claimed = FenceableManifest::init_writer_with_epoch( + stored, + timeout, + Arc::new(DefaultSystemClock::new()), + 12, + ) + .await + .unwrap(); + assert_eq!(claimed.local_epoch(), 12); + + for stale_epoch in [11, 12] { + let stored = StoredManifest::load(store.clone(), Arc::new(DefaultSystemClock::new())) + .await + .unwrap(); + let result = FenceableManifest::init_writer_with_epoch( + stored, + timeout, + Arc::new(DefaultSystemClock::new()), + stale_epoch, + ) + .await; + assert!(matches!(result, Err(SlateDBError::Fenced))); + let latest = store.read_latest_manifest().await.unwrap(); + assert_eq!(latest.manifest.writer_epoch, 12); + } + } + #[tokio::test] async fn test_should_fail_refresh_on_writer_fenced() { let ms = new_memory_manifest_store(); diff --git a/slatedb/src/replay_task_scope.rs b/slatedb/src/replay_task_scope.rs new file mode 100644 index 000000000..f80980b54 --- /dev/null +++ b/slatedb/src/replay_task_scope.rs @@ -0,0 +1,89 @@ +use std::fmt; +use std::future::Future; + +use tokio::task::JoinHandle; +use tokio_util::{sync::CancellationToken, task::TaskTracker}; + +/// Owns every spawned task for one reader replay operation. +/// +/// Cancellation stops the operation. `shutdown` additionally waits until all +/// tracked descendants have exited, which is the boundary required before a +/// reader may report that its database runtime is closed. +#[derive(Clone)] +pub(crate) struct ReplayTaskScope { + cancellation: CancellationToken, + tasks: TaskTracker, +} + +impl ReplayTaskScope { + pub(crate) fn new() -> Self { + Self { + cancellation: CancellationToken::new(), + tasks: TaskTracker::new(), + } + } + + pub(crate) fn child(&self) -> Self { + Self { + cancellation: self.cancellation.child_token(), + tasks: TaskTracker::new(), + } + } + + pub(crate) fn cancel(&self) { + self.cancellation.cancel(); + } + + pub(crate) async fn cancelled(&self) { + self.cancellation.cancelled().await; + } + + pub(crate) fn spawn(&self, task: F) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.tasks.spawn(task) + } + + pub(crate) async fn shutdown(&self) { + self.cancel(); + self.tasks.close(); + self.tasks.wait().await; + } + + #[cfg(test)] + pub(crate) fn active_tasks(&self) -> usize { + self.tasks.len() + } +} + +impl fmt::Debug for ReplayTaskScope { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ReplayTaskScope") + .field("cancelled", &self.cancellation.is_cancelled()) + .field("active_tasks", &self.tasks.len()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::ReplayTaskScope; + + #[tokio::test] + async fn shutdown_waits_until_tracked_tasks_exit() { + let scope = ReplayTaskScope::new(); + let cancellation = scope.clone(); + let task = scope.spawn(async move { + cancellation.cancelled().await; + }); + assert_eq!(scope.active_tasks(), 1); + + scope.shutdown().await; + + task.await.unwrap(); + assert_eq!(scope.active_tasks(), 0); + } +} diff --git a/slatedb/src/sst_builder.rs b/slatedb/src/sst_builder.rs index edc5853e8..e7574cdc0 100644 --- a/slatedb/src/sst_builder.rs +++ b/slatedb/src/sst_builder.rs @@ -1321,6 +1321,10 @@ mod tests { #[async_trait] impl BlockTransformer for XorTransformer { + fn max_decoded_len(&self, encoded_len: usize) -> Option { + Some(encoded_len) + } + async fn encode(&self, data: Bytes) -> Result { let transformed: Vec = data.iter().map(|b| b ^ self.key).collect(); Ok(Bytes::from(transformed)) diff --git a/slatedb/src/sst_iter.rs b/slatedb/src/sst_iter.rs index ff3855d20..ca60761e8 100644 --- a/slatedb/src/sst_iter.rs +++ b/slatedb/src/sst_iter.rs @@ -18,6 +18,7 @@ use crate::filter_policy::{FilterContext, FilterQuery, NamedFilter}; use crate::flatbuffer_types::SsTableIndexOwned; use crate::format::block::Block; use crate::prefix_extractor::PrefixTarget; +use crate::replay_task_scope::ReplayTaskScope; use crate::{ iter::{IterationOrder, RowEntryIterator}, partitioned_keyspace, @@ -31,6 +32,14 @@ enum FetchTask { Finished(VecDeque>), } +impl Drop for FetchTask { + fn drop(&mut self) { + if let Self::InFlight(handle) = self { + handle.abort(); + } + } +} + #[derive(Clone, Debug)] pub(crate) struct SstIteratorOptions { pub(crate) max_fetch_tasks: usize, @@ -279,6 +288,7 @@ pub(crate) struct InternalSstIterator<'a> { fetch_tasks: VecDeque, table_store: Arc, options: SstIteratorOptions, + replay_tasks: Option, /// Buffer for descending iteration to maintain correct sequence order within keys. descending_buffer: Option>, /// Pending entry that was read ahead but belongs to the next key group. @@ -291,6 +301,15 @@ impl<'a> InternalSstIterator<'a> { view: SstView<'a>, table_store: Arc, options: SstIteratorOptions, + ) -> Result { + Self::new_scoped(view, table_store, options, None) + } + + fn new_scoped( + view: SstView<'a>, + table_store: Arc, + options: SstIteratorOptions, + replay_tasks: Option, ) -> Result { assert!(options.max_fetch_tasks > 0); assert!(options.blocks_to_fetch > 0); @@ -309,6 +328,7 @@ impl<'a> InternalSstIterator<'a> { fetch_tasks: VecDeque::new(), table_store, options, + replay_tasks, descending_buffer, pending_entry: None, }) @@ -335,6 +355,20 @@ impl<'a> InternalSstIterator<'a> { Self::new(view, table_store, options).map(Some) } + fn new_owned_scoped>( + range: T, + table: SsTableView, + table_store: Arc, + options: SstIteratorOptions, + replay_tasks: ReplayTaskScope, + ) -> Result, SlateDBError> { + let Some(view_range) = table.calculate_view_range(BytesRange::from(range)) else { + return Ok(None); + }; + let view = SstView::Owned(Box::new(table), view_range); + Self::new_scoped(view, table_store, options, Some(replay_tasks)).map(Some) + } + fn new_borrowed>( range: T, table: &'a SsTableView, @@ -391,17 +425,22 @@ impl<'a> InternalSstIterator<'a> { let blocks_end = self.next_block_idx_to_fetch + blocks_to_fetch; let index = index.clone(); let cache_blocks = self.options.cache_blocks; - self.fetch_tasks - .push_back(FetchTask::InFlight(tokio::spawn(async move { - table_store - .read_blocks_using_index( - &table, - index, - blocks_start..blocks_end, - cache_blocks, - ) - .await - }))); + let task = async move { + table_store + .read_blocks_using_index( + &table, + index, + blocks_start..blocks_end, + cache_blocks, + ) + .await + }; + let task = if let Some(scope) = self.replay_tasks.as_ref() { + scope.spawn(task) + } else { + tokio::spawn(task) + }; + self.fetch_tasks.push_back(FetchTask::InFlight(task)); self.next_block_idx_to_fetch = blocks_end; } } @@ -420,17 +459,22 @@ impl<'a> InternalSstIterator<'a> { let blocks_start = blocks_end - blocks_to_fetch; let index = index.clone(); let cache_blocks = self.options.cache_blocks; - self.fetch_tasks - .push_back(FetchTask::InFlight(tokio::spawn(async move { - table_store - .read_blocks_using_index( - &table, - index, - blocks_start..blocks_end, - cache_blocks, - ) - .await - }))); + let task = async move { + table_store + .read_blocks_using_index( + &table, + index, + blocks_start..blocks_end, + cache_blocks, + ) + .await + }; + let task = if let Some(scope) = self.replay_tasks.as_ref() { + scope.spawn(task) + } else { + tokio::spawn(task) + }; + self.fetch_tasks.push_back(FetchTask::InFlight(task)); self.next_block_idx_to_fetch = blocks_start; } } @@ -916,6 +960,40 @@ impl<'a> SstIterator<'a> { Self::new_owned_initialized_with_stats(range, table, table_store, options, None).await } + pub(crate) async fn new_owned_initialized_scoped>( + range: T, + table: SsTableView, + table_store: Arc, + options: SstIteratorOptions, + replay_tasks: ReplayTaskScope, + ) -> Result, SlateDBError> { + let internal = InternalSstIterator::new_owned_scoped( + range, + table, + table_store, + options, + replay_tasks, + )?; + match internal { + Some(inner) => { + let mut iterator = Self::from_internal(inner, None); + match &mut iterator.delegate { + SstIteratorDelegate::Filter(filter_iter) => { + filter_iter.init().await?; + if filter_iter.is_filtered_out() { + return Ok(None); + } + } + SstIteratorDelegate::Direct(inner_iter) => { + inner_iter.init().await?; + } + } + Ok(Some(iterator)) + } + None => Ok(None), + } + } + pub(crate) async fn new_owned_initialized_with_stats>( range: T, table: SsTableView, @@ -1104,8 +1182,41 @@ mod tests { use slatedb_common::metrics::{ lookup_metric_with_labels, DefaultMetricsRecorder, MetricLevel, MetricsRecorderHelper, }; + use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; + struct FetchCancellationProbe(Arc); + + impl Drop for FetchCancellationProbe { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + + #[tokio::test] + async fn dropping_an_in_flight_fetch_aborts_and_joins_its_scoped_task() { + let scope = ReplayTaskScope::new(); + let cancelled = Arc::new(AtomicBool::new(false)); + let probe = Arc::clone(&cancelled); + let handle = scope.spawn(async move { + let _probe = FetchCancellationProbe(probe); + std::future::pending::>, SlateDBError>>().await + }); + tokio::task::yield_now().await; + assert_eq!(scope.active_tasks(), 1); + + drop(FetchTask::InFlight(handle)); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !cancelled.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropped block fetch remained detached"); + scope.shutdown().await; + assert_eq!(scope.active_tasks(), 0); + } + #[tokio::test] async fn test_one_block_sst_iter() { test_one_block_sst_iter_with_order(IterationOrder::Ascending).await; diff --git a/slatedb/src/sst_reader.rs b/slatedb/src/sst_reader.rs index 6d852fe1b..e818c1f73 100644 --- a/slatedb/src/sst_reader.rs +++ b/slatedb/src/sst_reader.rs @@ -525,6 +525,10 @@ mod tests { #[async_trait::async_trait] impl BlockTransformer for XorTransformer { + fn max_decoded_len(&self, encoded_len: usize) -> Option { + Some(encoded_len) + } + async fn encode(&self, data: Bytes) -> Result { let transformed: Vec = data.iter().map(|b| b ^ self.key).collect(); Ok(Bytes::from(transformed)) diff --git a/slatedb/src/tablestore.rs b/slatedb/src/tablestore.rs index 7a0717264..4e5b007be 100644 --- a/slatedb/src/tablestore.rs +++ b/slatedb/src/tablestore.rs @@ -16,16 +16,19 @@ use slatedb_common::ObjectMetadata; use tokio::io::AsyncWriteExt; use ulid::Ulid; -use crate::blob::ReadOnlyBlob; +use crate::blob::{BytesBlob, ReadOnlyBlob}; use crate::block_cache_policy::{should_cache_data_block, BlockCachePolicy}; use crate::db_cache::CacheTarget; use crate::db_cache::{CacheLoader, CachedEntry, CachedKey, DbCache, EncodedCachedFilter}; -use crate::db_state::{SsTableHandle, SsTableId, SstType}; +use crate::db_state::{SsTableHandle, SsTableId, SsTableInfo, SstType}; use crate::error::SlateDBError; use crate::filter_policy::NamedFilter; use crate::flatbuffer_types::SsTableIndexOwned; use crate::format::block::Block; -use crate::format::sst::{EncodedSsTable, EncodedSsTableBlock, SsTableFormat}; +use crate::format::sst::{ + EncodedSsTable, EncodedSsTableBlock, SsTableFormat, StagedSstInfoError, CHECKSUM_SIZE, + METADATA_OFFSET_SIZE, VERSION_SIZE, +}; use crate::object_store_tag::ObjectStoreCallTag; pub(crate) use crate::object_store_tag::TableStoreKind; use crate::object_stores::{ObjectStoreType, ObjectStores}; @@ -49,6 +52,38 @@ pub(crate) struct TableStore { kind: TableStoreKind, } +pub(crate) enum DecodedWalSst { + Fence, + Data(Box), +} + +pub(crate) struct DecodedWalSstData { + pub(crate) wal_id: u64, + pub(crate) format_version: u16, + pub(crate) object_bytes: Bytes, + pub(crate) info: SsTableInfo, + pub(crate) index: SsTableIndexOwned, + pub(crate) retained_decode_bytes: usize, +} + +pub(crate) enum RangedWalSst { + Fence, + Data(Box), +} + +pub(crate) struct RangedWalSstData { + pub(crate) wal_id: u64, + pub(crate) format_version: u16, + pub(crate) info: SsTableInfo, + pub(crate) index: SsTableIndexOwned, + pub(crate) retained_decode_bytes: usize, +} + +pub(crate) enum RuntimeWalOpenError { + MissingInitialObject(SlateDBError), + Replay(SlateDBError), +} + struct ReadOnlyObject { object_store: Arc, path: Path, @@ -167,6 +202,129 @@ impl TableStore { bytes.div_ceil(self.sst_format.block_size) } + pub(crate) fn validate_wal_sst_replay_memory( + &self, + encoded_sst: &EncodedSsTable, + metadata_memory_limit: usize, + block_memory_limit: usize, + ) -> Result<(), SlateDBError> { + let index_encoded_bytes = usize::try_from(encoded_sst.info.index_len).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + } + })?; + let metadata_encoded_bytes = encoded_sst + .footer + .len() + .checked_sub(METADATA_OFFSET_SIZE + VERSION_SIZE) + .and_then(|footer_bytes| footer_bytes.checked_sub(index_encoded_bytes)) + .ok_or(SlateDBError::InvalidDBState)?; + let retained_info_bytes = encoded_sst + .info + .first_entry + .as_ref() + .map_or(0, Bytes::len) + .checked_add(encoded_sst.info.last_entry.as_ref().map_or(0, Bytes::len)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + let metadata_required_bytes = metadata_encoded_bytes + .checked_add(retained_info_bytes) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + if metadata_required_bytes > metadata_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + required_bytes: metadata_required_bytes, + limit_bytes: metadata_memory_limit, + }); + } + + let index_payload_bytes = index_encoded_bytes + .checked_sub(CHECKSUM_SIZE) + .ok_or(SlateDBError::InvalidDBState)?; + let transformed_index_bytes = match &self.sst_format.block_transformer { + Some(transformer) => transformer.max_decoded_len(index_payload_bytes).ok_or( + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "transformed WAL index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + }, + )?, + None => 0, + }; + let index_required_bytes = retained_info_bytes + .checked_add(index_encoded_bytes) + .and_then(|bytes| bytes.checked_add(transformed_index_bytes)) + .and_then(|bytes| bytes.checked_add(encoded_sst.index.size())) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + if index_required_bytes > metadata_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + required_bytes: index_required_bytes, + limit_bytes: metadata_memory_limit, + }); + } + + for block in &encoded_sst.unconsumed_blocks { + let encoded_payload_bytes = block + .encoded_bytes + .len() + .checked_sub(CHECKSUM_SIZE) + .ok_or(SlateDBError::InvalidDBState)?; + let transformed_block_bytes = match &self.sst_format.block_transformer { + Some(transformer) => transformer.max_decoded_len(encoded_payload_bytes).ok_or( + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "transformed WAL block", + required_bytes: usize::MAX, + limit_bytes: block_memory_limit, + }, + )?, + None => 0, + }; + let offsets_bytes = block + .block + .offsets + .len() + .checked_mul(std::mem::size_of::()) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL block", + required_bytes: usize::MAX, + limit_bytes: block_memory_limit, + })?; + let required_bytes = block + .encoded_bytes + .len() + .checked_add(transformed_block_bytes) + .and_then(|bytes| bytes.checked_add(block.block.size())) + .and_then(|bytes| bytes.checked_add(offsets_bytes)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL block", + required_bytes: usize::MAX, + limit_bytes: block_memory_limit, + })?; + if required_bytes > block_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL block", + required_bytes, + limit_bytes: block_memory_limit, + }); + } + } + Ok(()) + } + /// Find the highest WAL SST id present in the object store at or above /// `start_after + 1`, returning `start_after` if none exist. /// @@ -320,6 +478,421 @@ impl TableStore { Ok(wal_list) } + /// Lists only the WAL suffix needed for one replay operation. + pub(crate) async fn list_wal_ssts_for_replay( + &self, + id_range: Range, + ) -> Result>, SlateDBError> { + if id_range.is_empty() { + return Ok(Vec::new()); + } + + let wal_path = self.path_resolver.wal_path(); + let object_store = self.object_stores.store_of(ObjectStoreType::Wal); + let mut files_stream = if id_range.start == 0 { + object_store.list(Some(&wal_path)) + } else { + let offset = self.path(&SsTableId::Wal(id_range.start - 1)); + object_store.list_with_offset(Some(&wal_path), &offset) + }; + let mut wal_list = Vec::new(); + + while let Some(file) = files_stream.next().await.transpose()? { + let Ok(Some(SsTableId::Wal(id))) = self.path_resolver.parse_table_id(&file.location) + else { + continue; + }; + if id >= id_range.end { + break; + } + if id >= id_range.start { + wal_list.push(IdentifiedObjectMetadata::from_object_meta( + SsTableId::Wal(id), + file, + )); + } + } + wal_list.sort_by_key(|metadata| metadata.id.unwrap_wal_id()); + Ok(wal_list) + } + + /// Fetches a WAL SST with one full-object GET. + pub(crate) async fn read_wal_sst_bytes( + &self, + wal_id: u64, + expected_size: Option, + max_object_bytes: usize, + ) -> Result { + self.read_wal_sst_bytes_with_retry(wal_id, expected_size, max_object_bytes, None) + .await + } + + async fn read_wal_sst_bytes_with_retry( + &self, + wal_id: u64, + expected_size: Option, + max_object_bytes: usize, + retry: Option, + ) -> Result { + let id = SsTableId::Wal(wal_id); + let obj = ReadOnlyObject { + object_store: self.object_stores.store_for(&id), + path: self.path(&id), + tag: ObjectStoreCallTag::new(self.kind, SstType::Wal), + }; + let mut tag = ObjectStoreCallTag::new(self.kind, SstType::Wal); + tag.retry = retry; + // Some object stores reject a body read for an existing zero-byte object. + // A tagged HEAD still proves that the exact fence object exists and has + // the size observed by the replay LIST. + let is_expected_fence = expected_size == Some(0); + let opts = GetOptions { + head: is_expected_fence, + extensions: tag.into(), + ..GetOptions::default() + }; + let result = obj + .object_store + .get_opts(&obj.path, opts) + .await + .map_err(SlateDBError::from) + .map_err(|err| err.with_path(&obj.path))?; + let response_size = usize::try_from(result.meta.size).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL object", + required_bytes: usize::MAX, + limit_bytes: max_object_bytes, + } + })?; + if response_size > max_object_bytes { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL object", + required_bytes: response_size, + limit_bytes: max_object_bytes, + }); + } + if expected_size.is_some_and(|expected| result.meta.size != expected) { + return Err(invalid_wal_size( + &obj.path, + wal_id, + expected_size, + result.meta.size, + )); + } + if is_expected_fence { + return Ok(Bytes::new()); + } + let response_size_u64 = result.meta.size; + let bytes = result + .bytes() + .await + .map_err(SlateDBError::from) + .map_err(|err| err.with_path(&obj.path))?; + let actual_size = u64::try_from(bytes.len()).map_err(|err| { + SlateDBError::WalDataError(Arc::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + err, + ))) + })?; + if actual_size != response_size_u64 { + return Err(invalid_wal_size( + &obj.path, + wal_id, + Some(response_size_u64), + actual_size, + )); + } + Ok(bytes) + } + + pub(crate) async fn refetch_wal_sst_after_validation( + &self, + wal_id: u64, + expected_size: usize, + max_object_bytes: usize, + working_memory_limit: usize, + reason: crate::error::RetryReason, + ) -> Result { + let expected_size = u64::try_from(expected_size).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL object", + required_bytes: expected_size, + limit_bytes: max_object_bytes, + } + })?; + let bytes = self + .read_wal_sst_bytes_with_retry( + wal_id, + Some(expected_size), + max_object_bytes, + Some(reason), + ) + .await?; + self.decode_wal_sst(wal_id, bytes, working_memory_limit) + .await + } + + /// Decodes the metadata and index of a fully fetched WAL locally. Data + /// blocks are decoded lazily by [`Self::decode_wal_block`]. + pub(crate) async fn decode_wal_sst( + &self, + wal_id: u64, + bytes: Bytes, + working_memory_limit: usize, + ) -> Result { + let metadata_memory_limit = working_memory_limit / 2; + let path = self.path(&SsTableId::Wal(wal_id)); + if bytes.is_empty() { + return Ok(DecodedWalSst::Fence); + } + let blob = BytesBlob::new(bytes.clone()); + let object_len = blob.len().await?; + let footer_size = METADATA_OFFSET_SIZE + VERSION_SIZE; + let footer_start = bytes + .len() + .checked_sub(footer_size) + .ok_or(SlateDBError::EmptySSTable)?; + let metadata_offset = u64::from_be_bytes( + bytes[footer_start..footer_start + METADATA_OFFSET_SIZE] + .try_into() + .map_err(|_| SlateDBError::EmptySSTable)?, + ); + let (info, format_version) = self + .sst_format + .read_info_and_version(&blob) + .await + .map_err(|error| error.with_path(&path))?; + validate_wal_sst_info_layout(&info, object_len, metadata_offset, &path)?; + let retained_info_bytes = info + .first_entry + .as_ref() + .map_or(0, Bytes::len) + .checked_add(info.last_entry.as_ref().map_or(0, Bytes::len)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + let index_memory_limit = metadata_memory_limit + .checked_sub(retained_info_bytes) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: retained_info_bytes, + limit_bytes: metadata_memory_limit, + })?; + let index = self + .sst_format + .read_wal_index_bounded(&info, &blob, index_memory_limit) + .await + .map_err(|error| error.with_path(&path))?; + validate_wal_sst_index_layout(&info, &index, &path)?; + let retained_decode_bytes = index + .size() + .checked_add(info.first_entry.as_ref().map_or(0, Bytes::len)) + .and_then(|size| size.checked_add(info.last_entry.as_ref().map_or(0, Bytes::len))) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + if retained_decode_bytes > metadata_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: retained_decode_bytes, + limit_bytes: metadata_memory_limit, + }); + } + Ok(DecodedWalSst::Data(Box::new(DecodedWalSstData { + wal_id, + format_version, + object_bytes: bytes, + info, + index, + retained_decode_bytes, + }))) + } + + /// Opens an oversized WAL without retaining its encoded object. Metadata, + /// the index, and each data block are fetched with bounded range reads. + pub(crate) async fn open_ranged_wal_sst( + &self, + wal_id: u64, + expected_size: u64, + working_memory_limit: usize, + ) -> Result { + let metadata_memory_limit = working_memory_limit / 2; + let id = SsTableId::Wal(wal_id); + read_obj!(self, &id, |obj| async { + let object_len = obj.len().await?; + if object_len != expected_size { + return Err(invalid_wal_size( + &obj.path, + wal_id, + Some(expected_size), + object_len, + )); + } + if object_len == 0 { + return Ok(RangedWalSst::Fence); + } + + let footer_size = u64::try_from(METADATA_OFFSET_SIZE + VERSION_SIZE) + .map_err(|_| SlateDBError::InvalidDBState)?; + let footer_start = object_len + .checked_sub(footer_size) + .ok_or(SlateDBError::EmptySSTable)?; + let footer = obj.read_range(footer_start..object_len).await?; + let metadata_offset = u64::from_be_bytes( + footer + .get(..METADATA_OFFSET_SIZE) + .ok_or(SlateDBError::EmptySSTable)? + .try_into() + .map_err(|_| SlateDBError::EmptySSTable)?, + ); + let version = u16::from_be_bytes( + footer + .get(METADATA_OFFSET_SIZE..METADATA_OFFSET_SIZE + VERSION_SIZE) + .ok_or(SlateDBError::EmptySSTable)? + .try_into() + .map_err(|_| SlateDBError::EmptySSTable)?, + ); + let (info, format_version) = self + .sst_format + .read_wal_info_and_version_bounded( + &obj, + object_len, + metadata_offset, + version, + metadata_memory_limit, + ) + .await?; + validate_wal_sst_info_layout(&info, object_len, metadata_offset, &obj.path)?; + let retained_info_bytes = info + .first_entry + .as_ref() + .map_or(0, Bytes::len) + .checked_add(info.last_entry.as_ref().map_or(0, Bytes::len)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + let index_memory_limit = metadata_memory_limit + .checked_sub(retained_info_bytes) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: retained_info_bytes, + limit_bytes: metadata_memory_limit, + })?; + let index = self + .sst_format + .read_ranged_wal_index_bounded(&info, &obj, index_memory_limit) + .await?; + validate_wal_sst_index_layout(&info, &index, &obj.path)?; + let retained_decode_bytes = index + .size() + .checked_add(info.first_entry.as_ref().map_or(0, Bytes::len)) + .and_then(|size| size.checked_add(info.last_entry.as_ref().map_or(0, Bytes::len))) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: usize::MAX, + limit_bytes: metadata_memory_limit, + })?; + if retained_decode_bytes > metadata_memory_limit { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata and index", + required_bytes: retained_decode_bytes, + limit_bytes: metadata_memory_limit, + }); + } + Ok(RangedWalSst::Data(Box::new(RangedWalSstData { + wal_id, + format_version, + info, + index, + retained_decode_bytes, + }))) + }) + .await + } + + pub(crate) async fn decode_wal_block( + &self, + wal: &DecodedWalSstData, + block_index: usize, + working_memory_limit: usize, + ) -> Result { + let path = self.path(&SsTableId::Wal(wal.wal_id)); + self.sst_format + .decode_wal_block_from_object( + &wal.info, + &wal.index, + block_index, + &wal.object_bytes, + working_memory_limit, + ) + .await + .map_err(|error| error.with_path(&path)) + } + + pub(crate) async fn read_ranged_wal_block( + &self, + wal: &RangedWalSstData, + block_index: usize, + working_memory_limit: usize, + ) -> Result { + let (start, end) = { + let index = wal.index.borrow(); + if block_index >= index.block_meta().len() { + return Err(SlateDBError::CorruptSst { + reason: "WAL block index is out of range", + path: None, + }); + } + let start = index.block_meta().get(block_index).offset(); + let end = if block_index + 1 < index.block_meta().len() { + index.block_meta().get(block_index + 1).offset() + } else { + wal.info.filter_offset + }; + (start, end) + }; + if start >= end { + return Err(SlateDBError::CorruptSst { + reason: "WAL block range is empty or reversed", + path: None, + }); + } + let expected_len = usize::try_from(end - start).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL block", + required_bytes: usize::MAX, + limit_bytes: working_memory_limit, + } + })?; + let decode_memory_limit = working_memory_limit.checked_sub(expected_len).ok_or( + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL block", + required_bytes: expected_len, + limit_bytes: working_memory_limit, + }, + )?; + let id = SsTableId::Wal(wal.wal_id); + read_obj!(self, &id, |obj| async { + let bytes = obj.read_range(start..end).await?; + if bytes.len() != expected_len { + return Err(SlateDBError::CorruptSst { + reason: "WAL block range returned an unexpected length", + path: None, + }); + } + self.sst_format + .decode_wal_block_bounded(bytes, wal.info.compression_codec, decode_memory_limit) + .await + }) + .await + } + pub(crate) async fn next_wal_sst_id( &self, wal_id_last_compacted: u64, @@ -618,6 +1191,52 @@ impl TableStore { Ok(SsTableHandle::new(*id, version, info)) } + pub(crate) async fn open_runtime_wal_sst( + &self, + wal_id: u64, + ) -> Result { + fail_point!(Arc::clone(&self.fp_registry), "runtime-wal-replay", |_| { + Err(RuntimeWalOpenError::Replay(SlateDBError::from( + std::io::Error::other("runtime WAL replay failpoint"), + ))) + }); + let id = SsTableId::Wal(wal_id); + let object_store = self.object_stores.store_for(&id); + let path = self.path(&id); + let mut tag = ObjectStoreCallTag::new(self.kind, SstType::Wal); + let mut validation_retries = 0_usize; + loop { + let object = ReadOnlyObject { + object_store: Arc::clone(&object_store), + path: path.clone(), + tag, + }; + match self.sst_format.read_info_and_version_staged(&object).await { + Ok((info, version)) => return Ok(SsTableHandle::new(id, version, info)), + Err(StagedSstInfoError::Initial(error)) if error.has_object_store_not_found() => { + return Err(RuntimeWalOpenError::MissingInitialObject( + error.with_path(&path), + )); + } + Err(StagedSstInfoError::Initial(error)) | Err(StagedSstInfoError::Later(error)) => { + let error = error.with_path(&path); + let Some(reason) = error.maybe_validation_retry_reason() else { + return Err(RuntimeWalOpenError::Replay(error)); + }; + if validation_retries >= MAX_VALIDATION_RETRIES { + return Err(RuntimeWalOpenError::Replay(error)); + } + validation_retries += 1; + tag.retry = Some(reason); + warn!( + "retrying runtime WAL read after validation failure [wal_id={}, reason={:?}, error={}]", + wal_id, reason, error + ); + } + } + } + } + #[cfg(test)] pub(crate) async fn read_sst_version(&self, id: &SsTableId) -> Result { let (_, version) = @@ -1109,6 +1728,177 @@ impl TableStore { } } +fn validate_wal_sst_info_layout( + info: &SsTableInfo, + object_len: u64, + metadata_offset: u64, + path: &Path, +) -> Result<(), SlateDBError> { + for (name, offset, len) in [ + ("index", info.index_offset, info.index_len), + ("filter", info.filter_offset, info.filter_len), + ("stats", info.stats_offset, info.stats_len), + ] { + let Some(end) = offset.checked_add(len) else { + return Err(invalid_wal_layout(path, format!("{name} range overflow"))); + }; + if end > metadata_offset { + return Err(invalid_wal_layout( + path, + format!("{name} range {offset}..{end} exceeds metadata offset {metadata_offset}"), + )); + } + } + if metadata_offset >= object_len { + return Err(invalid_wal_layout( + path, + format!("metadata offset {metadata_offset} reaches object length {object_len}"), + )); + } + + let filter_end = info + .filter_offset + .checked_add(info.filter_len) + .ok_or_else(|| invalid_wal_layout(path, "filter range overflow".to_string()))?; + if filter_end != info.index_offset { + return Err(invalid_wal_layout( + path, + format!( + "filter end {filter_end} does not equal index offset {}", + info.index_offset + ), + )); + } + let index_end = info + .index_offset + .checked_add(info.index_len) + .ok_or_else(|| invalid_wal_layout(path, "index range overflow".to_string()))?; + let checksum_size = CHECKSUM_SIZE as u64; + if info.index_len < checksum_size { + return Err(invalid_wal_layout( + path, + format!( + "index length {} is smaller than its checksum", + info.index_len + ), + )); + } + if info.filter_len > 0 && info.filter_len < checksum_size { + return Err(invalid_wal_layout( + path, + format!( + "filter length {} is smaller than its checksum", + info.filter_len + ), + )); + } + if info.stats_len > 0 { + if info.stats_len < checksum_size { + return Err(invalid_wal_layout( + path, + format!( + "stats length {} is smaller than its checksum", + info.stats_len + ), + )); + } + if info.stats_offset != index_end { + return Err(invalid_wal_layout( + path, + format!( + "index end {index_end} does not equal stats offset {}", + info.stats_offset + ), + )); + } + } + let data_end = if info.stats_len > 0 { + info.stats_offset + .checked_add(info.stats_len) + .ok_or_else(|| invalid_wal_layout(path, "stats range overflow".to_string()))? + } else { + index_end + }; + if data_end != metadata_offset { + return Err(invalid_wal_layout( + path, + format!("last WAL section ends at {data_end}, metadata starts at {metadata_offset}"), + )); + } + Ok(()) +} + +fn validate_wal_sst_index_layout( + info: &SsTableInfo, + index: &SsTableIndexOwned, + path: &Path, +) -> Result<(), SlateDBError> { + let block_meta = index.borrow().block_meta(); + if block_meta.is_empty() { + if info.filter_offset != 0 { + return Err(invalid_wal_layout( + path, + format!( + "empty block index has nonzero data length {}", + info.filter_offset + ), + )); + } + return Ok(()); + } + let mut previous_offset = None; + for block in 0..block_meta.len() { + let offset = block_meta.get(block).offset(); + if block == 0 && offset != 0 { + return Err(invalid_wal_layout( + path, + format!("first block offset {offset} is not zero"), + )); + } + if offset >= info.filter_offset { + return Err(invalid_wal_layout( + path, + format!( + "block {block} offset {offset} reaches data end {}", + info.filter_offset + ), + )); + } + if previous_offset.is_some_and(|previous| offset <= previous) { + return Err(invalid_wal_layout( + path, + format!("block {block} offset {offset} is not strictly increasing"), + )); + } + previous_offset = Some(offset); + } + Ok(()) +} + +fn invalid_wal_layout(path: &Path, reason: String) -> SlateDBError { + SlateDBError::WalDataError(Arc::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("invalid WAL SST layout at {path}: {reason}"), + ))) +} + +fn invalid_wal_size( + path: &Path, + wal_id: u64, + expected_size: Option, + actual_size: u64, +) -> SlateDBError { + SlateDBError::WalDataError(Arc::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!( + "WAL {wal_id} size changed at {path}: expected {}, got {actual_size}", + expected_size + .map(|size| size.to_string()) + .unwrap_or_else(|| "unknown".to_string()) + ), + ))) +} + async fn wal_object_exists( object_store: &Arc, path: &Path, @@ -1329,16 +2119,89 @@ mod tests { use crate::object_stores::ObjectStores; use crate::retrying_object_store::RetryingObjectStore; use crate::sst_iter::{SstIterator, SstIteratorOptions}; - use crate::tablestore::{TableStore, TableStoreKind}; + use crate::tablestore::{ + validate_wal_sst_index_layout, validate_wal_sst_info_layout, TableStore, TableStoreKind, + }; use crate::test_utils::FlakyObjectStore; use crate::test_utils::{assert_iterator, build_test_sst}; use crate::types::{RowEntry, ValueDeletable}; - use crate::{block_iterator::BlockIteratorLatest, db_state::SsTableId, iter::RowEntryIterator}; + use crate::{ + block_iterator::BlockIteratorLatest, + db_state::{SsTableId, SsTableInfo, SstType}, + iter::RowEntryIterator, + }; use slatedb_common::clock::DefaultSystemClock; use slatedb_common::DbRand; const ROOT: &str = "/root"; + #[test] + fn should_reject_invalid_wal_section_layouts_before_reading_index_bytes() { + let path = Path::from("/root/wal/00000000000000000001.sst"); + let valid = SsTableInfo { + filter_offset: 16, + filter_len: 4, + index_offset: 20, + index_len: 8, + sst_type: SstType::Wal, + ..SsTableInfo::default() + }; + assert!(validate_wal_sst_info_layout(&valid, 128, 28, &path).is_ok()); + + let invalid = [ + SsTableInfo { + index_offset: u64::MAX, + index_len: 8, + ..valid.clone() + }, + SsTableInfo { + index_offset: 21, + ..valid.clone() + }, + SsTableInfo { + index_len: 3, + ..valid.clone() + }, + SsTableInfo { + filter_len: 3, + index_offset: 19, + ..valid.clone() + }, + SsTableInfo { + stats_offset: 28, + stats_len: 3, + ..valid.clone() + }, + SsTableInfo { + stats_offset: 29, + stats_len: 4, + ..valid.clone() + }, + ]; + for info in invalid { + assert!( + validate_wal_sst_info_layout(&info, 128, 28, &path).is_err(), + "invalid WAL layout was accepted: {info:?}" + ); + } + assert!(validate_wal_sst_info_layout(&valid, 128, 29, &path).is_err()); + } + + #[tokio::test] + async fn should_reject_unindexed_wal_data() { + let path = Path::from("/root/wal/00000000000000000001.sst"); + let encoded = SsTableFormat::default() + .wal_table_builder() + .build() + .await + .unwrap(); + let mut info = encoded.info.clone(); + info.filter_offset = 1; + info.index_offset = 1; + + assert!(validate_wal_sst_index_layout(&info, &encoded.index, &path).is_err()); + } + /// Wraps an object store: counts range-bounded `get_opts` calls and pauses the first /// one until `release` is notified. Other methods just delegate. Shared by the /// concurrent-read dedup tests. @@ -3185,12 +4048,21 @@ mod tests { use crate::error::{RetryReason, SlateDBError}; use crate::format::sst::SsTableFormat; use crate::object_stores::ObjectStores; + use crate::replay_task_scope::ReplayTaskScope; use crate::tablestore::TableStore; use crate::test_utils::{build_test_sst, RecordingObjectStore}; use object_store::memory::InMemory; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; + struct RetryCancellationProbe(Arc); + + impl Drop for RetryCancellationProbe { + fn drop(&mut self) { + self.0.store(true, Ordering::Release); + } + } + fn format() -> SsTableFormat { SsTableFormat { block_size: 32, @@ -3372,6 +4244,48 @@ mod tests { ); } + #[tokio::test] + async fn validation_retry_is_cancelled_when_its_owner_is_aborted() { + let scope = ReplayTaskScope::new(); + let attempts = Arc::new(AtomicUsize::new(0)); + let retry_started = Arc::new(tokio::sync::Notify::new()); + let cancelled = Arc::new(AtomicBool::new(false)); + let task_attempts = Arc::clone(&attempts); + let task_retry_started = Arc::clone(&retry_started); + let task_cancelled = Arc::clone(&cancelled); + let task = scope.spawn(async move { + read_with_validation_retry( + ObjectStoreCallTag::new(TableStoreKind::Reader, SstType::Wal), + move |_| { + let attempt = task_attempts.fetch_add(1, Ordering::SeqCst); + let retry_started = Arc::clone(&task_retry_started); + let cancelled = Arc::clone(&task_cancelled); + async move { + if attempt == 0 { + return Err(SlateDBError::ChecksumMismatch { path: None }); + } + let _probe = RetryCancellationProbe(cancelled); + retry_started.notify_one(); + std::future::pending::>().await + } + }, + ) + .await + }); + tokio::time::timeout(std::time::Duration::from_secs(1), retry_started.notified()) + .await + .expect("validation retry did not start"); + + task.abort(); + let _ = task.await; + scope.shutdown().await; + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert!( + cancelled.load(Ordering::Acquire), + "validation retry future outlived its owner" + ); + } + #[tokio::test] async fn compacted_writes_carry_source_and_compacted_type() { let (recording, ts) = recording_store(TableStoreKind::Compactor); diff --git a/slatedb/src/test_utils.rs b/slatedb/src/test_utils.rs index df22e7009..9842c546d 100644 --- a/slatedb/src/test_utils.rs +++ b/slatedb/src/test_utils.rs @@ -16,8 +16,9 @@ use futures::stream::BoxStream; use futures::{stream, StreamExt}; use object_store::path::Path; use object_store::{ - CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, - PutMultipartOptions, PutOptions as OS_PutOptions, PutPayload, PutResult, RenameOptions, + CopyOptions, GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta, + ObjectStore, PutMultipartOptions, PutOptions as OS_PutOptions, PutPayload, PutResult, + RenameOptions, }; use rand::{Rng, RngCore}; use std::cmp::Ordering as CmpOrdering; @@ -1112,6 +1113,9 @@ pub(crate) struct GatedObjectStore { pub(crate) put_multipart_opts_gate: Gate, pub(crate) copy_gate: Gate, pub(crate) rename_gate: Gate, + /// Gates each path emitted by LIST operations. Per-item gating also lets + /// tests observe cancellation while a list stream is being consumed. + pub(crate) list_gate: Arc, /// Gates each path emitted by `delete_stream`. Per-item gating preserves /// the pre-0.13 single-call `delete` semantics now that callers go through /// `ObjectStoreExt::delete`, which fans out into `delete_stream`. @@ -1128,6 +1132,7 @@ impl GatedObjectStore { put_multipart_opts_gate: Gate::default(), copy_gate: Gate::default(), rename_gate: Gate::default(), + list_gate: Arc::new(Gate::default()), delete_stream_gate: Arc::new(Gate::default()), } } @@ -1192,7 +1197,17 @@ impl ObjectStore for GatedObjectStore { } fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { - self.inner.list(prefix) + let gate = Arc::clone(&self.list_gate); + self.inner + .list(prefix) + .then(move |result| { + let gate = Arc::clone(&gate); + async move { + gate.wait().await?; + result + } + }) + .boxed() } fn list_with_offset( @@ -1200,7 +1215,17 @@ impl ObjectStore for GatedObjectStore { prefix: Option<&Path>, offset: &Path, ) -> BoxStream<'static, object_store::Result> { - self.inner.list_with_offset(prefix, offset) + let gate = Arc::clone(&self.list_gate); + self.inner + .list_with_offset(prefix, offset) + .then(move |result| { + let gate = Arc::clone(&gate); + async move { + gate.wait().await?; + result + } + }) + .boxed() } async fn list_with_delimiter(&self, prefix: Option<&Path>) -> object_store::Result { @@ -1549,6 +1574,7 @@ mod tests { pub(crate) enum RecordedCall { Get { head: bool, + range_bytes: Option, kind: Option, sst_type: Option, retry: Option, @@ -1569,6 +1595,7 @@ pub(crate) enum RecordedCall { pub(crate) struct RecordingObjectStore { inner: Arc, calls: parking_lot::Mutex>, + list_calls: AtomicUsize, } impl RecordingObjectStore { @@ -1576,11 +1603,13 @@ impl RecordingObjectStore { Self { inner, calls: parking_lot::Mutex::new(Vec::new()), + list_calls: AtomicUsize::new(0), } } pub(crate) fn clear(&self) { self.calls.lock().clear(); + self.list_calls.store(0, Ordering::SeqCst); } pub(crate) fn get_kinds(&self, head: bool) -> Vec> { @@ -1618,6 +1647,24 @@ impl RecordingObjectStore { .collect() } + pub(crate) fn get_range_sizes(&self) -> Vec { + self.calls + .lock() + .iter() + .filter_map(|call| match call { + RecordedCall::Get { + range_bytes: Some(range_bytes), + .. + } => Some(*range_bytes), + _ => None, + }) + .collect() + } + + pub(crate) fn list_calls(&self) -> usize { + self.list_calls.load(Ordering::SeqCst) + } + pub(crate) fn write_kinds(&self) -> Vec> { self.calls .lock() @@ -1660,6 +1707,10 @@ impl ObjectStore for RecordingObjectStore { let tag = ObjectStoreCallTag::from_extensions(&options.extensions); self.calls.lock().push(RecordedCall::Get { head: options.head, + range_bytes: match &options.range { + Some(GetRange::Bounded(range)) => range.end.checked_sub(range.start), + _ => None, + }, kind: tag.map(|t| t.kind), sst_type: tag.map(|t| t.sst_type), retry: tag.and_then(|t| t.retry), @@ -1702,6 +1753,7 @@ impl ObjectStore for RecordingObjectStore { } fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, object_store::Result> { + self.list_calls.fetch_add(1, Ordering::SeqCst); self.inner.list(prefix) } @@ -1710,6 +1762,7 @@ impl ObjectStore for RecordingObjectStore { prefix: Option<&Path>, offset: &Path, ) -> BoxStream<'static, object_store::Result> { + self.list_calls.fetch_add(1, Ordering::SeqCst); self.inner.list_with_offset(prefix, offset) } diff --git a/slatedb/src/utils.rs b/slatedb/src/utils.rs index d89338250..3f5b90d88 100644 --- a/slatedb/src/utils.rs +++ b/slatedb/src/utils.rs @@ -639,6 +639,34 @@ pub(crate) fn decode_varint(buf: &mut impl Buf) -> u32 { result } +/// Decode a bounded u32 LEB128 value from untrusted storage bytes. +pub(crate) fn decode_varint_checked(buf: &mut impl Buf) -> Result { + let mut result = 0u32; + for byte_index in 0..5 { + if !buf.has_remaining() { + return Err(SlateDBError::CorruptSst { + reason: "truncated row varint", + path: None, + }); + } + let byte = buf.get_u8(); + if byte_index == 4 && (byte & 0xf0) != 0 { + return Err(SlateDBError::CorruptSst { + reason: "row varint overflows u32", + path: None, + }); + } + result |= u32::from(byte & 0x7f) << (byte_index * 7); + if byte & 0x80 == 0 { + return Ok(result); + } + } + Err(SlateDBError::CorruptSst { + reason: "row varint exceeds five bytes", + path: None, + }) +} + /// Calculate the encoded length of a u32 varint without actually encoding it. #[allow(dead_code)] pub(crate) fn varint_len(mut value: u32) -> usize { diff --git a/slatedb/src/wal/wal_sst_builder.rs b/slatedb/src/wal/wal_sst_builder.rs index 758a56984..595d2c0bf 100644 --- a/slatedb/src/wal/wal_sst_builder.rs +++ b/slatedb/src/wal/wal_sst_builder.rs @@ -902,6 +902,10 @@ mod tests { #[async_trait] impl BlockTransformer for MockBlockTransformer { + fn max_decoded_len(&self, encoded_len: usize) -> Option { + Some(encoded_len) + } + async fn encode(&self, data: Bytes) -> Result { self.encode_call_count.fetch_add(1, Ordering::SeqCst); Ok(data) diff --git a/slatedb/src/wal/writer_init.rs b/slatedb/src/wal/writer_init.rs index f98c85ff4..050d3b45c 100644 --- a/slatedb/src/wal/writer_init.rs +++ b/slatedb/src/wal/writer_init.rs @@ -15,6 +15,8 @@ use std::time::Duration; #[derive(Clone, Copy)] pub(crate) struct WalWriterInitOptions { max_wal_bytes_size: usize, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, max_flush_interval: Option, } @@ -22,6 +24,8 @@ impl From<&Settings> for WalWriterInitOptions { fn from(settings: &Settings) -> Self { Self { max_wal_bytes_size: settings.l0_sst_size_bytes, + max_replay_metadata_bytes: settings.wal_replay.metadata_working_memory_limit(), + max_replay_block_bytes: settings.wal_replay.block_working_memory_limit(), max_flush_interval: settings.flush_interval, } } @@ -32,6 +36,8 @@ pub(crate) struct WalWriterInit { recorder: MetricsRecorderHelper, table_store: Arc, max_wal_bytes_size: usize, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, max_flush_interval: Option, empty_wal_id: u64, task_executor: Arc, @@ -58,6 +64,8 @@ impl WalWriterInit { recorder, table_store, max_wal_bytes_size: options.max_wal_bytes_size, + max_replay_metadata_bytes: options.max_replay_metadata_bytes, + max_replay_block_bytes: options.max_replay_block_bytes, max_flush_interval: options.max_flush_interval, empty_wal_id, task_executor, @@ -115,6 +123,8 @@ impl wal::WriterInit for WalWriterInit { empty_wal_id, self.table_store.clone(), self.max_wal_bytes_size, + self.max_replay_metadata_bytes, + self.max_replay_block_bytes, self.max_flush_interval, self.task_executor.clone(), ) diff --git a/slatedb/src/wal_buffer.rs b/slatedb/src/wal_buffer.rs index a20020808..5c173a89d 100644 --- a/slatedb/src/wal_buffer.rs +++ b/slatedb/src/wal_buffer.rs @@ -51,6 +51,7 @@ pub(crate) struct WalBufferManager { stats: Arc, table_store: Arc, max_wal_bytes_size: usize, + max_replay_block_bytes: usize, /// The largest flush_epoch for which a size-triggered flush request has been /// sent. Compared against `flush_epoch` in the inner struct to avoid sending /// redundant flush requests for the same WAL. @@ -113,6 +114,8 @@ impl WalBufferManager { last_flushed_wal_id: u64, table_store: Arc, max_wal_bytes_size: usize, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, max_flush_interval: Option, task_executor: Arc, ) -> Result { @@ -133,6 +136,8 @@ impl WalBufferManager { let stats = Arc::new(WalBufferStats::new(recorder)); let wal_flush_handler = WalFlushHandler { max_flush_interval, + max_replay_metadata_bytes, + max_replay_block_bytes, inner: inner.clone(), table_store: table_store.clone(), stats: stats.clone(), @@ -149,6 +154,7 @@ impl WalBufferManager { stats, table_store, max_wal_bytes_size, + max_replay_block_bytes, last_flush_requested_epoch: AtomicU64::new(0), task_executor, }) @@ -210,6 +216,25 @@ impl WalWriter for WalBufferManager { /// Append row entries to the current WAL. Returns a watcher for durability notification. async fn append(&mut self, entries: &[RowEntry]) -> Result<(), WalError> { + for entry in entries { + let encoded_row_bytes = entry.encoded_size(0); + let required_bytes = encoded_row_bytes + .checked_mul(2) + .and_then(|size| size.checked_add(14)) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "decoded WAL row block", + required_bytes: usize::MAX, + limit_bytes: self.max_replay_block_bytes, + })?; + if required_bytes > self.max_replay_block_bytes { + return Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "decoded WAL row block", + required_bytes, + limit_bytes: self.max_replay_block_bytes, + } + .into()); + } + } self.inner.write().append(entries)?; self.maybe_trigger_flush()?; Ok(()) @@ -481,6 +506,8 @@ impl Debug for WalFlushWork { struct WalFlushHandler { max_flush_interval: Option, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, inner: Arc>, table_store: Arc, stats: Arc, @@ -540,6 +567,11 @@ impl WalFlushHandler { } let encoded_sst = sst_builder.build().await?; + self.table_store.validate_wal_sst_replay_memory( + &encoded_sst, + self.max_replay_metadata_bytes, + self.max_replay_block_bytes, + )?; let written_bytes = encoded_sst.remaining_len() as u64; self.table_store .write_sst(&SsTableId::Wal(wal_id), &encoded_sst) @@ -903,11 +935,55 @@ mod tests { Arc, Arc, Arc, + ) { + setup_wal_buffer_with_replay_limits( + flush_interval, + listener, + 32 * 1024 * 1024, + 32 * 1024 * 1024, + ) + .await + } + + async fn setup_wal_buffer_with_replay_limits( + flush_interval: Duration, + listener: wal::WalStatusListener, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, + ) -> ( + WalBufferManager, + Arc, + Arc, + Arc, + ) { + setup_wal_buffer_with_format_and_replay_limits( + SsTableFormat::default(), + 1000, + flush_interval, + listener, + max_replay_metadata_bytes, + max_replay_block_bytes, + ) + .await + } + + async fn setup_wal_buffer_with_format_and_replay_limits( + format: SsTableFormat, + max_wal_bytes_size: usize, + flush_interval: Duration, + listener: wal::WalStatusListener, + max_replay_metadata_bytes: usize, + max_replay_block_bytes: usize, + ) -> ( + WalBufferManager, + Arc, + Arc, + Arc, ) { let object_store: Arc = Arc::new(InMemory::new()); let table_store = Arc::new(TableStore::new( ObjectStores::new(object_store, None), - SsTableFormat::default(), + format, Path::from("/root"), None, TableStoreKind::Main, @@ -927,7 +1003,9 @@ mod tests { &helper, 0, // recent_flushed_wal_id table_store.clone(), - 1000, // max_wal_bytes_size + max_wal_bytes_size, + max_replay_metadata_bytes, + max_replay_block_bytes, Some(flush_interval), // max_flush_interval task_executor.clone(), ) @@ -1017,6 +1095,92 @@ mod tests { assert_eq!(wal_buffer.status().unwrap().last_flushed_wal_id, 1); } + #[tokio::test] + async fn oversized_single_entry_is_rejected_before_wal_append() { + let (mut wal_buffer, _, _, _) = setup_wal_buffer().await; + wal_buffer.max_replay_block_bytes = 1024; + let entry = make_entry("key", &"v".repeat(1024), 1, None); + + let error = wal_buffer.append(&[entry]).await.unwrap_err(); + + assert!(matches!(error, WalError::DataError(_))); + let status = wal_buffer.status().unwrap(); + assert_eq!(status.buffered_wal_entries_count, 0); + assert_eq!(status.last_flushed_wal_id, 0); + } + + #[tokio::test] + async fn wal_with_unreplayable_metadata_is_rejected_before_object_write() { + let (mut wal_buffer, table_store, _, _) = setup_wal_buffer_with_replay_limits( + Duration::MAX, + Arc::new(|_status| {}), + 1, + 32 * 1024 * 1024, + ) + .await; + wal_buffer + .append(&[make_entry("key", "value", 1, None)]) + .await + .unwrap(); + + let flush = wal_buffer.flush().await.unwrap(); + let error = flush.await.unwrap_err(); + + assert!(matches!(error, WalError::DataError(_))); + assert!(table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .is_empty()); + } + + #[tokio::test] + async fn many_block_wal_with_unreplayable_index_is_rejected_before_object_write() { + let metadata_memory_limit = 8 * 1024; + let format = SsTableFormat { + block_size: 128, + ..SsTableFormat::default() + }; + let (mut wal_buffer, table_store, _, _) = setup_wal_buffer_with_format_and_replay_limits( + format, + usize::MAX, + Duration::MAX, + Arc::new(|_status| {}), + metadata_memory_limit, + 32 * 1024 * 1024, + ) + .await; + let entries = (1..=1024) + .map(|seq| { + make_entry( + &format!("large-index-{seq:04}-{}", "k".repeat(64)), + &"v".repeat(128), + seq, + None, + ) + }) + .collect::>(); + let mut builder = table_store.wal_table_builder(); + for entry in &entries { + builder.add(entry.clone()).await.unwrap(); + } + let encoded = builder.build().await.unwrap(); + assert!(encoded.unconsumed_blocks.len() > 100); + assert!(usize::try_from(encoded.info.index_len).unwrap() > metadata_memory_limit); + + wal_buffer.append(&entries).await.unwrap(); + let flush = wal_buffer.flush().await.unwrap(); + let error = flush.await.unwrap_err(); + + assert!(matches!(error, WalError::DataError(_))); + assert!(error.to_string().contains("WAL index")); + assert!(table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .is_empty()); + } + #[tokio::test] async fn test_immutable_wal_reclaim() { let (mut wal_buffer, _, _, _) = setup_wal_buffer().await; diff --git a/slatedb/src/wal_replay.rs b/slatedb/src/wal_replay.rs index c98277ab9..032cd8d81 100644 --- a/slatedb/src/wal_replay.rs +++ b/slatedb/src/wal_replay.rs @@ -1,43 +1,128 @@ -use crate::db_state::SsTableId; +//! WAL replay has two deliberately separate discovery and I/O paths. +//! +//! Exact recovery (writer/reader open and reader checkpoint recovery) fixes a +//! finite range with LIST, then uses concurrent full-object GETs with bounded +//! oversized-object fallback. Runtime reader replay never LISTs: it restores +//! the bounded range-read iterator and probes only explicit WAL IDs. Keeping +//! these types separate prevents recurring reader refresh from accidentally +//! inheriting exact-open discovery cost. + +use crate::block_iterator::DataBlockIterator; +use crate::config::WalReplaySettings; use crate::error::SlateDBError; -use crate::iter::{EmptyIterator, RowEntryIterator}; -use crate::manifest::ManifestCore; -use crate::manifest::SsTableView; +use crate::iter::{EmptyIterator, IterationOrder, RowEntryIterator}; +use crate::manifest::{ManifestCore, SsTableView}; use crate::mem_table::WritableKVTable; +use crate::replay_task_scope::ReplayTaskScope; use crate::sst_iter::{SstIterator, SstIteratorOptions}; -use crate::tablestore::TableStore; +use crate::tablestore::{ + DecodedWalSst, DecodedWalSstData, RangedWalSst, RangedWalSstData, RuntimeWalOpenError, + TableStore, +}; +use crate::types::RowEntry; use crate::utils::panic_string; -use log::error; +use async_trait::async_trait; +use bytes::Bytes; +use log::{error, info}; use std::collections::VecDeque; +use std::num::NonZeroUsize; use std::ops::Range; use std::sync::Arc; -use tokio::task; -use tokio::task::JoinHandle; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::task::{self, JoinHandle}; pub(crate) struct WalReplayOptions { - /// The number of SSTs to preload while replaying - pub(crate) sst_batch_size: usize, + /// Limits concurrent full-object WAL prefetch. + pub(crate) prefetch: WalReplaySettings, /// The target maximum number of bytes in each returned table. WAL replay only /// splits between complete WAL SSTs, so a returned table may exceed this if a /// single WAL SST is larger. pub(crate) max_memtable_bytes: usize, - /// Options to pass through to underlying SST iterators - pub(crate) sst_iter_options: SstIteratorOptions, - /// The minimum seq number to replay. If unset, will replay all /// entries after `last_l0_seq` in the manifest. pub(crate) min_seq: Option, + + pub(crate) source: ExactWalReplaySource, + + /// Tracks spawned fetches for a live reader. Writer opening leaves this + /// unset because its open future already owns the entire replay lifetime. + pub(crate) task_scope: Option, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ExactWalReplaySource { + /// A writable database is recovering before it can serve writes. + WriterOpen, + /// A reader is recovering its initial exact view. + ReaderOpen, + /// A running reader lost or advanced beyond its checkpoint and must recover exactly. + CheckpointRecovery, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum RuntimeWalReplaySource { + /// Replay of the exact range declared by a newly applied manifest. + Manifest, + /// Recurring exact-next WAL tail while the reader remains open. + Tail, +} + +pub(crate) struct RuntimeWalReplayOptions { + /// The target maximum number of bytes in each returned table. Runtime + /// replay splits only between complete WAL SSTs. + pub(crate) max_memtable_bytes: usize, + + /// Options passed to the range-read SST iterators. + pub(crate) sst_iter_options: SstIteratorOptions, + + /// The minimum sequence number to replay. + pub(crate) min_seq: Option, + + /// Optional fairness boundary used by live tailing. Exact recovery and + /// manifest-range replay leave this unset to retain historical memtable + /// boundaries. + pub(crate) max_wals_per_batch: Option, + + pub(crate) source: RuntimeWalReplaySource, + + /// Exact owner for WAL-open and nested block-fetch tasks. + pub(crate) task_scope: Option, +} + +impl Default for RuntimeWalReplayOptions { + fn default() -> Self { + Self { + max_memtable_bytes: 64 * 1024 * 1024, + // Preserve the pre-full-object reader replay behavior: one fetch + // task eagerly reads up to 256 blocks through bounded range GETs. + sst_iter_options: SstIteratorOptions { + max_fetch_tasks: 1, + blocks_to_fetch: 256, + cache_blocks: true, + cache_metadata: false, + eager_spawn: true, + order: IterationOrder::Ascending, + prefix: None, + filter_context: None, + }, + min_seq: None, + max_wals_per_batch: None, + source: RuntimeWalReplaySource::Manifest, + task_scope: None, + } + } } impl Default for WalReplayOptions { fn default() -> Self { Self { - sst_batch_size: 4, + prefetch: WalReplaySettings::default(), max_memtable_bytes: 64 * 1024 * 1024, - sst_iter_options: SstIteratorOptions::default(), min_seq: None, + source: ExactWalReplaySource::WriterOpen, + task_scope: None, } } } @@ -52,6 +137,310 @@ pub(crate) struct ReplayedMemtable { struct WalIdAndIter { wal_id: u64, iter: Box, + _permit: Option, +} + +struct RuntimeWalIdAndIter { + wal_id: u64, + iter: Box, +} + +#[derive(Debug)] +pub(crate) enum RuntimeWalReplayError { + MissingInitialObject { wal_id: u64, source: SlateDBError }, + Replay(SlateDBError), +} + +impl From for RuntimeWalReplayError { + fn from(error: SlateDBError) -> Self { + Self::Replay(error) + } +} + +struct WalObjectPlan { + wal_id: u64, + expected_size: Option, +} + +enum FetchedWal { + Full { + wal_id: u64, + bytes: Bytes, + permit: OwnedSemaphorePermit, + }, + Ranged { + wal_id: u64, + expected_size: u64, + }, +} + +impl FetchedWal { + fn wal_id(&self) -> u64 { + match self { + Self::Full { wal_id, .. } | Self::Ranged { wal_id, .. } => *wal_id, + } + } +} + +struct PendingWal { + wal_id: u64, + handle: JoinHandle>, +} + +struct WalBlocksIterator { + table_store: Arc, + wal: Option, + next_block: usize, + current: Option, + encoded_byte_limit: usize, + working_memory_limit: usize, + block_working_memory_limit: usize, + validation_retried: bool, + yielded_rows: bool, + #[cfg(test)] + block_lifetime_probe: Option>, +} + +struct CurrentWalBlock { + iterator: DataBlockIterator, + #[cfg(test)] + lifetime_probe: Option>, +} + +impl CurrentWalBlock { + fn new( + iterator: DataBlockIterator, + #[cfg(test)] lifetime_probe: Option>, + ) -> Self { + #[cfg(test)] + if let Some(probe) = &lifetime_probe { + let previous = probe + .active + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + probe + .peak + .fetch_max(previous + 1, std::sync::atomic::Ordering::SeqCst); + } + Self { + iterator, + #[cfg(test)] + lifetime_probe, + } + } +} + +#[cfg(test)] +impl Drop for CurrentWalBlock { + fn drop(&mut self) { + if let Some(probe) = &self.lifetime_probe { + let previous = probe + .active + .fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + assert_eq!(previous, 1, "decoded WAL blocks overlapped"); + } + } +} + +#[cfg(test)] +#[derive(Default)] +struct BlockLifetimeProbe { + active: std::sync::atomic::AtomicUsize, + peak: std::sync::atomic::AtomicUsize, +} + +enum WalBlockSource { + Full(DecodedWalSstData), + Ranged(RangedWalSstData), +} + +impl WalBlockSource { + fn block_count(&self) -> usize { + match self { + Self::Full(wal) => wal.index.borrow().block_meta().len(), + Self::Ranged(wal) => wal.index.borrow().block_meta().len(), + } + } + + fn format_version(&self) -> u16 { + match self { + Self::Full(wal) => wal.format_version, + Self::Ranged(wal) => wal.format_version, + } + } + + fn retained_decode_bytes(&self) -> usize { + match self { + Self::Full(wal) => wal.retained_decode_bytes, + Self::Ranged(wal) => wal.retained_decode_bytes, + } + } +} + +impl WalBlocksIterator { + fn new( + table_store: Arc, + wal: DecodedWalSstData, + encoded_byte_limit: usize, + working_memory_limit: usize, + ) -> Result { + Self::from_source( + table_store, + WalBlockSource::Full(wal), + encoded_byte_limit, + working_memory_limit, + ) + } + + fn new_ranged( + table_store: Arc, + wal: RangedWalSstData, + encoded_byte_limit: usize, + working_memory_limit: usize, + ) -> Result { + Self::from_source( + table_store, + WalBlockSource::Ranged(wal), + encoded_byte_limit, + working_memory_limit, + ) + } + + fn from_source( + table_store: Arc, + wal: WalBlockSource, + encoded_byte_limit: usize, + working_memory_limit: usize, + ) -> Result { + let block_working_memory_limit = (working_memory_limit / 2).min( + working_memory_limit + .checked_sub(wal.retained_decode_bytes()) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata, index, and block", + required_bytes: wal.retained_decode_bytes(), + limit_bytes: working_memory_limit, + })?, + ); + Ok(Self { + table_store, + wal: Some(wal), + next_block: 0, + current: None, + encoded_byte_limit, + working_memory_limit, + block_working_memory_limit, + validation_retried: false, + yielded_rows: false, + #[cfg(test)] + block_lifetime_probe: None, + }) + } + + #[cfg(test)] + fn observe_block_lifetimes_with(mut self, probe: Arc) -> Self { + self.block_lifetime_probe = Some(probe); + self + } +} + +#[async_trait] +impl RowEntryIterator for WalBlocksIterator { + async fn init(&mut self) -> Result<(), SlateDBError> { + Ok(()) + } + + async fn next(&mut self) -> Result, SlateDBError> { + loop { + if let Some(current) = &mut self.current { + if let Some(entry) = current.iterator.next().await? { + self.yielded_rows = true; + return Ok(Some(entry)); + } + } + self.current.take(); + let wal = self.wal.as_ref().ok_or(SlateDBError::InvalidDBState)?; + let block_count = wal.block_count(); + if self.next_block >= block_count { + return Ok(None); + } + debug_assert!( + self.current.is_none(), + "exhausted WAL block must be released before decoding the next block" + ); + let decoded = match wal { + WalBlockSource::Full(wal) => { + self.table_store + .decode_wal_block(wal, self.next_block, self.block_working_memory_limit) + .await + } + WalBlockSource::Ranged(wal) => { + self.table_store + .read_ranged_wal_block( + wal, + self.next_block, + self.block_working_memory_limit, + ) + .await + } + }; + let block = match decoded { + Ok(block) => block, + Err(error) => { + let Some(reason) = error.maybe_validation_retry_reason() else { + return Err(error); + }; + if self.validation_retried || self.yielded_rows { + return Err(error); + } + let old_wal = self.wal.take().ok_or(SlateDBError::InvalidDBState)?; + let WalBlockSource::Full(old_wal) = old_wal else { + return Err(error); + }; + let wal_id = old_wal.wal_id; + let expected_size = old_wal.object_bytes.len(); + drop(old_wal); + self.validation_retried = true; + let replacement = self + .table_store + .refetch_wal_sst_after_validation( + wal_id, + expected_size, + self.encoded_byte_limit, + self.working_memory_limit, + reason, + ) + .await?; + let DecodedWalSst::Data(replacement) = replacement else { + return Err(SlateDBError::CorruptSst { + reason: "nonempty WAL changed into a fence during validation retry", + path: None, + }); + }; + self.block_working_memory_limit = (self.working_memory_limit / 2).min( + self.working_memory_limit + .checked_sub(replacement.retained_decode_bytes) + .ok_or(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "WAL metadata, index, and block", + required_bytes: replacement.retained_decode_bytes, + limit_bytes: self.working_memory_limit, + })?, + ); + self.wal = Some(WalBlockSource::Full(*replacement)); + continue; + } + }; + self.next_block += 1; + let wal = self.wal.as_ref().ok_or(SlateDBError::InvalidDBState)?; + self.current = Some(CurrentWalBlock::new( + DataBlockIterator::new(block, wal.format_version(), IterationOrder::Ascending)?, + #[cfg(test)] + self.block_lifetime_probe.clone(), + )); + } + } + + async fn seek(&mut self, _next_key: &[u8]) -> Result<(), SlateDBError> { + Err(SlateDBError::InvalidDBState) + } } struct IteratorHolder { @@ -82,43 +471,38 @@ impl IteratorHolder { } } -pub(crate) struct WalReplayIterator { - options: WalReplayOptions, +/// Replays a finite WAL range through the established bounded SST range-read +/// iterator. This path deliberately performs no object-store LIST and retains +/// only a small ordered prefetch window. +pub(crate) struct RuntimeWalReplayIterator { + options: RuntimeWalReplayOptions, wal_id_range: Range, table_store: Arc, - current_iter: IteratorHolder, - next_iters: VecDeque, SlateDBError>>>, + current_iter: IteratorHolder, + next_iters: VecDeque, RuntimeWalReplayError>>>, last_tick: i64, last_seq: u64, min_seq: u64, next_wal_id: u64, } -impl WalReplayIterator { - pub(crate) async fn range( +const RUNTIME_WAL_PREFETCH: usize = 4; + +impl RuntimeWalReplayIterator { + pub(crate) fn range( wal_id_range: Range, db_state: &ManifestCore, - options: WalReplayOptions, + options: RuntimeWalReplayOptions, table_store: Arc, ) -> Result { - let sst_batch_size = options.sst_batch_size; - if sst_batch_size < 1 { - return Err(SlateDBError::InvalidSSTBatchSize(sst_batch_size)); - } - - // load the last seq number from manifest, and use it as the starting seq number to avoid - // replaying the entries that are already in the L0 SST. while replaying the WALs, we'll - // update the last seq number to the max seq number, and this final `last_seq` will be passed - // to the db_state for the further writes. let min_seq = options.min_seq.unwrap_or(db_state.last_l0_seq); let last_seq = db_state.last_l0_seq; let last_tick = db_state.last_l0_clock_tick; let next_wal_id = wal_id_range.start; - - let mut replay_iter = WalReplayIterator { + let mut replay = Self { options, wal_id_range, - table_store: Arc::clone(&table_store), + table_store, current_iter: IteratorHolder::new(), next_iters: VecDeque::new(), last_tick, @@ -126,81 +510,101 @@ impl WalReplayIterator { min_seq, next_wal_id, }; - - for _ in 0..sst_batch_size { - if !replay_iter.maybe_load_next_iter() { - break; - } - } - - Ok(replay_iter) + while replay.maybe_load_next_iter() {} + info!( + "SlateDB runtime WAL replay initialized [source={:?}, replay_start_wal_id={}, replay_end_wal_id={}, prefetch={}]", + replay.options.source, + replay.wal_id_range.start, + replay.wal_id_range.end, + RUNTIME_WAL_PREFETCH, + ); + Ok(replay) } fn maybe_load_next_iter(&mut self) -> bool { if !self.wal_id_range.contains(&self.next_wal_id) - || self.next_iters.len() >= self.options.sst_batch_size + || self.next_iters.len() >= RUNTIME_WAL_PREFETCH { return false; } - let next_wal_id = self.next_wal_id; + let wal_id = self.next_wal_id; self.next_wal_id += 1; - - async fn load_iter( - wal_id: u64, - sst_iter_options: SstIteratorOptions, - table_store: Arc, - ) -> Result, SlateDBError> { - let sst = match table_store.open_sst(&SsTableId::Wal(wal_id)).await { + let table_store = Arc::clone(&self.table_store); + let sst_iter_options = self.options.sst_iter_options.clone(); + let task_scope = self.options.task_scope.clone(); + let nested_scope = task_scope.clone(); + let task = async move { + let sst = match table_store.open_runtime_wal_sst(wal_id).await { Ok(sst) => sst, - Err(SlateDBError::EmptySSTable) => { - // Zero-byte WAL files are fence markers; replay them as empty WALs - // so the last replayed WAL ID still advances past the marker. - return Ok(Some(WalIdAndIter { + Err(RuntimeWalOpenError::Replay(SlateDBError::EmptySSTable)) => { + return Ok(Some(RuntimeWalIdAndIter { wal_id, iter: Box::new(EmptyIterator::new()), })); } - Err(err) => return Err(err), + Err(RuntimeWalOpenError::MissingInitialObject(error)) => { + return Err(RuntimeWalReplayError::MissingInitialObject { + wal_id, + source: error, + }); + } + Err(RuntimeWalOpenError::Replay(error)) => { + return Err(RuntimeWalReplayError::Replay(error)); + } }; - let iter = SstIterator::new_owned_initialized( - .., - SsTableView::identity(sst), - Arc::clone(&table_store), - sst_iter_options, - ) - .await?; - Ok(iter.map(|iter| WalIdAndIter { + let iter = if let Some(scope) = nested_scope { + SstIterator::new_owned_initialized_scoped( + .., + SsTableView::identity(sst), + Arc::clone(&table_store), + sst_iter_options, + scope, + ) + .await + } else { + SstIterator::new_owned_initialized( + .., + SsTableView::identity(sst), + Arc::clone(&table_store), + sst_iter_options, + ) + .await + } + .map_err(RuntimeWalReplayError::Replay)?; + Ok(iter.map(|iter| RuntimeWalIdAndIter { wal_id, iter: Box::new(iter) as Box, })) - } - - let handle = task::spawn(load_iter( - next_wal_id, - self.options.sst_iter_options.clone(), - Arc::clone(&self.table_store), - )); + }; + let handle = if let Some(scope) = task_scope.as_ref() { + scope.spawn(task) + } else { + task::spawn(task) + }; self.next_iters.push_back(handle); true } - async fn advance_current_iter(&mut self) -> Result<(), SlateDBError> { + async fn advance_current_iter(&mut self) -> Result<(), RuntimeWalReplayError> { let next_iter = if let Some(join_handle) = self.next_iters.pop_front() { match join_handle.await { - Ok(Ok(sst_iter)) => sst_iter, - Ok(Err(slate_err)) => return Err(slate_err), - Err(join_err) => { - let task_name = format!("wal_replay[{:?}]", self.wal_id_range); - if let Ok(panic_err) = join_err.try_into_panic() { + Ok(result) => result?, + Err(join_error) => { + let task_name = format!("runtime_wal_replay[{:?}]", self.wal_id_range); + if let Ok(panic_error) = join_error.try_into_panic() { error!( - "wal_replay task panicked unexpectedly. [task_name={}, panic={}]", + "runtime WAL replay task panicked [task_name={}, panic={}]", task_name, - panic_string(&panic_err), + panic_string(&panic_error), ); - return Err(SlateDBError::BackgroundTaskPanic(task_name)); + return Err(RuntimeWalReplayError::Replay( + SlateDBError::BackgroundTaskPanic(task_name), + )); } - return Err(SlateDBError::BackgroundTaskCancelled(task_name)); + return Err(RuntimeWalReplayError::Replay( + SlateDBError::BackgroundTaskCancelled(task_name), + )); } } } else { @@ -210,219 +614,2934 @@ impl WalReplayIterator { Ok(()) } - /// Get the next table replayed from the WAL. Replay accumulates complete WAL - /// SSTs until the returned table reaches [`WalReplayOptions::max_memtable_bytes`], - /// unless it is the final table replayed from the WAL. The final table may even - /// be empty since writers use an empty WAL to fence zombie writers. The empty - /// table must still be returned so that replay logic can account for the latest - /// WAL ID. - /// - /// The returned table may exceed [`WalReplayOptions::max_memtable_bytes`] when - /// a complete WAL SST is larger than the configured target, because replay - /// must not split a WAL SST across replayed memtables. - pub(crate) async fn next(&mut self) -> Result, SlateDBError> { + pub(crate) async fn next(&mut self) -> Result, RuntimeWalReplayError> { if self.current_iter.is_finished() { return Ok(None); } + if !self.current_iter.initialized { + self.advance_current_iter().await?; + } let table = WritableKVTable::new(); let mut last_wal_id = 0; - + let mut replayed_wals = 0_usize; while !self.current_iter.is_finished() { - if let Some(wal_id_and_iter) = &mut self.current_iter.current_iter { - while let Some(row_entry) = wal_id_and_iter.iter.next().await? { - // skip the entries that are already in the L0 SST. - if row_entry.seq <= self.min_seq { + if let Some(wal) = &mut self.current_iter.current_iter { + while let Some(row) = wal + .iter + .next() + .await + .map_err(RuntimeWalReplayError::Replay)? + { + if row.seq <= self.min_seq { continue; } - - if let Some(ts) = row_entry.create_ts { - self.last_tick = self.last_tick.max(ts); + if let Some(timestamp) = row.create_ts { + self.last_tick = self.last_tick.max(timestamp); } - self.last_seq = self.last_seq.max(row_entry.seq); - table.put(row_entry); + self.last_seq = self.last_seq.max(row.seq); + table.put(row); } - last_wal_id = wal_id_and_iter.wal_id; + last_wal_id = wal.wal_id; + replayed_wals = replayed_wals.saturating_add(1); + let metadata = table.metadata(); + let estimated_bytes = self.table_store.estimate_encoded_size_compacted( + metadata.entry_num, + metadata.entries_size_in_bytes, + ); + if (!table.is_empty() && estimated_bytes >= self.options.max_memtable_bytes) + || self + .options + .max_wals_per_batch + .is_some_and(|limit| replayed_wals >= limit.get()) + { + self.current_iter.reset(); + break; + } + } + + self.maybe_load_next_iter(); + self.advance_current_iter().await?; + } + + Ok((last_wal_id > 0).then_some(ReplayedMemtable { + table, + last_tick: self.last_tick, + last_seq: self.last_seq, + last_wal_id, + })) + } +} + +impl Drop for RuntimeWalReplayIterator { + fn drop(&mut self) { + for pending in &self.next_iters { + pending.abort(); + } + } +} + +/// Replays one exact, contiguous WAL-ID range. +/// +/// Safety invariants: +/// - every ID is fetched directly, including IDs omitted by object listing; +/// - fetched objects are consumed strictly in WAL-ID order; +/// - encoded-byte permits remain owned until the corresponding WAL is consumed; +/// - an error permanently fails the iterator and aborts every outstanding fetch; +/// - the replay cursor is lazy, so memory does not grow with a sparse ID span. +pub(crate) struct ExactWalReplayIterator { + options: WalReplayOptions, + wal_id_range: Range, + table_store: Arc, + current_iter: IteratorHolder, + next_wal_id: u64, + listed_sizes: VecDeque<(u64, u64)>, + pending_fetches: VecDeque, + byte_semaphore: Arc, + encoded_byte_limit: usize, + working_memory_limit: usize, + last_tick: i64, + last_seq: u64, + min_seq: u64, + fetched_objects: u64, + fetched_bytes: u64, + decoded_objects: u64, + ranged_objects: u64, + peak_concurrent_objects: usize, + peak_reserved_bytes: usize, + completion_logged: bool, + failed: Option, +} + +impl ExactWalReplayIterator { + pub(crate) async fn range( + wal_id_range: Range, + db_state: &ManifestCore, + options: WalReplayOptions, + table_store: Arc, + ) -> Result { + options.prefetch.validate()?; + + // load the last seq number from manifest, and use it as the starting seq number to avoid + // replaying the entries that are already in the L0 SST. while replaying the WALs, we'll + // update the last seq number to the max seq number, and this final `last_seq` will be passed + // to the db_state for the further writes. + let min_seq = options.min_seq.unwrap_or(db_state.last_l0_seq); + let last_seq = db_state.last_l0_seq; + let last_tick = db_state.last_l0_clock_tick; + let listed = table_store + .list_wal_ssts_for_replay(wal_id_range.clone()) + .await?; + let listed_wal_count = listed.len(); + let mut listed_sizes = VecDeque::with_capacity(listed_wal_count); + let mut previous_listed_wal_id = None; + for metadata in listed { + let wal_id = metadata.id.unwrap_wal_id(); + if !wal_id_range.contains(&wal_id) + || previous_listed_wal_id.is_some_and(|previous| wal_id <= previous) + { + return Err(SlateDBError::InvalidDBState); + } + previous_listed_wal_id = Some(wal_id); + listed_sizes.push_back((wal_id, metadata.metadata.size)); + } + let working_memory_limit = options.prefetch.working_memory_limit(); + let encoded_byte_limit = options.prefetch.encoded_byte_limit()?; + let byte_semaphore = Arc::new(Semaphore::new(encoded_byte_limit)); + let next_wal_id = wal_id_range.start; + + let mut replay_iter = ExactWalReplayIterator { + options, + wal_id_range, + table_store: Arc::clone(&table_store), + current_iter: IteratorHolder::new(), + next_wal_id, + listed_sizes, + pending_fetches: VecDeque::new(), + byte_semaphore, + encoded_byte_limit, + working_memory_limit, + last_tick, + last_seq, + min_seq, + fetched_objects: 0, + fetched_bytes: 0, + decoded_objects: 0, + ranged_objects: 0, + peak_concurrent_objects: 0, + peak_reserved_bytes: 0, + completion_logged: false, + failed: None, + }; + + replay_iter.fill_prefetch(); + info!( + "SlateDB exact WAL replay initialized [source={:?}, replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}, listed_wal_count={}, missing_size_count={}, max_concurrent_objects={}, max_inflight_bytes={}]", + replay_iter.options.source, + replay_iter.wal_id_range.start, + replay_iter.wal_id_range.end, + replay_iter + .wal_id_range + .end + .saturating_sub(replay_iter.wal_id_range.start), + listed_wal_count, + replay_iter + .wal_id_range + .end + .saturating_sub(replay_iter.wal_id_range.start) + .saturating_sub(listed_wal_count as u64), + replay_iter.options.prefetch.max_concurrent_objects, + replay_iter.options.prefetch.max_inflight_bytes, + ); + + Ok(replay_iter) + } + + fn spawn_fetch(&self, task: F) -> JoinHandle + where + F: std::future::Future + Send + 'static, + T: Send + 'static, + { + if let Some(scope) = self.options.task_scope.as_ref() { + scope.spawn(task) + } else { + tokio::spawn(task) + } + } + + fn fill_prefetch(&mut self) { + while self.pending_fetches.len() < self.options.prefetch.max_concurrent_objects { + if self.next_wal_id >= self.wal_id_range.end { + break; + } + while self + .listed_sizes + .front() + .is_some_and(|(wal_id, _)| *wal_id < self.next_wal_id) + { + self.listed_sizes.pop_front(); + } + let plan = WalObjectPlan { + wal_id: self.next_wal_id, + expected_size: self + .listed_sizes + .front() + .filter(|(wal_id, _)| *wal_id == self.next_wal_id) + .map(|(_, size)| *size), + }; + let byte_limit = self.encoded_byte_limit; + let reserved_bytes = plan + .expected_size + .and_then(|size| usize::try_from(size).ok()) + .unwrap_or(byte_limit) + .max(1); + if reserved_bytes > byte_limit { + let wal_id = plan.wal_id; + let expected_size = plan + .expected_size + .expect("an oversized planned WAL must have a listed size"); + self.listed_sizes.pop_front(); + self.next_wal_id += 1; + let handle = self.spawn_fetch(async move { + Ok(FetchedWal::Ranged { + wal_id, + expected_size, + }) + }); + self.pending_fetches + .push_back(PendingWal { wal_id, handle }); + self.peak_concurrent_objects = + self.peak_concurrent_objects.max(self.pending_fetches.len()); + continue; + } + let permit_count = u32::try_from(reserved_bytes) + .expect("validated WAL replay byte limit must fit in u32"); + let Ok(permit) = Arc::clone(&self.byte_semaphore).try_acquire_many_owned(permit_count) + else { + break; + }; + if plan.expected_size.is_some() { + self.listed_sizes.pop_front(); + } + self.next_wal_id += 1; + let table_store = Arc::clone(&self.table_store); + let handle = self.spawn_fetch(async move { + match table_store + .read_wal_sst_bytes(plan.wal_id, plan.expected_size, byte_limit) + .await + { + Ok(bytes) => Ok(FetchedWal::Full { + wal_id: plan.wal_id, + bytes, + permit, + }), + Err(SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL object", + required_bytes, + limit_bytes, + }) if required_bytes > limit_bytes => { + drop(permit); + Ok(FetchedWal::Ranged { + wal_id: plan.wal_id, + expected_size: u64::try_from(required_bytes).map_err(|_| { + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded WAL object", + required_bytes, + limit_bytes, + } + })?, + }) + } + Err(error) => Err(error), + } + }); + self.pending_fetches.push_back(PendingWal { + wal_id: plan.wal_id, + handle, + }); + self.peak_concurrent_objects = + self.peak_concurrent_objects.max(self.pending_fetches.len()); + self.peak_reserved_bytes = self.peak_reserved_bytes.max( + self.encoded_byte_limit + .saturating_sub(self.byte_semaphore.available_permits()), + ); + } + } + + async fn advance_current_iter(&mut self) -> Result<(), SlateDBError> { + self.current_iter.current_iter.take(); + self.fill_prefetch(); + let next_iter = if let Some(pending) = self.pending_fetches.pop_front() { + let fetched = match pending.handle.await { + Ok(Ok(fetched)) => fetched, + Ok(Err(slate_err)) => return Err(self.fail(slate_err)), + Err(join_err) => { + let task_name = format!("wal_replay[{:?}]", self.wal_id_range); + if let Ok(panic_err) = join_err.try_into_panic() { + error!( + "wal_replay task panicked unexpectedly. [task_name={}, panic={}]", + task_name, + panic_string(&panic_err), + ); + return Err(self.fail(SlateDBError::BackgroundTaskPanic(task_name))); + } + return Err(self.fail(SlateDBError::BackgroundTaskCancelled(task_name))); + } + }; + if pending.wal_id != fetched.wal_id() { + return Err(self.fail(SlateDBError::InvalidDBState)); + } + let (wal_id, iter, permit): (_, Box, _) = match fetched + { + FetchedWal::Full { + wal_id, + bytes, + permit, + } => { + self.fetched_objects = self.fetched_objects.saturating_add(1); + self.fetched_bytes = self + .fetched_bytes + .saturating_add(u64::try_from(bytes.len()).unwrap_or(u64::MAX)); + let expected_size = bytes.len(); + let initial_decode = self + .table_store + .decode_wal_sst(wal_id, bytes, self.working_memory_limit) + .await; + let decoded = match initial_decode { + Ok(decoded) => decoded, + Err(error) => { + let Some(reason) = error.maybe_validation_retry_reason() else { + return Err(self.fail(error)); + }; + match self + .table_store + .refetch_wal_sst_after_validation( + wal_id, + expected_size, + self.encoded_byte_limit, + self.working_memory_limit, + reason, + ) + .await + { + Ok(decoded) => decoded, + Err(error) => return Err(self.fail(error)), + } + } + }; + let iter: Box = match decoded { + DecodedWalSst::Fence => Box::new(crate::iter::EmptyIterator::new()), + DecodedWalSst::Data(wal) => Box::new( + WalBlocksIterator::new( + Arc::clone(&self.table_store), + *wal, + self.encoded_byte_limit, + self.working_memory_limit, + ) + .map_err(|error| self.fail(error))?, + ), + }; + (wal_id, iter, Some(permit)) + } + FetchedWal::Ranged { + wal_id, + expected_size, + } => { + self.ranged_objects = self.ranged_objects.saturating_add(1); + let decoded = match self + .table_store + .open_ranged_wal_sst(wal_id, expected_size, self.working_memory_limit) + .await + { + Ok(decoded) => decoded, + Err(error) => return Err(self.fail(error)), + }; + let iter: Box = match decoded { + RangedWalSst::Fence => Box::new(crate::iter::EmptyIterator::new()), + RangedWalSst::Data(wal) => Box::new( + WalBlocksIterator::new_ranged( + Arc::clone(&self.table_store), + *wal, + self.encoded_byte_limit, + self.working_memory_limit, + ) + .map_err(|error| self.fail(error))?, + ), + }; + (wal_id, iter, None) + } + }; + self.decoded_objects = self.decoded_objects.saturating_add(1); + Some(WalIdAndIter { + wal_id, + iter, + _permit: permit, + }) + } else { + None + }; + self.current_iter.advance(next_iter); + self.fill_prefetch(); + Ok(()) + } + + /// Get the next table replayed from the WAL. Replay accumulates complete WAL + /// SSTs until the returned table reaches [`WalReplayOptions::max_memtable_bytes`], + /// unless it is the final table replayed from the WAL. The final table may even + /// be empty since writers use an empty WAL to fence zombie writers. The empty + /// table must still be returned so that replay logic can account for the latest + /// WAL ID. + /// + /// The returned table may exceed [`WalReplayOptions::max_memtable_bytes`] when + /// a complete WAL SST is larger than the configured target, because replay + /// must not split a WAL SST across replayed memtables. + pub(crate) async fn next(&mut self) -> Result, SlateDBError> { + if let Some(error) = &self.failed { + return Err(error.clone()); + } + if self.current_iter.is_finished() { + self.log_completion(); + return Ok(None); + } + + if !self.current_iter.initialized { + self.advance_current_iter().await?; + } + + let table = WritableKVTable::new(); + let mut last_wal_id = 0; + + while !self.current_iter.is_finished() { + if let Some(wal_id_and_iter) = &mut self.current_iter.current_iter { + loop { + let row_entry = match wal_id_and_iter.iter.next().await { + Ok(Some(row_entry)) => row_entry, + Ok(None) => break, + Err(error) => return Err(self.fail(error)), + }; + // skip the entries that are already in the L0 SST. + if row_entry.seq <= self.min_seq { + continue; + } + + if let Some(ts) = row_entry.create_ts { + self.last_tick = self.last_tick.max(ts); + } + self.last_seq = self.last_seq.max(row_entry.seq); + table.put(row_entry); + } + + last_wal_id = wal_id_and_iter.wal_id; + let replayed_wal_count = last_wal_id + .saturating_sub(self.wal_id_range.start) + .saturating_add(1); + if replayed_wal_count.is_multiple_of(128) { + let metadata = table.metadata(); + info!( + "SlateDB WAL replay progress [replay_start_wal_id={}, replay_end_wal_id={}, replay_wal_count={}, last_replayed_wal_id={}, replayed_wal_count={}, replayed_entries={}, replayed_bytes={}]", + self.wal_id_range.start, + self.wal_id_range.end, + self.wal_id_range + .end + .saturating_sub(self.wal_id_range.start), + last_wal_id, + replayed_wal_count, + metadata.entry_num, + metadata.entries_size_in_bytes + ); + } + + let meta = table.metadata(); + let estimated_bytes = self + .table_store + .estimate_encoded_size_compacted(meta.entry_num, meta.entries_size_in_bytes); + if !table.is_empty() && estimated_bytes >= self.options.max_memtable_bytes { + self.current_iter.reset(); + break; + } + } + + self.advance_current_iter().await? + } + + if last_wal_id > 0 { + Ok(Some(ReplayedMemtable { + table, + last_tick: self.last_tick, + last_seq: self.last_seq, + last_wal_id, + })) + } else { + self.log_completion(); + Ok(None) + } + } + + fn log_completion(&mut self) { + if self.completion_logged { + return; + } + self.completion_logged = true; + info!( + "SlateDB exact WAL replay completed [source={:?}, replay_start_wal_id={}, replay_end_wal_id={}, fetched_objects={}, fetched_bytes={}, ranged_objects={}, decoded_objects={}, peak_concurrent_objects={}, peak_reserved_bytes={}]", + self.options.source, + self.wal_id_range.start, + self.wal_id_range.end, + self.fetched_objects, + self.fetched_bytes, + self.ranged_objects, + self.decoded_objects, + self.peak_concurrent_objects, + self.peak_reserved_bytes, + ); + } + + fn fail(&mut self, error: SlateDBError) -> SlateDBError { + for pending in self.pending_fetches.drain(..) { + pending.handle.abort(); + } + self.current_iter.current_iter.take(); + self.failed = Some(error.clone()); + error + } +} + +impl Drop for ExactWalReplayIterator { + fn drop(&mut self) { + for pending in &self.pending_fetches { + pending.handle.abort(); + } + } +} + +#[cfg(test)] +mod tests { + use super::{ + ExactWalReplayIterator, RuntimeWalReplayError, RuntimeWalReplayIterator, + RuntimeWalReplayOptions, WalReplayOptions, + }; + use crate::block_cache_policy::BlockCachePolicy; + use crate::bytes_range::BytesRange; + use crate::config::WalReplaySettings; + use crate::db_state::SsTableId; + use crate::format::sst::{BlockTransformer, SsTableFormat}; + use crate::iter::{IterationOrder, RowEntryIterator}; + use crate::manifest::ManifestCore; + use crate::mem_table::WritableKVTable; + use crate::object_stores::ObjectStores; + use crate::proptest_util::{rng, sample}; + use crate::tablestore::{DecodedWalSst, TableStore, TableStoreKind}; + use crate::test_utils::{FlakyObjectStore, GatedObjectStore, RecordingObjectStore}; + use crate::types::{RowEntry, ValueDeletable}; + use crate::{error::SlateDBError, test_utils}; + use async_trait::async_trait; + use bytes::Bytes; + use futures::stream::BoxStream; + use futures::FutureExt; + use object_store::memory::InMemory; + use object_store::path::Path; + use object_store::{ + CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, + ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, RenameOptions, + }; + use proptest::prelude::*; + use proptest::test_runner::TestRng; + use rand::Rng; + use std::cmp::min; + use std::collections::btree_map::Iter; + use std::collections::BTreeMap; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + use std::time::Duration; + use tokio::sync::Notify; + + struct HistoricalXorTransformer; + + #[async_trait] + impl BlockTransformer for HistoricalXorTransformer { + fn max_decoded_len(&self, encoded_len: usize) -> Option { + Some(encoded_len) + } + + async fn encode(&self, data: Bytes) -> Result { + Ok(Bytes::from( + data.iter().map(|byte| byte ^ 0xa5).collect::>(), + )) + } + + async fn decode(&self, data: Bytes) -> Result { + self.encode(data).await + } + } + + struct UnboundedTransformer { + decode_called: Arc, + } + + #[async_trait] + impl BlockTransformer for UnboundedTransformer { + fn max_decoded_len(&self, _encoded_len: usize) -> Option { + None + } + + async fn encode(&self, data: Bytes) -> Result { + Ok(data) + } + + async fn decode(&self, data: Bytes) -> Result { + self.decode_called.store(true, Ordering::SeqCst); + Ok(data) + } + } + + #[cfg(test)] + use crate::sst_builder::BlockFormat; + + #[derive(Debug)] + struct PausedFirstWalGetStore { + inner: Arc, + first_path: Path, + panic_path: Option, + first_started: Arc, + later_get_returned: Arc, + first_cancelled: Arc, + release_first: Arc, + } + + impl PausedFirstWalGetStore { + fn new(inner: Arc, first_path: Path) -> Self { + Self { + inner, + first_path, + panic_path: None, + first_started: Arc::new(AtomicBool::new(false)), + later_get_returned: Arc::new(AtomicBool::new(false)), + first_cancelled: Arc::new(AtomicBool::new(false)), + release_first: Arc::new(Notify::new()), + } + } + + fn panic_on(mut self, path: Path) -> Self { + self.panic_path = Some(path); + self + } + + async fn wait_until(flag: &AtomicBool) { + while !flag.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + } + } + + impl std::fmt::Display for PausedFirstWalGetStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "PausedFirstWalGetStore({})", self.inner) + } + } + + struct PendingGetGuard { + cancelled: Arc, + completed: bool, + } + + impl Drop for PendingGetGuard { + fn drop(&mut self) { + if !self.completed { + self.cancelled.store(true, Ordering::Release); + } + } + } + + #[async_trait] + impl ObjectStore for PausedFirstWalGetStore { + async fn get_opts( + &self, + location: &Path, + options: GetOptions, + ) -> object_store::Result { + if !options.head && self.panic_path.as_ref() == Some(location) { + panic!("injected WAL fetch panic for {location}"); + } + if !options.head && location == &self.first_path { + let mut guard = PendingGetGuard { + cancelled: Arc::clone(&self.first_cancelled), + completed: false, + }; + self.first_started.store(true, Ordering::Release); + self.release_first.notified().await; + let result = self.inner.get_opts(location, options).await; + guard.completed = true; + return result; + } + let result = self.inner.get_opts(location, options).await; + if location != &self.first_path { + self.later_get_returned.store(true, Ordering::Release); + } + result + } + + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + options: PutOptions, + ) -> object_store::Result { + self.inner.put_opts(location, payload, options).await + } + + async fn put_multipart_opts( + &self, + location: &Path, + options: PutMultipartOptions, + ) -> object_store::Result> { + self.inner.put_multipart_opts(location, options).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, object_store::Result>, + ) -> BoxStream<'static, object_store::Result> { + self.inner.delete_stream(locations) + } + + fn list( + &self, + prefix: Option<&Path>, + ) -> BoxStream<'static, object_store::Result> { + self.inner.list(prefix) + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, object_store::Result> { + self.inner.list_with_offset(prefix, offset) + } + + async fn list_with_delimiter( + &self, + prefix: Option<&Path>, + ) -> object_store::Result { + self.inner.list_with_delimiter(prefix).await + } + + async fn copy_opts( + &self, + from: &Path, + to: &Path, + options: CopyOptions, + ) -> object_store::Result<()> { + self.inner.copy_opts(from, to, options).await + } + + async fn rename_opts( + &self, + from: &Path, + to: &Path, + options: RenameOptions, + ) -> object_store::Result<()> { + self.inner.rename_opts(from, to, options).await + } + } + + impl ExactWalReplayIterator { + async fn all_wal_ids( + db_state: &ManifestCore, + options: WalReplayOptions, + table_store: Arc, + ) -> Result { + let wal_id_start = db_state.replay_after_wal_id + 1; + let wal_id_end = table_store + .last_seen_wal_id(db_state.replay_after_wal_id) + .await?; + let wal_id_range = wal_id_start..(wal_id_end + 1); + Self::range(wal_id_range, db_state, options, table_store).await + } + } + + #[tokio::test] + async fn runtime_replay_uses_range_reads_without_listing() { + let recording = Arc::new(RecordingObjectStore::new(Arc::new(InMemory::new()))); + let object_store: Arc = recording.clone(); + let table_store = test_table_store_with_object_store(object_store); + write_empty_wal(1, Arc::clone(&table_store)).await.unwrap(); + + let mut runtime = RuntimeWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + RuntimeWalReplayOptions::default(), + Arc::clone(&table_store), + ) + .unwrap(); + assert!(runtime.next().await.unwrap().is_some()); + assert!(runtime.next().await.unwrap().is_none()); + assert_eq!(recording.list_calls(), 0); + assert!(!recording.get_range_sizes().is_empty()); + + let mut exact = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + assert!(exact.next().await.unwrap().is_some()); + assert!(recording.list_calls() > 0); + } + + #[tokio::test] + async fn runtime_replay_treats_metadata_not_found_as_corruption_not_caught_up() { + let inner: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner)); + let object_store: Arc = gated.clone(); + let table_store = test_table_store_with_object_store(object_store); + let mut rng = rng::new_test_rng(None); + let entries = sample::table(&mut rng, 1, 16); + let mut entries = entries.iter(); + write_wal(1, 1, &mut entries, 1, Arc::clone(&table_store)) + .await + .unwrap(); + + // The initial HEAD and footer read establish that the exact object + // exists. A later metadata 404 must remain a replay error. + let prior_range_arrivals = gated.get_opts_gate.arrivals(); + gated.get_opts_gate.close(); + gated.get_opts_gate.admit(1); + let mut runtime = RuntimeWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + RuntimeWalReplayOptions::default(), + table_store, + ) + .unwrap(); + let replay = tokio::spawn(async move { runtime.next().await }); + tokio::time::timeout( + Duration::from_secs(2), + gated + .get_opts_gate + .wait_for_arrivals(prior_range_arrivals + 2), + ) + .await + .unwrap_or_else(|_| { + panic!( + "only {} range reads arrived", + gated.get_opts_gate.arrivals() + ) + }); + gated + .get_opts_gate + .set_error(|| object_store::Error::NotFound { + path: "injected-metadata-404".to_string(), + source: Box::new(std::io::Error::other("injected metadata 404")), + }); + gated.get_opts_gate.release(); + + let replay = tokio::time::timeout(Duration::from_secs(2), replay) + .await + .unwrap_or_else(|_| { + panic!( + "runtime replay remained blocked after {} range reads", + gated.get_opts_gate.arrivals() + ) + }) + .unwrap(); + let error = match replay { + Ok(_) => panic!("metadata 404 unexpectedly replayed"), + Err(error) => error, + }; + let RuntimeWalReplayError::Replay(error) = error else { + panic!("metadata 404 was incorrectly classified as caught up"); + }; + assert!(error.has_object_store_not_found()); + } + + #[tokio::test] + async fn runtime_replay_returns_corruption_for_a_truncated_footer_without_panicking() { + let inner: Arc = Arc::new(InMemory::new()); + let flaky = Arc::new(FlakyObjectStore::new(inner, 0).with_truncate_get_range_bytes(1, 2)); + let object_store: Arc = flaky; + let table_store = test_table_store_with_object_store(object_store); + let entries = BTreeMap::from([(Bytes::from_static(b"key"), Bytes::from_static(b"value"))]); + let mut entries = entries.iter(); + write_wal(1, 1, &mut entries, 1, Arc::clone(&table_store)) + .await + .unwrap(); + let mut runtime = RuntimeWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + RuntimeWalReplayOptions::default(), + table_store, + ) + .unwrap(); + + let replay = std::panic::AssertUnwindSafe(runtime.next()) + .catch_unwind() + .await; + let result = replay.expect("malformed object-store bytes must never panic"); + let error = match result { + Ok(_) => panic!("truncated footer unexpectedly replayed"), + Err(error) => error, + }; + let RuntimeWalReplayError::Replay(SlateDBError::CorruptSst { .. }) = error else { + panic!("truncated footer did not return typed corruption"); + }; + } + + #[tokio::test] + async fn should_replay_empty_wal() { + let table_store = test_table_store(); + write_empty_wal(1, Arc::clone(&table_store)).await.unwrap(); + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + + let Some(table) = replay_iter.next().await.unwrap() else { + panic!("Expected empty table to be returned from iterator") + }; + + assert_eq!(table.last_wal_id, 1); + assert_eq!(table.last_seq, 0); + assert!(table.table.is_empty()); + assert_eq!(table.last_tick, i64::MIN); + assert!(replay_iter.next().await.unwrap().is_none()); + } + + #[tokio::test] + async fn should_replay_zero_byte_wal_fence() { + let table_store = test_table_store(); + table_store.write_wal_fence(1).await.unwrap(); + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + + let Some(table) = replay_iter.next().await.unwrap() else { + panic!("Expected empty table to be returned from iterator") + }; + + assert_eq!(table.last_wal_id, 1); + assert_eq!(table.last_seq, 0); + assert!(table.table.is_empty()); + assert_eq!(table.last_tick, i64::MIN); + assert!(replay_iter.next().await.unwrap().is_none()); + } + + #[tokio::test] + async fn should_replay_zero_byte_wal_fence_before_real_wal() { + let table_store = test_table_store(); + table_store.write_wal_fence(1).await.unwrap(); + + let row = RowEntry::new_value(b"key", b"value", 1); + let mut builder = table_store.wal_table_builder(); + builder.add(row.clone()).await.unwrap(); + let encoded_sst = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(2), &encoded_sst) + .await + .unwrap(); + + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + + let Some(replayed_table) = replay_iter.next().await.unwrap() else { + panic!("Expected table to be returned from iterator") + }; + assert_eq!(replayed_table.last_wal_id, 2); + assert_eq!(replayed_table.last_seq, 1); + + let mut iter = replayed_table.table.table().iter(); + test_utils::assert_iterator(&mut iter, vec![row]).await; + assert!(replay_iter.next().await.unwrap().is_none()); + } + + #[tokio::test] + async fn should_replay_zero_byte_wal_fences_between_and_after_data_wals() { + let table_store = test_table_store(); + let first = RowEntry::new_value(b"first", b"value-1", 1); + let second = RowEntry::new_value(b"second", b"value-2", 2); + + for (wal_id, row) in [(1, first.clone()), (3, second.clone())] { + let mut builder = table_store.wal_table_builder(); + builder.add(row).await.unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + table_store.write_wal_fence(2).await.unwrap(); + table_store.write_wal_fence(4).await.unwrap(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..5, + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + + let Some(replayed) = replay_iter.next().await.unwrap() else { + panic!("expected replayed WALs"); + }; + assert_eq!(replayed.last_wal_id, 4); + assert_eq!(replayed.last_seq, 2); + let mut iter = replayed.table.table().iter(); + test_utils::assert_iterator(&mut iter, vec![first, second]).await; + assert!(replay_iter.next().await.unwrap().is_none()); + } + + #[tokio::test] + async fn should_replay_legacy_v1_wal_encoding() { + let format = SsTableFormat { + block_format: Some(BlockFormat::V1), + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let rows = vec![ + RowEntry::new_value(b"first", b"value-1", 1), + RowEntry::new_value(b"second", b"value-2", 2), + ]; + let mut builder = table_store.table_builder(); + for row in &rows { + builder.add(row.clone()).await.unwrap(); + } + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + let Some(replayed) = replay_iter.next().await.unwrap() else { + panic!("expected legacy WAL to replay"); + }; + let mut iter = replayed.table.table().iter(); + test_utils::assert_iterator(&mut iter, rows).await; + assert_eq!(replayed.last_wal_id, 1); + assert_eq!(replayed.last_seq, 2); + } + + #[tokio::test] + async fn historical_wal_goldens_match_range_and_full_object_decoders() { + for (name, bytes) in historical_wal_fixtures() { + let mut format = SsTableFormat::default(); + if name.ends_with("-xor") { + format.block_transformer = Some(Arc::new(HistoricalXorTransformer)); + } + let store: Arc = Arc::new(InMemory::new()); + let table_store = + test_table_store_with_format_and_object_store(format, Arc::clone(&store)); + let path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + store.put(&path, bytes.clone().into()).await.unwrap(); + + let handle = table_store.open_sst(&SsTableId::Wal(1)).await.unwrap(); + let index = table_store.read_index(&handle, false).await.unwrap(); + let block_count = index.borrow().block_meta().len(); + let range_blocks = table_store + .read_blocks_using_index(&handle, index, 0..block_count, false) + .await + .unwrap(); + let mut range_rows = Vec::new(); + for block in range_blocks { + let mut iterator = crate::block_iterator::DataBlockIterator::new( + block, + handle.format_version, + IterationOrder::Ascending, + ) + .unwrap(); + while let Some(row) = iterator.next().await.unwrap() { + range_rows.push(row); + } + } + + let decoded = table_store + .decode_wal_sst(1, bytes, 64 * 1024 * 1024) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = decoded else { + panic!("historical fixture {name} decoded as a fence"); + }; + let mut iterator = super::WalBlocksIterator::new( + Arc::clone(&table_store), + *wal, + 192 * 1024 * 1024, + 64 * 1024 * 1024, + ) + .unwrap(); + let mut full_rows = Vec::new(); + while let Some(row) = iterator.next().await.unwrap() { + full_rows.push(row); + } + + let ranged = table_store + .open_ranged_wal_sst(1, store.head(&path).await.unwrap().size, 64 * 1024 * 1024) + .await + .unwrap(); + let super::RangedWalSst::Data(wal) = ranged else { + panic!("historical fixture {name} range-decoded as a fence"); + }; + let mut iterator = super::WalBlocksIterator::new_ranged( + Arc::clone(&table_store), + *wal, + 192 * 1024 * 1024, + 64 * 1024 * 1024, + ) + .unwrap(); + let mut bounded_range_rows = Vec::new(); + while let Some(row) = iterator.next().await.unwrap() { + bounded_range_rows.push(row); + } + + assert_eq!(full_rows, range_rows, "fixture {name}"); + assert_eq!(bounded_range_rows, range_rows, "fixture {name}"); + assert_historical_rows(name, &full_rows); + + let expected_filtered = range_rows + .iter() + .filter(|row| row.seq > 8) + .cloned() + .collect::>(); + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + min_seq: Some(8), + ..WalReplayOptions::default() + }, + Arc::clone(&table_store), + ) + .await + .unwrap(); + let replayed = replay.next().await.unwrap().unwrap(); + let mut replayed_rows = Vec::new(); + let mut replayed_iter = replayed.table.table().iter(); + while let Some(row) = replayed_iter.next().await.unwrap() { + replayed_rows.push(row); + } + assert_eq!(replayed_rows, expected_filtered, "filtered fixture {name}"); + assert!(replay.next().await.unwrap().is_none()); + + let mut runtime_replay = RuntimeWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + RuntimeWalReplayOptions { + min_seq: Some(8), + ..RuntimeWalReplayOptions::default() + }, + Arc::clone(&table_store), + ) + .unwrap(); + let runtime_table = runtime_replay.next().await.unwrap().unwrap(); + let mut runtime_rows = Vec::new(); + let mut runtime_iter = runtime_table.table.table().iter(); + while let Some(row) = runtime_iter.next().await.unwrap() { + runtime_rows.push(row); + } + assert_eq!(runtime_rows, expected_filtered, "runtime fixture {name}"); + assert!(runtime_replay.next().await.unwrap().is_none()); + } + } + + #[tokio::test] + async fn historical_wal_goldens_report_identical_checksum_errors() { + for (name, bytes) in historical_wal_fixtures() { + let mut format = SsTableFormat::default(); + if name.ends_with("-xor") { + format.block_transformer = Some(Arc::new(HistoricalXorTransformer)); + } + let store: Arc = Arc::new(InMemory::new()); + let table_store = + test_table_store_with_format_and_object_store(format, Arc::clone(&store)); + let path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + let mut corrupted = bytes.to_vec(); + corrupted[0] ^= 1; + let corrupted = Bytes::from(corrupted); + store.put(&path, corrupted.clone().into()).await.unwrap(); + + let handle = table_store.open_sst(&SsTableId::Wal(1)).await.unwrap(); + let index = table_store.read_index(&handle, false).await.unwrap(); + let range_error = table_store + .read_blocks_using_index( + &handle, + index.clone(), + 0..index.borrow().block_meta().len(), + false, + ) + .await; + let Err(range_error) = range_error else { + panic!("corrupted historical fixture {name} passed range decoding"); + }; + + let decoded = table_store + .decode_wal_sst(1, corrupted, 64 * 1024 * 1024) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = decoded else { + panic!("corrupted historical fixture {name} decoded as a fence"); + }; + let mut full_iterator = super::WalBlocksIterator::new( + Arc::clone(&table_store), + *wal, + 192 * 1024 * 1024, + 64 * 1024 * 1024, + ) + .unwrap(); + let full_error = full_iterator.next().await.unwrap_err(); + + let mut runtime_replay = RuntimeWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + RuntimeWalReplayOptions::default(), + Arc::clone(&table_store), + ) + .unwrap(); + let runtime_error = match runtime_replay.next().await { + Ok(_) => panic!("corrupted historical fixture {name} passed runtime replay"), + Err(error) => error, + }; + let RuntimeWalReplayError::Replay(runtime_error) = runtime_error else { + panic!("corrupted historical fixture {name} looked missing"); + }; + + assert_eq!( + full_error.maybe_validation_retry_reason(), + range_error.maybe_validation_retry_reason(), + "fixture {name}" + ); + assert_eq!( + full_error.maybe_validation_retry_reason(), + Some(crate::error::RetryReason::CrcMismatch), + "fixture {name}" + ); + assert_eq!( + runtime_error.maybe_validation_retry_reason(), + full_error.maybe_validation_retry_reason(), + "runtime fixture {name}" + ); + } + } + + #[cfg(feature = "snappy")] + #[tokio::test] + async fn should_replay_compressed_wal() { + let format = SsTableFormat { + compression_codec: Some(crate::config::CompressionCodec::Snappy), + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let row = RowEntry::new_value(b"key", &[b'x'; 16 * 1024], 1); + let mut builder = table_store.wal_table_builder(); + builder.add(row.clone()).await.unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + let Some(replayed) = replay_iter.next().await.unwrap() else { + panic!("expected compressed WAL to replay"); + }; + let mut iter = replayed.table.table().iter(); + test_utils::assert_iterator(&mut iter, vec![row]).await; + } + + #[tokio::test] + async fn should_finish_empty_replay_range_without_object_io() { + let recording = Arc::new(RecordingObjectStore::new(Arc::new(InMemory::new()))); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut replay_iter = ExactWalReplayIterator::range( + 7..7, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + + assert!(replay_iter.next().await.unwrap().is_none()); + assert!(recording.get_kinds(false).is_empty()); + assert!(recording.get_kinds(true).is_empty()); + } + + #[tokio::test] + async fn should_replay_all_entries() { + let table_store = test_table_store(); + let mut rng = rng::new_test_rng(None); + let entries = sample::table(&mut rng, 1000, 10); + let next_wal_id = write_wals(&entries, 1, &mut rng, 200, Arc::clone(&table_store)) + .await + .unwrap(); + + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( + &ManifestCore::new(), + WalReplayOptions::default(), + Arc::clone(&table_store), + ) + .await + .unwrap(); + + let Some(replayed_table) = replay_iter.next().await.unwrap() else { + panic!("Expected table to be returned from iterator") + }; + assert_eq!(replayed_table.last_wal_id + 1, next_wal_id); + + let mut imm_table_iter = replayed_table.table.table().iter(); + test_utils::assert_ranged_kv_scan( + &entries, + &BytesRange::from(..), + IterationOrder::Ascending, + &mut imm_table_iter, + ) + .await; + assert!(replay_iter.next().await.unwrap().is_none()); + } + + #[tokio::test] + async fn should_issue_one_full_get_per_wal() { + let inner: Arc = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let table_store = test_table_store_with_object_store(recording.clone()); + for wal_id in 1..=3 { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + recording.clear(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..4, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + while replay_iter.next().await.unwrap().is_some() {} + + assert_eq!(recording.get_kinds(false).len(), 3); + assert!(recording.get_kinds(true).is_empty()); + assert_eq!( + recording.get_sst_types(false), + vec![Some(crate::db_state::SstType::Wal); 3] + ); + assert_eq!(recording.get_retries(false), vec![None; 3]); + } + + #[tokio::test] + async fn should_fetch_replay_objects_only_from_dedicated_wal_store() { + let main = Arc::new(RecordingObjectStore::new(Arc::new(InMemory::new()))); + let wal = Arc::new(RecordingObjectStore::new(Arc::new(InMemory::new()))); + let table_store = Arc::new(TableStore::new( + ObjectStores::new(main.clone(), Some(wal.clone())), + SsTableFormat::default(), + Path::from("/tmp/test_kv_store"), + None, + TableStoreKind::Main, + BlockCachePolicy::default(), + )); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + main.clear(); + wal.clear(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + while replay_iter.next().await.unwrap().is_some() {} + + assert!(main.get_kinds(false).is_empty()); + assert!(main.get_kinds(true).is_empty()); + assert_eq!(wal.get_kinds(false), vec![Some(TableStoreKind::Main)]); + assert!(wal.get_kinds(true).is_empty()); + } + + #[tokio::test] + async fn should_refetch_the_full_wal_once_after_checksum_failure() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner.clone())); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + let metadata = table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .pop() + .unwrap(); + let mut corrupted = inner + .get(&metadata.metadata.location) + .await + .unwrap() + .bytes() + .await + .unwrap() + .to_vec(); + corrupted[0] ^= 1; + inner + .put(&metadata.metadata.location, corrupted.into()) + .await + .unwrap(); + recording.clear(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + assert!(matches!( + replay_iter.next().await, + Err(SlateDBError::ChecksumMismatch { .. }) + )); + assert_eq!(recording.get_kinds(false).len(), 2); + assert_eq!(recording.get_retries(false)[0], None); + assert!(recording.get_retries(false)[1].is_some()); + } + + #[tokio::test] + async fn should_recover_when_validation_refetch_returns_valid_wal() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + let valid = encoded.remaining_as_bytes(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + recording.clear(); + + let mut corrupted = valid.to_vec(); + corrupted[0] ^= 1; + let initially_decoded = table_store + .decode_wal_sst(1, Bytes::from(corrupted), 64 * 1024 * 1024) + .await + .unwrap(); + let DecodedWalSst::Data(initially_decoded) = initially_decoded else { + panic!("nonempty WAL decoded as a fence"); + }; + let Err(initial_error) = table_store + .decode_wal_block(&initially_decoded, 0, 64 * 1024 * 1024) + .await + else { + panic!("corrupt initial WAL block unexpectedly decoded"); + }; + drop(initially_decoded); + let reason = initial_error + .maybe_validation_retry_reason() + .expect("corruption must be classified for a validation retry"); + let decoded = table_store + .refetch_wal_sst_after_validation( + 1, + valid.len(), + 192 * 1024 * 1024, + 64 * 1024 * 1024, + reason, + ) + .await + .unwrap(); + + assert!(matches!(decoded, DecodedWalSst::Data(_))); + assert_eq!(recording.get_kinds(false).len(), 1); + assert!(recording.get_retries(false)[0].is_some()); + } + + #[tokio::test] + async fn should_fail_without_panicking_for_every_truncated_wal_prefix() { + let inner = Arc::new(InMemory::new()); + let table_store = test_table_store_with_object_store(inner.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=8 { + builder + .add(RowEntry::new_value( + format!("key-{seq}").as_bytes(), + &[b'x'; 128], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + let valid = encoded.remaining_as_bytes(); + let location = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + + for prefix_len in 1..valid.len() { + let truncated = valid.slice(..prefix_len); + inner + .put(&location, truncated.clone().into()) + .await + .unwrap(); + let result = std::panic::AssertUnwindSafe(table_store.decode_wal_sst( + 1, + truncated, + 64 * 1024 * 1024, + )) + .catch_unwind() + .await; + let decode = result + .unwrap_or_else(|_| panic!("WAL decoder panicked for prefix length {prefix_len}")); + assert!( + decode.is_err(), + "truncated WAL prefix {prefix_len} decoded successfully" + ); + } + } + + #[tokio::test] + async fn should_not_skip_a_missing_wal_id() { + let table_store = test_table_store(); + for wal_id in [1, 3] { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + + let mut replay_iter = ExactWalReplayIterator::range( + 1..4, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + let Err(error) = replay_iter.next().await else { + panic!("missing WAL must fail replay"); + }; + assert!(error.has_object_store_not_found()); + } + + #[tokio::test] + async fn should_remain_failed_after_a_replay_error() { + let table_store = test_table_store(); + for wal_id in [1, 3] { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + + let mut replay_iter = ExactWalReplayIterator::range( + 1..4, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + let Err(first) = replay_iter.next().await else { + panic!("missing WAL must fail replay"); + }; + let Err(second) = replay_iter.next().await else { + panic!("failed replay iterator must remain failed"); + }; + assert!(first.has_object_store_not_found()); + assert!(second.has_object_store_not_found()); + } + + #[tokio::test] + async fn should_construct_huge_sparse_range_without_materializing_every_wal_id() { + let inner: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner)); + gated.get_opts_gate.close(); + let table_store = test_table_store_with_object_store(gated.clone()); + + let replay_iter = tokio::time::timeout( + Duration::from_secs(1), + ExactWalReplayIterator::range( + 1..u64::MAX, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: 2, + }, + ..WalReplayOptions::default() + }, + table_store, + ), + ) + .await + .expect("replay range construction allocated proportional to the WAL ID span") + .unwrap(); + + gated.get_opts_gate.wait_for_arrivals(1).await; + assert_eq!(gated.get_opts_gate.arrivals(), 1); + drop(replay_iter); + } + + #[tokio::test] + async fn should_fail_if_wal_size_changes_after_listing() { + let inner = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner.clone())); + let table_store = test_table_store_with_object_store(gated.clone()); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + gated.get_opts_gate.close(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + gated.get_opts_gate.wait_for_arrivals(1).await; + let location = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + inner + .put(&location, Bytes::from_static(b"changed").into()) + .await + .unwrap(); + gated.get_opts_gate.release(); + + let Err(error) = replay_iter.next().await else { + panic!("changed WAL size must fail replay"); + }; + assert!(matches!(error, SlateDBError::WalDataError(_))); + } + + #[tokio::test] + async fn should_prefetch_up_to_the_object_limit() { + let inner: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner)); + let table_store = test_table_store_with_object_store(gated.clone()); + for wal_id in 1..=4 { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + gated.get_opts_gate.close(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..5, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: 1024 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + gated.get_opts_gate.wait_for_arrivals(2).await; + assert_eq!(gated.get_opts_gate.arrivals(), 2); + + gated.get_opts_gate.release(); + while replay_iter.next().await.unwrap().is_some() {} + assert_eq!(gated.get_opts_gate.arrivals(), 4); + } + + #[cfg(any( + feature = "snappy", + feature = "zlib", + feature = "lz4", + feature = "zstd" + ))] + #[tokio::test] + async fn compressed_wals_cannot_expand_past_total_replay_memory_budget() { + let mut codecs = Vec::new(); + #[cfg(feature = "snappy")] + codecs.push(crate::config::CompressionCodec::Snappy); + #[cfg(feature = "zlib")] + codecs.push(crate::config::CompressionCodec::Zlib); + #[cfg(feature = "lz4")] + codecs.push(crate::config::CompressionCodec::Lz4); + #[cfg(feature = "zstd")] + codecs.push(crate::config::CompressionCodec::Zstd); + + for codec in codecs { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let format = SsTableFormat { + compression_codec: Some(codec), + ..SsTableFormat::default() + }; + let table_store = + test_table_store_with_format_and_object_store(format, recording.clone()); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"large", &vec![b'x'; 1024 * 1024], 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + recording.clear(); + + let total_budget = 256 * 1024; + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 64, + max_inflight_bytes: total_budget, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + assert!(replay.peak_reserved_bytes + replay.working_memory_limit <= total_budget); + let Err(first_error) = replay.next().await else { + panic!("{codec:?} WAL expanded past its memory budget"); + }; + let Err(second_error) = replay.next().await else { + panic!("{codec:?} memory-limit failure was not sticky"); + }; + assert!(matches!( + first_error, + SlateDBError::WalReplayMemoryLimitExceeded { .. } + )); + assert!(matches!( + second_error, + SlateDBError::WalReplayMemoryLimitExceeded { .. } + )); + assert_eq!(recording.get_kinds(false).len(), 1, "{codec:?}"); + } + } + + #[tokio::test] + async fn oversized_encoded_wal_uses_bounded_range_replay() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=64 { + builder + .add(RowEntry::new_value( + format!("large-{seq:03}").as_bytes(), + &[b'x'; 2048], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + recording.clear(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 64, + max_inflight_bytes: 64 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + let replayed = replay + .next() + .await + .expect("oversized WAL should remain recoverable") + .expect("oversized WAL should produce a replayed table"); + let mut rows = replayed.table.table().iter(); + let mut row_count = 0; + while let Some(row) = rows.next().await.unwrap() { + assert_eq!(row.value.len(), 2048); + row_count += 1; + } + assert_eq!(row_count, 64); + assert!(replay.next().await.unwrap().is_none()); + assert!(recording.get_kinds(false).len() > 1); + assert!( + recording + .get_range_sizes() + .into_iter() + .all(|range_bytes| range_bytes <= 8 * 1024), + "range replay must not fetch a section larger than its working partition" + ); + } + + #[tokio::test] + async fn oversized_range_replay_rejects_large_metadata_before_fetching_it() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner.clone())); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=64 { + builder + .add(RowEntry::new_value( + format!("large-metadata-{seq:03}").as_bytes(), + &[b'x'; 2048], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + let metadata = table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .pop() + .unwrap() + .metadata; + let mut bytes = inner + .get(&metadata.location) + .await + .unwrap() + .bytes() + .await + .unwrap() + .to_vec(); + let footer_start = bytes.len() - 10; + bytes[footer_start..footer_start + 8].copy_from_slice(&0_u64.to_be_bytes()); + inner + .put(&metadata.location, Bytes::from(bytes).into()) + .await + .unwrap(); + recording.clear(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 4, + max_inflight_bytes: 64 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + let Err(error) = replay.next().await else { + panic!("oversized metadata was not rejected"); + }; + + assert!(matches!( + error, + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL metadata", + .. + } + )); + assert!( + recording + .get_range_sizes() + .into_iter() + .all(|range_bytes| range_bytes <= 8 * 1024), + "metadata validation must reject before an oversized range GET" + ); + } + + #[tokio::test] + async fn oversized_range_replay_rejects_large_index_before_fetching_it() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let format = SsTableFormat { + block_size: 128, + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format_and_object_store(format, recording.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=1024 { + builder + .add(RowEntry::new_value( + format!("large-index-{seq:04}-{}", "k".repeat(64)).as_bytes(), + &[b'x'; 128], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + let index_len = usize::try_from(encoded.info.index_len).unwrap(); + let total_budget = 64 * 1024; + let metadata_memory_limit = total_budget / 8; + assert!(encoded.remaining_len() > total_budget * 3 / 4); + assert!(index_len > metadata_memory_limit); + assert!(encoded.unconsumed_blocks.len() > 100); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + recording.clear(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 4, + max_inflight_bytes: total_budget, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + let Err(error) = replay.next().await else { + panic!("oversized WAL index was not rejected"); + }; + + assert!(matches!( + error, + SlateDBError::WalReplayMemoryLimitExceeded { + kind: "encoded and decoded WAL index", + .. + } + )); + assert!( + recording + .get_range_sizes() + .into_iter() + .all(|range_bytes| range_bytes <= metadata_memory_limit as u64), + "index validation must reject before the oversized index range GET" + ); + } + + #[tokio::test] + async fn oversized_range_replay_supports_codecs_and_block_transformations() { + let mut codecs = Vec::with_capacity(5); + codecs.push(None); + #[cfg(feature = "snappy")] + codecs.push(Some(crate::config::CompressionCodec::Snappy)); + #[cfg(feature = "zlib")] + codecs.push(Some(crate::config::CompressionCodec::Zlib)); + #[cfg(feature = "lz4")] + codecs.push(Some(crate::config::CompressionCodec::Lz4)); + #[cfg(feature = "zstd")] + codecs.push(Some(crate::config::CompressionCodec::Zstd)); + + for codec in codecs { + for transformed in [false, true] { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner)); + let format = SsTableFormat { + compression_codec: codec, + block_transformer: transformed + .then(|| Arc::new(HistoricalXorTransformer) as Arc), + ..SsTableFormat::default() + }; + let table_store = + test_table_store_with_format_and_object_store(format, recording.clone()); + let mut builder = table_store.wal_table_builder(); + let mut test_rng = rng::new_test_rng(None); + for seq in 1..=64 { + let mut value = vec![0_u8; 2048]; + test_rng.fill(value.as_mut_slice()); + builder + .add(RowEntry::new_value( + format!("codec-{seq:03}").as_bytes(), + &value, + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + assert!(encoded.remaining_len() > 96 * 1024, "{codec:?}"); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + recording.clear(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 4, + max_inflight_bytes: 128 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + let replayed = replay.next().await.unwrap().unwrap(); + assert_eq!(replayed.table.metadata().entry_num, 64, "{codec:?}"); + assert!(replay.next().await.unwrap().is_none()); + assert!( + recording.get_kinds(false).len() > 1, + "{codec:?}, transformed={transformed}" + ); + } + } + } + + #[tokio::test] + async fn exhausted_block_is_dropped_before_decoding_the_next_block() { + let format = SsTableFormat { + block_size: 64, + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=3 { + builder + .add(RowEntry::new_value( + format!("key-{seq}").as_bytes(), + &[b'x'; 48], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + let bytes = table_store + .read_wal_sst_bytes(1, None, encoded.remaining_len()) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = table_store + .decode_wal_sst(1, bytes, 1024 * 1024) + .await + .unwrap() + else { + panic!("data WAL decoded as a fence"); + }; + let probe = Arc::new(super::BlockLifetimeProbe::default()); + let mut iterator = + super::WalBlocksIterator::new(Arc::clone(&table_store), *wal, 1024 * 1024, 1024 * 1024) + .unwrap() + .observe_block_lifetimes_with(Arc::clone(&probe)); + + while iterator.next().await.unwrap().is_some() {} + + assert_eq!(probe.active.load(Ordering::SeqCst), 0); + assert_eq!(probe.peak.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn cancelling_wal_block_iteration_releases_the_current_block() { + let format = SsTableFormat { + block_size: 64, + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", &[b'x'; 48], 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + let bytes = table_store + .read_wal_sst_bytes(1, None, encoded.remaining_len()) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = table_store + .decode_wal_sst(1, bytes, 1024 * 1024) + .await + .unwrap() + else { + panic!("data WAL decoded as a fence"); + }; + let probe = Arc::new(super::BlockLifetimeProbe::default()); + let mut iterator = + super::WalBlocksIterator::new(table_store, *wal, 1024 * 1024, 1024 * 1024) + .unwrap() + .observe_block_lifetimes_with(Arc::clone(&probe)); + assert!(iterator.next().await.unwrap().is_some()); + assert_eq!(probe.active.load(Ordering::SeqCst), 1); + + drop(iterator); + + assert_eq!(probe.active.load(Ordering::SeqCst), 0); + assert_eq!(probe.peak.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn decoding_failure_drops_the_exhausted_block_first() { + let format = SsTableFormat { + block_size: 64, + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=3 { + builder + .add(RowEntry::new_value( + format!("key-{seq}").as_bytes(), + &[b'x'; 48], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + let original = encoded.remaining_as_bytes(); + let DecodedWalSst::Data(original_wal) = table_store + .decode_wal_sst(1, original.clone(), 1024 * 1024) + .await + .unwrap() + else { + panic!("data WAL decoded as a fence"); + }; + let second_block_offset = original_wal.index.borrow().block_meta().get(1).offset() as usize; + let mut corrupted = original.to_vec(); + corrupted[second_block_offset] ^= 1; + let DecodedWalSst::Data(wal) = table_store + .decode_wal_sst(1, Bytes::from(corrupted), 1024 * 1024) + .await + .unwrap() + else { + panic!("data WAL decoded as a fence"); + }; + let probe = Arc::new(super::BlockLifetimeProbe::default()); + let mut iterator = + super::WalBlocksIterator::new(table_store, *wal, 1024 * 1024, 1024 * 1024) + .unwrap() + .observe_block_lifetimes_with(Arc::clone(&probe)); + + assert!(iterator.next().await.unwrap().is_some()); + let Err(error) = iterator.next().await else { + panic!("corrupt second block did not fail"); + }; + + assert!(matches!(error, SlateDBError::ChecksumMismatch { .. })); + assert_eq!(probe.active.load(Ordering::SeqCst), 0); + assert_eq!(probe.peak.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn oversized_range_replay_never_exposes_a_partial_wal_after_late_corruption() { + let inner = Arc::new(InMemory::new()); + let table_store = test_table_store_with_object_store(inner.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=64 { + builder + .add(RowEntry::new_value( + format!("corrupt-{seq:03}").as_bytes(), + &[b'x'; 2048], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + assert!(encoded.remaining_len() > 96 * 1024); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + let handle = table_store.open_sst(&SsTableId::Wal(1)).await.unwrap(); + let index = table_store.read_index(&handle, false).await.unwrap(); + let second_block_offset = index.borrow().block_meta().get(1).offset() as usize; + let metadata = table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .pop() + .unwrap() + .metadata; + let mut bytes = inner + .get(&metadata.location) + .await + .unwrap() + .bytes() + .await + .unwrap() + .to_vec(); + bytes[second_block_offset] ^= 1; + inner + .put(&metadata.location, Bytes::from(bytes).into()) + .await + .unwrap(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 4, + max_inflight_bytes: 128 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + let Err(first) = replay.next().await else { + panic!("corrupt oversized WAL was partially exposed"); + }; + let Err(second) = replay.next().await else { + panic!("corrupt oversized WAL failure was not sticky"); + }; + assert!(matches!(first, SlateDBError::ChecksumMismatch { .. })); + assert!(matches!(second, SlateDBError::ChecksumMismatch { .. })); + } + + #[tokio::test] + async fn unbounded_transform_is_rejected_before_decode() { + let decode_called = Arc::new(AtomicBool::new(false)); + let format = SsTableFormat { + block_transformer: Some(Arc::new(UnboundedTransformer { + decode_called: Arc::clone(&decode_called), + })), + ..SsTableFormat::default() + }; + let table_store = test_table_store_with_format(format); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + assert!(matches!( + replay.next().await, + Err(SlateDBError::BlockTransformError) + )); + assert!(!decode_called.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn later_checksum_valid_structural_corruption_returns_no_partial_wal() { + let inner = Arc::new(InMemory::new()); + let recording = Arc::new(RecordingObjectStore::new(inner.clone())); + let table_store = test_table_store_with_object_store(recording.clone()); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=12 { + builder + .add(RowEntry::new_value( + format!("key-{seq:02}").as_bytes(), + &[b'x'; 1024], + seq, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + let mut bytes = encoded.remaining_as_bytes().to_vec(); + let decoded = table_store + .decode_wal_sst(1, Bytes::copy_from_slice(&bytes), 64 * 1024 * 1024) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = decoded else { + panic!("data WAL decoded as a fence"); + }; + let block_meta = wal.index.borrow().block_meta(); + assert!(block_meta.len() >= 2); + let block_index = block_meta.len() - 1; + let start = block_meta.get(block_index).offset() as usize; + let end = wal.info.filter_offset as usize; + drop(wal); + let payload_end = end - crate::format::sst::CHECKSUM_SIZE; + bytes[payload_end - 2..payload_end].copy_from_slice(&0u16.to_be_bytes()); + let checksum = crc32fast::hash(&bytes[start..payload_end]); + bytes[payload_end..end].copy_from_slice(&checksum.to_be_bytes()); + let path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + inner.put(&path, Bytes::from(bytes).into()).await.unwrap(); + recording.clear(); + + let mut replay = ExactWalReplayIterator::range( + 1..2, + &ManifestCore::new(), + WalReplayOptions::default(), + table_store, + ) + .await + .unwrap(); + assert!(replay.next().await.is_err()); + assert!(replay.next().await.is_err()); + assert_eq!(recording.get_kinds(false).len(), 1); + } + + #[tokio::test] + async fn checksum_valid_metadata_index_block_and_row_mutations_never_panic() { + let table_store = test_table_store(); + let mut builder = table_store.wal_table_builder(); + for seq in 1..=16 { + builder + .add(RowEntry::new( + format!("key-{seq:02}").into(), + if seq % 3 == 0 { + ValueDeletable::Merge(Bytes::from_static(b"merge")) + } else if seq % 5 == 0 { + ValueDeletable::Tombstone + } else { + ValueDeletable::Value(Bytes::from(vec![seq as u8; 256])) + }, + seq, + (seq % 2 == 0).then_some(seq as i64), + (seq % 4 == 0).then_some((seq * 10) as i64), + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap().remaining_as_bytes(); + let decoded = table_store + .decode_wal_sst(1, encoded.clone(), 64 * 1024 * 1024) + .await + .unwrap(); + let DecodedWalSst::Data(wal) = decoded else { + panic!("data WAL decoded as a fence"); + }; + let metadata_start = u64::from_be_bytes( + encoded[encoded.len() - 10..encoded.len() - 2] + .try_into() + .unwrap(), + ) as usize; + let metadata_end = encoded.len() - 10; + let index_start = wal.info.index_offset as usize; + let index_end = index_start + wal.info.index_len as usize; + let block_meta = wal.index.borrow().block_meta(); + let block_start = block_meta.get(0).offset() as usize; + let block_end = if block_meta.len() > 1 { + block_meta.get(1).offset() as usize + } else { + wal.info.filter_offset as usize + }; + drop(wal); - let meta = table.metadata(); - let estimated_bytes = self - .table_store - .estimate_encoded_size_compacted(meta.entry_num, meta.entries_size_in_bytes); - if !table.is_empty() && estimated_bytes >= self.options.max_memtable_bytes { - self.current_iter.reset(); - break; + let sections = [ + ("metadata", metadata_start, metadata_end), + ("index", index_start, index_end), + ("block-and-row", block_start, block_end), + ]; + for (section_name, start, end) in sections { + let payload_end = end - crate::format::sst::CHECKSUM_SIZE; + assert!(payload_end > start); + let mut structural_errors = 0; + for mutation in 0..256usize { + let mut mutated = encoded.to_vec(); + let position = start + mutation % (payload_end - start); + mutated[position] ^= 1u8 << (mutation % 8); + let checksum = crc32fast::hash(&mutated[start..payload_end]); + mutated[payload_end..end].copy_from_slice(&checksum.to_be_bytes()); + let table_store = Arc::clone(&table_store); + let decode = std::panic::AssertUnwindSafe(async move { + let decoded = table_store + .decode_wal_sst(1, Bytes::from(mutated), 64 * 1024 * 1024) + .await?; + let DecodedWalSst::Data(wal) = decoded else { + return Err(SlateDBError::CorruptSst { + reason: "mutated data WAL decoded as a fence", + path: None, + }); + }; + let mut iterator = super::WalBlocksIterator::new( + Arc::clone(&table_store), + *wal, + 192 * 1024 * 1024, + 64 * 1024 * 1024, + )?; + while iterator.next().await?.is_some() {} + Ok::<(), SlateDBError>(()) + }) + .catch_unwind() + .await; + let result = decode.unwrap_or_else(|_| { + panic!("checksum-valid {section_name} mutation {mutation} panicked") + }); + if result.is_err() { + structural_errors += 1; } } - - self.maybe_load_next_iter(); - self.advance_current_iter().await? + assert!( + structural_errors > 0, + "{section_name} mutations never reached a structural error" + ); } + } - if last_wal_id > 0 { - Ok(Some(ReplayedMemtable { - table, - last_tick: self.last_tick, - last_seq: self.last_seq, - last_wal_id, - })) - } else { - Ok(None) + proptest! { + #![proptest_config(ProptestConfig::with_cases(512))] + + #[test] + fn checksum_valid_wal_structural_fuzz_never_panics( + section in 0usize..3, + mutations in prop::collection::vec((any::(), any::()), 1..=16), + ) { + let (_, encoded) = historical_wal_fixtures() + .into_iter() + .find(|(name, _)| *name == "v2-none") + .unwrap(); + let table_store = test_table_store(); + let decoded = futures::executor::block_on( + table_store.decode_wal_sst(1, encoded.clone(), 64 * 1024 * 1024), + ) + .unwrap(); + let DecodedWalSst::Data(wal) = decoded else { + panic!("historical data WAL decoded as a fence"); + }; + let metadata_start = u64::from_be_bytes( + encoded[encoded.len() - 10..encoded.len() - 2] + .try_into() + .unwrap(), + ) as usize; + let metadata_end = encoded.len() - 10; + let index_start = wal.info.index_offset as usize; + let index_end = index_start + wal.info.index_len as usize; + let block_meta = wal.index.borrow().block_meta(); + let block_start = block_meta.get(0).offset() as usize; + let block_end = if block_meta.len() > 1 { + block_meta.get(1).offset() as usize + } else { + wal.info.filter_offset as usize + }; + let sections = [ + (metadata_start, metadata_end), + (index_start, index_end), + (block_start, block_end), + ]; + let (start, end) = sections[section]; + let payload_end = end - crate::format::sst::CHECKSUM_SIZE; + let mut mutated = encoded.to_vec(); + for (offset, value) in mutations { + let position = start + offset % (payload_end - start); + mutated[position] ^= value.max(1); + } + let checksum = crc32fast::hash(&mutated[start..payload_end]); + mutated[payload_end..end].copy_from_slice(&checksum.to_be_bytes()); + + let decode = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + futures::executor::block_on(async { + let decoded = table_store + .decode_wal_sst(1, Bytes::from(mutated), 64 * 1024 * 1024) + .await?; + let DecodedWalSst::Data(wal) = decoded else { + return Err(SlateDBError::CorruptSst { + reason: "mutated data WAL decoded as a fence", + path: None, + }); + }; + let mut iterator = super::WalBlocksIterator::new( + Arc::clone(&table_store), + *wal, + 192 * 1024 * 1024, + 64 * 1024 * 1024, + )?; + while iterator.next().await?.is_some() {} + Ok::<(), SlateDBError>(()) + }) + })); + prop_assert!(decode.is_ok()); } } -} -#[cfg(test)] -mod tests { - use super::{WalReplayIterator, WalReplayOptions}; - use crate::block_cache_policy::BlockCachePolicy; - use crate::bytes_range::BytesRange; - use crate::db_state::SsTableId; - use crate::format::sst::SsTableFormat; - use crate::iter::{IterationOrder, RowEntryIterator}; - use crate::manifest::ManifestCore; - use crate::mem_table::WritableKVTable; - use crate::object_stores::ObjectStores; - use crate::proptest_util::{rng, sample}; - use crate::tablestore::{TableStore, TableStoreKind}; - use crate::types::RowEntry; - use crate::{error::SlateDBError, test_utils}; - use bytes::Bytes; - use object_store::memory::InMemory; - use object_store::path::Path; - use object_store::ObjectStore; - use proptest::test_runner::TestRng; - use rand::Rng; - use std::cmp::min; - use std::collections::btree_map::Iter; - use std::collections::BTreeMap; - use std::sync::Arc; + #[tokio::test] + async fn should_apply_wals_in_id_order_when_later_get_returns_first() { + let inner: Arc = Arc::new(InMemory::new()); + let first_path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + let reordered = Arc::new(PausedFirstWalGetStore::new(inner, first_path)); + let table_store = test_table_store_with_object_store(reordered.clone()); + for wal_id in 1..=2 { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + &[b'x'; 128], + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } - impl WalReplayIterator { - async fn all_wal_ids( - db_state: &ManifestCore, - options: WalReplayOptions, - table_store: Arc, - ) -> Result { - let wal_id_start = db_state.replay_after_wal_id + 1; - let wal_id_end = table_store - .last_seen_wal_id(db_state.replay_after_wal_id) - .await?; - let wal_id_range = wal_id_start..(wal_id_end + 1); - Self::range(wal_id_range, db_state, options, table_store).await + let mut replay_iter = ExactWalReplayIterator::range( + 1..3, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: 1024 * 1024, + }, + max_memtable_bytes: 1, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + PausedFirstWalGetStore::wait_until(&reordered.first_started).await; + PausedFirstWalGetStore::wait_until(&reordered.later_get_returned).await; + + let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel(); + let replay_task = tokio::spawn(async move { + while let Some(replayed) = replay_iter.next().await.unwrap() { + result_tx.send(replayed.last_wal_id).unwrap(); + } + }); + for _ in 0..100 { + tokio::task::yield_now().await; } + assert!(matches!( + result_rx.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + + reordered.release_first.notify_one(); + assert_eq!(result_rx.recv().await, Some(1)); + assert_eq!(result_rx.recv().await, Some(2)); + replay_task.await.unwrap(); } #[tokio::test] - async fn should_replay_empty_wal() { - let table_store = test_table_store(); - write_empty_wal(1, Arc::clone(&table_store)).await.unwrap(); - let mut replay_iter = WalReplayIterator::all_wal_ids( + async fn should_abort_pending_fetches_when_replay_iterator_is_dropped() { + let inner: Arc = Arc::new(InMemory::new()); + let first_path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + let paused = Arc::new(PausedFirstWalGetStore::new(inner, first_path)); + let table_store = test_table_store_with_object_store(paused.clone()); + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(1), &encoded) + .await + .unwrap(); + + let replay_iter = ExactWalReplayIterator::range( + 1..2, &ManifestCore::new(), WalReplayOptions::default(), - Arc::clone(&table_store), + table_store, ) .await .unwrap(); + PausedFirstWalGetStore::wait_until(&paused.first_started).await; + drop(replay_iter); - let Some(table) = replay_iter.next().await.unwrap() else { - panic!("Expected empty table to be returned from iterator") - }; + tokio::time::timeout( + Duration::from_secs(1), + PausedFirstWalGetStore::wait_until(&paused.first_cancelled), + ) + .await + .expect("pending WAL fetch was detached instead of aborted"); + } - assert_eq!(table.last_wal_id, 1); - assert_eq!(table.last_seq, 0); - assert!(table.table.is_empty()); - assert_eq!(table.last_tick, i64::MIN); - assert!(replay_iter.next().await.unwrap().is_none()); + #[tokio::test] + async fn should_abort_later_fetches_after_a_replay_error() { + let inner = Arc::new(InMemory::new()); + let paused_path = Path::from("/tmp/test_kv_store/wal/00000000000000000002.sst"); + let paused = Arc::new(PausedFirstWalGetStore::new(inner.clone(), paused_path)); + let table_store = test_table_store_with_object_store(paused.clone()); + let mut first_wal_location = None; + for wal_id in 1..=2 { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + if wal_id == 1 { + first_wal_location = Some( + table_store + .list_wal_ssts_for_replay(1..2) + .await + .unwrap() + .pop() + .unwrap() + .metadata + .location, + ); + } + } + let first_wal_location = first_wal_location.unwrap(); + let mut corrupted = inner + .get(&first_wal_location) + .await + .unwrap() + .bytes() + .await + .unwrap() + .to_vec(); + corrupted[0] ^= 1; + inner + .put(&first_wal_location, corrupted.into()) + .await + .unwrap(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..3, + &ManifestCore::new(), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: 1024 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, + ) + .await + .unwrap(); + PausedFirstWalGetStore::wait_until(&paused.first_started).await; + + assert!(matches!( + replay_iter.next().await, + Err(SlateDBError::ChecksumMismatch { .. }) + )); + tokio::time::timeout( + Duration::from_secs(1), + PausedFirstWalGetStore::wait_until(&paused.first_cancelled), + ) + .await + .expect("a replay error did not abort a later in-flight WAL fetch"); } #[tokio::test] - async fn should_replay_zero_byte_wal_fence() { - let table_store = test_table_store(); - table_store.write_wal_fence(1).await.unwrap(); - let mut replay_iter = WalReplayIterator::all_wal_ids( + async fn should_fail_terminally_and_abort_later_fetches_after_a_fetch_panic() { + let inner: Arc = Arc::new(InMemory::new()); + let first_path = Path::from("/tmp/test_kv_store/wal/00000000000000000001.sst"); + let paused_path = Path::from("/tmp/test_kv_store/wal/00000000000000000002.sst"); + let store = Arc::new(PausedFirstWalGetStore::new(inner, paused_path).panic_on(first_path)); + let table_store = test_table_store_with_object_store(store.clone()); + for wal_id in 1..=2 { + let mut builder = table_store.wal_table_builder(); + builder + .add(RowEntry::new_value( + format!("key-{wal_id}").as_bytes(), + b"value", + wal_id, + )) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + + let mut replay_iter = ExactWalReplayIterator::range( + 1..3, &ManifestCore::new(), - WalReplayOptions::default(), - Arc::clone(&table_store), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: 1024 * 1024, + }, + ..WalReplayOptions::default() + }, + table_store, ) .await .unwrap(); + PausedFirstWalGetStore::wait_until(&store.first_started).await; - let Some(table) = replay_iter.next().await.unwrap() else { - panic!("Expected empty table to be returned from iterator") + let Err(first) = replay_iter.next().await else { + panic!("panicking WAL fetch must fail replay"); }; - - assert_eq!(table.last_wal_id, 1); - assert_eq!(table.last_seq, 0); - assert!(table.table.is_empty()); - assert_eq!(table.last_tick, i64::MIN); - assert!(replay_iter.next().await.unwrap().is_none()); + let Err(second) = replay_iter.next().await else { + panic!("failed replay iterator must remain failed"); + }; + assert!(matches!(first, SlateDBError::BackgroundTaskPanic(_))); + assert!(matches!(second, SlateDBError::BackgroundTaskPanic(_))); + tokio::time::timeout( + Duration::from_secs(1), + PausedFirstWalGetStore::wait_until(&store.first_cancelled), + ) + .await + .expect("fetch panic did not abort a later in-flight WAL fetch"); } #[tokio::test] - async fn should_replay_zero_byte_wal_fence_before_real_wal() { - let table_store = test_table_store(); - table_store.write_wal_fence(1).await.unwrap(); - - let row = RowEntry::new_value(b"key", b"value", 1); + async fn should_fail_terminally_if_a_fetch_task_is_cancelled() { + let inner: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner)); + let table_store = test_table_store_with_object_store(gated.clone()); let mut builder = table_store.wal_table_builder(); - builder.add(row.clone()).await.unwrap(); - let encoded_sst = builder.build().await.unwrap(); + builder + .add(RowEntry::new_value(b"key", b"value", 1)) + .await + .unwrap(); + let encoded = builder.build().await.unwrap(); table_store - .write_sst(&SsTableId::Wal(2), &encoded_sst) + .write_sst(&SsTableId::Wal(1), &encoded) .await .unwrap(); + gated.get_opts_gate.close(); - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::range( + 1..2, &ManifestCore::new(), WalReplayOptions::default(), - Arc::clone(&table_store), + table_store, ) .await .unwrap(); + gated.get_opts_gate.wait_for_arrivals(1).await; + replay_iter.pending_fetches.front().unwrap().handle.abort(); - let Some(replayed_table) = replay_iter.next().await.unwrap() else { - panic!("Expected table to be returned from iterator") + let Err(first) = replay_iter.next().await else { + panic!("cancelled WAL fetch must fail replay"); }; - assert_eq!(replayed_table.last_wal_id, 2); - assert_eq!(replayed_table.last_seq, 1); - - let mut iter = replayed_table.table.table().iter(); - test_utils::assert_iterator(&mut iter, vec![row]).await; - assert!(replay_iter.next().await.unwrap().is_none()); + let Err(second) = replay_iter.next().await else { + panic!("failed replay iterator must remain failed"); + }; + assert!(matches!(first, SlateDBError::BackgroundTaskCancelled(_))); + assert!(matches!(second, SlateDBError::BackgroundTaskCancelled(_))); } #[tokio::test] - async fn should_replay_all_entries() { - let table_store = test_table_store(); - let mut rng = rng::new_test_rng(None); - let entries = sample::table(&mut rng, 1000, 10); - let next_wal_id = write_wals(&entries, 1, &mut rng, 200, Arc::clone(&table_store)) - .await + async fn should_hold_byte_permits_until_the_wal_is_consumed() { + let inner: Arc = Arc::new(InMemory::new()); + let gated = Arc::new(GatedObjectStore::new(inner)); + let table_store = test_table_store_with_format_and_object_store( + SsTableFormat { + block_size: 256, + ..SsTableFormat::default() + }, + gated.clone(), + ); + for wal_id in 1..=2 { + let mut builder = table_store.wal_table_builder(); + for row_id in 0..16 { + builder + .add(RowEntry::new_value( + format!("key-{wal_id}-{row_id:02}").as_bytes(), + &[b'x'; 128], + wal_id * 16 + row_id, + )) + .await + .unwrap(); + } + let encoded = builder.build().await.unwrap(); + table_store + .write_sst(&SsTableId::Wal(wal_id), &encoded) + .await + .unwrap(); + } + let listed = table_store.list_wal_ssts_for_replay(1..3).await.unwrap(); + let largest_object = listed + .iter() + .map(|entry| usize::try_from(entry.metadata.size).unwrap()) + .max() .unwrap(); - - let mut replay_iter = WalReplayIterator::all_wal_ids( + let total_budget = largest_object.checked_mul(2).unwrap(); + let working_budget = total_budget / 4; + let encoded_budget = total_budget - working_budget; + assert!(encoded_budget >= largest_object); + assert!(encoded_budget < largest_object * 2); + gated.get_opts_gate.close(); + + let mut replay_iter = ExactWalReplayIterator::range( + 1..3, &ManifestCore::new(), - WalReplayOptions::default(), - Arc::clone(&table_store), + WalReplayOptions { + prefetch: WalReplaySettings { + max_concurrent_objects: 2, + max_inflight_bytes: total_budget, + }, + max_memtable_bytes: 1, + ..WalReplayOptions::default() + }, + table_store, ) .await .unwrap(); + gated.get_opts_gate.wait_for_arrivals(1).await; + for _ in 0..100 { + tokio::task::yield_now().await; + } + assert_eq!(gated.get_opts_gate.arrivals(), 1); - let Some(replayed_table) = replay_iter.next().await.unwrap() else { - panic!("Expected table to be returned from iterator") - }; - assert_eq!(replayed_table.last_wal_id + 1, next_wal_id); - - let mut imm_table_iter = replayed_table.table.table().iter(); - test_utils::assert_ranged_kv_scan( - &entries, - &BytesRange::from(..), - IterationOrder::Ascending, - &mut imm_table_iter, - ) - .await; + gated.get_opts_gate.release(); + assert!(replay_iter.next().await.unwrap().is_some()); + assert!(replay_iter.next().await.unwrap().is_some()); assert!(replay_iter.next().await.unwrap().is_none()); + assert_eq!(gated.get_opts_gate.arrivals(), 2); } #[tokio::test] @@ -436,7 +3555,7 @@ mod tests { .unwrap(); let max_memtable_bytes = 1024; - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &ManifestCore::new(), WalReplayOptions { max_memtable_bytes, @@ -506,7 +3625,7 @@ mod tests { .unwrap(); } - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &ManifestCore::new(), WalReplayOptions { max_memtable_bytes, @@ -574,7 +3693,7 @@ mod tests { // Replay the single WAL SST into in-memory tables. If the replay code // can split a single commit sequence, it will do so here. - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &ManifestCore::new(), WalReplayOptions { max_memtable_bytes, @@ -632,7 +3751,7 @@ mod tests { .unwrap(); // Replay the single WAL SST into in-memory tables. - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &ManifestCore::new(), WalReplayOptions { max_memtable_bytes, @@ -694,7 +3813,7 @@ mod tests { db_state.replay_after_wal_id = replay_after_wal_id; db_state.next_wal_sst_id = replay_after_wal_id + 1; - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &db_state, WalReplayOptions::default(), Arc::clone(&table_store), @@ -733,7 +3852,7 @@ mod tests { db_state.last_l0_seq = min_seq; db_state.last_l0_clock_tick = 0; - let mut replay_iter = WalReplayIterator::all_wal_ids( + let mut replay_iter = ExactWalReplayIterator::all_wal_ids( &db_state, WalReplayOptions::default(), Arc::clone(&table_store), @@ -758,12 +3877,88 @@ mod tests { assert_eq!(total, 500); } + fn historical_wal_fixtures() -> Vec<(&'static str, Bytes)> { + let source = include_str!("../testdata/historical-wal-fixtures.hex"); + source + .lines() + .filter(|line| !line.is_empty() && !line.starts_with('#')) + .filter_map(|line| { + let (name, encoded) = line.split_once(' ')?; + // Each recognized suffix has a different compile-time value; + // this cannot be reduced to one `matches!` expression. + #[allow(clippy::match_like_matches_macro)] + let feature_enabled = match name.rsplit_once('-').map(|(_, suffix)| suffix) { + Some("snappy") => cfg!(feature = "snappy"), + Some("zlib") => cfg!(feature = "zlib"), + Some("lz4") => cfg!(feature = "lz4"), + Some("zstd") => cfg!(feature = "zstd"), + _ => true, + }; + feature_enabled.then(|| (name, Bytes::from(decode_hex(encoded)))) + }) + .collect() + } + + fn decode_hex(encoded: &str) -> Vec { + assert!(encoded.len().is_multiple_of(2)); + encoded + .as_bytes() + .chunks_exact(2) + .map(|pair| { + let high = (pair[0] as char).to_digit(16).unwrap(); + let low = (pair[1] as char).to_digit(16).unwrap(); + ((high << 4) | low) as u8 + }) + .collect() + } + + fn assert_historical_rows(name: &str, rows: &[RowEntry]) { + assert_eq!(rows.len(), 5, "fixture {name}"); + assert_eq!( + rows.iter().map(|row| row.seq).collect::>(), + vec![7, 8, 9, 10, 11] + ); + assert_eq!(rows[0].key, Bytes::from_static(b"alpha")); + assert_eq!( + rows[0].value, + ValueDeletable::Value(Bytes::from_static(b"value-a")) + ); + assert_eq!( + (rows[0].create_ts, rows[0].expire_ts), + (Some(100), Some(200)) + ); + assert_eq!(rows[1].key, Bytes::from_static(b"dup")); + assert_eq!( + rows[1].value, + ValueDeletable::Merge(Bytes::from_static(b"merge-1")) + ); + assert_eq!(rows[2].key, Bytes::from_static(b"dup")); + assert_eq!(rows[3].value, ValueDeletable::Tombstone); + assert_eq!(rows[4].value.len(), 20); + } + fn test_table_store() -> Arc { let object_store: Arc = Arc::new(InMemory::new()); + test_table_store_with_object_store(object_store) + } + + fn test_table_store_with_object_store(object_store: Arc) -> Arc { + test_table_store_with_format_and_object_store(SsTableFormat::default(), object_store) + } + + fn test_table_store_with_format(format: SsTableFormat) -> Arc { + let object_store: Arc = Arc::new(InMemory::new()); + test_table_store_with_format_and_object_store(format, object_store) + } + + fn test_table_store_with_format_and_object_store( + format: SsTableFormat, + object_store: Arc, + ) -> Arc { let path = Path::from("/tmp/test_kv_store"); Arc::new(TableStore::new( - ObjectStores::new(object_store.clone(), None), - SsTableFormat::default(), + ObjectStores::new(object_store, None), + format, path, None, TableStoreKind::Main, diff --git a/slatedb/testdata/historical-wal-fixtures.hex b/slatedb/testdata/historical-wal-fixtures.hex new file mode 100644 index 000000000..b214d3645 --- /dev/null +++ b/slatedb/testdata/historical-wal-fixtures.hex @@ -0,0 +1,17 @@ +# Immutable WAL SST fixtures generated by historical SlateDB code. +# V1: commit ecb57a3 (before V2 became the default). +# V2: commit 35a8059 (the V2 default introduction). +# Each fixture contains values, a merge operand, a tombstone, timestamps, +# expiry, duplicate keys, and multiple blocks. +v1-none 00000005616c70686100000000000000070600000000000000c800000000000000640000000776616c75652d6100000001933ac4c30000000364757000000000000000080a000000000000012c000000076d657267652d31000000017c9e71520000000364757000000000000000090400000000000001900000000776616c75652d320000000128e2c2fd000000056f6d656761000000000000000a0500000000000001f4000000016aa3888d000000047a657461000000000000000b00000000147a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a000000012139482010000000000000000000060008000400060000000400000005000000a800000080000000540000002800000004000000c0ffffff10000000ad000000000000000000000008000000000000000000000be0ffffff100000008b000000000000000000000008000000000000000000000a080014000800040008000000100000006000000000000000000000000800000000000000000000090800100008000400080000000c000000350000000000000008000000000000000000000808000800000004000800000004000000080000000000000000000007562e8673100000000c00240004000800100018000c00000020000000de00000000000000dc00000000000000de0000000000000000000000080000000000000000000007021fc50000000000000001ba0001 +v1-snappy 312400000005616c70686100090104070609080400c80d083c640000000776616c75652d6100000001e8f8e866271c0000000364757000090104080a090840012c000000076d657267652d310000000130f40ee8271c0000000364757000090104090409084001900000000776616c75652d320000000167d42bd71e24000000056f6d656761000901340a0500000000000001f4000000017088516e2d20000000047a657461000901000b010704147a4a01000c00000001d1ed9ecbd801041000110178060008000400060000000400000005000000a80000008000000054000000280d180cc0ffffff0134009c11370c0000000801050901040be00d20007b091005011d200c0a080014056801140148015c110101140d010c090800101528100c0000002d0d191d24040808050e00040d063400000800000000000000000000073dbac8b0140000001000280008000c0014001c0000000700100000000000000120000000bc000000000000008a00000000000000bc00000000000000000000000800000000000000000000075ba739e100000000000001460001 +v1-zlib 789c2d8ab10900200c04df429dc6a11e142c52d89819335608c9377f1c07a053de2572731458fd0ea9947f5634cd01891f05f8fc683eb3789c636060604e292d6080000e2e08cda80324d873538bd253750d417c0045b804024e29b76a789c636060604e292d6080004e1608cd380148b09725e694a6ea1a81f8004bfb046fc381c753789c63606060cdcf4d4d4f6480002e5608cdf80544000037370314fe2c0505789c63606060a94a2d496480006e102152850500c51901d1ac0b61c49228ca789c6d8dbb0d80300c44af08603e450a86600316a165172436a1620c46611373891c4528b1f4ce9f3bc91ea95a081c1554a0213739c84616bb3faaead92fe4923c8eaff967dd1f0473fc136e21b7d773bdd04db989ac6546246ecef4e7751ff4070c3ae308b800140000001000280008000c0014001c0000000700100000000000000220000000c1000000000000006700000000000000c1000000000000000000000008000000000000000000000799eb90f400000000000001280001 +v1-lz4 31000000a200000005616c70686100010022070608002300c80800f001640000000776616c75652d6100000001ac6ccd3427000000820000000364757000010022080a0800f002012c000000076d657267652d31000000013d0fae442700000082000000036475700001002209040800f00201900000000776616c75652d3200000001fe09755f1e000000a2000000056f6d656761000100220a0508006001f40000000126315b562d00000092000000047a657461000100100b07002d147a0100607a7a0000000140ca2ae2d80000002410000100f310060008000400060000000400000005000000a8000000800000005400000028180040c0ffffff3400001c00040200140809004300000be0200010870e00030200072000410a0800146800001400004800135c240000020000140003020045090800102800530c0000003119000724002108080e00024c0000b400000e00800000000000000007d3cd2547140000001000280008000c0014001c0000000700100000000000000320000000cc000000000000009500000000000000cc000000000000000000000008000000000000000000000785655ddb00000000000001610001 +v1-zstd 28b52ffd20313d0100f800000005616c70686100070600c800640000000776616c75652d610000000103140003782001b9f1252828b52ffd2027250100e00000000364757000080a00012c000000076d657267652d31000000010200408c017001208677bf28b52ffd2027250100e0000000036475700009040001900000000776616c75652d32000000010200408c017001d1b8fa5028b52ffd201edd000098000000056f6d656761000a050001f4000000010200408c01200186bc599728b52ffd202dd5000080000000047a657461000b147a000000010300b8c494c0062007d30d070f28b52ffd20d805030032840d1ab035c3180cc59262883b77c5c55bd700824eec1c112149fbb82902da921bea86e538a20914f773b5d6022057cedc7970f414dd540510003f0fa4e48d69f8e21880d34b69047a4b88b81a600d17017587caa3a58cadca79677604c002a8a84a4e140000001000280008000c0014001c0000000700100000000000000420000000e5000000000000006d00000000000000e50000000000000000000000080000000000000000000007484276d800000000000001520001 +v1-xor a5a5a5a0c4c9d5cdc4a5a5a5a5a5a5a5a2a3a5a5a5a5a5a5a56da5a5a5a5a5a5a5c1a5a5a5a2d3c4c9d0c088c4a5a5a5a4bd8b0552a5a5a5a6c1d0d5a5a5a5a5a5a5a5adafa5a5a5a5a5a5a489a5a5a5a2c8c0d7c2c08894a5a5a5a43530ed80a5a5a5a6c1d0d5a5a5a5a5a5a5a5aca1a5a5a5a5a5a5a435a5a5a5a2d3c4c9d0c08897a5a5a5a4614c5e2fa5a5a5a0cac8c0c2c4a5a5a5a5a5a5a5afa0a5a5a5a5a5a5a451a5a5a5a4781b61cfa5a5a5a1dfc0d1c4a5a5a5a5a5a5a5aea5a5a5a5b1dfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfa5a5a5a4991a1ea0b5a5a5a5a5a5a5a5a5a5a3a5ada5a1a5a3a5a5a5a1a5a5a5a0a5a5a50da5a5a525a5a5a5f1a5a5a58da5a5a5a1a5a5a5655a5a5ab5a5a5a508a5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5ae455a5a5ab5a5a5a52ea5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5afada5b1a5ada5a1a5ada5a5a5b5a5a5a5c5a5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5acada5b5a5ada5a1a5ada5a5a5a9a5a5a590a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5adada5ada5a5a5a1a5ada5a5a5a1a5a5a5ada5a5a5a5a5a5a5a5a5a5a2e569645a100000000c00240004000800100018000c00000020000000de00000000000000dc00000000000000de0000000000000000000000080000000000000000000007021fc50000000000000001ba0001 +v2-none 000507616c70686176616c75652d6100000000000000070600000000000000c8000000000000006400000001f8f3fcc70003076475706d657267652d3100000000000000080a000000000000012c00000001accf424b00030764757076616c75652d320000000000000009040000000000000190000000015ad218c40005006f6d656761000000000000000a0500000000000001f400000001695c0a730004147a6574617a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a7a000000000000000b0000000001c6dd813810000000000000000000060008000400060000000400000005000000a800000080000000540000002800000004000000c0ffffff100000009d000000000000000000000008000000000000000000000be0ffffff100000007c000000000000000000000008000000000000000000000a080014000800040008000000100000005600000000000000000000000800000000000000000000090800100008000400080000000c000000300000000000000008000000000000000000000808000800000004000800000004000000080000000000000000000007db1cfa7b100000000c00240004000800100018000c00000020000000c900000000000000dc00000000000000c9000000000000000000000008000000000000000000000742b67de300000000000001a50002 +v2-snappy 2c3c000507616c70686176616c75652d6100090104070609083400c800000000000000640000000153b8823722340003076475706d657267652d3100090134080a000000000000012c00000001715f75c7223400030764757076616c75652d32000901340904000000000000019000000001238653b91d200005006f6d656761000901340a0500000000000001f40000000169df6b3d281c0004147a6574617a4a010030000000000000000b0000000001da5124dfd801041000110178060008000400060000000400000005000000a80000008000000054000000280d180cc0ffffff0134009611370c0000000801050901040be00d200076091005011d200c0a08001405680114014800510522090101140d010c090800101528100c0000002c0d191d24040808050e00040d06340000080000000000000000000007bf5f8165140000001000280008000c0014001c0000000700100000000000000120000000b5000000000000008c00000000000000b5000000000000000000000008000000000000000000000794f764ff00000000000001410002 +v2-zlib 789c6360654fcc29c8482c4bcc294dd54d648000763628e304944e01624600b70505f8f3d7afa9789c6360664f292dc84d2d4a4fd5356480000e2e08cda803220067ab04029528b008789c6360664f292d284bcc294dd5356280004e1608cd38014400006ac7046ff9edaa34789c636065c8cf4d4d4f6480002e5608cdf805440000373b031471936ce1789c636011a94a2d49acc2021820801b44300200209e0b61069c6582789c6d8dbb0d80300c44af08603e450a768055e8198321a899843118854dcc257214a1c4d23b7fee247ba46a217054508186dce4243b59edfea8aa67bf904bf238bee61f757f10ccf14fb885dc56cff54237e526b2941991b839d39fd77dcfa70be7f2baaf9f140000001000280008000c0014001c0000000700100000000000000220000000a6000000000000006600000000000000a600000000000000000000000800000000000000000000072cdf5776000000000000010c0002 +v2-lz4 2c000000f201000507616c70686176616c75652d6100010022070608002200c80800600064000000014056bbfa22000000e20003076475706d657267652d3100010022080a080060012c0000000170a7681122000000e200030764757076616c75652d32000100220904080060019000000001f6e2ce861d000000920005006f6d656761000100220a0508006001f400000001a5d32e79280000008f0004147a6574617a01000012000100600b0000000001a8e33ad6d80000002410000100f310060008000400060000000400000005000000a8000000800000005400000028180040c0ffffff34001497370041000000080600010200230be0200011770f00020200072000410a08001468000014000048001252230001020000140003020045090800102800530c0000002d19000724002108087a00024c0000b400021c00600000000000078252952a140000001000280008000c0014001c0000000700100000000000000320000000b6000000000000009500000000000000b60000000000000000000000080000000000000000000007759f7491000000000000014b0002 +v2-zstd 28b52ffd202c150100d0000507616c70686176616c75652d6100070600c800640000000103140003785802840694bd28b52ffd2022fd0000b80003076475706d657267652d3100080a00012c000000010200400c0350023a01f62228b52ffd2022fd0000b800030764757076616c75652d32000904000190000000010200400c035002da39d8a928b52ffd201dd50000900005006f6d656761000a050001f4000000010200408c011c0187f0af2628b52ffd2028bd0000780004147a6574617a000b00000000010200601009e0020f8e9bc428b52ffd20d805030032840d1ab035c3180cc592c2c511a6d64bbabea4eb4bdc14112149fbb82902da921bea86e538a20914f773b5d6022057cedc7970f414dd540510003f0fa4e48d69f8e21880d34b69047a4b88b81a600d17017587caa3a58cadca79677604c0028bbbbbbd140000001000280008000c0014001c0000000700100000000000000420000000d2000000000000006d00000000000000d20000000000000000000000080000000000000000000007c48231fc000000000000013f0002 +v2-xor a5a0a2c4c9d5cdc4d3c4c9d0c088c4a5a5a5a5a5a5a5a2a3a5a5a5a5a5a5a56da5a5a5a5a5a5a5c1a5a5a5a466007355a5a6a2c1d0d5c8c0d7c2c08894a5a5a5a5a5a5a5adafa5a5a5a5a5a5a489a5a5a5a4348e055da5a6a2c1d0d5d3c4c9d0c08897a5a5a5a5a5a5a5aca1a5a5a5a5a5a5a435a5a5a5a4c2935fd2a5a0a5cac8c0c2c4a5a5a5a5a5a5a5afa0a5a5a5a5a5a5a451a5a5a5a4665c188fa5a1b1dfc0d1c4dfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfdfa5a5a5a5a5a5a5aea5a5a5a5a408f5cb3bb5a5a5a5a5a5a5a5a5a5a3a5ada5a1a5a3a5a5a5a1a5a5a5a0a5a5a50da5a5a525a5a5a5f1a5a5a58da5a5a5a1a5a5a5655a5a5ab5a5a5a538a5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5ae455a5a5ab5a5a5a5d9a5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5afada5b1a5ada5a1a5ada5a5a5b5a5a5a5f3a5a5a5a5a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5acada5b5a5ada5a1a5ada5a5a5a9a5a5a595a5a5a5a5a5a5a5ada5a5a5a5a5a5a5a5a5a5adada5ada5a5a5a1a5ada5a5a5a1a5a5a5ada5a5a5a5a5a5a5a5a5a5a2685b1852100000000c00240004000800100018000c00000020000000c900000000000000dc00000000000000c9000000000000000000000008000000000000000000000742b67de300000000000001a50002