Skip to main content

hashtree_core/
store.rs

1//! Content-addressed key-value store interfaces and implementations
2
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::{Arc, RwLock};
6
7use crate::types::{to_hex, Hash};
8
9/// Return a byte range from an in-memory blob.
10///
11/// `end_inclusive` follows HTTP/storage range conventions. Out-of-bounds ranges
12/// return an empty or clamped slice instead of an error.
13pub fn slice_blob_range(
14    data: &[u8],
15    start: u64,
16    end_inclusive: u64,
17) -> Result<Vec<u8>, StoreError> {
18    if data.is_empty() || end_inclusive < start {
19        return Ok(Vec::new());
20    }
21
22    let len = data.len() as u64;
23    if start >= len {
24        return Ok(Vec::new());
25    }
26
27    let actual_end = end_inclusive.min(len - 1);
28    let start = usize::try_from(start)
29        .map_err(|_| StoreError::Other("blob range start is too large".to_string()))?;
30    let end_exclusive = usize::try_from(actual_end.saturating_add(1))
31        .map_err(|_| StoreError::Other("blob range end is too large".to_string()))?;
32
33    Ok(data[start..end_exclusive].to_vec())
34}
35
36/// Storage statistics
37#[derive(Debug, Clone, Default)]
38pub struct StoreStats {
39    /// Number of items in store
40    pub count: u64,
41    /// Total bytes stored
42    pub bytes: u64,
43    /// Number of pinned items
44    pub pinned_count: u64,
45    /// Bytes used by pinned items
46    pub pinned_bytes: u64,
47}
48
49/// Result of a batch insert.
50#[derive(Debug, Clone, Default, PartialEq, Eq)]
51pub struct PutManyReport {
52    /// Number of candidate items passed by the caller, including duplicates.
53    pub total: usize,
54    /// Number of blobs newly inserted into the store.
55    pub inserted: usize,
56    /// Logical bytes newly inserted into the store.
57    pub inserted_bytes: u64,
58    /// Hashes that were newly inserted, in insertion order.
59    pub inserted_hashes: Vec<Hash>,
60}
61
62/// Content-addressed key-value store interface
63#[async_trait]
64pub trait Store: Send + Sync {
65    /// Store data by its hash
66    /// Returns true if newly stored, false if already existed
67    async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError>;
68
69    /// Store multiple blobs.
70    /// Returns the number of newly stored items.
71    async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
72        let mut inserted = 0usize;
73        for (hash, data) in items {
74            if self.put(hash, data).await? {
75                inserted += 1;
76            }
77        }
78        Ok(inserted)
79    }
80
81    /// Retrieve data by hash
82    /// Returns data or None if not found
83    async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError>;
84
85    /// Retrieve a byte range from a blob by hash.
86    ///
87    /// `end_inclusive` is inclusive, matching HTTP range requests and the sync
88    /// filesystem/LMDB storage primitives. Implementations may override this to
89    /// avoid loading whole blobs.
90    async fn get_range(
91        &self,
92        hash: &Hash,
93        start: u64,
94        end_inclusive: u64,
95    ) -> Result<Option<Vec<u8>>, StoreError> {
96        let Some(data) = self.get(hash).await? else {
97            return Ok(None);
98        };
99        Ok(Some(slice_blob_range(&data, start, end_inclusive)?))
100    }
101
102    /// Return the stored blob size in bytes.
103    ///
104    /// The default implementation reads the blob. Backends with metadata should
105    /// override this to answer without loading payload bytes.
106    async fn blob_size(&self, hash: &Hash) -> Result<Option<u64>, StoreError> {
107        Ok(self.get(hash).await?.map(|data| data.len() as u64))
108    }
109
110    /// Check if hash exists
111    async fn has(&self, hash: &Hash) -> Result<bool, StoreError>;
112
113    /// Delete by hash
114    /// Returns true if deleted, false if didn't exist
115    async fn delete(&self, hash: &Hash) -> Result<bool, StoreError>;
116
117    // ========================================================================
118    // Optional: Storage limits and eviction (default no-op implementations)
119    // ========================================================================
120
121    /// Set maximum storage size in bytes. 0 = unlimited.
122    fn set_max_bytes(&self, _max: u64) {}
123
124    /// Get maximum storage size. None = unlimited.
125    fn max_bytes(&self) -> Option<u64> {
126        None
127    }
128
129    /// Get storage statistics
130    async fn stats(&self) -> StoreStats {
131        StoreStats::default()
132    }
133
134    /// Evict unpinned items if over storage limit.
135    /// Returns number of bytes freed.
136    async fn evict_if_needed(&self) -> Result<u64, StoreError> {
137        Ok(0)
138    }
139
140    // ========================================================================
141    // Optional: Pinning (default no-op implementations)
142    // ========================================================================
143
144    /// Pin a hash (increment ref count). Pinned items are not evicted.
145    async fn pin(&self, _hash: &Hash) -> Result<(), StoreError> {
146        Ok(())
147    }
148
149    /// Unpin a hash (decrement ref count). Item can be evicted when count reaches 0.
150    async fn unpin(&self, _hash: &Hash) -> Result<(), StoreError> {
151        Ok(())
152    }
153
154    /// Get pin count for a hash. 0 = not pinned.
155    fn pin_count(&self, _hash: &Hash) -> u32 {
156        0
157    }
158
159    /// Check if hash is pinned (pin count > 0)
160    fn is_pinned(&self, hash: &Hash) -> bool {
161        self.pin_count(hash) > 0
162    }
163}
164
165/// Store error type
166#[derive(Debug, thiserror::Error)]
167pub enum StoreError {
168    #[error("IO error: {0}")]
169    Io(#[from] std::io::Error),
170    #[error("Store error: {0}")]
171    Other(String),
172}
173
174#[derive(Debug, Default)]
175struct BufferedStoreInner {
176    pending: HashMap<Hash, Vec<u8>>,
177    order: Vec<Hash>,
178}
179
180#[derive(Debug, Clone, Copy)]
181struct BufferedStoreOptions {
182    check_base_on_put: bool,
183}
184
185impl Default for BufferedStoreOptions {
186    fn default() -> Self {
187        Self {
188            check_base_on_put: true,
189        }
190    }
191}
192
193/// Buffered overlay store that keeps writes in memory until flushed.
194#[derive(Debug, Clone)]
195pub struct BufferedStore<S: Store> {
196    base: Arc<S>,
197    inner: Arc<RwLock<BufferedStoreInner>>,
198    options: BufferedStoreOptions,
199}
200
201impl<S: Store> BufferedStore<S> {
202    pub fn new(base: Arc<S>) -> Self {
203        Self::with_options(base, BufferedStoreOptions::default())
204    }
205
206    pub fn new_optimistic(base: Arc<S>) -> Self {
207        Self::with_options(
208            base,
209            BufferedStoreOptions {
210                check_base_on_put: false,
211            },
212        )
213    }
214
215    fn with_options(base: Arc<S>, options: BufferedStoreOptions) -> Self {
216        Self {
217            base,
218            inner: Arc::new(RwLock::new(BufferedStoreInner::default())),
219            options,
220        }
221    }
222
223    pub async fn flush(&self) -> Result<usize, StoreError> {
224        let items = {
225            let mut inner = self.inner.write().unwrap();
226            if inner.order.is_empty() {
227                return Ok(0);
228            }
229
230            let order = std::mem::take(&mut inner.order);
231            let mut items = Vec::with_capacity(order.len());
232            for hash in order {
233                if let Some(data) = inner.pending.remove(&hash) {
234                    items.push((hash, data));
235                }
236            }
237            items
238        };
239
240        self.base.put_many(items).await
241    }
242}
243
244#[async_trait]
245impl<S: Store> Store for BufferedStore<S> {
246    async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
247        {
248            let inner = self.inner.read().unwrap();
249            if inner.pending.contains_key(&hash) {
250                return Ok(false);
251            }
252        }
253
254        if self.options.check_base_on_put && self.base.has(&hash).await? {
255            return Ok(false);
256        }
257
258        let mut inner = self.inner.write().unwrap();
259        if inner.pending.contains_key(&hash) {
260            return Ok(false);
261        }
262        inner.order.push(hash);
263        inner.pending.insert(hash, data);
264        Ok(true)
265    }
266
267    async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
268        let mut inserted = 0usize;
269        for (hash, data) in items {
270            if self.put(hash, data).await? {
271                inserted += 1;
272            }
273        }
274        Ok(inserted)
275    }
276
277    async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
278        if let Some(data) = self.inner.read().unwrap().pending.get(hash).cloned() {
279            return Ok(Some(data));
280        }
281        self.base.get(hash).await
282    }
283
284    async fn get_range(
285        &self,
286        hash: &Hash,
287        start: u64,
288        end_inclusive: u64,
289    ) -> Result<Option<Vec<u8>>, StoreError> {
290        let pending = {
291            let inner = self.inner.read().unwrap();
292            inner
293                .pending
294                .get(hash)
295                .map(|data| slice_blob_range(data, start, end_inclusive))
296                .transpose()?
297        };
298        if pending.is_some() {
299            return Ok(pending);
300        }
301        self.base.get_range(hash, start, end_inclusive).await
302    }
303
304    async fn blob_size(&self, hash: &Hash) -> Result<Option<u64>, StoreError> {
305        if let Some(size) = self
306            .inner
307            .read()
308            .unwrap()
309            .pending
310            .get(hash)
311            .map(|data| data.len() as u64)
312        {
313            return Ok(Some(size));
314        }
315        self.base.blob_size(hash).await
316    }
317
318    async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
319        if self.inner.read().unwrap().pending.contains_key(hash) {
320            return Ok(true);
321        }
322        self.base.has(hash).await
323    }
324
325    async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
326        let removed = {
327            let mut inner = self.inner.write().unwrap();
328            let removed = inner.pending.remove(hash).is_some();
329            if removed {
330                inner.order.retain(|queued| queued != hash);
331            }
332            removed
333        };
334
335        if removed {
336            return Ok(true);
337        }
338
339        self.base.delete(hash).await
340    }
341
342    async fn stats(&self) -> StoreStats {
343        let mut stats = self.base.stats().await;
344        let pending_bytes = self
345            .inner
346            .read()
347            .unwrap()
348            .pending
349            .values()
350            .map(|data| data.len() as u64)
351            .sum::<u64>();
352        stats.count += self.inner.read().unwrap().pending.len() as u64;
353        stats.bytes += pending_bytes;
354        stats
355    }
356
357    async fn evict_if_needed(&self) -> Result<u64, StoreError> {
358        self.base.evict_if_needed().await
359    }
360
361    async fn pin(&self, hash: &Hash) -> Result<(), StoreError> {
362        self.base.pin(hash).await
363    }
364
365    async fn unpin(&self, hash: &Hash) -> Result<(), StoreError> {
366        self.base.unpin(hash).await
367    }
368
369    fn pin_count(&self, hash: &Hash) -> u32 {
370        self.base.pin_count(hash)
371    }
372}
373
374/// Entry in the memory store with metadata for LRU
375#[derive(Debug, Clone)]
376struct MemoryEntry {
377    data: Vec<u8>,
378    /// Insertion order for LRU (lower = older)
379    order: u64,
380}
381
382/// Internal state for MemoryStore
383#[derive(Debug, Default)]
384struct MemoryStoreInner {
385    data: HashMap<String, MemoryEntry>,
386    pins: HashMap<String, u32>,
387    next_order: u64,
388    max_bytes: Option<u64>,
389}
390
391/// In-memory content-addressed store with LRU eviction and pinning
392#[derive(Debug, Clone, Default)]
393pub struct MemoryStore {
394    inner: Arc<RwLock<MemoryStoreInner>>,
395}
396
397impl MemoryStore {
398    pub fn new() -> Self {
399        Self {
400            inner: Arc::new(RwLock::new(MemoryStoreInner::default())),
401        }
402    }
403
404    /// Create a new store with a maximum size limit
405    pub fn with_max_bytes(max_bytes: u64) -> Self {
406        Self {
407            inner: Arc::new(RwLock::new(MemoryStoreInner {
408                max_bytes: if max_bytes > 0 { Some(max_bytes) } else { None },
409                ..Default::default()
410            })),
411        }
412    }
413
414    /// Get number of stored items
415    pub fn size(&self) -> usize {
416        self.inner.read().unwrap().data.len()
417    }
418
419    /// Get total bytes stored
420    pub fn total_bytes(&self) -> usize {
421        self.inner
422            .read()
423            .unwrap()
424            .data
425            .values()
426            .map(|e| e.data.len())
427            .sum()
428    }
429
430    /// Clear all data (but not pins)
431    pub fn clear(&self) {
432        self.inner.write().unwrap().data.clear();
433    }
434
435    /// List all hashes
436    pub fn keys(&self) -> Vec<Hash> {
437        self.inner
438            .read()
439            .unwrap()
440            .data
441            .keys()
442            .filter_map(|hex| {
443                let bytes = hex::decode(hex).ok()?;
444                if bytes.len() != 32 {
445                    return None;
446                }
447                let mut hash = [0u8; 32];
448                hash.copy_from_slice(&bytes);
449                Some(hash)
450            })
451            .collect()
452    }
453
454    /// Evict oldest unpinned entries until under target bytes
455    fn evict_to_target(&self, target_bytes: u64) -> u64 {
456        let mut inner = self.inner.write().unwrap();
457
458        let current_bytes: u64 = inner.data.values().map(|e| e.data.len() as u64).sum();
459        if current_bytes <= target_bytes {
460            return 0;
461        }
462
463        // Collect unpinned entries sorted by order (oldest first)
464        let mut unpinned: Vec<(String, u64, u64)> = inner
465            .data
466            .iter()
467            .filter(|(key, _)| inner.pins.get(*key).copied().unwrap_or(0) == 0)
468            .map(|(key, entry)| (key.clone(), entry.order, entry.data.len() as u64))
469            .collect();
470
471        unpinned.sort_by_key(|(_, order, _)| *order);
472
473        let mut freed = 0u64;
474        let to_free = current_bytes - target_bytes;
475
476        for (key, _, size) in unpinned {
477            if freed >= to_free {
478                break;
479            }
480            inner.data.remove(&key);
481            freed += size;
482        }
483
484        freed
485    }
486}
487
488#[async_trait]
489impl Store for MemoryStore {
490    async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError> {
491        let key = to_hex(&hash);
492        let mut inner = self.inner.write().unwrap();
493        if inner.data.contains_key(&key) {
494            return Ok(false);
495        }
496        let order = inner.next_order;
497        inner.next_order += 1;
498        inner.data.insert(key, MemoryEntry { data, order });
499        Ok(true)
500    }
501
502    async fn put_many(&self, items: Vec<(Hash, Vec<u8>)>) -> Result<usize, StoreError> {
503        let mut inserted = 0usize;
504        let mut inner = self.inner.write().unwrap();
505        for (hash, data) in items {
506            let key = to_hex(&hash);
507            if inner.data.contains_key(&key) {
508                continue;
509            }
510            let order = inner.next_order;
511            inner.next_order += 1;
512            inner.data.insert(key, MemoryEntry { data, order });
513            inserted += 1;
514        }
515        Ok(inserted)
516    }
517
518    async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError> {
519        let key = to_hex(hash);
520        let inner = self.inner.read().unwrap();
521        Ok(inner.data.get(&key).map(|e| e.data.clone()))
522    }
523
524    async fn get_range(
525        &self,
526        hash: &Hash,
527        start: u64,
528        end_inclusive: u64,
529    ) -> Result<Option<Vec<u8>>, StoreError> {
530        let key = to_hex(hash);
531        let inner = self.inner.read().unwrap();
532        inner
533            .data
534            .get(&key)
535            .map(|entry| slice_blob_range(&entry.data, start, end_inclusive))
536            .transpose()
537    }
538
539    async fn blob_size(&self, hash: &Hash) -> Result<Option<u64>, StoreError> {
540        let key = to_hex(hash);
541        let inner = self.inner.read().unwrap();
542        Ok(inner.data.get(&key).map(|entry| entry.data.len() as u64))
543    }
544
545    async fn has(&self, hash: &Hash) -> Result<bool, StoreError> {
546        let key = to_hex(hash);
547        Ok(self.inner.read().unwrap().data.contains_key(&key))
548    }
549
550    async fn delete(&self, hash: &Hash) -> Result<bool, StoreError> {
551        let key = to_hex(hash);
552        let mut inner = self.inner.write().unwrap();
553        // Also remove pin entry if exists
554        inner.pins.remove(&key);
555        Ok(inner.data.remove(&key).is_some())
556    }
557
558    fn set_max_bytes(&self, max: u64) {
559        self.inner.write().unwrap().max_bytes = if max > 0 { Some(max) } else { None };
560    }
561
562    fn max_bytes(&self) -> Option<u64> {
563        self.inner.read().unwrap().max_bytes
564    }
565
566    async fn stats(&self) -> StoreStats {
567        let inner = self.inner.read().unwrap();
568        let mut count = 0u64;
569        let mut bytes = 0u64;
570        let mut pinned_count = 0u64;
571        let mut pinned_bytes = 0u64;
572
573        for (key, entry) in &inner.data {
574            count += 1;
575            bytes += entry.data.len() as u64;
576            if inner.pins.get(key).copied().unwrap_or(0) > 0 {
577                pinned_count += 1;
578                pinned_bytes += entry.data.len() as u64;
579            }
580        }
581
582        StoreStats {
583            count,
584            bytes,
585            pinned_count,
586            pinned_bytes,
587        }
588    }
589
590    async fn evict_if_needed(&self) -> Result<u64, StoreError> {
591        let max = match self.inner.read().unwrap().max_bytes {
592            Some(m) => m,
593            None => return Ok(0), // No limit set
594        };
595
596        let current: u64 = self
597            .inner
598            .read()
599            .unwrap()
600            .data
601            .values()
602            .map(|e| e.data.len() as u64)
603            .sum();
604
605        if current <= max {
606            return Ok(0);
607        }
608
609        // Evict to 90% of max to avoid frequent evictions
610        let target = max * 9 / 10;
611        Ok(self.evict_to_target(target))
612    }
613
614    async fn pin(&self, hash: &Hash) -> Result<(), StoreError> {
615        let key = to_hex(hash);
616        let mut inner = self.inner.write().unwrap();
617        *inner.pins.entry(key).or_insert(0) += 1;
618        Ok(())
619    }
620
621    async fn unpin(&self, hash: &Hash) -> Result<(), StoreError> {
622        let key = to_hex(hash);
623        let mut inner = self.inner.write().unwrap();
624        if let Some(count) = inner.pins.get_mut(&key) {
625            if *count > 0 {
626                *count -= 1;
627            }
628            if *count == 0 {
629                inner.pins.remove(&key);
630            }
631        }
632        Ok(())
633    }
634
635    fn pin_count(&self, hash: &Hash) -> u32 {
636        let key = to_hex(hash);
637        self.inner
638            .read()
639            .unwrap()
640            .pins
641            .get(&key)
642            .copied()
643            .unwrap_or(0)
644    }
645}
646
647#[cfg(test)]
648mod tests {
649    use super::*;
650    use crate::hash::sha256;
651
652    #[tokio::test]
653    async fn test_put_returns_true_for_new() {
654        let store = MemoryStore::new();
655        let data = vec![1u8, 2, 3];
656        let hash = sha256(&data);
657
658        let result = store.put(hash, data).await.unwrap();
659        assert!(result);
660    }
661
662    #[tokio::test]
663    async fn test_put_returns_false_for_duplicate() {
664        let store = MemoryStore::new();
665        let data = vec![1u8, 2, 3];
666        let hash = sha256(&data);
667
668        store.put(hash, data.clone()).await.unwrap();
669        let result = store.put(hash, data).await.unwrap();
670        assert!(!result);
671    }
672
673    #[tokio::test]
674    async fn test_put_many_counts_only_new_items() {
675        let store = MemoryStore::new();
676        let data1 = vec![1u8, 2, 3];
677        let data2 = vec![4u8, 5, 6];
678        let hash1 = sha256(&data1);
679        let hash2 = sha256(&data2);
680
681        store.put(hash1, data1.clone()).await.unwrap();
682        let inserted = store
683            .put_many(vec![(hash1, data1), (hash2, data2.clone())])
684            .await
685            .unwrap();
686
687        assert_eq!(inserted, 1);
688        assert_eq!(store.get(&hash2).await.unwrap(), Some(data2));
689    }
690
691    #[tokio::test]
692    async fn test_buffered_store_flushes_pending_writes() {
693        let base = std::sync::Arc::new(MemoryStore::new());
694        let buffered = BufferedStore::new(std::sync::Arc::clone(&base));
695        let data = vec![9u8, 8, 7];
696        let hash = sha256(&data);
697
698        assert!(buffered.put(hash, data.clone()).await.unwrap());
699        assert_eq!(buffered.get(&hash).await.unwrap(), Some(data.clone()));
700        assert_eq!(base.get(&hash).await.unwrap(), None);
701
702        let flushed = buffered.flush().await.unwrap();
703
704        assert_eq!(flushed, 1);
705        assert_eq!(base.get(&hash).await.unwrap(), Some(data));
706    }
707
708    #[tokio::test]
709    async fn test_optimistic_buffered_store_avoids_base_probe_but_preserves_contents() {
710        let base = std::sync::Arc::new(MemoryStore::new());
711        let buffered = BufferedStore::new_optimistic(std::sync::Arc::clone(&base));
712        let data = vec![4u8, 5, 6];
713        let hash = sha256(&data);
714
715        base.put(hash, data.clone()).await.unwrap();
716
717        assert!(buffered.put(hash, data.clone()).await.unwrap());
718        assert_eq!(buffered.get(&hash).await.unwrap(), Some(data.clone()));
719
720        let flushed = buffered.flush().await.unwrap();
721
722        assert_eq!(flushed, 0);
723        assert_eq!(base.get(&hash).await.unwrap(), Some(data));
724    }
725
726    #[tokio::test]
727    async fn test_get_returns_data() {
728        let store = MemoryStore::new();
729        let data = vec![1u8, 2, 3];
730        let hash = sha256(&data);
731
732        store.put(hash, data.clone()).await.unwrap();
733        let result = store.get(&hash).await.unwrap();
734
735        assert_eq!(result, Some(data));
736    }
737
738    #[tokio::test]
739    async fn test_get_returns_none_for_missing() {
740        let store = MemoryStore::new();
741        let hash = [0u8; 32];
742
743        let result = store.get(&hash).await.unwrap();
744        assert!(result.is_none());
745    }
746
747    #[tokio::test]
748    async fn test_has_returns_true() {
749        let store = MemoryStore::new();
750        let data = vec![1u8, 2, 3];
751        let hash = sha256(&data);
752
753        store.put(hash, data).await.unwrap();
754        assert!(store.has(&hash).await.unwrap());
755    }
756
757    #[tokio::test]
758    async fn test_has_returns_false() {
759        let store = MemoryStore::new();
760        let hash = [0u8; 32];
761
762        assert!(!store.has(&hash).await.unwrap());
763    }
764
765    #[tokio::test]
766    async fn test_delete_returns_true() {
767        let store = MemoryStore::new();
768        let data = vec![1u8, 2, 3];
769        let hash = sha256(&data);
770
771        store.put(hash, data).await.unwrap();
772        let result = store.delete(&hash).await.unwrap();
773
774        assert!(result);
775        assert!(!store.has(&hash).await.unwrap());
776    }
777
778    #[tokio::test]
779    async fn test_delete_returns_false() {
780        let store = MemoryStore::new();
781        let hash = [0u8; 32];
782
783        let result = store.delete(&hash).await.unwrap();
784        assert!(!result);
785    }
786
787    #[tokio::test]
788    async fn test_size() {
789        let store = MemoryStore::new();
790        assert_eq!(store.size(), 0);
791
792        let data1 = vec![1u8];
793        let data2 = vec![2u8];
794        let hash1 = sha256(&data1);
795        let hash2 = sha256(&data2);
796
797        store.put(hash1, data1).await.unwrap();
798        store.put(hash2, data2).await.unwrap();
799
800        assert_eq!(store.size(), 2);
801    }
802
803    #[tokio::test]
804    async fn test_total_bytes() {
805        let store = MemoryStore::new();
806        assert_eq!(store.total_bytes(), 0);
807
808        let data1 = vec![1u8, 2, 3];
809        let data2 = vec![4u8, 5];
810        let hash1 = sha256(&data1);
811        let hash2 = sha256(&data2);
812
813        store.put(hash1, data1).await.unwrap();
814        store.put(hash2, data2).await.unwrap();
815
816        assert_eq!(store.total_bytes(), 5);
817    }
818
819    #[tokio::test]
820    async fn test_clear() {
821        let store = MemoryStore::new();
822        let data = vec![1u8, 2, 3];
823        let hash = sha256(&data);
824
825        store.put(hash, data).await.unwrap();
826        store.clear();
827
828        assert_eq!(store.size(), 0);
829        assert!(!store.has(&hash).await.unwrap());
830    }
831
832    #[tokio::test]
833    async fn test_keys() {
834        let store = MemoryStore::new();
835        assert!(store.keys().is_empty());
836
837        let data1 = vec![1u8];
838        let data2 = vec![2u8];
839        let hash1 = sha256(&data1);
840        let hash2 = sha256(&data2);
841
842        store.put(hash1, data1).await.unwrap();
843        store.put(hash2, data2).await.unwrap();
844
845        let keys = store.keys();
846        assert_eq!(keys.len(), 2);
847
848        let mut hex_keys: Vec<_> = keys.iter().map(to_hex).collect();
849        hex_keys.sort();
850        let mut expected: Vec<_> = vec![to_hex(&hash1), to_hex(&hash2)];
851        expected.sort();
852        assert_eq!(hex_keys, expected);
853    }
854
855    #[tokio::test]
856    async fn test_pin_and_unpin() {
857        let store = MemoryStore::new();
858        let data = vec![1u8, 2, 3];
859        let hash = sha256(&data);
860
861        store.put(hash, data).await.unwrap();
862
863        // Initially not pinned
864        assert!(!store.is_pinned(&hash));
865        assert_eq!(store.pin_count(&hash), 0);
866
867        // Pin
868        store.pin(&hash).await.unwrap();
869        assert!(store.is_pinned(&hash));
870        assert_eq!(store.pin_count(&hash), 1);
871
872        // Unpin
873        store.unpin(&hash).await.unwrap();
874        assert!(!store.is_pinned(&hash));
875        assert_eq!(store.pin_count(&hash), 0);
876    }
877
878    #[tokio::test]
879    async fn test_pin_count_ref_counting() {
880        let store = MemoryStore::new();
881        let data = vec![1u8, 2, 3];
882        let hash = sha256(&data);
883
884        store.put(hash, data).await.unwrap();
885
886        // Pin multiple times
887        store.pin(&hash).await.unwrap();
888        store.pin(&hash).await.unwrap();
889        store.pin(&hash).await.unwrap();
890        assert_eq!(store.pin_count(&hash), 3);
891
892        // Unpin once
893        store.unpin(&hash).await.unwrap();
894        assert_eq!(store.pin_count(&hash), 2);
895        assert!(store.is_pinned(&hash));
896
897        // Unpin remaining
898        store.unpin(&hash).await.unwrap();
899        store.unpin(&hash).await.unwrap();
900        assert_eq!(store.pin_count(&hash), 0);
901        assert!(!store.is_pinned(&hash));
902
903        // Extra unpin shouldn't go negative
904        store.unpin(&hash).await.unwrap();
905        assert_eq!(store.pin_count(&hash), 0);
906    }
907
908    #[tokio::test]
909    async fn test_stats() {
910        let store = MemoryStore::new();
911
912        let data1 = vec![1u8, 2, 3]; // 3 bytes
913        let data2 = vec![4u8, 5]; // 2 bytes
914        let hash1 = sha256(&data1);
915        let hash2 = sha256(&data2);
916
917        store.put(hash1, data1).await.unwrap();
918        store.put(hash2, data2).await.unwrap();
919
920        // Pin one item
921        store.pin(&hash1).await.unwrap();
922
923        let stats = store.stats().await;
924        assert_eq!(stats.count, 2);
925        assert_eq!(stats.bytes, 5);
926        assert_eq!(stats.pinned_count, 1);
927        assert_eq!(stats.pinned_bytes, 3);
928    }
929
930    #[tokio::test]
931    async fn test_max_bytes() {
932        let store = MemoryStore::new();
933        assert!(store.max_bytes().is_none());
934
935        store.set_max_bytes(1000);
936        assert_eq!(store.max_bytes(), Some(1000));
937
938        // 0 means unlimited
939        store.set_max_bytes(0);
940        assert!(store.max_bytes().is_none());
941    }
942
943    #[tokio::test]
944    async fn test_with_max_bytes() {
945        let store = MemoryStore::with_max_bytes(500);
946        assert_eq!(store.max_bytes(), Some(500));
947
948        let store_unlimited = MemoryStore::with_max_bytes(0);
949        assert!(store_unlimited.max_bytes().is_none());
950    }
951
952    #[tokio::test]
953    async fn test_eviction_respects_pins() {
954        // Store with 10 byte limit
955        let store = MemoryStore::with_max_bytes(10);
956
957        // Insert 3 items: 3 + 3 + 3 = 9 bytes
958        let data1 = vec![1u8, 1, 1]; // oldest
959        let data2 = vec![2u8, 2, 2];
960        let data3 = vec![3u8, 3, 3]; // newest
961        let hash1 = sha256(&data1);
962        let hash2 = sha256(&data2);
963        let hash3 = sha256(&data3);
964
965        store.put(hash1, data1).await.unwrap();
966        store.put(hash2, data2).await.unwrap();
967        store.put(hash3, data3).await.unwrap();
968
969        // Pin the oldest item
970        store.pin(&hash1).await.unwrap();
971
972        // Add more data to exceed limit: 9 + 3 = 12 bytes > 10
973        let data4 = vec![4u8, 4, 4];
974        let hash4 = sha256(&data4);
975        store.put(hash4, data4).await.unwrap();
976
977        // Evict - should remove hash2 (oldest unpinned)
978        let freed = store.evict_if_needed().await.unwrap();
979        assert!(freed > 0);
980
981        // hash1 should still exist (pinned)
982        assert!(store.has(&hash1).await.unwrap());
983        // hash2 should be gone (oldest unpinned)
984        assert!(!store.has(&hash2).await.unwrap());
985        // hash3 and hash4 should exist
986        assert!(store.has(&hash3).await.unwrap());
987        assert!(store.has(&hash4).await.unwrap());
988    }
989
990    #[tokio::test]
991    async fn test_eviction_lru_order() {
992        // Store with 15 byte limit
993        let store = MemoryStore::with_max_bytes(15);
994
995        // Insert items in order (oldest first)
996        let data1 = vec![1u8; 5]; // oldest
997        let data2 = vec![2u8; 5];
998        let data3 = vec![3u8; 5];
999        let data4 = vec![4u8; 5]; // newest
1000        let hash1 = sha256(&data1);
1001        let hash2 = sha256(&data2);
1002        let hash3 = sha256(&data3);
1003        let hash4 = sha256(&data4);
1004
1005        store.put(hash1, data1).await.unwrap();
1006        store.put(hash2, data2).await.unwrap();
1007        store.put(hash3, data3).await.unwrap();
1008        store.put(hash4, data4).await.unwrap();
1009
1010        // Now at 20 bytes, limit is 15
1011        assert_eq!(store.total_bytes(), 20);
1012
1013        // Evict - should remove oldest items first
1014        let freed = store.evict_if_needed().await.unwrap();
1015        assert!(freed >= 5); // At least one item evicted
1016
1017        // Oldest should be gone
1018        assert!(!store.has(&hash1).await.unwrap());
1019        // Newest should still exist
1020        assert!(store.has(&hash4).await.unwrap());
1021    }
1022
1023    #[tokio::test]
1024    async fn test_no_eviction_when_under_limit() {
1025        let store = MemoryStore::with_max_bytes(100);
1026
1027        let data = vec![1u8, 2, 3];
1028        let hash = sha256(&data);
1029        store.put(hash, data).await.unwrap();
1030
1031        let freed = store.evict_if_needed().await.unwrap();
1032        assert_eq!(freed, 0);
1033        assert!(store.has(&hash).await.unwrap());
1034    }
1035
1036    #[tokio::test]
1037    async fn test_no_eviction_without_limit() {
1038        let store = MemoryStore::new();
1039
1040        // Add lots of data
1041        for i in 0..100u8 {
1042            let data = vec![i; 100];
1043            let hash = sha256(&data);
1044            store.put(hash, data).await.unwrap();
1045        }
1046
1047        let freed = store.evict_if_needed().await.unwrap();
1048        assert_eq!(freed, 0);
1049        assert_eq!(store.size(), 100);
1050    }
1051
1052    #[tokio::test]
1053    async fn test_delete_removes_pin() {
1054        let store = MemoryStore::new();
1055        let data = vec![1u8, 2, 3];
1056        let hash = sha256(&data);
1057
1058        store.put(hash, data).await.unwrap();
1059        store.pin(&hash).await.unwrap();
1060        assert!(store.is_pinned(&hash));
1061
1062        store.delete(&hash).await.unwrap();
1063        // Pin should be gone after delete
1064        assert_eq!(store.pin_count(&hash), 0);
1065    }
1066}