1use async_trait::async_trait;
4use hashtree_core::store::{Store, StoreError, StoreStats};
5use hashtree_core::types::Hash;
6use heed::types::*;
7use heed::{Database, EnvOpenOptions, Error as HeedError, MdbError, PutFlags};
8use std::path::Path;
9use std::sync::atomic::{AtomicU64, Ordering};
10
11pub use hashtree_core::hash::sha256 as compute_sha256;
13
14#[cfg(target_pointer_width = "64")]
15const DEFAULT_MAP_SIZE: usize = 10 * 1024 * 1024 * 1024;
16#[cfg(target_pointer_width = "32")]
17const DEFAULT_MAP_SIZE: usize = 1024 * 1024 * 1024;
18const DEFAULT_MAX_READERS: u32 = 1024;
19const DATABASE_COUNT: u32 = 5;
20const REOPEN_HEADROOM_BYTES: u64 = 64 * 1024 * 1024;
21const EVICTION_BATCH_TARGET_BYTES: u64 = 256 * 1024 * 1024;
22const EVICTION_BATCH_MAX_ITEMS: usize = 4096;
23const BLOB_META_BYTES: usize = 16;
24const ORDER_KEY_BYTES: usize = 40;
25const PIN_COUNT_BYTES: usize = 4;
26const STORE_TOTALS_BYTES: usize = 32;
27const STORE_TOTALS_KEY: &[u8] = b"totals";
28
29#[derive(Debug, Clone, Copy)]
30struct BlobMeta {
31 order: u64,
32 size: u64,
33}
34
35#[derive(Debug, Default, Clone, Copy)]
36struct StoreTotals {
37 count: u64,
38 total_bytes: u64,
39 pinned_count: u64,
40 pinned_bytes: u64,
41}
42
43pub struct LmdbBlobStore {
45 env: heed::Env,
46 blobs: Database<Bytes, Bytes>,
48 metadata: Database<Bytes, Bytes>,
50 eviction_order: Database<Bytes, Unit>,
52 pins: Database<Bytes, Bytes>,
54 stats: Database<Bytes, Bytes>,
56 max_bytes: AtomicU64,
57 current_bytes: AtomicU64,
58 next_order: AtomicU64,
59}
60
61impl LmdbBlobStore {
62 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, StoreError> {
64 Self::with_map_size(path, DEFAULT_MAP_SIZE)
65 }
66
67 pub fn with_max_bytes<P: AsRef<Path>>(path: P, max_bytes: u64) -> Result<Self, StoreError> {
69 let path_ref = path.as_ref();
70 let existing_map_size = std::fs::metadata(path_ref.join("data.mdb"))
71 .map(|metadata| metadata.len())
72 .unwrap_or(0);
73 let existing_headroom = if existing_map_size == 0 {
74 0
75 } else {
76 existing_map_size
77 .saturating_div(10)
78 .max(REOPEN_HEADROOM_BYTES)
79 };
80 let requested_map_size = usize::try_from(Self::align_map_size_bytes(
81 existing_map_size
82 .saturating_add(existing_headroom)
83 .max(max_bytes),
84 ))
85 .unwrap_or(usize::MAX)
86 .max(DEFAULT_MAP_SIZE);
87 let store = Self::with_map_size(path_ref, requested_map_size)?;
88 store.max_bytes.store(max_bytes, Ordering::Relaxed);
89 let current = store.current_bytes.load(Ordering::Relaxed);
90 if max_bytes > 0 && current > max_bytes {
91 let target = max_bytes.saturating_mul(9) / 10;
92 store.evict_to_target(current, target)?;
93 }
94 Ok(store)
95 }
96
97 fn align_map_size_bytes(bytes: u64) -> u64 {
98 if bytes == 0 {
99 return 0;
100 }
101 let page_size = (page_size::get() as u64).max(4096);
102 let remainder = bytes % page_size;
103 if remainder == 0 {
104 bytes
105 } else {
106 bytes.saturating_add(page_size - remainder)
107 }
108 }
109
110 pub fn with_map_size<P: AsRef<Path>>(path: P, map_size: usize) -> Result<Self, StoreError> {
112 let path_ref = path.as_ref();
113 std::fs::create_dir_all(path_ref).map_err(StoreError::Io)?;
114 let existing_map_size = std::fs::metadata(path_ref.join("data.mdb"))
115 .map(|metadata| metadata.len())
116 .unwrap_or(0);
117 let existing_headroom = if existing_map_size == 0 {
118 0
119 } else {
120 existing_map_size
121 .saturating_div(10)
122 .max(REOPEN_HEADROOM_BYTES)
123 };
124 let requested_map_size = u64::try_from(map_size).unwrap_or(u64::MAX);
125 let map_size = usize::try_from(Self::align_map_size_bytes(
126 requested_map_size.max(existing_map_size.saturating_add(existing_headroom)),
127 ))
128 .unwrap_or(usize::MAX);
129
130 let env = unsafe {
131 EnvOpenOptions::new()
132 .map_size(map_size)
133 .max_dbs(DATABASE_COUNT)
134 .max_readers(DEFAULT_MAX_READERS)
135 .open(path_ref)
136 .map_err(map_heed_error)?
137 };
138 let _ = env.clear_stale_readers();
139 if env.info().map_size < map_size {
140 unsafe { env.resize(map_size) }.map_err(map_heed_error)?;
141 }
142
143 let mut wtxn = env.write_txn().map_err(map_heed_error)?;
144 let blobs = env
145 .create_database(&mut wtxn, Some("blobs"))
146 .map_err(map_heed_error)?;
147 let metadata = env
148 .create_database(&mut wtxn, Some("metadata"))
149 .map_err(map_heed_error)?;
150 let eviction_order = env
151 .create_database(&mut wtxn, Some("eviction_order"))
152 .map_err(map_heed_error)?;
153 let pins = env
154 .create_database(&mut wtxn, Some("pins"))
155 .map_err(map_heed_error)?;
156 let stats = env
157 .create_database(&mut wtxn, Some("stats"))
158 .map_err(map_heed_error)?;
159 wtxn.commit().map_err(map_heed_error)?;
160
161 let next_order = {
162 let rtxn = env.read_txn().map_err(map_heed_error)?;
163 let next = eviction_order
164 .iter(&rtxn)
165 .map_err(map_heed_error)?
166 .last()
167 .transpose()
168 .map_err(map_heed_error)?
169 .map(|(key, _)| {
170 Self::decode_order_from_order_key(key).map(|order| order.saturating_add(1))
171 })
172 .transpose()?
173 .unwrap_or(0);
174 next
175 };
176 let totals = Self::load_or_rebuild_totals(&env, metadata, pins, stats)?;
177
178 Ok(Self {
179 env,
180 blobs,
181 metadata,
182 eviction_order,
183 pins,
184 stats,
185 max_bytes: AtomicU64::new(0),
186 current_bytes: AtomicU64::new(totals.total_bytes),
187 next_order: AtomicU64::new(next_order),
188 })
189 }
190
191 fn load_or_rebuild_totals(
192 env: &heed::Env,
193 metadata: Database<Bytes, Bytes>,
194 pins: Database<Bytes, Bytes>,
195 stats: Database<Bytes, Bytes>,
196 ) -> Result<StoreTotals, StoreError> {
197 {
198 let rtxn = env.read_txn().map_err(map_heed_error)?;
199 if let Some(totals) = Self::read_store_totals(stats, &rtxn)? {
200 return Ok(totals);
201 }
202 }
203
204 let mut totals = StoreTotals::default();
205 {
206 let rtxn = env.read_txn().map_err(map_heed_error)?;
207 for item in metadata.iter(&rtxn).map_err(map_heed_error)? {
208 let (_, bytes) = item.map_err(map_heed_error)?;
209 let meta = Self::decode_blob_meta(bytes)?;
210 totals.count = totals.count.saturating_add(1);
211 totals.total_bytes = totals.total_bytes.saturating_add(meta.size);
212 }
213 for item in pins.iter(&rtxn).map_err(map_heed_error)? {
214 let (hash, pin_bytes) = item.map_err(map_heed_error)?;
215 let count = Self::decode_pin_count(pin_bytes)?;
216 if count == 0 {
217 continue;
218 }
219 let Some(meta_bytes) = metadata.get(&rtxn, hash).map_err(map_heed_error)? else {
220 continue;
221 };
222 let meta = Self::decode_blob_meta(meta_bytes)?;
223 totals.pinned_count = totals.pinned_count.saturating_add(1);
224 totals.pinned_bytes = totals.pinned_bytes.saturating_add(meta.size);
225 }
226 }
227
228 let mut wtxn = env.write_txn().map_err(map_heed_error)?;
229 Self::write_store_totals(stats, &mut wtxn, totals).map_err(map_heed_error)?;
230 wtxn.commit().map_err(map_heed_error)?;
231 Ok(totals)
232 }
233
234 fn read_store_totals(
235 stats: Database<Bytes, Bytes>,
236 txn: &heed::RoTxn,
237 ) -> Result<Option<StoreTotals>, StoreError> {
238 stats
239 .get(txn, STORE_TOTALS_KEY)
240 .map_err(map_heed_error)?
241 .map(Self::decode_store_totals)
242 .transpose()
243 }
244
245 fn read_store_totals_lossy(
246 stats: Database<Bytes, Bytes>,
247 txn: &heed::RoTxn,
248 ) -> std::result::Result<StoreTotals, HeedError> {
249 Ok(stats
250 .get(txn, STORE_TOTALS_KEY)?
251 .and_then(Self::decode_store_totals_lossy)
252 .unwrap_or_default())
253 }
254
255 fn write_store_totals(
256 stats: Database<Bytes, Bytes>,
257 txn: &mut heed::RwTxn,
258 totals: StoreTotals,
259 ) -> std::result::Result<(), HeedError> {
260 let encoded = Self::encode_store_totals(totals);
261 stats.put(txn, STORE_TOTALS_KEY, &encoded)
262 }
263
264 fn increment_totals_in_txn(
265 &self,
266 txn: &mut heed::RwTxn,
267 count: u64,
268 total_bytes: u64,
269 pinned_count: u64,
270 pinned_bytes: u64,
271 ) -> std::result::Result<(), HeedError> {
272 let mut totals = Self::read_store_totals_lossy(self.stats, txn)?;
273 totals.count = totals.count.saturating_add(count);
274 totals.total_bytes = totals.total_bytes.saturating_add(total_bytes);
275 totals.pinned_count = totals.pinned_count.saturating_add(pinned_count);
276 totals.pinned_bytes = totals.pinned_bytes.saturating_add(pinned_bytes);
277 Self::write_store_totals(self.stats, txn, totals)
278 }
279
280 fn decrement_totals_in_txn(
281 &self,
282 txn: &mut heed::RwTxn,
283 count: u64,
284 total_bytes: u64,
285 pinned_count: u64,
286 pinned_bytes: u64,
287 ) -> std::result::Result<(), HeedError> {
288 let mut totals = Self::read_store_totals_lossy(self.stats, txn)?;
289 totals.count = totals.count.saturating_sub(count);
290 totals.total_bytes = totals.total_bytes.saturating_sub(total_bytes);
291 totals.pinned_count = totals.pinned_count.saturating_sub(pinned_count);
292 totals.pinned_bytes = totals.pinned_bytes.saturating_sub(pinned_bytes);
293 Self::write_store_totals(self.stats, txn, totals)
294 }
295
296 pub fn exists(&self, hash: &Hash) -> Result<bool, StoreError> {
298 let rtxn = self
299 .env
300 .read_txn()
301 .map_err(|e| StoreError::Other(e.to_string()))?;
302
303 Ok(self
304 .blobs
305 .get(&rtxn, hash)
306 .map_err(|e| StoreError::Other(e.to_string()))?
307 .is_some())
308 }
309
310 pub fn map_size_bytes(&self) -> usize {
311 self.env.info().map_size
312 }
313
314 fn evict_before_write(&self, incoming_bytes: u64) -> Result<u64, StoreError> {
315 let max = self.max_bytes.load(Ordering::Relaxed);
316 if max == 0 {
317 return Ok(0);
318 }
319
320 let current = self.current_bytes.load(Ordering::Relaxed);
321 if current.saturating_add(incoming_bytes) <= max {
322 return Ok(0);
323 }
324
325 let target = if incoming_bytes >= max {
326 0
327 } else {
328 (max.saturating_mul(9) / 10).min(max.saturating_sub(incoming_bytes))
329 };
330 self.evict_to_target(current, target)
331 }
332
333 fn evict_for_write_pressure(&self, incoming_bytes: u64) -> Result<u64, StoreError> {
334 let current = self.current_bytes.load(Ordering::Relaxed);
335 if current == 0 {
336 return Ok(0);
337 }
338
339 let headroom = incoming_bytes.max(current / 10).max(1);
340 let target = current.saturating_sub(headroom);
341 self.evict_to_target(current, target)
342 }
343
344 fn is_map_full_error(err: &HeedError) -> bool {
345 matches!(err, HeedError::Mdb(MdbError::MapFull))
346 }
347
348 fn put_sync_attempt(&self, hash: Hash, data: &[u8]) -> std::result::Result<bool, HeedError> {
349 let mut wtxn = self.env.write_txn()?;
350 let inserted =
351 match self
352 .blobs
353 .put_with_flags(&mut wtxn, PutFlags::NO_OVERWRITE, &hash, data)
354 {
355 Ok(()) => true,
356 Err(HeedError::Mdb(MdbError::KeyExist)) => false,
357 Err(err) => return Err(err),
358 };
359
360 if inserted {
361 let order = self.next_order.fetch_add(1, Ordering::Relaxed);
362 let meta = Self::encode_blob_meta(BlobMeta {
363 order,
364 size: data.len() as u64,
365 });
366 let order_key = Self::encode_order_key(order, &hash);
367 self.metadata.put(&mut wtxn, &hash, &meta)?;
368 self.eviction_order.put(&mut wtxn, &order_key, &())?;
369 let pin_count = self.read_pin_count_lossy(&wtxn, &hash)?;
370 self.increment_totals_in_txn(
371 &mut wtxn,
372 1,
373 data.len() as u64,
374 u64::from(pin_count > 0),
375 if pin_count > 0 { data.len() as u64 } else { 0 },
376 )?;
377 }
378
379 wtxn.commit()?;
380 Ok(inserted)
381 }
382
383 fn put_many_sync_attempt(
384 &self,
385 items: &[(Hash, Vec<u8>)],
386 ) -> std::result::Result<(usize, u64), HeedError> {
387 let mut wtxn = self.env.write_txn()?;
388 let mut inserted = 0usize;
389 let mut inserted_bytes = 0u64;
390 let mut inserted_pinned = 0u64;
391 let mut inserted_pinned_bytes = 0u64;
392
393 for (hash, data) in items {
394 let inserted_blob = match self.blobs.put_with_flags(
395 &mut wtxn,
396 PutFlags::NO_OVERWRITE,
397 hash,
398 data.as_slice(),
399 ) {
400 Ok(()) => true,
401 Err(HeedError::Mdb(MdbError::KeyExist)) => false,
402 Err(err) => return Err(err),
403 };
404
405 if !inserted_blob {
406 continue;
407 }
408
409 let order = self.next_order.fetch_add(1, Ordering::Relaxed);
410 let meta = Self::encode_blob_meta(BlobMeta {
411 order,
412 size: data.len() as u64,
413 });
414 let order_key = Self::encode_order_key(order, hash);
415 self.metadata.put(&mut wtxn, hash, &meta)?;
416 self.eviction_order.put(&mut wtxn, &order_key, &())?;
417 inserted += 1;
418 inserted_bytes = inserted_bytes.saturating_add(data.len() as u64);
419 if self.read_pin_count_lossy(&wtxn, hash)? > 0 {
420 inserted_pinned = inserted_pinned.saturating_add(1);
421 inserted_pinned_bytes = inserted_pinned_bytes.saturating_add(data.len() as u64);
422 }
423 }
424
425 if inserted > 0 {
426 self.increment_totals_in_txn(
427 &mut wtxn,
428 inserted as u64,
429 inserted_bytes,
430 inserted_pinned,
431 inserted_pinned_bytes,
432 )?;
433 }
434
435 wtxn.commit()?;
436 Ok((inserted, inserted_bytes))
437 }
438
439 pub fn stats(&self) -> Result<LmdbStats, StoreError> {
441 let rtxn = self
442 .env
443 .read_txn()
444 .map_err(|e| StoreError::Other(e.to_string()))?;
445 let totals = Self::read_store_totals(self.stats, &rtxn)?.unwrap_or_default();
446
447 Ok(LmdbStats {
448 count: totals.count as usize,
449 total_bytes: totals.total_bytes,
450 pinned_count: totals.pinned_count as usize,
451 pinned_bytes: totals.pinned_bytes,
452 })
453 }
454
455 pub fn list(&self) -> Result<Vec<Hash>, StoreError> {
457 let rtxn = self
458 .env
459 .read_txn()
460 .map_err(|e| StoreError::Other(e.to_string()))?;
461
462 let mut hashes = Vec::new();
463 for item in self
464 .blobs
465 .iter(&rtxn)
466 .map_err(|e| StoreError::Other(e.to_string()))?
467 {
468 let (hash, _) = item.map_err(|e| StoreError::Other(e.to_string()))?;
469 let hash_arr: Hash = hash
470 .try_into()
471 .map_err(|_| StoreError::Other("invalid hash length".into()))?;
472 hashes.push(hash_arr);
473 }
474
475 Ok(hashes)
476 }
477
478 pub fn put_sync(&self, hash: Hash, data: &[u8]) -> Result<bool, StoreError> {
480 let incoming_bytes = data.len() as u64;
481 self.evict_before_write(incoming_bytes)?;
482
483 let mut retried_after_eviction = false;
484 loop {
485 match self.put_sync_attempt(hash, data) {
486 Ok(inserted) => {
487 if inserted {
488 self.current_bytes
489 .fetch_add(incoming_bytes, Ordering::Relaxed);
490 }
491 return Ok(inserted);
492 }
493 Err(err) if Self::is_map_full_error(&err) && !retried_after_eviction => {
494 let freed = self.evict_for_write_pressure(incoming_bytes)?;
495 if freed == 0 {
496 return Err(StoreError::Other(err.to_string()));
497 }
498 retried_after_eviction = true;
499 }
500 Err(err) => return Err(StoreError::Other(err.to_string())),
501 }
502 }
503 }
504
505 pub fn put_many_sync(&self, items: &[(Hash, Vec<u8>)]) -> Result<usize, StoreError> {
507 if items.is_empty() {
508 return Ok(0);
509 }
510
511 let incoming_bytes = items
512 .iter()
513 .map(|(_, data)| data.len() as u64)
514 .fold(0u64, |total, size| total.saturating_add(size));
515 self.evict_before_write(incoming_bytes)?;
516
517 let mut retried_after_eviction = false;
518 loop {
519 match self.put_many_sync_attempt(items) {
520 Ok((inserted, inserted_bytes)) => {
521 if inserted_bytes > 0 {
522 self.current_bytes
523 .fetch_add(inserted_bytes, Ordering::Relaxed);
524 }
525 return Ok(inserted);
526 }
527 Err(err) if Self::is_map_full_error(&err) && !retried_after_eviction => {
528 let freed = self.evict_for_write_pressure(incoming_bytes)?;
529 if freed == 0 {
530 return Err(StoreError::Other(err.to_string()));
531 }
532 retried_after_eviction = true;
533 }
534 Err(err) => return Err(StoreError::Other(err.to_string())),
535 }
536 }
537 }
538
539 pub fn get_sync(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
541 let rtxn = self
542 .env
543 .read_txn()
544 .map_err(|e| StoreError::Other(e.to_string()))?;
545
546 Ok(self
547 .blobs
548 .get(&rtxn, hash)
549 .map_err(|e| StoreError::Other(e.to_string()))?
550 .map(|b| b.to_vec()))
551 }
552
553 pub fn delete_sync(&self, hash: &Hash) -> Result<bool, StoreError> {
555 let mut wtxn = self
556 .env
557 .write_txn()
558 .map_err(|e| StoreError::Other(e.to_string()))?;
559 let (existed, freed) = self.delete_blob_in_txn(&mut wtxn, hash)?;
560
561 wtxn.commit()
562 .map_err(|e| StoreError::Other(e.to_string()))?;
563 if freed > 0 {
564 self.current_bytes.fetch_sub(freed, Ordering::Relaxed);
565 }
566
567 Ok(existed)
568 }
569
570 fn pin_sync(&self, hash: &Hash) -> Result<(), StoreError> {
571 let mut wtxn = self
572 .env
573 .write_txn()
574 .map_err(|e| StoreError::Other(e.to_string()))?;
575 let previous = self.read_pin_count(&wtxn, hash)?;
576 let count = previous.saturating_add(1);
577 let encoded = count.to_be_bytes();
578 self.pins
579 .put(&mut wtxn, hash, &encoded)
580 .map_err(|e| StoreError::Other(e.to_string()))?;
581 if previous == 0 {
582 if let Some(size) = self.blob_size_in_txn(&wtxn, hash)? {
583 self.increment_totals_in_txn(&mut wtxn, 0, 0, 1, size)
584 .map_err(|e| StoreError::Other(e.to_string()))?;
585 }
586 }
587 wtxn.commit()
588 .map_err(|e| StoreError::Other(e.to_string()))?;
589 Ok(())
590 }
591
592 fn unpin_sync(&self, hash: &Hash) -> Result<(), StoreError> {
593 let mut wtxn = self
594 .env
595 .write_txn()
596 .map_err(|e| StoreError::Other(e.to_string()))?;
597 let count = self.read_pin_count(&wtxn, hash)?;
598 if count == 1 {
599 if let Some(size) = self.blob_size_in_txn(&wtxn, hash)? {
600 self.decrement_totals_in_txn(&mut wtxn, 0, 0, 1, size)
601 .map_err(|e| StoreError::Other(e.to_string()))?;
602 }
603 }
604 if count <= 1 {
605 let _ = self
606 .pins
607 .delete(&mut wtxn, hash)
608 .map_err(|e| StoreError::Other(e.to_string()))?;
609 } else {
610 let encoded = (count - 1).to_be_bytes();
611 self.pins
612 .put(&mut wtxn, hash, &encoded)
613 .map_err(|e| StoreError::Other(e.to_string()))?;
614 }
615 wtxn.commit()
616 .map_err(|e| StoreError::Other(e.to_string()))?;
617 Ok(())
618 }
619
620 fn evict_to_target(&self, current_bytes: u64, target_bytes: u64) -> Result<u64, StoreError> {
621 if current_bytes <= target_bytes {
622 return Ok(0);
623 }
624
625 let rtxn = self
626 .env
627 .read_txn()
628 .map_err(|e| StoreError::Other(e.to_string()))?;
629 let order_keys: Vec<Vec<u8>> = self
630 .eviction_order
631 .iter(&rtxn)
632 .map_err(|e| StoreError::Other(e.to_string()))?
633 .map(|item| {
634 item.map(|(key, _)| key.to_vec())
635 .map_err(|e| StoreError::Other(e.to_string()))
636 })
637 .collect::<Result<_, _>>()?;
638 drop(rtxn);
639
640 let mut freed_total = 0u64;
641 let to_free = current_bytes - target_bytes;
642 let mut index = 0usize;
643
644 while freed_total < to_free && index < order_keys.len() {
645 let mut wtxn = self
646 .env
647 .write_txn()
648 .map_err(|e| StoreError::Other(e.to_string()))?;
649 let mut batch_freed = 0u64;
650 let mut batch_items = 0usize;
651
652 while freed_total + batch_freed < to_free && index < order_keys.len() {
653 let order_key = &order_keys[index];
654 index += 1;
655
656 let hash = Self::decode_hash_from_order_key(order_key)?;
657 if self.read_pin_count(&wtxn, &hash)? > 0 {
658 continue;
659 }
660
661 let (_, bytes_freed) = self.delete_blob_in_txn(&mut wtxn, &hash)?;
662 if bytes_freed == 0 {
663 let _ = self
664 .eviction_order
665 .delete(&mut wtxn, order_key)
666 .map_err(|e| StoreError::Other(e.to_string()))?;
667 continue;
668 }
669
670 batch_freed = batch_freed.saturating_add(bytes_freed);
671 batch_items += 1;
672
673 if batch_freed >= EVICTION_BATCH_TARGET_BYTES
674 || batch_items >= EVICTION_BATCH_MAX_ITEMS
675 {
676 break;
677 }
678 }
679
680 wtxn.commit()
681 .map_err(|e| StoreError::Other(e.to_string()))?;
682 if batch_freed > 0 {
683 self.current_bytes.fetch_sub(batch_freed, Ordering::Relaxed);
684 freed_total = freed_total.saturating_add(batch_freed);
685 }
686 }
687
688 Ok(freed_total)
689 }
690
691 fn delete_blob_in_txn(
692 &self,
693 wtxn: &mut heed::RwTxn,
694 hash: &Hash,
695 ) -> Result<(bool, u64), StoreError> {
696 let pin_count = self.read_pin_count(wtxn, hash)?;
697 let data_len = self
698 .blobs
699 .get(wtxn, hash)
700 .map_err(|e| StoreError::Other(e.to_string()))?
701 .map(|data| data.len() as u64);
702 let meta = self
703 .metadata
704 .get(wtxn, hash)
705 .map_err(|e| StoreError::Other(e.to_string()))?
706 .map(Self::decode_blob_meta)
707 .transpose()?;
708
709 let existed = self
710 .blobs
711 .delete(wtxn, hash)
712 .map_err(|e| StoreError::Other(e.to_string()))?;
713 let _ = self
714 .metadata
715 .delete(wtxn, hash)
716 .map_err(|e| StoreError::Other(e.to_string()))?;
717 let _ = self
718 .pins
719 .delete(wtxn, hash)
720 .map_err(|e| StoreError::Other(e.to_string()))?;
721 if let Some(meta) = meta {
722 let order_key = Self::encode_order_key(meta.order, hash);
723 let _ = self
724 .eviction_order
725 .delete(wtxn, &order_key)
726 .map_err(|e| StoreError::Other(e.to_string()))?;
727 }
728 let bytes_freed = data_len.or(meta.map(|m| m.size)).unwrap_or(0);
729 if existed || meta.is_some() {
730 self.decrement_totals_in_txn(
731 wtxn,
732 1,
733 bytes_freed,
734 u64::from(pin_count > 0),
735 if pin_count > 0 { bytes_freed } else { 0 },
736 )
737 .map_err(|e| StoreError::Other(e.to_string()))?;
738 }
739
740 Ok((existed || meta.is_some(), bytes_freed))
741 }
742
743 fn blob_size_in_txn(&self, txn: &heed::RoTxn, hash: &Hash) -> Result<Option<u64>, StoreError> {
744 if let Some(meta) = self
745 .metadata
746 .get(txn, hash)
747 .map_err(|e| StoreError::Other(e.to_string()))?
748 .map(Self::decode_blob_meta)
749 .transpose()?
750 {
751 return Ok(Some(meta.size));
752 }
753
754 self.blobs
755 .get(txn, hash)
756 .map_err(|e| StoreError::Other(e.to_string()))
757 .map(|data| data.map(|bytes| bytes.len() as u64))
758 }
759
760 fn read_pin_count(&self, txn: &heed::RoTxn, hash: &[u8]) -> Result<u32, StoreError> {
761 self.pins
762 .get(txn, hash)
763 .map_err(|e| StoreError::Other(e.to_string()))?
764 .map(Self::decode_pin_count)
765 .transpose()?
766 .map_or(Ok(0), Ok)
767 }
768
769 fn read_pin_count_lossy(
770 &self,
771 txn: &heed::RoTxn,
772 hash: &[u8],
773 ) -> std::result::Result<u32, HeedError> {
774 Ok(self
775 .pins
776 .get(txn, hash)?
777 .and_then(Self::decode_pin_count_lossy)
778 .unwrap_or(0))
779 }
780
781 fn encode_blob_meta(meta: BlobMeta) -> [u8; BLOB_META_BYTES] {
782 let mut encoded = [0u8; BLOB_META_BYTES];
783 encoded[..8].copy_from_slice(&meta.order.to_be_bytes());
784 encoded[8..].copy_from_slice(&meta.size.to_be_bytes());
785 encoded
786 }
787
788 fn decode_blob_meta(bytes: &[u8]) -> Result<BlobMeta, StoreError> {
789 if bytes.len() != BLOB_META_BYTES {
790 return Err(StoreError::Other(format!(
791 "invalid blob metadata length: {}",
792 bytes.len()
793 )));
794 }
795 Ok(BlobMeta {
796 order: Self::decode_order(&bytes[..8])?,
797 size: u64::from_be_bytes(
798 bytes[8..16]
799 .try_into()
800 .map_err(|_| StoreError::Other("invalid blob size bytes".into()))?,
801 ),
802 })
803 }
804
805 fn encode_store_totals(totals: StoreTotals) -> [u8; STORE_TOTALS_BYTES] {
806 let mut encoded = [0u8; STORE_TOTALS_BYTES];
807 encoded[0..8].copy_from_slice(&totals.count.to_be_bytes());
808 encoded[8..16].copy_from_slice(&totals.total_bytes.to_be_bytes());
809 encoded[16..24].copy_from_slice(&totals.pinned_count.to_be_bytes());
810 encoded[24..32].copy_from_slice(&totals.pinned_bytes.to_be_bytes());
811 encoded
812 }
813
814 fn decode_store_totals(bytes: &[u8]) -> Result<StoreTotals, StoreError> {
815 Self::decode_store_totals_lossy(bytes).ok_or_else(|| {
816 StoreError::Other(format!("invalid store totals length: {}", bytes.len()))
817 })
818 }
819
820 fn decode_store_totals_lossy(bytes: &[u8]) -> Option<StoreTotals> {
821 if bytes.len() != STORE_TOTALS_BYTES {
822 return None;
823 }
824 Some(StoreTotals {
825 count: u64::from_be_bytes(bytes[0..8].try_into().ok()?),
826 total_bytes: u64::from_be_bytes(bytes[8..16].try_into().ok()?),
827 pinned_count: u64::from_be_bytes(bytes[16..24].try_into().ok()?),
828 pinned_bytes: u64::from_be_bytes(bytes[24..32].try_into().ok()?),
829 })
830 }
831
832 fn encode_order_key(order: u64, hash: &Hash) -> [u8; ORDER_KEY_BYTES] {
833 let mut key = [0u8; ORDER_KEY_BYTES];
834 key[..8].copy_from_slice(&order.to_be_bytes());
835 key[8..].copy_from_slice(hash);
836 key
837 }
838
839 fn decode_order(bytes: &[u8]) -> Result<u64, StoreError> {
840 if bytes.len() != 8 {
841 return Err(StoreError::Other(format!(
842 "invalid order length: {}",
843 bytes.len()
844 )));
845 }
846 Ok(u64::from_be_bytes(bytes.try_into().map_err(|_| {
847 StoreError::Other("invalid order bytes".into())
848 })?))
849 }
850
851 fn decode_hash_from_order_key(bytes: &[u8]) -> Result<Hash, StoreError> {
852 if bytes.len() != ORDER_KEY_BYTES {
853 return Err(StoreError::Other(format!(
854 "invalid order key length: {}",
855 bytes.len()
856 )));
857 }
858 let mut hash = [0u8; 32];
859 hash.copy_from_slice(&bytes[8..]);
860 Ok(hash)
861 }
862
863 fn decode_order_from_order_key(bytes: &[u8]) -> Result<u64, StoreError> {
864 if bytes.len() != ORDER_KEY_BYTES {
865 return Err(StoreError::Other(format!(
866 "invalid order key length: {}",
867 bytes.len()
868 )));
869 }
870 Self::decode_order(&bytes[..8])
871 }
872
873 fn decode_pin_count(bytes: &[u8]) -> Result<u32, StoreError> {
874 Self::decode_pin_count_lossy(bytes)
875 .ok_or_else(|| StoreError::Other(format!("invalid pin count length: {}", bytes.len())))
876 }
877
878 fn decode_pin_count_lossy(bytes: &[u8]) -> Option<u32> {
879 if bytes.len() != PIN_COUNT_BYTES {
880 return None;
881 }
882 Some(u32::from_be_bytes(bytes.try_into().ok()?))
883 }
884}
885
886fn map_heed_error(error: HeedError) -> StoreError {
887 match error {
888 HeedError::Io(io_error) => StoreError::Io(io_error),
889 other => StoreError::Other(other.to_string()),
890 }
891}
892
893#[derive(Debug, Clone)]
894pub struct LmdbStats {
895 pub count: usize,
896 pub total_bytes: u64,
897 pub pinned_count: usize,
898 pub pinned_bytes: u64,
899}
900
901#[async_trait]
902impl Store for LmdbBlobStore {
903 async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
904 self.put_sync(hash, &data)
905 }
906
907 async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
908 self.put_many_sync(&items)
909 }
910
911 async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
912 self.get_sync(hash)
913 }
914
915 async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
916 self.exists(hash)
917 }
918
919 async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
920 self.delete_sync(hash)
921 }
922
923 fn set_max_bytes(&self, max: u64) {
924 self.max_bytes.store(max, Ordering::Relaxed);
925 }
926
927 fn max_bytes(&self) -> Option<u64> {
928 let max = self.max_bytes.load(Ordering::Relaxed);
929 if max > 0 {
930 Some(max)
931 } else {
932 None
933 }
934 }
935
936 async fn stats(&self) -> StoreStats {
937 match self.stats() {
938 Ok(stats) => StoreStats {
939 count: stats.count as u64,
940 bytes: stats.total_bytes,
941 pinned_count: stats.pinned_count as u64,
942 pinned_bytes: stats.pinned_bytes,
943 },
944 Err(_) => StoreStats::default(),
945 }
946 }
947
948 async fn evict_if_needed(&self) -> Result<u64, StoreError> {
949 let max = self.max_bytes.load(Ordering::Relaxed);
950 if max == 0 {
951 return Ok(0);
952 }
953
954 let current = self.current_bytes.load(Ordering::Relaxed);
955 if current <= max {
956 return Ok(0);
957 }
958
959 let target = max * 9 / 10;
960 self.evict_to_target(current, target)
961 }
962
963 async fn pin(&self, hash: &Hash) -> Result<(), StoreError> {
964 self.pin_sync(hash)
965 }
966
967 async fn unpin(&self, hash: &Hash) -> Result<(), StoreError> {
968 self.unpin_sync(hash)
969 }
970
971 fn pin_count(&self, hash: &Hash) -> u32 {
972 let Ok(rtxn) = self.env.read_txn() else {
973 return 0;
974 };
975 self.read_pin_count(&rtxn, hash).unwrap_or(0)
976 }
977}
978
979#[cfg(test)]
980mod tests {
981 use super::*;
982 use hashtree_core::sha256;
983 use heed::EnvOpenOptions;
984 #[cfg(unix)]
985 use std::path::{Path, PathBuf};
986 #[cfg(unix)]
987 use std::process::Command;
988 use std::sync::{
989 atomic::{AtomicBool, Ordering},
990 Arc, Barrier,
991 };
992 use std::time::Duration;
993 use tempfile::TempDir;
994
995 #[cfg(unix)]
996 const STALE_READER_HELPER_ENV: &str = "HASHTREE_LMDB_STALE_READER_HELPER";
997 #[cfg(unix)]
998 const STALE_READER_HELPER_MODE_ENV: &str = "HASHTREE_LMDB_STALE_READER_HELPER_MODE";
999 #[cfg(unix)]
1000 const STALE_READER_DB_PATH_ENV: &str = "HASHTREE_LMDB_STALE_READER_DB_PATH";
1001 #[cfg(unix)]
1002 const STALE_READER_MARKER_PATH_ENV: &str = "HASHTREE_LMDB_STALE_READER_MARKER_PATH";
1003 #[cfg(unix)]
1004 const TEST_MAX_READERS: u32 = 4;
1005
1006 #[cfg(unix)]
1007 fn run_helper(mode: &str, path: &Path, marker: &Path) {
1008 let output = Command::new(std::env::current_exe().expect("test binary path"))
1009 .arg("--ignored")
1010 .arg("--exact")
1011 .arg("tests::lmdb_stale_reader_helper")
1012 .env(STALE_READER_HELPER_ENV, "1")
1013 .env(STALE_READER_HELPER_MODE_ENV, mode)
1014 .env(STALE_READER_DB_PATH_ENV, path)
1015 .env(STALE_READER_MARKER_PATH_ENV, marker)
1016 .env("RUST_TEST_THREADS", "1")
1017 .output()
1018 .expect("spawn stale-reader helper");
1019
1020 assert!(
1021 output.status.success(),
1022 "stale-reader helper failed: stdout={} stderr={}",
1023 String::from_utf8_lossy(&output.stdout),
1024 String::from_utf8_lossy(&output.stderr)
1025 );
1026 assert!(
1027 marker.exists(),
1028 "stale-reader helper did not run: stdout={} stderr={}",
1029 String::from_utf8_lossy(&output.stdout),
1030 String::from_utf8_lossy(&output.stderr)
1031 );
1032 }
1033
1034 #[tokio::test]
1035 async fn test_put_get() -> Result<(), StoreError> {
1036 let temp = TempDir::new().unwrap();
1037 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1038
1039 let data = b"hello lmdb";
1040 let hash = sha256(data);
1041 store.put(hash, data.to_vec()).await?;
1042
1043 assert!(store.has(&hash).await?);
1044 assert_eq!(store.get(&hash).await?, Some(data.to_vec()));
1045
1046 Ok(())
1047 }
1048
1049 #[tokio::test]
1050 async fn test_delete() -> Result<(), StoreError> {
1051 let temp = TempDir::new().unwrap();
1052 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1053
1054 let data = b"delete me";
1055 let hash = sha256(data);
1056 store.put(hash, data.to_vec()).await?;
1057 assert!(store.has(&hash).await?);
1058
1059 assert!(store.delete(&hash).await?);
1060 assert!(!store.has(&hash).await?);
1061 assert!(!store.delete(&hash).await?);
1062
1063 Ok(())
1064 }
1065
1066 #[tokio::test]
1067 async fn test_list() -> Result<(), StoreError> {
1068 let temp = TempDir::new().unwrap();
1069 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1070
1071 let d1 = b"one";
1072 let d2 = b"two";
1073 let d3 = b"three";
1074 let h1 = sha256(d1);
1075 let h2 = sha256(d2);
1076 let h3 = sha256(d3);
1077
1078 store.put(h1, d1.to_vec()).await?;
1079 store.put(h2, d2.to_vec()).await?;
1080 store.put(h3, d3.to_vec()).await?;
1081
1082 let hashes = store.list()?;
1083 assert_eq!(hashes.len(), 3);
1084 assert!(hashes.contains(&h1));
1085 assert!(hashes.contains(&h2));
1086 assert!(hashes.contains(&h3));
1087
1088 Ok(())
1089 }
1090
1091 #[tokio::test]
1092 async fn test_stats() -> Result<(), StoreError> {
1093 let temp = TempDir::new().unwrap();
1094 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1095
1096 let d1 = b"hello";
1097 let d2 = b"world";
1098 store.put(sha256(d1), d1.to_vec()).await?;
1099 store.put(sha256(d2), d2.to_vec()).await?;
1100
1101 let stats = store.stats()?;
1102 assert_eq!(stats.count, 2);
1103 assert_eq!(stats.total_bytes, 10);
1104
1105 Ok(())
1106 }
1107
1108 #[tokio::test]
1109 async fn test_stats_persist_across_reopen_and_mutations() -> Result<(), StoreError> {
1110 let temp = TempDir::new().unwrap();
1111 let path = temp.path().join("blobs");
1112 let h1 = sha256(b"hello");
1113 let h2 = sha256(b"world!");
1114 let h3 = sha256(b"prepin");
1115
1116 {
1117 let store = LmdbBlobStore::new(&path)?;
1118 store.put(h1, b"hello".to_vec()).await?;
1119 store.put(h2, b"world!".to_vec()).await?;
1120 store.pin(&h1).await?;
1121 store.pin(&h3).await?;
1122 store.put(h3, b"prepin".to_vec()).await?;
1123
1124 let stats = store.stats()?;
1125 assert_eq!(stats.count, 3);
1126 assert_eq!(stats.total_bytes, 17);
1127 assert_eq!(stats.pinned_count, 2);
1128 assert_eq!(stats.pinned_bytes, 11);
1129 }
1130
1131 {
1132 let reopened = LmdbBlobStore::new(&path)?;
1133 let stats = reopened.stats()?;
1134 assert_eq!(stats.count, 3);
1135 assert_eq!(stats.total_bytes, 17);
1136 assert_eq!(stats.pinned_count, 2);
1137 assert_eq!(stats.pinned_bytes, 11);
1138
1139 assert!(reopened.delete(&h1).await?);
1140 let stats = reopened.stats()?;
1141 assert_eq!(stats.count, 2);
1142 assert_eq!(stats.total_bytes, 12);
1143 assert_eq!(stats.pinned_count, 1);
1144 assert_eq!(stats.pinned_bytes, 6);
1145 }
1146
1147 let reopened = LmdbBlobStore::new(&path)?;
1148 let stats = reopened.stats()?;
1149 assert_eq!(stats.count, 2);
1150 assert_eq!(stats.total_bytes, 12);
1151 assert_eq!(stats.pinned_count, 1);
1152 assert_eq!(stats.pinned_bytes, 6);
1153
1154 Ok(())
1155 }
1156
1157 #[tokio::test]
1158 async fn test_deduplication() -> Result<(), StoreError> {
1159 let temp = TempDir::new().unwrap();
1160 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1161
1162 let data = b"same";
1163 let hash = sha256(data);
1164 assert!(store.put(hash, data.to_vec()).await?); assert!(!store.put(hash, data.to_vec()).await?); assert_eq!(store.list()?.len(), 1);
1168
1169 Ok(())
1170 }
1171
1172 #[tokio::test]
1173 async fn test_max_bytes() -> Result<(), StoreError> {
1174 let temp = TempDir::new().unwrap();
1175 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1176
1177 assert!(store.max_bytes().is_none());
1178
1179 store.set_max_bytes(1000);
1180 assert_eq!(store.max_bytes(), Some(1000));
1181
1182 store.set_max_bytes(0);
1183 assert!(store.max_bytes().is_none());
1184
1185 Ok(())
1186 }
1187
1188 #[test]
1189 fn test_with_max_bytes_expands_lmdb_map_size() -> Result<(), StoreError> {
1190 let temp = TempDir::new().unwrap();
1191 let requested = (DEFAULT_MAP_SIZE as u64) + 64 * 1024 * 1024;
1192 let store = LmdbBlobStore::with_max_bytes(temp.path().join("blobs"), requested)?;
1193
1194 assert!(
1195 store.map_size_bytes() as u64 >= requested,
1196 "expected LMDB map to grow to at least {requested} bytes, got {}",
1197 store.map_size_bytes()
1198 );
1199 assert_eq!(store.max_bytes(), Some(requested));
1200
1201 Ok(())
1202 }
1203
1204 #[tokio::test]
1205 async fn test_eviction_over_limit() -> Result<(), StoreError> {
1206 let temp = TempDir::new().unwrap();
1207 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1208 store.set_max_bytes(25);
1209
1210 let h1 = sha256(b"aaaaaaaaaa");
1211 let h2 = sha256(b"bbbbbbbbbb");
1212 let h3 = sha256(b"cccccccccc");
1213
1214 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
1215 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
1216 store.put(h3, b"cccccccccc".to_vec()).await?;
1217
1218 let freed = store.evict_if_needed().await?;
1219 assert_eq!(
1220 freed, 0,
1221 "write path should have already evicted stale blobs"
1222 );
1223
1224 assert!(
1225 !store.has(&h1).await?,
1226 "oldest blob should be evicted before the third write"
1227 );
1228 assert!(store.has(&h2).await?);
1229 assert!(store.has(&h3).await?);
1230
1231 let stats = store.stats()?;
1232 assert!(
1233 stats.total_bytes <= 22,
1234 "store should be reduced to 90% target"
1235 );
1236
1237 Ok(())
1238 }
1239
1240 #[tokio::test]
1241 async fn test_eviction_respects_pins() -> Result<(), StoreError> {
1242 let temp = TempDir::new().unwrap();
1243 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
1244 store.set_max_bytes(25);
1245
1246 let h1 = sha256(b"aaaaaaaaaa");
1247 let h2 = sha256(b"bbbbbbbbbb");
1248 let h3 = sha256(b"cccccccccc");
1249
1250 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
1251 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
1252 store.pin(&h1).await?;
1253 store.put(h3, b"cccccccccc".to_vec()).await?;
1254
1255 let freed = store.evict_if_needed().await?;
1256 assert_eq!(
1257 freed, 0,
1258 "write path should have already evicted stale blobs"
1259 );
1260
1261 assert!(store.has(&h1).await?, "pinned blob must not be evicted");
1262 assert!(
1263 !store.has(&h2).await?,
1264 "oldest unpinned blob should be evicted before the third write"
1265 );
1266 assert!(store.has(&h3).await?);
1267
1268 Ok(())
1269 }
1270
1271 #[tokio::test]
1272 async fn test_reopen_with_existing_eviction_order() -> Result<(), StoreError> {
1273 let temp = TempDir::new().unwrap();
1274 let path = temp.path().join("blobs");
1275
1276 {
1277 let store = LmdbBlobStore::new(&path)?;
1278 let h1 = sha256(b"aaaaaaaaaa");
1279 let h2 = sha256(b"bbbbbbbbbb");
1280 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
1281 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
1282 }
1283
1284 let reopened = LmdbBlobStore::new(&path)?;
1285 let h3 = sha256(b"cccccccccc");
1286 assert!(reopened.put(h3, b"cccccccccc".to_vec()).await?);
1287 assert!(reopened.has(&h3).await?);
1288
1289 Ok(())
1290 }
1291
1292 #[tokio::test]
1293 async fn test_reopen_with_lower_max_bytes_evicts_existing_blobs() -> Result<(), StoreError> {
1294 let temp = TempDir::new().unwrap();
1295 let path = temp.path().join("blobs");
1296
1297 {
1298 let store = LmdbBlobStore::new(&path)?;
1299 let h1 = sha256(b"aaaaaaaaaa");
1300 let h2 = sha256(b"bbbbbbbbbb");
1301 let h3 = sha256(b"cccccccccc");
1302 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
1303 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
1304 store.put(h3, b"cccccccccc".to_vec()).await?;
1305 }
1306
1307 let reopened = LmdbBlobStore::with_max_bytes(&path, 25)?;
1308 let h1 = sha256(b"aaaaaaaaaa");
1309 let h2 = sha256(b"bbbbbbbbbb");
1310 let h3 = sha256(b"cccccccccc");
1311
1312 assert!(
1313 !reopened.has(&h1).await?,
1314 "oldest blob should be evicted when reopening over the new cap"
1315 );
1316 assert!(reopened.has(&h2).await?);
1317 assert!(reopened.has(&h3).await?);
1318
1319 let stats = reopened.stats()?;
1320 assert!(
1321 stats.total_bytes <= 22,
1322 "reopened store should be reduced to the 90% target"
1323 );
1324
1325 Ok(())
1326 }
1327
1328 #[test]
1329 fn test_supports_many_concurrent_readers() -> Result<(), Box<dyn std::error::Error>> {
1330 const READER_THREADS: usize = 160;
1331
1332 let temp = TempDir::new()?;
1333 let store = Arc::new(LmdbBlobStore::new(temp.path().join("blobs"))?);
1334 let hash = sha256(b"many readers");
1335 store.put_sync(hash, b"many readers")?;
1336
1337 let start = Arc::new(Barrier::new(READER_THREADS + 1));
1338 let release = Arc::new(AtomicBool::new(false));
1339 let mut handles = Vec::with_capacity(READER_THREADS);
1340
1341 for _ in 0..READER_THREADS {
1342 let env = store.env.clone();
1343 let start = Arc::clone(&start);
1344 let release = Arc::clone(&release);
1345 handles.push(std::thread::spawn(move || -> Result<(), String> {
1346 start.wait();
1347 let _rtxn = env.read_txn().map_err(|err| err.to_string())?;
1348 while !release.load(Ordering::Relaxed) {
1349 std::thread::sleep(Duration::from_millis(1));
1350 }
1351 Ok(())
1352 }));
1353 }
1354
1355 start.wait();
1356 std::thread::sleep(Duration::from_millis(50));
1357 release.store(true, Ordering::Relaxed);
1358
1359 let results: Vec<Result<(), String>> = handles
1360 .into_iter()
1361 .map(|handle| handle.join().expect("reader thread panicked"))
1362 .collect();
1363
1364 let failures: Vec<String> = results.into_iter().filter_map(Result::err).collect();
1365 assert!(
1366 failures.is_empty(),
1367 "concurrent reader failures: {}",
1368 failures.join(" | ")
1369 );
1370 assert!(store.exists(&hash)?);
1371
1372 Ok(())
1373 }
1374
1375 #[cfg(unix)]
1376 #[test]
1377 fn test_reclaims_stale_reader_slots() -> Result<(), Box<dyn std::error::Error>> {
1378 let temp = TempDir::new()?;
1379 let path = temp.path().join("blobs");
1380 let data = b"hello stale readers";
1381 let hash = sha256(data);
1382
1383 run_helper("setup", &path, &temp.path().join("setup.marker"));
1384
1385 for index in 0..TEST_MAX_READERS {
1386 let marker = temp.path().join(format!("helper-{index}.marker"));
1387 run_helper("stale", &path, &marker);
1388 }
1389
1390 let store = LmdbBlobStore::with_map_size(&path, 1024 * 1024)?;
1391 assert!(store.exists(&hash)?);
1392
1393 Ok(())
1394 }
1395
1396 #[cfg(unix)]
1397 #[test]
1398 fn test_reopens_existing_env_with_larger_map_size() -> Result<(), Box<dyn std::error::Error>> {
1399 let temp = TempDir::new()?;
1400 let path = temp.path().join("blobs");
1401 run_helper("small-map", &path, &temp.path().join("small-map.marker"));
1402
1403 let reopened = LmdbBlobStore::with_map_size(&path, 8 * 1024 * 1024)?;
1404 assert!(reopened.map_size_bytes() >= 8 * 1024 * 1024);
1405
1406 Ok(())
1407 }
1408
1409 #[cfg(unix)]
1410 #[test]
1411 fn test_reopens_existing_env_with_smaller_requested_map_size(
1412 ) -> Result<(), Box<dyn std::error::Error>> {
1413 let temp = TempDir::new()?;
1414 let path = temp.path().join("blobs");
1415 let hash = sha256(b"large existing blob");
1416 run_helper("large-map", &path, &temp.path().join("large-map.marker"));
1417
1418 let existing_size = std::fs::metadata(path.join("data.mdb"))?.len();
1419 assert!(
1420 existing_size > 1024 * 1024,
1421 "test setup should create an environment larger than the reopen request"
1422 );
1423
1424 let reopened = LmdbBlobStore::with_map_size(&path, 1024 * 1024)?;
1425 assert!(
1426 reopened.map_size_bytes() as u64 >= existing_size,
1427 "expected map size to cover existing data.mdb size {existing_size}, got {}",
1428 reopened.map_size_bytes()
1429 );
1430 assert!(reopened.exists(&hash)?);
1431
1432 Ok(())
1433 }
1434
1435 #[cfg(unix)]
1436 #[test]
1437 #[ignore = "used as a subprocess helper by test_reclaims_stale_reader_slots"]
1438 fn lmdb_stale_reader_helper() {
1439 let Some(db_path) = std::env::var_os(STALE_READER_DB_PATH_ENV) else {
1440 return;
1441 };
1442 let marker_path =
1443 PathBuf::from(std::env::var_os(STALE_READER_MARKER_PATH_ENV).expect("marker path"));
1444 std::fs::write(&marker_path, b"started").expect("write helper marker");
1445
1446 let _env_flag = std::env::var_os(STALE_READER_HELPER_ENV).expect("helper mode enabled");
1447 let mode = std::env::var(STALE_READER_HELPER_MODE_ENV).expect("helper mode");
1448 let db_path = PathBuf::from(db_path);
1449 std::fs::create_dir_all(&db_path).expect("create helper db dir");
1450 let helper_map_size = if mode == "large-map" {
1451 8 * 1024 * 1024
1452 } else {
1453 1024 * 1024
1454 };
1455 let env = unsafe {
1456 EnvOpenOptions::new()
1457 .map_size(helper_map_size)
1458 .max_dbs(DATABASE_COUNT)
1459 .max_readers(TEST_MAX_READERS)
1460 .open(&db_path)
1461 .expect("open lmdb env")
1462 };
1463 match mode.as_str() {
1464 "setup" => {
1465 let mut wtxn = env.write_txn().expect("open write txn");
1466 let blobs: Database<Bytes, Bytes> = env
1467 .create_database(&mut wtxn, Some("blobs"))
1468 .expect("create blobs database");
1469 let data = b"hello stale readers";
1470 let hash = sha256(data);
1471 blobs.put(&mut wtxn, &hash, data).expect("seed blob");
1472 wtxn.commit().expect("commit setup txn");
1473 std::process::exit(0);
1474 }
1475 "stale" => {
1476 let _rtxn = env.read_txn().expect("open read txn");
1477 std::process::exit(0);
1478 }
1479 "small-map" => {
1480 let mut wtxn = env.write_txn().expect("open write txn");
1481 let _blobs: Database<Bytes, Bytes> = env
1482 .create_database(&mut wtxn, Some("blobs"))
1483 .expect("create blobs database");
1484 let _metadata: Database<Bytes, Bytes> = env
1485 .create_database(&mut wtxn, Some("metadata"))
1486 .expect("create metadata database");
1487 let _eviction_order: Database<Bytes, Unit> = env
1488 .create_database(&mut wtxn, Some("eviction_order"))
1489 .expect("create eviction_order database");
1490 let _pins: Database<Bytes, Bytes> = env
1491 .create_database(&mut wtxn, Some("pins"))
1492 .expect("create pins database");
1493 wtxn.commit().expect("commit small-map setup txn");
1494 std::process::exit(0);
1495 }
1496 "large-map" => {
1497 let mut wtxn = env.write_txn().expect("open write txn");
1498 let blobs: Database<Bytes, Bytes> = env
1499 .create_database(&mut wtxn, Some("blobs"))
1500 .expect("create blobs database");
1501 let _metadata: Database<Bytes, Bytes> = env
1502 .create_database(&mut wtxn, Some("metadata"))
1503 .expect("create metadata database");
1504 let _eviction_order: Database<Bytes, Unit> = env
1505 .create_database(&mut wtxn, Some("eviction_order"))
1506 .expect("create eviction_order database");
1507 let _pins: Database<Bytes, Bytes> = env
1508 .create_database(&mut wtxn, Some("pins"))
1509 .expect("create pins database");
1510 let data = vec![42u8; 3 * 1024 * 1024];
1511 let hash = sha256(b"large existing blob");
1512 blobs.put(&mut wtxn, &hash, &data).expect("seed blob");
1513 wtxn.commit().expect("commit large-map setup txn");
1514 std::process::exit(0);
1515 }
1516 other => panic!("unknown helper mode: {other}"),
1517 }
1518 }
1519}