hashtree_lmdb/
lib.rs

1//! LMDB-backed content-addressed blob storage.
2
3use async_trait::async_trait;
4use heed::types::*;
5use heed::{Database, EnvOpenOptions};
6use hashtree_core::store::{Store, StoreError};
7use hashtree_core::types::Hash;
8use std::path::Path;
9
10// Re-export sha256 for convenience
11pub use hashtree_core::hash::sha256 as compute_sha256;
12
13/// LMDB-backed blob store implementing hashtree's Store trait.
14pub struct LmdbBlobStore {
15    env: heed::Env,
16    /// Maps SHA256 hash (32 bytes) → blob data
17    blobs: Database<Bytes, Bytes>,
18}
19
20impl LmdbBlobStore {
21    /// Open or create an LMDB blob store at the given path.
22    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, StoreError> {
23        Self::with_map_size(path, 10 * 1024 * 1024 * 1024) // 10GB default
24    }
25
26    /// Open or create with custom map size.
27    pub fn with_map_size<P: AsRef<Path>>(path: P, map_size: usize) -> Result<Self, StoreError> {
28        std::fs::create_dir_all(&path).map_err(StoreError::Io)?;
29
30        let env = unsafe {
31            EnvOpenOptions::new()
32                .map_size(map_size)
33                .max_dbs(1)
34                .open(path)
35                .map_err(|e| StoreError::Other(e.to_string()))?
36        };
37
38        let mut wtxn = env
39            .write_txn()
40            .map_err(|e| StoreError::Other(e.to_string()))?;
41        let blobs = env
42            .create_database(&mut wtxn, Some("blobs"))
43            .map_err(|e| StoreError::Other(e.to_string()))?;
44        wtxn.commit()
45            .map_err(|e| StoreError::Other(e.to_string()))?;
46
47        Ok(Self { env, blobs })
48    }
49
50    /// Check if a hash exists (sync version for internal use).
51    pub fn exists(&self, hash: &Hash) -> Result<bool, StoreError> {
52        let rtxn = self
53            .env
54            .read_txn()
55            .map_err(|e| StoreError::Other(e.to_string()))?;
56
57        Ok(self
58            .blobs
59            .get(&rtxn, hash)
60            .map_err(|e| StoreError::Other(e.to_string()))?
61            .is_some())
62    }
63
64    /// Get storage statistics.
65    pub fn stats(&self) -> Result<LmdbStats, StoreError> {
66        let rtxn = self
67            .env
68            .read_txn()
69            .map_err(|e| StoreError::Other(e.to_string()))?;
70
71        let count = self
72            .blobs
73            .len(&rtxn)
74            .map_err(|e| StoreError::Other(e.to_string()))?
75            as usize;
76
77        let mut total_bytes = 0u64;
78        for item in self
79            .blobs
80            .iter(&rtxn)
81            .map_err(|e| StoreError::Other(e.to_string()))?
82        {
83            let (_, data) = item.map_err(|e| StoreError::Other(e.to_string()))?;
84            total_bytes += data.len() as u64;
85        }
86
87        Ok(LmdbStats { count, total_bytes })
88    }
89
90    /// List all hashes in the store.
91    pub fn list(&self) -> Result<Vec<Hash>, StoreError> {
92        let rtxn = self
93            .env
94            .read_txn()
95            .map_err(|e| StoreError::Other(e.to_string()))?;
96
97        let mut hashes = Vec::new();
98        for item in self
99            .blobs
100            .iter(&rtxn)
101            .map_err(|e| StoreError::Other(e.to_string()))?
102        {
103            let (hash, _) = item.map_err(|e| StoreError::Other(e.to_string()))?;
104            let hash_arr: Hash = hash
105                .try_into()
106                .map_err(|_| StoreError::Other("invalid hash length".into()))?;
107            hashes.push(hash_arr);
108        }
109
110        Ok(hashes)
111    }
112
113    /// Sync put operation (for use in sync contexts).
114    pub fn put_sync(&self, hash: Hash, data: &[u8]) -> Result<bool, StoreError> {
115        let mut wtxn = self
116            .env
117            .write_txn()
118            .map_err(|e| StoreError::Other(e.to_string()))?;
119
120        let existed = self
121            .blobs
122            .get(&wtxn, &hash)
123            .map_err(|e| StoreError::Other(e.to_string()))?
124            .is_some();
125
126        if !existed {
127            self.blobs
128                .put(&mut wtxn, &hash, data)
129                .map_err(|e| StoreError::Other(e.to_string()))?;
130        }
131
132        wtxn.commit()
133            .map_err(|e| StoreError::Other(e.to_string()))?;
134
135        Ok(!existed)
136    }
137
138    /// Sync get operation (for use in sync contexts).
139    pub fn get_sync(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
140        let rtxn = self
141            .env
142            .read_txn()
143            .map_err(|e| StoreError::Other(e.to_string()))?;
144
145        Ok(self
146            .blobs
147            .get(&rtxn, hash)
148            .map_err(|e| StoreError::Other(e.to_string()))?
149            .map(|b| b.to_vec()))
150    }
151
152    /// Sync delete operation (for use in sync contexts).
153    pub fn delete_sync(&self, hash: &Hash) -> Result<bool, StoreError> {
154        let mut wtxn = self
155            .env
156            .write_txn()
157            .map_err(|e| StoreError::Other(e.to_string()))?;
158
159        let existed = self
160            .blobs
161            .delete(&mut wtxn, hash)
162            .map_err(|e| StoreError::Other(e.to_string()))?;
163
164        wtxn.commit()
165            .map_err(|e| StoreError::Other(e.to_string()))?;
166
167        Ok(existed)
168    }
169}
170
171#[derive(Debug, Clone)]
172pub struct LmdbStats {
173    pub count: usize,
174    pub total_bytes: u64,
175}
176
177#[async_trait]
178impl Store for LmdbBlobStore {
179    async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
180        self.put_sync(hash, &data)
181    }
182
183    async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
184        self.get_sync(hash)
185    }
186
187    async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
188        self.exists(hash)
189    }
190
191    async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
192        self.delete_sync(hash)
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199    use hashtree_core::sha256;
200    use tempfile::TempDir;
201
202    #[tokio::test]
203    async fn test_put_get() -> Result<(), StoreError> {
204        let temp = TempDir::new().unwrap();
205        let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
206
207        let data = b"hello lmdb";
208        let hash = sha256(data);
209        store.put(hash, data.to_vec()).await?;
210
211        assert!(store.has(&hash).await?);
212        assert_eq!(store.get(&hash).await?, Some(data.to_vec()));
213
214        Ok(())
215    }
216
217    #[tokio::test]
218    async fn test_delete() -> Result<(), StoreError> {
219        let temp = TempDir::new().unwrap();
220        let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
221
222        let data = b"delete me";
223        let hash = sha256(data);
224        store.put(hash, data.to_vec()).await?;
225        assert!(store.has(&hash).await?);
226
227        assert!(store.delete(&hash).await?);
228        assert!(!store.has(&hash).await?);
229        assert!(!store.delete(&hash).await?);
230
231        Ok(())
232    }
233
234    #[tokio::test]
235    async fn test_list() -> Result<(), StoreError> {
236        let temp = TempDir::new().unwrap();
237        let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
238
239        let d1 = b"one";
240        let d2 = b"two";
241        let d3 = b"three";
242        let h1 = sha256(d1);
243        let h2 = sha256(d2);
244        let h3 = sha256(d3);
245
246        store.put(h1, d1.to_vec()).await?;
247        store.put(h2, d2.to_vec()).await?;
248        store.put(h3, d3.to_vec()).await?;
249
250        let hashes = store.list()?;
251        assert_eq!(hashes.len(), 3);
252        assert!(hashes.contains(&h1));
253        assert!(hashes.contains(&h2));
254        assert!(hashes.contains(&h3));
255
256        Ok(())
257    }
258
259    #[tokio::test]
260    async fn test_stats() -> Result<(), StoreError> {
261        let temp = TempDir::new().unwrap();
262        let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
263
264        let d1 = b"hello";
265        let d2 = b"world";
266        store.put(sha256(d1), d1.to_vec()).await?;
267        store.put(sha256(d2), d2.to_vec()).await?;
268
269        let stats = store.stats()?;
270        assert_eq!(stats.count, 2);
271        assert_eq!(stats.total_bytes, 10);
272
273        Ok(())
274    }
275
276    #[tokio::test]
277    async fn test_deduplication() -> Result<(), StoreError> {
278        let temp = TempDir::new().unwrap();
279        let store = LmdbBlobStore::new(temp.path().join("blobs"))?;
280
281        let data = b"same";
282        let hash = sha256(data);
283        assert!(store.put(hash, data.to_vec()).await?); // Returns true (newly stored)
284        assert!(!store.put(hash, data.to_vec()).await?); // Returns false (already existed)
285
286        assert_eq!(store.list()?.len(), 1);
287
288        Ok(())
289    }
290}
hashtree_lmdb/lib.rs

hashtree_lmdb/
lib.rs
