rpocket/store/
mod.rs

1use crate::error::RPocketError;
2use async_trait::async_trait;
3
4#[async_trait]
5pub trait Storage {
6    async fn get(&self, key: &str) -> Result<Option<String>, RPocketError>;
7    async fn set(&self, key: &str, value: &str) -> Result<(), RPocketError>;
8    async fn delete(&self, key: &str) -> Result<(), RPocketError>;
9}
10
11/// MemoryStorage is a simple implementation of Storage.
12pub struct MemoryStorage {
13    pub data: std::sync::RwLock<std::collections::HashMap<String, std::sync::RwLock<String>>>,
14}
15
16impl MemoryStorage {
17    /// create a new MemoryStorage.
18    pub fn new() -> Self {
19        MemoryStorage {
20            data: std::sync::RwLock::new(std::collections::HashMap::new()),
21        }
22    }
23}
24
25impl Default for MemoryStorage {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31#[async_trait]
32impl Storage for MemoryStorage {
33    /// get the value of a key.
34    async fn get(&self, key: &str) -> Result<Option<String>, RPocketError> {
35        let data = self.data.read().map_err(|_| RPocketError::MutexError)?;
36        let value = data.get(key);
37
38        return match value {
39            Some(value) => {
40                let value = value.read().map_err(|_| RPocketError::MutexError)?;
41                return Ok(Some(value.to_string()));
42            }
43            None => Ok(None),
44        };
45    }
46
47    /// set the value of a key.
48    async fn set(&self, key: &str, value: &str) -> Result<(), RPocketError> {
49        let data = self.data.read().map_err(|_| RPocketError::MutexError)?;
50        let inner_entry = data.get(key);
51
52        match inner_entry {
53            Some(entry) => {
54                let mut inner_value = entry.write().map_err(|_| RPocketError::MutexError)?;
55                *inner_value = value.to_string();
56            }
57            None => {
58                drop(data);
59                let mut data = self.data.write().map_err(|_| RPocketError::MutexError)?;
60                data.insert(key.to_string(), std::sync::RwLock::new(value.to_string()));
61            }
62        }
63
64        return Ok(());
65    }
66
67    /// delete a key.
68    async fn delete(&self, key: &str) -> Result<(), RPocketError> {
69        let data = self.data.read().map_err(|_| RPocketError::MutexError)?;
70        let inner_entry = data.get(key);
71
72        if let Some(..) = inner_entry {
73            drop(data);
74            let mut data = self.data.write().map_err(|_| RPocketError::MutexError)?;
75            data.remove(key);
76        }
77
78        return Ok(());
79    }
80}
81
82#[cfg(test)]
83mod test {
84    use super::*;
85
86    #[tokio::test]
87    async fn test_memory_storage() {
88        let storage = MemoryStorage::new();
89        storage.set("key", "value").await.unwrap();
90        assert_eq!(storage.get("key").await.unwrap().unwrap(), "value");
91
92        storage.set("key", "value2").await.unwrap();
93        assert_eq!(storage.get("key").await.unwrap().unwrap(), "value2");
94
95        storage.delete("key").await.unwrap();
96        assert_eq!(storage.get("key").await.unwrap(), None);
97    }
98}
rpocket/store/mod.rs

rpocket/store/
mod.rs
