1use async_trait::async_trait;
4use std::collections::HashMap;
5use std::sync::{Arc, RwLock};
6
7use crate::types::{to_hex, Hash};
8
9pub 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#[derive(Debug, Clone, Default)]
38pub struct StoreStats {
39 pub count: u64,
41 pub bytes: u64,
43 pub pinned_count: u64,
45 pub pinned_bytes: u64,
47}
48
49#[derive(Debug, Clone, Default, PartialEq, Eq)]
51pub struct PutManyReport {
52 pub total: usize,
54 pub inserted: usize,
56 pub inserted_bytes: u64,
58 pub inserted_hashes: Vec<Hash>,
60}
61
62#[async_trait]
64pub trait Store: Send + Sync {
65 async fn put(&self, hash: Hash, data: Vec<u8>) -> Result<bool, StoreError>;
68
69 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 async fn get(&self, hash: &Hash) -> Result<Option<Vec<u8>>, StoreError>;
84
85 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 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 async fn has(&self, hash: &Hash) -> Result<bool, StoreError>;
112
113 async fn delete(&self, hash: &Hash) -> Result<bool, StoreError>;
116
117 fn set_max_bytes(&self, _max: u64) {}
123
124 fn max_bytes(&self) -> Option<u64> {
126 None
127 }
128
129 async fn stats(&self) -> StoreStats {
131 StoreStats::default()
132 }
133
134 async fn evict_if_needed(&self) -> Result<u64, StoreError> {
137 Ok(0)
138 }
139
140 async fn pin(&self, _hash: &Hash) -> Result<(), StoreError> {
146 Ok(())
147 }
148
149 async fn unpin(&self, _hash: &Hash) -> Result<(), StoreError> {
151 Ok(())
152 }
153
154 fn pin_count(&self, _hash: &Hash) -> u32 {
156 0
157 }
158
159 fn is_pinned(&self, hash: &Hash) -> bool {
161 self.pin_count(hash) > 0
162 }
163}
164
165#[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#[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#[derive(Debug, Clone)]
376struct MemoryEntry {
377 data: Vec<u8>,
378 order: u64,
380}
381
382#[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#[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 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 pub fn size(&self) -> usize {
416 self.inner.read().unwrap().data.len()
417 }
418
419 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 pub fn clear(&self) {
432 self.inner.write().unwrap().data.clear();
433 }
434
435 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 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 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 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), };
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 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 assert!(!store.is_pinned(&hash));
865 assert_eq!(store.pin_count(&hash), 0);
866
867 store.pin(&hash).await.unwrap();
869 assert!(store.is_pinned(&hash));
870 assert_eq!(store.pin_count(&hash), 1);
871
872 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 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 store.unpin(&hash).await.unwrap();
894 assert_eq!(store.pin_count(&hash), 2);
895 assert!(store.is_pinned(&hash));
896
897 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 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]; let data2 = vec![4u8, 5]; 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 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 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 let store = MemoryStore::with_max_bytes(10);
956
957 let data1 = vec![1u8, 1, 1]; let data2 = vec![2u8, 2, 2];
960 let data3 = vec![3u8, 3, 3]; 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 store.pin(&hash1).await.unwrap();
971
972 let data4 = vec![4u8, 4, 4];
974 let hash4 = sha256(&data4);
975 store.put(hash4, data4).await.unwrap();
976
977 let freed = store.evict_if_needed().await.unwrap();
979 assert!(freed > 0);
980
981 assert!(store.has(&hash1).await.unwrap());
983 assert!(!store.has(&hash2).await.unwrap());
985 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 let store = MemoryStore::with_max_bytes(15);
994
995 let data1 = vec![1u8; 5]; let data2 = vec![2u8; 5];
998 let data3 = vec![3u8; 5];
999 let data4 = vec![4u8; 5]; 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 assert_eq!(store.total_bytes(), 20);
1012
1013 let freed = store.evict_if_needed().await.unwrap();
1015 assert!(freed >= 5); assert!(!store.has(&hash1).await.unwrap());
1019 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 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 assert_eq!(store.pin_count(&hash), 0);
1065 }
1066}