1use crate::{map_heed_error, LmdbBlobReader};
2use hashtree_core::store::StoreError;
3use hashtree_core::types::Hash;
4use heed::types::Bytes;
5use heed::{Database, RoTxn};
6use std::collections::BTreeMap;
7use std::ops::Bound;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::sync::{mpsc, Arc, Mutex};
10use std::thread::{self, JoinHandle};
11
12const HASH_RANGE_BITS: u16 = 12;
13const HASH_RANGE_COUNT: u16 = 1 << HASH_RANGE_BITS;
14const DEFAULT_RANGE_CHUNK_KEYS: usize = 32 * 1024;
15const SNAPSHOT_OPEN_ATTEMPTS: usize = 64;
16
17pub const MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY: usize = 16;
22
23#[derive(Debug, Clone, Copy)]
24struct RangeTask {
25 range: u16,
26 after: Option<Hash>,
27}
28
29struct RangePage {
30 range: u16,
31 hashes: Vec<Hash>,
32 complete: bool,
33}
34
35struct RangeResult {
36 range: u16,
37 result: Result<RangePage, StoreError>,
38}
39
40struct CurrentChunk {
41 hashes: Vec<Hash>,
42 offset: usize,
43}
44
45pub struct LmdbParallelRawKeyScanner {
53 task_tx: Option<mpsc::SyncSender<RangeTask>>,
54 result_rx: mpsc::Receiver<RangeResult>,
55 workers: Vec<JoinHandle<()>>,
56 cancelled: Arc<AtomicBool>,
57 concurrency: usize,
58 next_range_to_schedule: u16,
59 next_range_to_emit: u16,
60 initial_after: Option<Hash>,
61 in_flight: usize,
62 pending: BTreeMap<u16, RangePage>,
63 current_chunk: Option<CurrentChunk>,
64 previous_chunk_hash: Option<Hash>,
65 exhausted: bool,
66 snapshot_transaction_id: usize,
67}
68
69impl LmdbParallelRawKeyScanner {
70 fn new(
71 reader: &LmdbBlobReader,
72 after: Option<Hash>,
73 requested_concurrency: usize,
74 range_chunk_keys: usize,
75 ) -> Result<Self, StoreError> {
76 if requested_concurrency == 0 {
77 return Err(StoreError::Other(
78 "parallel raw-key scan concurrency must be non-zero".into(),
79 ));
80 }
81 if range_chunk_keys == 0 {
82 return Err(StoreError::Other(
83 "parallel raw-key scan range chunk must be non-zero".into(),
84 ));
85 }
86 let concurrency = requested_concurrency.min(MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY);
87 let transactions = open_coherent_snapshot_transactions(reader, concurrency)?;
88 let snapshot_transaction_id = transactions[0].id();
89 let blobs = reader.store.blobs;
90 let (task_tx, task_rx) = mpsc::sync_channel(concurrency);
91 let task_rx = Arc::new(Mutex::new(task_rx));
92 let (result_tx, result_rx) = mpsc::channel();
93 let cancelled = Arc::new(AtomicBool::new(false));
94 let mut workers = Vec::with_capacity(concurrency);
95
96 for (worker_index, transaction) in transactions.into_iter().enumerate() {
97 let task_rx = Arc::clone(&task_rx);
98 let result_tx = result_tx.clone();
99 let cancelled = Arc::clone(&cancelled);
100 let worker = thread::Builder::new()
101 .name(format!("lmdb-key-scan-{worker_index}"))
102 .spawn(move || {
103 raw_key_scan_worker(
104 blobs,
105 transaction,
106 task_rx,
107 result_tx,
108 cancelled,
109 range_chunk_keys,
110 );
111 })
112 .map_err(StoreError::Io)?;
113 workers.push(worker);
114 }
115 drop(result_tx);
116
117 let first_range = after.as_ref().map(hash_range).unwrap_or(0);
118 let mut scanner = Self {
119 task_tx: Some(task_tx),
120 result_rx,
121 workers,
122 cancelled,
123 concurrency,
124 next_range_to_schedule: first_range,
125 next_range_to_emit: first_range,
126 initial_after: after,
127 in_flight: 0,
128 pending: BTreeMap::new(),
129 current_chunk: None,
130 previous_chunk_hash: after,
131 exhausted: false,
132 snapshot_transaction_id,
133 };
134 scanner.fill_window()?;
135 Ok(scanner)
136 }
137
138 pub fn snapshot_transaction_id(&self) -> usize {
140 self.snapshot_transaction_id
141 }
142
143 pub fn next_page(&mut self, limit: usize) -> Result<Vec<Hash>, StoreError> {
149 if limit == 0 || self.exhausted {
150 return Ok(Vec::new());
151 }
152 let mut page = Vec::with_capacity(limit);
153 while page.len() < limit {
154 if self
155 .current_chunk
156 .as_ref()
157 .is_none_or(|chunk| chunk.offset == chunk.hashes.len())
158 {
159 self.current_chunk = self
160 .next_ordered_chunk()?
161 .map(|hashes| CurrentChunk { hashes, offset: 0 });
162 if self.current_chunk.is_none() {
163 self.exhausted = true;
164 break;
165 }
166 }
167 let chunk = self
168 .current_chunk
169 .as_mut()
170 .expect("parallel key scanner loaded a current chunk");
171 let take = (limit - page.len()).min(chunk.hashes.len() - chunk.offset);
172 page.extend_from_slice(&chunk.hashes[chunk.offset..chunk.offset + take]);
173 chunk.offset += take;
174 }
175 Ok(page)
176 }
177
178 fn next_ordered_chunk(&mut self) -> Result<Option<Vec<Hash>>, StoreError> {
179 loop {
180 if self.next_range_to_emit == HASH_RANGE_COUNT {
181 if self.in_flight != 0 || !self.pending.is_empty() {
182 return Err(StoreError::Other(
183 "parallel raw-key scan exhausted with outstanding ranges".into(),
184 ));
185 }
186 return Ok(None);
187 }
188
189 if let Some(range_page) = self.pending.remove(&self.next_range_to_emit) {
190 if range_page.complete {
191 self.next_range_to_emit = self
192 .next_range_to_emit
193 .checked_add(1)
194 .expect("12-bit hash range index overflow");
195 self.fill_window()?;
196 } else {
197 let after = range_page.hashes.last().copied().ok_or_else(|| {
198 StoreError::Other(
199 "parallel raw-key range continuation made no progress".into(),
200 )
201 })?;
202 self.send_task(RangeTask {
203 range: range_page.range,
204 after: Some(after),
205 })?;
206 }
207
208 if range_page.hashes.is_empty() {
209 continue;
210 }
211 if range_page.hashes.windows(2).any(|pair| pair[0] >= pair[1])
212 || self
213 .previous_chunk_hash
214 .zip(range_page.hashes.first().copied())
215 .is_some_and(|(previous, first)| previous >= first)
216 {
217 return Err(StoreError::Other(
218 "parallel raw-key scan produced non-increasing hashes".into(),
219 ));
220 }
221 self.previous_chunk_hash = range_page.hashes.last().copied();
222 return Ok(Some(range_page.hashes));
223 }
224
225 let received = self.result_rx.recv().map_err(|_| {
226 StoreError::Other("parallel raw-key scan workers stopped early".into())
227 })?;
228 self.in_flight = self.in_flight.checked_sub(1).ok_or_else(|| {
229 StoreError::Other("parallel raw-key scan in-flight count underflow".into())
230 })?;
231 let range_page = received.result?;
232 if range_page.range != received.range
233 || self.pending.insert(received.range, range_page).is_some()
234 {
235 return Err(StoreError::Other(
236 "parallel raw-key scan returned duplicate or mismatched range".into(),
237 ));
238 }
239 }
240 }
241
242 fn fill_window(&mut self) -> Result<(), StoreError> {
243 while self.in_flight + self.pending.len() < self.concurrency
244 && self.next_range_to_schedule < HASH_RANGE_COUNT
245 {
246 let range = self.next_range_to_schedule;
247 self.next_range_to_schedule = self
248 .next_range_to_schedule
249 .checked_add(1)
250 .expect("12-bit hash range index overflow");
251 let after = if range == self.next_range_to_emit {
252 self.initial_after.take()
253 } else {
254 None
255 };
256 self.send_task(RangeTask { range, after })?;
257 }
258 Ok(())
259 }
260
261 fn send_task(&mut self, task: RangeTask) -> Result<(), StoreError> {
262 self.task_tx
263 .as_ref()
264 .ok_or_else(|| StoreError::Other("parallel raw-key scanner is closed".into()))?
265 .send(task)
266 .map_err(|_| StoreError::Other("parallel raw-key scan workers stopped early".into()))?;
267 self.in_flight = self
268 .in_flight
269 .checked_add(1)
270 .ok_or_else(|| StoreError::Other("parallel raw-key scan count overflow".into()))?;
271 Ok(())
272 }
273}
274
275impl Drop for LmdbParallelRawKeyScanner {
276 fn drop(&mut self) {
277 self.cancelled.store(true, Ordering::Release);
278 self.task_tx.take();
279 for worker in self.workers.drain(..) {
280 let _ = worker.join();
281 }
282 }
283}
284
285impl LmdbBlobReader {
286 pub fn parallel_raw_key_scanner(
288 &self,
289 after: Option<Hash>,
290 concurrency: usize,
291 ) -> Result<LmdbParallelRawKeyScanner, StoreError> {
292 LmdbParallelRawKeyScanner::new(self, after, concurrency, DEFAULT_RANGE_CHUNK_KEYS)
293 }
294
295 #[cfg(test)]
296 fn parallel_raw_key_scanner_with_chunk_limit(
297 &self,
298 after: Option<Hash>,
299 concurrency: usize,
300 range_chunk_keys: usize,
301 ) -> Result<LmdbParallelRawKeyScanner, StoreError> {
302 LmdbParallelRawKeyScanner::new(self, after, concurrency, range_chunk_keys)
303 }
304}
305
306fn open_coherent_snapshot_transactions(
307 reader: &LmdbBlobReader,
308 concurrency: usize,
309) -> Result<Vec<RoTxn<'static>>, StoreError> {
310 let env = heed::Env::clone(&reader.store.env);
311 for _ in 0..SNAPSHOT_OPEN_ATTEMPTS {
312 let transactions = hashtree_core::lmdb_runtime::with_lmdb_lock_acquisition(|| {
313 (0..concurrency)
314 .map(|_| env.clone().static_read_txn())
315 .collect::<Result<Vec<_>, _>>()
316 })
317 .map_err(map_heed_error)?;
318 let snapshot = transactions[0].id();
319 if transactions
320 .iter()
321 .all(|transaction| transaction.id() == snapshot)
322 {
323 return Ok(transactions);
324 }
325 drop(transactions);
326 thread::yield_now();
327 }
328 Err(StoreError::Other(format!(
329 "could not open {concurrency} parallel raw-key readers on one LMDB snapshot"
330 )))
331}
332
333fn raw_key_scan_worker(
334 blobs: Database<Bytes, Bytes>,
335 transaction: RoTxn<'static>,
336 task_rx: Arc<Mutex<mpsc::Receiver<RangeTask>>>,
337 result_tx: mpsc::Sender<RangeResult>,
338 cancelled: Arc<AtomicBool>,
339 range_chunk_keys: usize,
340) {
341 while !cancelled.load(Ordering::Acquire) {
342 let task = {
343 let receiver = match task_rx.lock() {
344 Ok(receiver) => receiver,
345 Err(_) => return,
346 };
347 match receiver.recv() {
348 Ok(task) => task,
349 Err(_) => return,
350 }
351 };
352 if cancelled.load(Ordering::Acquire) {
353 return;
354 }
355 let result = scan_raw_key_range(blobs, &transaction, task, range_chunk_keys).map(
356 |(hashes, complete)| RangePage {
357 range: task.range,
358 hashes,
359 complete,
360 },
361 );
362 if result_tx
363 .send(RangeResult {
364 range: task.range,
365 result,
366 })
367 .is_err()
368 {
369 return;
370 }
371 }
372}
373
374fn scan_raw_key_range(
375 blobs: Database<Bytes, Bytes>,
376 transaction: &RoTxn<'_>,
377 task: RangeTask,
378 range_chunk_keys: usize,
379) -> Result<(Vec<Hash>, bool), StoreError> {
380 let lower = hash_range_boundary(task.range);
381 let upper = (task.range + 1 < HASH_RANGE_COUNT).then(|| hash_range_boundary(task.range + 1));
382 let start = match task.after.as_ref() {
383 Some(after) => {
384 if hash_range(after) != task.range {
385 return Err(StoreError::Other(
386 "parallel raw-key continuation left its 12-bit range".into(),
387 ));
388 }
389 Bound::Excluded(after.as_slice())
390 }
391 None => Bound::Included(lower.as_slice()),
392 };
393 let (keys, complete) = blobs
394 .raw_keys_in_range(
395 transaction,
396 start,
397 upper.as_ref().map(|bound| bound.as_slice()),
398 range_chunk_keys,
399 )
400 .map_err(map_heed_error)?;
401 let hashes = keys
402 .into_iter()
403 .map(|key| {
404 let hash: Hash = key
405 .try_into()
406 .map_err(|_| StoreError::Other("invalid hash length".into()))?;
407 if hash_range(&hash) != task.range {
408 return Err(StoreError::Other(
409 "parallel raw-key cursor escaped its 12-bit range".into(),
410 ));
411 }
412 Ok(hash)
413 })
414 .collect::<Result<Vec<_>, StoreError>>()?;
415 Ok((hashes, complete))
416}
417
418fn hash_range(hash: &Hash) -> u16 {
419 (u16::from(hash[0]) << 4) | u16::from(hash[1] >> 4)
420}
421
422fn hash_range_boundary(range: u16) -> [u8; 2] {
423 debug_assert!(range < HASH_RANGE_COUNT);
424 [(range >> 4) as u8, ((range & 0x0f) << 4) as u8]
425}
426
427#[cfg(test)]
428mod tests {
429 use super::*;
430 use crate::LmdbBlobStore;
431 use tempfile::TempDir;
432
433 fn hash_in_range(range: u16, suffix: u64) -> Hash {
434 let boundary = hash_range_boundary(range);
435 let mut hash = [0u8; 32];
436 hash[..2].copy_from_slice(&boundary);
437 hash[24..].copy_from_slice(&suffix.to_be_bytes());
438 hash
439 }
440
441 fn insert_raw(store: &LmdbBlobStore, hash: Hash, value: &[u8]) -> Result<(), StoreError> {
442 let mut wtxn = store.env.write_txn().map_err(map_heed_error)?;
443 store
444 .blobs
445 .put(&mut wtxn, &hash, value)
446 .map_err(map_heed_error)?;
447 wtxn.commit().map_err(map_heed_error)
448 }
449
450 fn collect(
451 scanner: &mut LmdbParallelRawKeyScanner,
452 page_size: usize,
453 ) -> Result<Vec<Hash>, StoreError> {
454 let mut hashes = Vec::new();
455 loop {
456 let page = scanner.next_page(page_size)?;
457 if page.is_empty() {
458 break;
459 }
460 assert!(page.len() <= page_size);
461 hashes.extend(page);
462 }
463 Ok(hashes)
464 }
465
466 #[test]
467 fn parallel_raw_key_scan_is_ordered_resumable_and_value_free() -> Result<(), StoreError> {
468 let temp = TempDir::new().map_err(StoreError::Io)?;
469 let store = LmdbBlobStore::new(temp.path().join("source"))?;
470 let mut expected = Vec::new();
471 for (range, count) in [(0, 5), (1, 3), (17, 7), (2048, 4), (4095, 6)] {
472 for suffix in 0..count {
473 let hash = hash_in_range(range, suffix);
474 insert_raw(&store, hash, &vec![range as u8; 128 * 1024])?;
477 expected.push(hash);
478 }
479 }
480 for (index, prefix) in [[0x00, 0x0f], [0x00, 0x10], [0xff, 0xef], [0xff, 0xf0]]
483 .into_iter()
484 .enumerate()
485 {
486 let mut hash = [0u8; 32];
487 hash[..2].copy_from_slice(&prefix);
488 hash[2] = 0xaa;
489 hash[31] = index as u8;
490 insert_raw(&store, hash, &vec![index as u8; 128 * 1024])?;
491 expected.push(hash);
492 }
493 expected.sort_unstable();
494 let reader = LmdbBlobReader {
495 store,
496 external_read_concurrency: 1,
497 };
498
499 let mut scanner = reader.parallel_raw_key_scanner_with_chunk_limit(None, 4, 2)?;
500 assert_eq!(
501 scanner.snapshot_transaction_id(),
502 reader.environment_generation().last_txn_id as usize,
503 );
504 assert_eq!(collect(&mut scanner, 3)?, expected);
505 drop(scanner);
506
507 let resume = expected[9];
508 let mut scanner = reader.parallel_raw_key_scanner_with_chunk_limit(Some(resume), 4, 2)?;
509 assert_eq!(collect(&mut scanner, 4)?, expected[10..]);
510 drop(scanner);
511
512 let final_hash = *expected.last().expect("nonempty source");
513 let mut scanner =
514 reader.parallel_raw_key_scanner_with_chunk_limit(Some(final_hash), 4, 2)?;
515 assert!(
516 scanner.next_page(4)?.is_empty(),
517 "resuming after the final raw key must reach exact EOF"
518 );
519 Ok(())
520 }
521
522 #[test]
523 fn parallel_raw_key_workers_share_one_snapshot() -> Result<(), StoreError> {
524 let temp = TempDir::new().map_err(StoreError::Io)?;
525 let store = LmdbBlobStore::new(temp.path().join("source"))?;
526 let original = hash_in_range(2, 1);
527 let committed_later = hash_in_range(3000, 1);
528 insert_raw(&store, original, b"original")?;
529 let reader = LmdbBlobReader {
530 store,
531 external_read_concurrency: 1,
532 };
533
534 let mut old_snapshot = reader.parallel_raw_key_scanner_with_chunk_limit(None, 8, 2)?;
535 insert_raw(&reader.store, committed_later, b"committed later")?;
536 assert_eq!(collect(&mut old_snapshot, 1)?, vec![original]);
537 drop(old_snapshot);
538
539 let mut new_snapshot = reader.parallel_raw_key_scanner_with_chunk_limit(None, 8, 2)?;
540 assert_eq!(
541 collect(&mut new_snapshot, 1)?,
542 vec![original, committed_later],
543 );
544 Ok(())
545 }
546
547 #[test]
548 fn parallel_raw_key_scan_caps_worker_and_buffer_concurrency() -> Result<(), StoreError> {
549 let temp = TempDir::new().map_err(StoreError::Io)?;
550 let store = LmdbBlobStore::new(temp.path().join("source"))?;
551 insert_raw(&store, hash_in_range(0, 1), b"one")?;
552 let reader = LmdbBlobReader {
553 store,
554 external_read_concurrency: 1,
555 };
556 let scanner = reader.parallel_raw_key_scanner_with_chunk_limit(None, 64, 1)?;
557 assert_eq!(scanner.workers.len(), MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY);
558 assert!(scanner.in_flight + scanner.pending.len() <= MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY);
559 assert_eq!(
560 scanner.in_flight, MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY,
561 "the initial prefetch window must be bounded and fully scheduled"
562 );
563 let drop_started = std::time::Instant::now();
564 drop(scanner);
565 assert!(
566 drop_started.elapsed() < std::time::Duration::from_secs(5),
567 "dropping a scanner with outstanding tasks must cancel and join promptly"
568 );
569 Ok(())
570 }
571
572 #[cfg(target_os = "macos")]
573 #[test]
574 fn parallel_raw_key_scanners_share_darwin_lock_acquisition_gate() -> Result<(), StoreError> {
575 const SCANNERS: usize = 12;
576
577 let temp = TempDir::new().map_err(StoreError::Io)?;
578 let mut readers = Vec::with_capacity(SCANNERS);
579 for index in 0..SCANNERS {
580 let store = LmdbBlobStore::new(temp.path().join(format!("source-{index}")))?;
581 let hash = hash_in_range(index as u16, index as u64);
582 insert_raw(&store, hash, b"one")?;
583 readers.push((
584 LmdbBlobReader {
585 store,
586 external_read_concurrency: 1,
587 },
588 hash,
589 ));
590 }
591
592 let start = Arc::new(std::sync::Barrier::new(SCANNERS));
593 let handles = readers
594 .into_iter()
595 .map(|(reader, expected)| {
596 let start = Arc::clone(&start);
597 thread::spawn(move || -> Result<(), String> {
598 start.wait();
599 let mut scanner = reader
600 .parallel_raw_key_scanner(None, MAX_PARALLEL_RAW_KEY_SCAN_CONCURRENCY)
601 .map_err(|error| error.to_string())?;
602 let hashes = collect(&mut scanner, 1).map_err(|error| error.to_string())?;
603 if hashes != [expected] {
604 return Err("parallel Darwin scanner returned the wrong keyset".into());
605 }
606 Ok(())
607 })
608 })
609 .collect::<Vec<_>>();
610 let failures = handles
611 .into_iter()
612 .filter_map(|handle| handle.join().expect("scanner thread panicked").err())
613 .collect::<Vec<_>>();
614 assert!(
615 failures.is_empty(),
616 "Darwin LMDB lock acquisition rejected parallel scanners: {}",
617 failures.join(" | ")
618 );
619 Ok(())
620 }
621
622 #[test]
623 #[ignore = "real ephemeral LMDB throughput benchmark"]
624 fn parallel_raw_key_scan_benchmark() -> Result<(), StoreError> {
625 const ENTRIES: usize = 250_000;
626 const PAGE_SIZE: usize = 4_096;
627 let temp = TempDir::new().map_err(StoreError::Io)?;
628 let store = LmdbBlobStore::new(temp.path().join("source"))?;
629 let value = vec![0x5a; 512];
630 let mut wtxn = store.env.write_txn().map_err(map_heed_error)?;
631 for index in 0..ENTRIES {
632 let hash = hash_in_range((index % HASH_RANGE_COUNT as usize) as u16, index as u64);
635 store
636 .blobs
637 .put(&mut wtxn, &hash, &value)
638 .map_err(map_heed_error)?;
639 }
640 wtxn.commit().map_err(map_heed_error)?;
641 store.force_sync()?;
642 let reader = LmdbBlobReader {
643 store,
644 external_read_concurrency: 1,
645 };
646
647 let sequential_started = std::time::Instant::now();
648 let mut sequential_count = 0usize;
649 let mut cursor = None;
650 loop {
651 let hashes = reader.scan_hashes_after(cursor, PAGE_SIZE)?;
652 if hashes.is_empty() {
653 break;
654 }
655 sequential_count += hashes.len();
656 cursor = hashes.last().copied();
657 }
658 let sequential_elapsed = sequential_started.elapsed();
659
660 let parallel_started = std::time::Instant::now();
661 let mut scanner = reader.parallel_raw_key_scanner(None, 16)?;
662 let parallel_count = collect(&mut scanner, PAGE_SIZE)?.len();
663 let parallel_elapsed = parallel_started.elapsed();
664 assert_eq!(sequential_count, ENTRIES);
665 assert_eq!(parallel_count, ENTRIES);
666 eprintln!(
667 "real ephemeral LMDB raw-key scan: entries={ENTRIES} sequential_ms={} parallel_ms={} sequential_keys_s={:.0} parallel_keys_s={:.0}",
668 sequential_elapsed.as_millis(),
669 parallel_elapsed.as_millis(),
670 ENTRIES as f64 / sequential_elapsed.as_secs_f64(),
671 ENTRIES as f64 / parallel_elapsed.as_secs_f64(),
672 );
673 Ok(())
674 }
675}