Skip to main content

hashtree_lmdb/
lib.rs

1//! LMDB-backed content-addressed blob storage.
2
3use 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
12// Re-export sha256 for convenience
13pub 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
46/// LMDB-backed blob store implementing hashtree's Store trait.
47pub struct LmdbBlobStore {
48    env: heed::Env,
49    /// Maps SHA256 hash (32 bytes) → blob data
50    blobs: Database<Bytes, Bytes>,
51    /// Maps SHA256 hash (32 bytes) → [order: u64][size: u64]
52    metadata: Database<Bytes, Bytes>,
53    /// Maps [order: u64][hash: 32 bytes] → ()
54    eviction_order: Database<Bytes, Unit>,
55    /// Maps SHA256 hash (32 bytes) → pin count (u32)
56    pins: Database<Bytes, Bytes>,
57    /// Small aggregate counters used by quota checks and stats.
58    stats: Database<Bytes, Bytes>,
59    max_bytes: AtomicU64,
60    current_bytes: AtomicU64,
61    next_order: AtomicU64,
62}
63
64impl LmdbBlobStore {
65    /// Open or create an LMDB blob store at the given path.
66    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, StoreError> {
67        Self::with_map_size(path, DEFAULT_MAP_SIZE)
68    }
69
70    /// Open or create with a maximum logical storage size.
71    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    /// Open or create with custom map size.
114    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    /// Check if a hash exists (sync version for internal use).
300    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    /// Mark which sorted hashes already exist using bounded LMDB range scans.
348    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    /// Get storage statistics.
506    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    /// List all hashes in the store.
522    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    /// Sync put operation (for use in sync contexts).
555    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    /// Sync batch put operation (for use in sync contexts).
584    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    /// Sync get operation (for use in sync contexts).
618    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    /// Sync delete operation (for use in sync contexts).
758    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?); // Returns true (newly stored)
1461        assert!(!store.put(hash, data.to_vec()).await?); // Returns false (already existed)
1462
1463        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}