rings_core/storage/
sled.rs

1#![warn(missing_docs)]
2
3//! Persistence Storage for default, use `sled` as backend db.
4
5use async_trait::async_trait;
6use itertools::Itertools;
7use serde::de::DeserializeOwned;
8use serde::Serialize;
9
10use crate::error::Error;
11use crate::error::Result;
12use crate::storage::KvStorageInterface;
13
14/// StorageInstance struct
15#[allow(dead_code)]
16pub struct SledStorage {
17    db: sled::Db,
18    cap: u32,
19    path: String,
20}
21
22impl SledStorage {
23    /// New SledStorage
24    /// * cap: max_size in bytes
25    /// * path: db file location
26    pub async fn new_with_cap_and_path<P>(cap: u32, path: P) -> Result<Self>
27    where P: AsRef<std::path::Path> {
28        let db = sled::Config::new()
29            .path(path.as_ref())
30            .mode(sled::Mode::HighThroughput)
31            .cache_capacity(cap as u64)
32            .open()
33            .map_err(Error::SledError)?;
34        Ok(Self {
35            db,
36            cap,
37            path: path.as_ref().to_string_lossy().to_string(),
38        })
39    }
40}
41
42#[async_trait]
43impl<V> KvStorageInterface<V> for SledStorage
44where V: Serialize + DeserializeOwned + Sync
45{
46    async fn get(&self, key: &str) -> Result<Option<V>> {
47        let v = self.db.get(key).map_err(Error::SledError)?;
48        if let Some(v) = v {
49            let v = v.as_ref();
50            return bincode::deserialize(v)
51                .map_err(Error::BincodeDeserialize)
52                .map(|r| Some(r));
53        }
54        Ok(None)
55    }
56
57    async fn put(&self, key: &str, value: &V) -> Result<()> {
58        let data = bincode::serialize(&value).map_err(Error::BincodeSerialize)?;
59        tracing::debug!("Try inserting key: {:?}", key);
60        self.db.insert(key, data).map_err(Error::SledError)?;
61        Ok(())
62    }
63
64    async fn get_all(&self) -> Result<Vec<(String, V)>> {
65        let iter = self.db.iter();
66        Ok(iter
67            .flatten()
68            .flat_map(|(k, v)| {
69                Some((
70                    std::str::from_utf8(k.as_ref()).ok()?.to_string(),
71                    bincode::deserialize(v.as_ref()).ok()?,
72                ))
73            })
74            .collect_vec())
75    }
76
77    async fn remove(&self, key: &str) -> Result<()> {
78        self.db
79            .remove(key.to_string().as_bytes())
80            .map_err(Error::SledError)?;
81        Ok(())
82    }
83
84    async fn clear(&self) -> Result<()> {
85        self.db.clear().map_err(Error::SledError)?;
86        Ok(())
87    }
88
89    async fn count(&self) -> Result<u32> {
90        Ok(self.db.len() as u32)
91    }
92}
93
94impl std::fmt::Debug for SledStorage {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("SledStorage")
97            .field("cap", &self.cap)
98            .field("path", &self.path)
99            .finish()
100    }
101}
102
103#[cfg(test)]
104mod test {
105    use serde::Deserialize;
106    use serde::Serialize;
107
108    use super::*;
109
110    #[derive(Debug, Serialize, Deserialize)]
111    struct TestStorageStruct {
112        content: String,
113    }
114
115    #[tokio::test]
116    async fn test_kv_storage_put_delete() {
117        let storage = SledStorage::new_with_cap_and_path(4096, "tmp/test_db")
118            .await
119            .unwrap();
120        let key1 = "test1".to_owned();
121        let data1 = TestStorageStruct {
122            content: "test1".to_string(),
123        };
124        storage.put(&key1, &data1).await.unwrap();
125        let count1 =
126            <SledStorage as KvStorageInterface<TestStorageStruct>>::count::<'_, '_>(&storage)
127                .await
128                .unwrap();
129        assert!(count1 == 1, "expect count1.1 is {}, got {}", 1, count1);
130        let got_v1: TestStorageStruct = storage.get(&key1).await.unwrap().unwrap();
131        assert!(
132            got_v1.content.eq(data1.content.as_str()),
133            "expect value1 is {}, got {}",
134            data1.content,
135            got_v1.content
136        );
137
138        let key2 = "test2".to_owned();
139        let data2 = TestStorageStruct {
140            content: "test2".to_string(),
141        };
142
143        storage.put(&key2, &data2).await.unwrap();
144
145        let count_got_2 =
146            <SledStorage as KvStorageInterface<TestStorageStruct>>::count::<'_, '_>(&storage)
147                .await
148                .unwrap();
149        assert!(count_got_2 == 2, "expect count 2, got {count_got_2}");
150
151        let all_entries: Vec<(String, TestStorageStruct)> = storage.get_all().await.unwrap();
152        assert!(
153            all_entries.len() == 2,
154            "all_entries len expect 2, got {}",
155            all_entries.len()
156        );
157
158        let keys = [key1, key2];
159        let values = [data1.content, data2.content];
160
161        assert!(
162            all_entries
163                .iter()
164                .any(|(k, v)| { keys.contains(k) && values.contains(&v.content) }),
165            "not found items"
166        );
167        let data3: u64 = 101;
168        let key3 = "key3".to_owned();
169        storage.put(&key3, &data3).await.unwrap();
170        let got_d3: u64 = storage.get(&key3).await.unwrap().unwrap();
171        assert!(data3 == got_d3, "expect {data3}, got {got_d3}");
172
173        // Clear full db and check if it's count is zero now.
174        <SledStorage as KvStorageInterface<TestStorageStruct>>::clear::<'_, '_>(&storage)
175            .await
176            .unwrap();
177        let count1 =
178            <SledStorage as KvStorageInterface<TestStorageStruct>>::count::<'_, '_>(&storage)
179                .await
180                .unwrap();
181        assert!(count1 == 0, "expect count1 is 0, got {count1}");
182        storage.db.flush_async().await.unwrap();
183
184        drop(storage)
185    }
186}
rings_core/storage/sled.rs

rings_core/storage/
sled.rs
