nectar_primitives/store/
memory.rs

1//! In-memory chunk storage.
2
3use std::collections::HashMap;
4
5use parking_lot::RwLock;
6
7use crate::bmt::DEFAULT_BODY_SIZE;
8use crate::chunk::{AnyChunk, ChunkAddress};
9
10use super::ChunkStoreError;
11use super::typed::{ChunkGet, ChunkHas, ChunkPut};
12
13/// In-memory chunk storage using a `RwLock<HashMap>`.
14///
15/// Uses interior mutability so `ChunkPut::put(&self)` works without
16/// external synchronization.
17#[derive(Debug)]
18pub struct MemoryStore<const BODY_SIZE: usize = DEFAULT_BODY_SIZE> {
19    chunks: RwLock<HashMap<ChunkAddress, AnyChunk<BODY_SIZE>>>,
20}
21
22impl<const BODY_SIZE: usize> Clone for MemoryStore<BODY_SIZE> {
23    fn clone(&self) -> Self {
24        Self {
25            chunks: RwLock::new(self.chunks.read().clone()),
26        }
27    }
28}
29
30impl<const BODY_SIZE: usize> Default for MemoryStore<BODY_SIZE> {
31    fn default() -> Self {
32        Self::new()
33    }
34}
35
36impl<const BODY_SIZE: usize> MemoryStore<BODY_SIZE> {
37    /// Create an empty memory store.
38    pub fn new() -> Self {
39        Self {
40            chunks: RwLock::new(HashMap::new()),
41        }
42    }
43
44    /// Build a store from a collection of chunks, keyed by address.
45    pub fn from_chunks(chunks: impl IntoIterator<Item = AnyChunk<BODY_SIZE>>) -> Self {
46        Self {
47            chunks: RwLock::new(chunks.into_iter().map(|c| (*c.address(), c)).collect()),
48        }
49    }
50
51    /// Get a cloned chunk by address.
52    pub fn get(&self, address: &ChunkAddress) -> Option<AnyChunk<BODY_SIZE>> {
53        self.chunks.read().get(address).cloned()
54    }
55
56    /// Number of stored chunks.
57    pub fn len(&self) -> usize {
58        self.chunks.read().len()
59    }
60
61    /// Whether the store is empty.
62    pub fn is_empty(&self) -> bool {
63        self.chunks.read().is_empty()
64    }
65
66    /// Consume the store and return all chunks.
67    pub fn into_chunks(self) -> HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
68        self.chunks.into_inner()
69    }
70}
71
72impl<const BODY_SIZE: usize> ChunkPut<BODY_SIZE> for MemoryStore<BODY_SIZE> {
73    type Error = std::convert::Infallible;
74
75    async fn put(&self, chunk: AnyChunk<BODY_SIZE>) -> Result<(), Self::Error> {
76        self.chunks.write().insert(*chunk.address(), chunk);
77        Ok(())
78    }
79}
80
81impl<const BODY_SIZE: usize> ChunkGet<BODY_SIZE> for MemoryStore<BODY_SIZE> {
82    type Error = ChunkStoreError;
83
84    async fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
85        self.chunks
86            .read()
87            .get(address)
88            .cloned()
89            .ok_or_else(|| ChunkStoreError::not_found(address))
90    }
91}
92
93impl<const BODY_SIZE: usize> ChunkHas<BODY_SIZE> for MemoryStore<BODY_SIZE> {
94    async fn has(&self, address: &ChunkAddress) -> bool {
95        self.chunks.read().contains_key(address)
96    }
97}
98
99impl<const BODY_SIZE: usize> ChunkGet<BODY_SIZE> for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
100    type Error = ChunkStoreError;
101
102    async fn get(&self, address: &ChunkAddress) -> Result<AnyChunk<BODY_SIZE>, Self::Error> {
103        self.get(address)
104            .cloned()
105            .ok_or_else(|| ChunkStoreError::not_found(address))
106    }
107}
108
109impl<const BODY_SIZE: usize> ChunkHas<BODY_SIZE> for HashMap<ChunkAddress, AnyChunk<BODY_SIZE>> {
110    async fn has(&self, address: &ChunkAddress) -> bool {
111        self.contains_key(address)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118    use crate::chunk::{Chunk, ContentChunk};
119    use futures::executor::block_on;
120
121    #[test]
122    fn test_memory_store() {
123        let store = MemoryStore::<DEFAULT_BODY_SIZE>::new();
124        assert!(store.is_empty());
125
126        let chunk = ContentChunk::new(b"hello".as_slice()).unwrap();
127        let addr = *chunk.address();
128        let any: AnyChunk = chunk.into();
129
130        block_on(ChunkPut::put(&store, any.clone())).unwrap();
131        assert_eq!(store.len(), 1);
132        assert!(block_on(ChunkHas::has(&store, &addr)));
133        assert_eq!(store.get(&addr), Some(any));
134    }
135}
nectar_primitives/store/memory.rs

nectar_primitives/store/
memory.rs
