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