From 7a55cdafdbd2cb1cd5275b0739542bab88e6efe4 Mon Sep 17 00:00:00 2001 From: trinity-1686a Date: Thu, 8 Oct 2026 13:30:14 +0200 Subject: [PATCH] add test for block download coalescing and skipping --- sstable/Cargo.toml | 1 + sstable/src/dictionary.rs | 120 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/sstable/Cargo.toml b/sstable/Cargo.toml index 8c774ab00..10ca4fc9b 100644 --- a/sstable/Cargo.toml +++ b/sstable/Cargo.toml @@ -26,6 +26,7 @@ proptest = "1" criterion = { version = "0.8", default-features = false } names = "0.14" rand = "0.9" +async-trait = "0.1" [[bench]] name = "stream_bench" diff --git a/sstable/src/dictionary.rs b/sstable/src/dictionary.rs index 8296c56c0..7169f9f8c 100644 --- a/sstable/src/dictionary.rs +++ b/sstable/src/dictionary.rs @@ -676,6 +676,7 @@ mod tests { struct PermissionedHandle { bytes: OwnedBytes, allowed_range: Mutex>, + read_ranges: Mutex>>, } impl PermissionedHandle { @@ -683,6 +684,7 @@ mod tests { let bytes = OwnedBytes::new(bytes); PermissionedHandle { allowed_range: Mutex::new(0..bytes.len()), + read_ranges: Mutex::new(Vec::new()), bytes, } } @@ -690,6 +692,11 @@ mod tests { fn restrict(&self, range: Range) { *self.allowed_range.lock().unwrap() = range; } + + /// The ranges that were successfully read so far. + fn read_ranges(&self) -> Vec> { + self.read_ranges.lock().unwrap().clone() + } } impl common::HasLen for PermissionedHandle { @@ -698,6 +705,7 @@ mod tests { } } + #[async_trait::async_trait] impl common::file_slice::FileHandle for PermissionedHandle { fn read_bytes(&self, range: Range) -> std::io::Result { let allowed_range = self.allowed_range.lock().unwrap(); @@ -707,8 +715,13 @@ mod tests { ))); } + self.read_ranges.lock().unwrap().push(range.clone()); Ok(self.bytes.slice(range)) } + + async fn read_bytes_async(&self, range: Range) -> std::io::Result { + self.read_bytes(range) + } } fn make_test_sstable() -> (Dictionary, Arc) { @@ -1174,4 +1187,111 @@ mod tests { } assert_eq!(seen, 3); // 0FFFF, 1FFFF, 2FFFF — the dictionary stops below 3FFFF } + + /// Minimal executor: the futures used here (in-memory async reads) never actually + /// yield, so spinning on `Poll::Pending` is enough. + fn block_on(fut: F) -> F::Output { + use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + + fn waker_clone(_: *const ()) -> RawWaker { + noop_raw_waker() + } + fn noop(_: *const ()) {} + fn noop_raw_waker() -> RawWaker { + static VTABLE: RawWakerVTable = RawWakerVTable::new(waker_clone, noop, noop, noop); + RawWaker::new(std::ptr::null(), &VTABLE) + } + + let waker = unsafe { Waker::from_raw(noop_raw_waker()) }; + let mut cx = Context::from_waker(&waker); + let mut fut = std::pin::pin!(fut); + loop { + if let Poll::Ready(output) = fut.as_mut().poll(&mut cx) { + return output; + } + std::thread::yield_now(); + } + } + + #[test] + fn test_merged_block_group_skips_non_matching_blocks() { + // define an sstable with + let elems_per_region: u64 = 10_000; + let (dict, handle) = { + let mut builder = Dictionary::::builder(Vec::new()).unwrap(); + let mut ord = 0u64; + for prefix in ['A', 'B', 'C'] { + for elem in 0..elems_per_region { + builder + .insert(format!("{prefix}{elem:05}").as_bytes(), &ord) + .unwrap(); + ord += 1; + } + } + let table = Arc::new(PermissionedHandle::new(builder.finish().unwrap())); + let dict = Dictionary::::open(common::file_slice::FileSlice::new( + table.clone(), + )) + .unwrap(); + assert!(dict.sstable_index.locate_with_ord(u64::MAX) > 2); + (dict, table) + }; + + let pattern = tantivy_fst::Regex::new("[^B].*").unwrap(); + + // verify there is a block that doesn't match the automaton at all + // (if we have only 2 blocks, one of A+B and one of B+C, this test doesn't work, + // we need at least one block of full B + let matching_ranges: Vec<_> = dict + .sstable_index + .get_block_for_automaton(&pattern) + .map(|(_, block_addr)| block_addr.byte_range) + .collect(); + let middle_block = dict.sstable_index.get_block_with_key(b"B05000").unwrap(); + assert!(!matching_ranges.is_empty()); + assert!(!matching_ranges.contains(&middle_block.byte_range)); + + // count what should be decoded + let num_blocks = dict.sstable_index.locate_with_ord(u64::MAX); + let mut expected_seen = 0usize; + for block_id in 0..=num_blocks { + let block = dict.sstable_index.get_block(block_id).unwrap(); + let end_ord = dict + .sstable_index + .get_block(block_id + 1) + .map(|block| block.first_ordinal) + .unwrap_or(dict.num_terms() as u64); + if matching_ranges.contains(&block.byte_range) { + expected_seen += (end_ord - block.first_ordinal) as usize; + } + } + assert!(expected_seen > 20_000); + assert!(expected_seen < 30_000); + + // a merge threshold big enough to coalesce the 'A' and 'C' block groups into a + // single read, downloading the non-matching 'B' blocks in between. + let merge_holes_under_bytes = 1 << 30; + + let mut delta_reader = block_on(dict.sstable_delta_reader_for_key_range_async( + .., + None, + &pattern, + merge_holes_under_bytes, + )) + .unwrap(); + + // check one read went through the whole table + let first_matching_block = dict.sstable_index.get_block_with_key(b"A00000").unwrap(); + let last_matching_block = dict.sstable_index.get_block_with_key(b"C09999").unwrap(); + let expected_range = + first_matching_block.byte_range.start..last_matching_block.byte_range.end; + assert!(handle.read_ranges().contains(&expected_range)); + + // check the delta-reader only decodes matching blocks + let mut seen = 0; + while delta_reader.advance().unwrap() { + seen += 1; + } + assert_eq!(seen, expected_seen); + } }