Skip to main content

rings_core/storage/sled/
mod.rs

1#![deny(missing_docs)]
2
3//! Persistent native key-value storage.
4
5use std::path::Path;
6use std::path::PathBuf;
7use std::sync::RwLock;
8
9use async_trait::async_trait;
10use itertools::Itertools;
11use serde::de::DeserializeOwned;
12use serde::Serialize;
13use sha1::Digest;
14use sha1::Sha1;
15
16use crate::error::Error;
17use crate::error::Result;
18use crate::storage::KvStorageInterface;
19
20/// StorageInstance struct
21#[allow(dead_code)]
22pub struct SledStorage {
23    root: PathBuf,
24    lock: RwLock<()>,
25    cap: u32,
26    path: String,
27}
28
29impl SledStorage {
30    /// New SledStorage
31    /// * cap: max_size in bytes
32    /// * path: db file location
33    pub async fn new_with_cap_and_path<P>(cap: u32, path: P) -> Result<Self>
34    where P: AsRef<std::path::Path> {
35        std::fs::create_dir_all(path.as_ref()).map_err(Error::ServiceIOError)?;
36        Ok(Self {
37            root: path.as_ref().to_path_buf(),
38            lock: RwLock::new(()),
39            cap,
40            path: path.as_ref().to_string_lossy().to_string(),
41        })
42    }
43
44    fn key_path(&self, key: &str) -> PathBuf {
45        let mut hasher = Sha1::new();
46        hasher.update(key.as_bytes());
47        self.root.join(hex::encode(hasher.finalize()))
48    }
49
50    fn entries(&self) -> Result<Vec<PathBuf>> {
51        match std::fs::read_dir(&self.root) {
52            Ok(entries) => Ok(entries
53                .flatten()
54                .map(|entry| entry.path())
55                .filter(|path| is_entry_path(path))
56                .collect_vec()),
57            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(Vec::new()),
58            Err(error) => Err(Error::ServiceIOError(error)),
59        }
60    }
61}
62
63#[async_trait]
64impl<V> KvStorageInterface<V> for SledStorage
65where V: Serialize + DeserializeOwned + Sync
66{
67    async fn get(&self, key: &str) -> Result<Option<V>> {
68        let _guard = self.lock.read().map_err(|_| Error::DHTSyncLockError)?;
69        match std::fs::read(self.key_path(key)) {
70            Ok(data) => {
71                let (stored_key, value): (String, V) =
72                    rings_codec::deserialize(&data).map_err(Error::CodecDeserialize)?;
73                if stored_key == key {
74                    Ok(Some(value))
75                } else {
76                    Ok(None)
77                }
78            }
79            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
80            Err(error) => Err(Error::ServiceIOError(error)),
81        }
82    }
83
84    async fn put(&self, key: &str, value: &V) -> Result<()> {
85        let _guard = self.lock.write().map_err(|_| Error::DHTSyncLockError)?;
86        std::fs::create_dir_all(&self.root).map_err(Error::ServiceIOError)?;
87        let data = rings_codec::serialize(&(key, value)).map_err(Error::CodecSerialize)?;
88        tracing::debug!("Try inserting key: {:?}", key);
89        let path = self.key_path(key);
90        let tmp_path = path.with_extension("tmp");
91        std::fs::write(&tmp_path, data).map_err(Error::ServiceIOError)?;
92        std::fs::rename(tmp_path, path).map_err(Error::ServiceIOError)?;
93        Ok(())
94    }
95
96    async fn get_all(&self) -> Result<Vec<(String, V)>> {
97        let _guard = self.lock.read().map_err(|_| Error::DHTSyncLockError)?;
98        Ok(self
99            .entries()?
100            .into_iter()
101            .flat_map(|path| {
102                let data = std::fs::read(path).ok()?;
103                rings_codec::deserialize::<(String, V)>(&data).ok()
104            })
105            .collect_vec())
106    }
107
108    async fn remove(&self, key: &str) -> Result<()> {
109        let _guard = self.lock.write().map_err(|_| Error::DHTSyncLockError)?;
110        match std::fs::remove_file(self.key_path(key)) {
111            Ok(()) => Ok(()),
112            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
113            Err(error) => Err(Error::ServiceIOError(error)),
114        }
115    }
116
117    async fn clear(&self) -> Result<()> {
118        let _guard = self.lock.write().map_err(|_| Error::DHTSyncLockError)?;
119        for path in self.entries()? {
120            std::fs::remove_file(path).map_err(Error::ServiceIOError)?;
121        }
122        Ok(())
123    }
124
125    async fn count(&self) -> Result<u32> {
126        let _guard = self.lock.read().map_err(|_| Error::DHTSyncLockError)?;
127        Ok(self.entries()?.len() as u32)
128    }
129}
130
131fn is_entry_path(path: &Path) -> bool {
132    let Some(file_name) = path.file_name().and_then(|name| name.to_str()) else {
133        return false;
134    };
135    file_name.len() == 40 && file_name.as_bytes().iter().all(u8::is_ascii_hexdigit)
136}
137
138impl std::fmt::Debug for SledStorage {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("SledStorage")
141            .field("cap", &self.cap)
142            .field("path", &self.path)
143            .finish()
144    }
145}
146
147#[cfg(test)]
148mod test_sled;