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 = 4;
20const BLOB_META_BYTES: usize = 16;
21const ORDER_KEY_BYTES: usize = 40;
22const PIN_COUNT_BYTES: usize = 4;
23
24#[derive(Debug, Clone, Copy)]
25struct BlobMeta {
26 order: u64,
27 size: u64,
28}
29
30pub struct LmdbBlobStore {
32 env: heed::Env,
33 blobs: Database<Bytes, Bytes>,
35 metadata: Database<Bytes, Bytes>,
37 eviction_order: Database<Bytes, Unit>,
39 pins: Database<Bytes, Bytes>,
41 max_bytes: AtomicU64,
42 next_order: AtomicU64,
43}
44
45impl LmdbBlobStore {
46 pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, StoreError> {
48 Self::with_map_size(path, DEFAULT_MAP_SIZE)
49 }
50
51 pub fn with_max_bytes<P: AsRef<Path>>(path: P, max_bytes: u64) -> Result<Self, StoreError> {
53 let store = Self::new(path)?;
54 store.max_bytes.store(max_bytes, Ordering::Relaxed);
55 Ok(store)
56 }
57
58 pub fn with_map_size<P: AsRef<Path>>(path: P, map_size: usize) -> Result<Self, StoreError> {
60 std::fs::create_dir_all(&path).map_err(StoreError::Io)?;
61
62 let env = unsafe {
63 EnvOpenOptions::new()
64 .map_size(map_size)
65 .max_dbs(DATABASE_COUNT)
66 .max_readers(DEFAULT_MAX_READERS)
67 .open(path)
68 .map_err(|e| StoreError::Other(e.to_string()))?
69 };
70 let _ = env.clear_stale_readers();
71 if env.info().map_size < map_size {
72 unsafe { env.resize(map_size) }.map_err(|e| StoreError::Other(e.to_string()))?;
73 }
74
75 let mut wtxn = env
76 .write_txn()
77 .map_err(|e| StoreError::Other(e.to_string()))?;
78 let blobs = env
79 .create_database(&mut wtxn, Some("blobs"))
80 .map_err(|e| StoreError::Other(e.to_string()))?;
81 let metadata = env
82 .create_database(&mut wtxn, Some("metadata"))
83 .map_err(|e| StoreError::Other(e.to_string()))?;
84 let eviction_order = env
85 .create_database(&mut wtxn, Some("eviction_order"))
86 .map_err(|e| StoreError::Other(e.to_string()))?;
87 let pins = env
88 .create_database(&mut wtxn, Some("pins"))
89 .map_err(|e| StoreError::Other(e.to_string()))?;
90 wtxn.commit()
91 .map_err(|e| StoreError::Other(e.to_string()))?;
92
93 let next_order = {
94 let rtxn = env
95 .read_txn()
96 .map_err(|e| StoreError::Other(e.to_string()))?;
97 let next = eviction_order
98 .iter(&rtxn)
99 .map_err(|e| StoreError::Other(e.to_string()))?
100 .last()
101 .transpose()
102 .map_err(|e| StoreError::Other(e.to_string()))?
103 .map(|(key, _)| {
104 Self::decode_order_from_order_key(key).map(|order| order.saturating_add(1))
105 })
106 .transpose()?
107 .unwrap_or(0);
108 next
109 };
110
111 Ok(Self {
112 env,
113 blobs,
114 metadata,
115 eviction_order,
116 pins,
117 max_bytes: AtomicU64::new(0),
118 next_order: AtomicU64::new(next_order),
119 })
120 }
121
122 pub fn exists(&self, hash: &Hash) -> Result<bool, StoreError> {
124 let rtxn = self
125 .env
126 .read_txn()
127 .map_err(|e| StoreError::Other(e.to_string()))?;
128
129 Ok(self
130 .blobs
131 .get(&rtxn, hash)
132 .map_err(|e| StoreError::Other(e.to_string()))?
133 .is_some())
134 }
135
136 pub fn map_size_bytes(&self) -> usize {
137 self.env.info().map_size
138 }
139
140 pub fn stats(&self) -> Result<LmdbStats, StoreError> {
142 let rtxn = self
143 .env
144 .read_txn()
145 .map_err(|e| StoreError::Other(e.to_string()))?;
146
147 let count = self
148 .blobs
149 .len(&rtxn)
150 .map_err(|e| StoreError::Other(e.to_string()))? as usize;
151
152 let mut total_bytes = 0u64;
153 let mut pinned_count = 0usize;
154 let mut pinned_bytes = 0u64;
155 for item in self
156 .blobs
157 .iter(&rtxn)
158 .map_err(|e| StoreError::Other(e.to_string()))?
159 {
160 let (hash, data) = item.map_err(|e| StoreError::Other(e.to_string()))?;
161 let size = data.len() as u64;
162 total_bytes += size;
163 if self.read_pin_count(&rtxn, hash)? > 0 {
164 pinned_count += 1;
165 pinned_bytes += size;
166 }
167 }
168
169 Ok(LmdbStats {
170 count,
171 total_bytes,
172 pinned_count,
173 pinned_bytes,
174 })
175 }
176
177 pub fn list(&self) -> Result<Vec<Hash>, StoreError> {
179 let rtxn = self
180 .env
181 .read_txn()
182 .map_err(|e| StoreError::Other(e.to_string()))?;
183
184 let mut hashes = Vec::new();
185 for item in self
186 .blobs
187 .iter(&rtxn)
188 .map_err(|e| StoreError::Other(e.to_string()))?
189 {
190 let (hash, _) = item.map_err(|e| StoreError::Other(e.to_string()))?;
191 let hash_arr: Hash = hash
192 .try_into()
193 .map_err(|_| StoreError::Other("invalid hash length".into()))?;
194 hashes.push(hash_arr);
195 }
196
197 Ok(hashes)
198 }
199
200 pub fn put_sync(&self, hash: Hash, data: &[u8]) -> Result<bool, StoreError> {
202 let mut wtxn = self
203 .env
204 .write_txn()
205 .map_err(|e| StoreError::Other(e.to_string()))?;
206 let inserted =
207 match self
208 .blobs
209 .put_with_flags(&mut wtxn, PutFlags::NO_OVERWRITE, &hash, data)
210 {
211 Ok(()) => true,
212 Err(HeedError::Mdb(MdbError::KeyExist)) => false,
213 Err(err) => return Err(StoreError::Other(err.to_string())),
214 };
215
216 if inserted {
217 let order = self.next_order.fetch_add(1, Ordering::Relaxed);
218 let meta = Self::encode_blob_meta(BlobMeta {
219 order,
220 size: data.len() as u64,
221 });
222 let order_key = Self::encode_order_key(order, &hash);
223 self.metadata
224 .put(&mut wtxn, &hash, &meta)
225 .map_err(|e| StoreError::Other(e.to_string()))?;
226 self.eviction_order
227 .put(&mut wtxn, &order_key, &())
228 .map_err(|e| StoreError::Other(e.to_string()))?;
229 }
230
231 wtxn.commit()
232 .map_err(|e| StoreError::Other(e.to_string()))?;
233
234 Ok(inserted)
235 }
236
237 pub fn put_many_sync(&self, items: &[(Hash, Vec<u8>)]) -> Result<usize, StoreError> {
239 if items.is_empty() {
240 return Ok(0);
241 }
242
243 let mut wtxn = self
244 .env
245 .write_txn()
246 .map_err(|e| StoreError::Other(e.to_string()))?;
247 let mut inserted = 0usize;
248
249 for (hash, data) in items {
250 let inserted_blob = match self.blobs.put_with_flags(
251 &mut wtxn,
252 PutFlags::NO_OVERWRITE,
253 hash,
254 data.as_slice(),
255 ) {
256 Ok(()) => true,
257 Err(HeedError::Mdb(MdbError::KeyExist)) => false,
258 Err(err) => return Err(StoreError::Other(err.to_string())),
259 };
260
261 if !inserted_blob {
262 continue;
263 }
264
265 let order = self.next_order.fetch_add(1, Ordering::Relaxed);
266 let meta = Self::encode_blob_meta(BlobMeta {
267 order,
268 size: data.len() as u64,
269 });
270 let order_key = Self::encode_order_key(order, hash);
271 self.blobs
272 .put(&mut wtxn, hash, data.as_slice())
273 .map_err(|e| StoreError::Other(e.to_string()))?;
274 self.metadata
275 .put(&mut wtxn, hash, &meta)
276 .map_err(|e| StoreError::Other(e.to_string()))?;
277 self.eviction_order
278 .put(&mut wtxn, &order_key, &())
279 .map_err(|e| StoreError::Other(e.to_string()))?;
280 inserted += 1;
281 }
282
283 wtxn.commit()
284 .map_err(|e| StoreError::Other(e.to_string()))?;
285
286 Ok(inserted)
287 }
288
289 pub fn get_sync(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
291 let rtxn = self
292 .env
293 .read_txn()
294 .map_err(|e| StoreError::Other(e.to_string()))?;
295
296 Ok(self
297 .blobs
298 .get(&rtxn, hash)
299 .map_err(|e| StoreError::Other(e.to_string()))?
300 .map(|b| b.to_vec()))
301 }
302
303 pub fn delete_sync(&self, hash: &Hash) -> Result<bool, StoreError> {
305 let mut wtxn = self
306 .env
307 .write_txn()
308 .map_err(|e| StoreError::Other(e.to_string()))?;
309 let (existed, _) = self.delete_blob_in_txn(&mut wtxn, hash)?;
310
311 wtxn.commit()
312 .map_err(|e| StoreError::Other(e.to_string()))?;
313
314 Ok(existed)
315 }
316
317 fn pin_sync(&self, hash: &Hash) -> Result<(), StoreError> {
318 let mut wtxn = self
319 .env
320 .write_txn()
321 .map_err(|e| StoreError::Other(e.to_string()))?;
322 let count = self.read_pin_count(&wtxn, hash)?.saturating_add(1);
323 let encoded = count.to_be_bytes();
324 self.pins
325 .put(&mut wtxn, hash, &encoded)
326 .map_err(|e| StoreError::Other(e.to_string()))?;
327 wtxn.commit()
328 .map_err(|e| StoreError::Other(e.to_string()))?;
329 Ok(())
330 }
331
332 fn unpin_sync(&self, hash: &Hash) -> Result<(), StoreError> {
333 let mut wtxn = self
334 .env
335 .write_txn()
336 .map_err(|e| StoreError::Other(e.to_string()))?;
337 let count = self.read_pin_count(&wtxn, hash)?;
338 if count <= 1 {
339 let _ = self
340 .pins
341 .delete(&mut wtxn, hash)
342 .map_err(|e| StoreError::Other(e.to_string()))?;
343 } else {
344 let encoded = (count - 1).to_be_bytes();
345 self.pins
346 .put(&mut wtxn, hash, &encoded)
347 .map_err(|e| StoreError::Other(e.to_string()))?;
348 }
349 wtxn.commit()
350 .map_err(|e| StoreError::Other(e.to_string()))?;
351 Ok(())
352 }
353
354 fn evict_to_target(&self, current_bytes: u64, target_bytes: u64) -> Result<u64, StoreError> {
355 if current_bytes <= target_bytes {
356 return Ok(0);
357 }
358
359 let mut wtxn = self
360 .env
361 .write_txn()
362 .map_err(|e| StoreError::Other(e.to_string()))?;
363 let order_keys: Vec<Vec<u8>> = self
364 .eviction_order
365 .iter(&wtxn)
366 .map_err(|e| StoreError::Other(e.to_string()))?
367 .map(|item| {
368 item.map(|(key, _)| key.to_vec())
369 .map_err(|e| StoreError::Other(e.to_string()))
370 })
371 .collect::<Result<_, _>>()?;
372
373 let mut freed = 0u64;
374 let to_free = current_bytes - target_bytes;
375
376 for order_key in order_keys {
377 if freed >= to_free {
378 break;
379 }
380
381 let hash = Self::decode_hash_from_order_key(&order_key)?;
382 if self.read_pin_count(&wtxn, &hash)? > 0 {
383 continue;
384 }
385
386 let (_, bytes_freed) = self.delete_blob_in_txn(&mut wtxn, &hash)?;
387 if bytes_freed == 0 {
388 let _ = self
389 .eviction_order
390 .delete(&mut wtxn, &order_key)
391 .map_err(|e| StoreError::Other(e.to_string()))?;
392 continue;
393 }
394 freed += bytes_freed;
395 }
396
397 wtxn.commit()
398 .map_err(|e| StoreError::Other(e.to_string()))?;
399 Ok(freed)
400 }
401
402 fn delete_blob_in_txn(
403 &self,
404 wtxn: &mut heed::RwTxn,
405 hash: &Hash,
406 ) -> Result<(bool, u64), StoreError> {
407 let data_len = self
408 .blobs
409 .get(wtxn, hash)
410 .map_err(|e| StoreError::Other(e.to_string()))?
411 .map(|data| data.len() as u64);
412 let meta = self
413 .metadata
414 .get(wtxn, hash)
415 .map_err(|e| StoreError::Other(e.to_string()))?
416 .map(Self::decode_blob_meta)
417 .transpose()?;
418
419 let existed = self
420 .blobs
421 .delete(wtxn, hash)
422 .map_err(|e| StoreError::Other(e.to_string()))?;
423 let _ = self
424 .metadata
425 .delete(wtxn, hash)
426 .map_err(|e| StoreError::Other(e.to_string()))?;
427 let _ = self
428 .pins
429 .delete(wtxn, hash)
430 .map_err(|e| StoreError::Other(e.to_string()))?;
431 if let Some(meta) = meta {
432 let order_key = Self::encode_order_key(meta.order, hash);
433 let _ = self
434 .eviction_order
435 .delete(wtxn, &order_key)
436 .map_err(|e| StoreError::Other(e.to_string()))?;
437 }
438
439 Ok((
440 existed || meta.is_some(),
441 data_len.or(meta.map(|m| m.size)).unwrap_or(0),
442 ))
443 }
444
445 fn read_pin_count(&self, txn: &heed::RoTxn, hash: &[u8]) -> Result<u32, StoreError> {
446 self.pins
447 .get(txn, hash)
448 .map_err(|e| StoreError::Other(e.to_string()))?
449 .map(Self::decode_pin_count)
450 .transpose()?
451 .map_or(Ok(0), Ok)
452 }
453
454 fn encode_blob_meta(meta: BlobMeta) -> [u8; BLOB_META_BYTES] {
455 let mut encoded = [0u8; BLOB_META_BYTES];
456 encoded[..8].copy_from_slice(&meta.order.to_be_bytes());
457 encoded[8..].copy_from_slice(&meta.size.to_be_bytes());
458 encoded
459 }
460
461 fn decode_blob_meta(bytes: &[u8]) -> Result<BlobMeta, StoreError> {
462 if bytes.len() != BLOB_META_BYTES {
463 return Err(StoreError::Other(format!(
464 "invalid blob metadata length: {}",
465 bytes.len()
466 )));
467 }
468 Ok(BlobMeta {
469 order: Self::decode_order(&bytes[..8])?,
470 size: u64::from_be_bytes(
471 bytes[8..16]
472 .try_into()
473 .map_err(|_| StoreError::Other("invalid blob size bytes".into()))?,
474 ),
475 })
476 }
477
478 fn encode_order_key(order: u64, hash: &Hash) -> [u8; ORDER_KEY_BYTES] {
479 let mut key = [0u8; ORDER_KEY_BYTES];
480 key[..8].copy_from_slice(&order.to_be_bytes());
481 key[8..].copy_from_slice(hash);
482 key
483 }
484
485 fn decode_order(bytes: &[u8]) -> Result<u64, StoreError> {
486 if bytes.len() != 8 {
487 return Err(StoreError::Other(format!(
488 "invalid order length: {}",
489 bytes.len()
490 )));
491 }
492 Ok(u64::from_be_bytes(bytes.try_into().map_err(|_| {
493 StoreError::Other("invalid order bytes".into())
494 })?))
495 }
496
497 fn decode_hash_from_order_key(bytes: &[u8]) -> Result<Hash, StoreError> {
498 if bytes.len() != ORDER_KEY_BYTES {
499 return Err(StoreError::Other(format!(
500 "invalid order key length: {}",
501 bytes.len()
502 )));
503 }
504 let mut hash = [0u8; 32];
505 hash.copy_from_slice(&bytes[8..]);
506 Ok(hash)
507 }
508
509 fn decode_order_from_order_key(bytes: &[u8]) -> Result<u64, StoreError> {
510 if bytes.len() != ORDER_KEY_BYTES {
511 return Err(StoreError::Other(format!(
512 "invalid order key length: {}",
513 bytes.len()
514 )));
515 }
516 Self::decode_order(&bytes[..8])
517 }
518
519 fn decode_pin_count(bytes: &[u8]) -> Result<u32, StoreError> {
520 if bytes.len() != PIN_COUNT_BYTES {
521 return Err(StoreError::Other(format!(
522 "invalid pin count length: {}",
523 bytes.len()
524 )));
525 }
526 Ok(u32::from_be_bytes(bytes.try_into().map_err(|_| {
527 StoreError::Other("invalid pin count bytes".into())
528 })?))
529 }
530}
531
532#[derive(Debug, Clone)]
533pub struct LmdbStats {
534 pub count: usize,
535 pub total_bytes: u64,
536 pub pinned_count: usize,
537 pub pinned_bytes: u64,
538}
539
540#[async_trait]
541impl Store for LmdbBlobStore {
542 async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
543 self.put_sync(hash, &data)
544 }
545
546 async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
547 self.put_many_sync(&items)
548 }
549
550 async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
551 self.get_sync(hash)
552 }
553
554 async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
555 self.exists(hash)
556 }
557
558 async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
559 self.delete_sync(hash)
560 }
561
562 fn set_max_bytes(&self, max: u64) {
563 self.max_bytes.store(max, Ordering::Relaxed);
564 }
565
566 fn max_bytes(&self) -> Option<u64> {
567 let max = self.max_bytes.load(Ordering::Relaxed);
568 if max > 0 {
569 Some(max)
570 } else {
571 None
572 }
573 }
574
575 async fn stats(&self) -> StoreStats {
576 match self.stats() {
577 Ok(stats) => StoreStats {
578 count: stats.count as u64,
579 bytes: stats.total_bytes,
580 pinned_count: stats.pinned_count as u64,
581 pinned_bytes: stats.pinned_bytes,
582 },
583 Err(_) => StoreStats::default(),
584 }
585 }
586
587 async fn evict_if_needed(&self) -> Result<u64, StoreError> {
588 let max = self.max_bytes.load(Ordering::Relaxed);
589 if max == 0 {
590 return Ok(0);
591 }
592
593 let current = self.stats()?.total_bytes;
594 if current <= max {
595 return Ok(0);
596 }
597
598 let target = max * 9 / 10;
599 self.evict_to_target(current, target)
600 }
601
602 async fn pin(&self, hash: &Hash) -> Result<(), StoreError> {
603 self.pin_sync(hash)
604 }
605
606 async fn unpin(&self, hash: &Hash) -> Result<(), StoreError> {
607 self.unpin_sync(hash)
608 }
609
610 fn pin_count(&self, hash: &Hash) -> u32 {
611 let Ok(rtxn) = self.env.read_txn() else {
612 return 0;
613 };
614 self.read_pin_count(&rtxn, hash).unwrap_or(0)
615 }
616}
617
618#[cfg(test)]
619mod tests {
620 use super::*;
621 use hashtree_core::sha256;
622 use heed::EnvOpenOptions;
623 #[cfg(unix)]
624 use std::path::{Path, PathBuf};
625 #[cfg(unix)]
626 use std::process::Command;
627 use std::sync::{
628 atomic::{AtomicBool, Ordering},
629 Arc, Barrier,
630 };
631 use std::time::Duration;
632 use tempfile::TempDir;
633
634 #[cfg(unix)]
635 const STALE_READER_HELPER_ENV: &str = "HASHTREE_LMDB_STALE_READER_HELPER";
636 #[cfg(unix)]
637 const STALE_READER_HELPER_MODE_ENV: &str = "HASHTREE_LMDB_STALE_READER_HELPER_MODE";
638 #[cfg(unix)]
639 const STALE_READER_DB_PATH_ENV: &str = "HASHTREE_LMDB_STALE_READER_DB_PATH";
640 #[cfg(unix)]
641 const STALE_READER_MARKER_PATH_ENV: &str = "HASHTREE_LMDB_STALE_READER_MARKER_PATH";
642 #[cfg(unix)]
643 const TEST_MAX_READERS: u32 = 4;
644
645 #[cfg(unix)]
646 fn run_helper(mode: &str, path: &Path, marker: &Path) {
647 let output = Command::new(std::env::current_exe().expect("test binary path"))
648 .arg("--ignored")
649 .arg("--exact")
650 .arg("tests::lmdb_stale_reader_helper")
651 .env(STALE_READER_HELPER_ENV, "1")
652 .env(STALE_READER_HELPER_MODE_ENV, mode)
653 .env(STALE_READER_DB_PATH_ENV, path)
654 .env(STALE_READER_MARKER_PATH_ENV, marker)
655 .env("RUST_TEST_THREADS", "1")
656 .output()
657 .expect("spawn stale-reader helper");
658
659 assert!(
660 output.status.success(),
661 "stale-reader helper failed: stdout={} stderr={}",
662 String::from_utf8_lossy(&output.stdout),
663 String::from_utf8_lossy(&output.stderr)
664 );
665 assert!(
666 marker.exists(),
667 "stale-reader helper did not run: stdout={} stderr={}",
668 String::from_utf8_lossy(&output.stdout),
669 String::from_utf8_lossy(&output.stderr)
670 );
671 }
672
673 #[tokio::test]
674 async fn test_put_get() -> Result<(), StoreError> {
675 let temp = TempDir::new().unwrap();
676 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
677
678 let data = b"hello lmdb";
679 let hash = sha256(data);
680 store.put(hash, data.to_vec()).await?;
681
682 assert!(store.has(&hash).await?);
683 assert_eq!(store.get(&hash).await?, Some(data.to_vec()));
684
685 Ok(())
686 }
687
688 #[tokio::test]
689 async fn test_delete() -> Result<(), StoreError> {
690 let temp = TempDir::new().unwrap();
691 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
692
693 let data = b"delete me";
694 let hash = sha256(data);
695 store.put(hash, data.to_vec()).await?;
696 assert!(store.has(&hash).await?);
697
698 assert!(store.delete(&hash).await?);
699 assert!(!store.has(&hash).await?);
700 assert!(!store.delete(&hash).await?);
701
702 Ok(())
703 }
704
705 #[tokio::test]
706 async fn test_list() -> Result<(), StoreError> {
707 let temp = TempDir::new().unwrap();
708 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
709
710 let d1 = b"one";
711 let d2 = b"two";
712 let d3 = b"three";
713 let h1 = sha256(d1);
714 let h2 = sha256(d2);
715 let h3 = sha256(d3);
716
717 store.put(h1, d1.to_vec()).await?;
718 store.put(h2, d2.to_vec()).await?;
719 store.put(h3, d3.to_vec()).await?;
720
721 let hashes = store.list()?;
722 assert_eq!(hashes.len(), 3);
723 assert!(hashes.contains(&h1));
724 assert!(hashes.contains(&h2));
725 assert!(hashes.contains(&h3));
726
727 Ok(())
728 }
729
730 #[tokio::test]
731 async fn test_stats() -> Result<(), StoreError> {
732 let temp = TempDir::new().unwrap();
733 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
734
735 let d1 = b"hello";
736 let d2 = b"world";
737 store.put(sha256(d1), d1.to_vec()).await?;
738 store.put(sha256(d2), d2.to_vec()).await?;
739
740 let stats = store.stats()?;
741 assert_eq!(stats.count, 2);
742 assert_eq!(stats.total_bytes, 10);
743
744 Ok(())
745 }
746
747 #[tokio::test]
748 async fn test_deduplication() -> Result<(), StoreError> {
749 let temp = TempDir::new().unwrap();
750 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
751
752 let data = b"same";
753 let hash = sha256(data);
754 assert!(store.put(hash, data.to_vec()).await?); assert!(!store.put(hash, data.to_vec()).await?); assert_eq!(store.list()?.len(), 1);
758
759 Ok(())
760 }
761
762 #[tokio::test]
763 async fn test_max_bytes() -> Result<(), StoreError> {
764 let temp = TempDir::new().unwrap();
765 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
766
767 assert!(store.max_bytes().is_none());
768
769 store.set_max_bytes(1000);
770 assert_eq!(store.max_bytes(), Some(1000));
771
772 store.set_max_bytes(0);
773 assert!(store.max_bytes().is_none());
774
775 Ok(())
776 }
777
778 #[tokio::test]
779 async fn test_eviction_over_limit() -> Result<(), StoreError> {
780 let temp = TempDir::new().unwrap();
781 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
782 store.set_max_bytes(25);
783
784 let h1 = sha256(b"aaaaaaaaaa");
785 let h2 = sha256(b"bbbbbbbbbb");
786 let h3 = sha256(b"cccccccccc");
787
788 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
789 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
790 store.put(h3, b"cccccccccc".to_vec()).await?;
791
792 let freed = store.evict_if_needed().await?;
793 assert!(freed >= 10, "expected eviction to free at least one blob");
794
795 assert!(
796 !store.has(&h1).await?,
797 "oldest blob should be evicted first"
798 );
799 assert!(store.has(&h2).await?);
800 assert!(store.has(&h3).await?);
801
802 let stats = store.stats()?;
803 assert!(
804 stats.total_bytes <= 22,
805 "store should be reduced to 90% target"
806 );
807
808 Ok(())
809 }
810
811 #[tokio::test]
812 async fn test_eviction_respects_pins() -> Result<(), StoreError> {
813 let temp = TempDir::new().unwrap();
814 let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
815 store.set_max_bytes(25);
816
817 let h1 = sha256(b"aaaaaaaaaa");
818 let h2 = sha256(b"bbbbbbbbbb");
819 let h3 = sha256(b"cccccccccc");
820
821 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
822 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
823 store.put(h3, b"cccccccccc".to_vec()).await?;
824 store.pin(&h1).await?;
825
826 let freed = store.evict_if_needed().await?;
827 assert!(freed >= 10, "expected eviction to free at least one blob");
828
829 assert!(store.has(&h1).await?, "pinned blob must not be evicted");
830 assert!(
831 !store.has(&h2).await?,
832 "oldest unpinned blob should be evicted"
833 );
834 assert!(store.has(&h3).await?);
835
836 Ok(())
837 }
838
839 #[tokio::test]
840 async fn test_reopen_with_existing_eviction_order() -> Result<(), StoreError> {
841 let temp = TempDir::new().unwrap();
842 let path = temp.path().join("blobs");
843
844 {
845 let store = LmdbBlobStore::new(&path)?;
846 let h1 = sha256(b"aaaaaaaaaa");
847 let h2 = sha256(b"bbbbbbbbbb");
848 store.put(h1, b"aaaaaaaaaa".to_vec()).await?;
849 store.put(h2, b"bbbbbbbbbb".to_vec()).await?;
850 }
851
852 let reopened = LmdbBlobStore::new(&path)?;
853 let h3 = sha256(b"cccccccccc");
854 assert!(reopened.put(h3, b"cccccccccc".to_vec()).await?);
855 assert!(reopened.has(&h3).await?);
856
857 Ok(())
858 }
859
860 #[test]
861 fn test_supports_many_concurrent_readers() -> Result<(), Box<dyn std::error::Error>> {
862 const READER_THREADS: usize = 160;
863
864 let temp = TempDir::new()?;
865 let store = Arc::new(LmdbBlobStore::new(temp.path().join("blobs"))?);
866 let hash = sha256(b"many readers");
867 store.put_sync(hash, b"many readers")?;
868
869 let start = Arc::new(Barrier::new(READER_THREADS + 1));
870 let release = Arc::new(AtomicBool::new(false));
871 let mut handles = Vec::with_capacity(READER_THREADS);
872
873 for _ in 0..READER_THREADS {
874 let env = store.env.clone();
875 let start = Arc::clone(&start);
876 let release = Arc::clone(&release);
877 handles.push(std::thread::spawn(move || -> Result<(), String> {
878 start.wait();
879 let _rtxn = env.read_txn().map_err(|err| err.to_string())?;
880 while !release.load(Ordering::Relaxed) {
881 std::thread::sleep(Duration::from_millis(1));
882 }
883 Ok(())
884 }));
885 }
886
887 start.wait();
888 std::thread::sleep(Duration::from_millis(50));
889 release.store(true, Ordering::Relaxed);
890
891 let results: Vec<Result<(), String>> = handles
892 .into_iter()
893 .map(|handle| handle.join().expect("reader thread panicked"))
894 .collect();
895
896 let failures: Vec<String> = results.into_iter().filter_map(Result::err).collect();
897 assert!(
898 failures.is_empty(),
899 "concurrent reader failures: {}",
900 failures.join(" | ")
901 );
902 assert!(store.exists(&hash)?);
903
904 Ok(())
905 }
906
907 #[cfg(unix)]
908 #[test]
909 fn test_reclaims_stale_reader_slots() -> Result<(), Box<dyn std::error::Error>> {
910 let temp = TempDir::new()?;
911 let path = temp.path().join("blobs");
912 let data = b"hello stale readers";
913 let hash = sha256(data);
914
915 run_helper("setup", &path, &temp.path().join("setup.marker"));
916
917 for index in 0..TEST_MAX_READERS {
918 let marker = temp.path().join(format!("helper-{index}.marker"));
919 run_helper("stale", &path, &marker);
920 }
921
922 let store = LmdbBlobStore::with_map_size(&path, 1024 * 1024)?;
923 assert!(store.exists(&hash)?);
924
925 Ok(())
926 }
927
928 #[cfg(unix)]
929 #[test]
930 fn test_reopens_existing_env_with_larger_map_size() -> Result<(), Box<dyn std::error::Error>> {
931 let temp = TempDir::new()?;
932 let path = temp.path().join("blobs");
933 run_helper("small-map", &path, &temp.path().join("small-map.marker"));
934
935 let reopened = LmdbBlobStore::with_map_size(&path, 8 * 1024 * 1024)?;
936 assert!(reopened.map_size_bytes() >= 8 * 1024 * 1024);
937
938 Ok(())
939 }
940
941 #[cfg(unix)]
942 #[test]
943 #[ignore = "used as a subprocess helper by test_reclaims_stale_reader_slots"]
944 fn lmdb_stale_reader_helper() {
945 let Some(db_path) = std::env::var_os(STALE_READER_DB_PATH_ENV) else {
946 return;
947 };
948 let marker_path =
949 PathBuf::from(std::env::var_os(STALE_READER_MARKER_PATH_ENV).expect("marker path"));
950 std::fs::write(&marker_path, b"started").expect("write helper marker");
951
952 let _env_flag = std::env::var_os(STALE_READER_HELPER_ENV).expect("helper mode enabled");
953 let mode = std::env::var(STALE_READER_HELPER_MODE_ENV).expect("helper mode");
954 let db_path = PathBuf::from(db_path);
955 std::fs::create_dir_all(&db_path).expect("create helper db dir");
956 let env = unsafe {
957 EnvOpenOptions::new()
958 .map_size(1024 * 1024)
959 .max_dbs(DATABASE_COUNT)
960 .max_readers(TEST_MAX_READERS)
961 .open(&db_path)
962 .expect("open lmdb env")
963 };
964 match mode.as_str() {
965 "setup" => {
966 let mut wtxn = env.write_txn().expect("open write txn");
967 let blobs: Database<Bytes, Bytes> = env
968 .create_database(&mut wtxn, Some("blobs"))
969 .expect("create blobs database");
970 let data = b"hello stale readers";
971 let hash = sha256(data);
972 blobs.put(&mut wtxn, &hash, data).expect("seed blob");
973 wtxn.commit().expect("commit setup txn");
974 std::process::exit(0);
975 }
976 "stale" => {
977 let _rtxn = env.read_txn().expect("open read txn");
978 std::process::exit(0);
979 }
980 "small-map" => {
981 let mut wtxn = env.write_txn().expect("open write txn");
982 let _blobs: Database<Bytes, Bytes> = env
983 .create_database(&mut wtxn, Some("blobs"))
984 .expect("create blobs database");
985 let _metadata: Database<Bytes, Bytes> = env
986 .create_database(&mut wtxn, Some("metadata"))
987 .expect("create metadata database");
988 let _eviction_order: Database<Bytes, Unit> = env
989 .create_database(&mut wtxn, Some("eviction_order"))
990 .expect("create eviction_order database");
991 let _pins: Database<Bytes, Bytes> = env
992 .create_database(&mut wtxn, Some("pins"))
993 .expect("create pins database");
994 wtxn.commit().expect("commit small-map setup txn");
995 std::process::exit(0);
996 }
997 other => panic!("unknown helper mode: {other}"),
998 }
999 }
1000}