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