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