tribbler/
storage.rs

1#![allow(dead_code)]
2//! module containing Tribbler storage-related structs and implementations
3use async_trait::async_trait;
4use std::{collections::HashMap, sync::RwLock};
5
6use crate::err::TribResult;
7
8#[derive(Debug, Clone)]
9
10/// A type comprising key-value pair
11pub struct KeyValue {
12    /// the key
13    pub key: String,
14    /// the value
15    pub value: String,
16}
17
18impl KeyValue {
19    /// convenience method for constructing a [KeyValue] from two `&str`s
20    pub fn new(key: &str, value: &str) -> KeyValue {
21        KeyValue {
22            key: key.to_string(),
23            value: value.to_string(),
24        }
25    }
26}
27
28#[derive(Debug, Clone, Default)]
29/// A type which represents a pattern that can be used to match on a String.
30pub struct Pattern {
31    /// exact-match string prefix
32    pub prefix: String,
33    /// exact-match string suffix
34    pub suffix: String,
35}
36
37impl Pattern {
38    /// this function returns true the provided string matches the prefix and
39    /// suffix of the given pattern
40    pub fn matches(&self, k: &str) -> bool {
41        k.starts_with(&self.prefix) && k.ends_with(&self.suffix)
42    }
43}
44
45#[derive(Debug, Clone)]
46/// A wrapper type around a [Vec<String>]
47pub struct List(pub Vec<String>);
48
49#[async_trait]
50/// Key-value pair interfaces
51/// Default value for all keys is empty string
52pub trait KeyString {
53    /// Gets a value. If no value set, return [None]
54    async fn get(&self, key: &str) -> TribResult<Option<String>>;
55
56    /// Set kv.key to kv.value. return true when no error.
57    async fn set(&self, kv: &KeyValue) -> TribResult<bool>;
58
59    /// List all the keys of non-empty pairs where the key matches
60    /// the given pattern.
61    async fn keys(&self, p: &Pattern) -> TribResult<List>;
62}
63
64#[async_trait]
65/// Key-list interfaces
66pub trait KeyList {
67    /// Get the list. Empty if not set.
68    async fn list_get(&self, key: &str) -> TribResult<List>;
69
70    /// Append a string to the list. return true when no error.
71    async fn list_append(&self, kv: &KeyValue) -> TribResult<bool>;
72
73    /// Removes all elements that are equal to `kv.value` in list `kv.key`
74    /// returns the number of elements removed.
75    async fn list_remove(&self, kv: &KeyValue) -> TribResult<u32>;
76
77    /// List all the keys of non-empty lists, where the key matches
78    /// the given pattern.
79    async fn list_keys(&self, p: &Pattern) -> TribResult<List>;
80}
81
82#[async_trait]
83/// A trait representing a storage interface
84/// The trait bounds for [KeyString] and [KeyList] respectively represent
85/// the functions requires for the single key-value and key-list parts of the
86/// storage interface.
87pub trait Storage: KeyString + KeyList + Send + Sync {
88    /// Returns an auto-incrementing clock. The returned value of each call will
89    /// be unique, no smaller than `at_least`, and strictly larger than the
90    /// value returned last time, unless it was [u64::MAX]
91    async fn clock(&self, at_least: u64) -> TribResult<u64>;
92}
93
94/// This is a toy implementation of a backend storage service.
95/// The trait definition requires this to be safe to utilize across threads
96/// because mutating methods (e.g. [KeyString::set] take `&self` instead of
97/// `&mut self`)
98#[derive(Debug, Default)]
99pub struct MemStorage {
100    kvs: RwLock<HashMap<String, String>>,
101    kv_list: RwLock<HashMap<String, List>>,
102    clock: RwLock<u64>,
103}
104
105impl MemStorage {
106    /// Creates a new instance of [MemStorage]
107    pub fn new() -> MemStorage {
108        MemStorage::default()
109    }
110}
111
112#[async_trait]
113impl KeyString for MemStorage {
114    async fn get(&self, key: &str) -> TribResult<Option<String>> {
115        match self.kvs.read().map_err(|e| e.to_string())?.get(key) {
116            Some(v) => Ok(Some(v.to_string())),
117            None => Ok(None),
118        }
119    }
120
121    async fn set(&self, kv: &KeyValue) -> TribResult<bool> {
122        let mut entry = self.kvs.write().map_err(|e| e.to_string())?;
123        if kv.value.is_empty() {
124            entry.remove(&kv.key);
125        } else {
126            entry.insert(kv.key.clone(), kv.value.clone());
127        }
128        Ok(true)
129    }
130
131    async fn keys(&self, p: &Pattern) -> TribResult<List> {
132        let result = self
133            .kvs
134            .read()
135            .map_err(|e| e.to_string())?
136            .iter()
137            .filter(|(k, _)| p.matches(*k))
138            .map(|(k, _)| k.to_string())
139            .collect::<Vec<String>>();
140        Ok(List(result))
141    }
142}
143
144#[async_trait]
145impl KeyList for MemStorage {
146    async fn list_get(&self, key: &str) -> TribResult<List> {
147        match self.kv_list.read().map_err(|e| e.to_string())?.get(key) {
148            Some(l) => Ok(l.clone()),
149            None => Ok(List(vec![])),
150        }
151    }
152
153    async fn list_append(&self, kv: &KeyValue) -> TribResult<bool> {
154        let mut kvl = self.kv_list.write().map_err(|e| e.to_string())?;
155        match kvl.get_mut(&kv.key) {
156            Some(list) => {
157                list.0.push(kv.value.clone());
158                Ok(true)
159            }
160            None => {
161                let list = vec![kv.value.clone()];
162                kvl.insert(kv.key.clone(), List(list));
163                Ok(true)
164            }
165        }
166    }
167
168    async fn list_remove(&self, kv: &KeyValue) -> TribResult<u32> {
169        let mut removed = 0;
170        let mut kvl = self.kv_list.write().map_err(|e| e.to_string())?;
171        kvl.entry(kv.key.clone()).and_modify(|list| {
172            let begin_size = list.0.len();
173            *list = List(
174                list.0
175                    .iter()
176                    .filter(|val| **val != kv.value)
177                    .map(String::from)
178                    .collect::<Vec<String>>(),
179            );
180            let end_size = list.0.len();
181            removed = begin_size - end_size;
182        });
183        if let Some(x) = kvl.get(&kv.key) {
184            if x.0.is_empty() {
185                kvl.remove(&kv.key);
186            }
187        };
188
189        Ok(removed as u32)
190    }
191
192    async fn list_keys(&self, p: &Pattern) -> TribResult<List> {
193        let mut result = vec![];
194        self.kv_list
195            .read()
196            .map_err(|e| e.to_string())?
197            .iter()
198            .filter(|(k, _)| p.matches(*k))
199            .for_each(|(v, _)| result.push((*v).clone()));
200        result.sort();
201        Ok(List(result))
202    }
203}
204
205#[async_trait]
206impl Storage for MemStorage {
207    async fn clock(&self, at_least: u64) -> TribResult<u64> {
208        let mut clk = self.clock.write().map_err(|e| e.to_string())?;
209        if *clk < at_least {
210            *clk = at_least
211        }
212
213        let ret = *clk;
214
215        if *clk < u64::MAX {
216            *clk += 1;
217        }
218        Ok(ret)
219    }
220}
221
222#[async_trait]
223/// Bin Storage interface
224pub trait BinStorage: Send + Sync {
225    /// Fetch a [Storage] bin based on the given bin name.
226    async fn bin(&self, name: &str) -> TribResult<Box<dyn Storage>>;
227}
228
229#[cfg(test)]
230mod test {
231    use crate::{
232        err::TribResult,
233        storage::{KeyValue, Pattern, Storage},
234    };
235
236    use super::{KeyList, KeyString, MemStorage};
237
238    async fn setup_test_storage() -> MemStorage {
239        let storage = MemStorage::new();
240        storage
241            .set(&KeyValue {
242                key: "test".to_string(),
243                value: "test-value".to_string(),
244            })
245            .await
246            .unwrap();
247        storage
248            .list_append(&KeyValue {
249                key: "test".to_string(),
250                value: "test-value".to_string(),
251            })
252            .await
253            .unwrap();
254        storage
255    }
256
257    #[tokio::test]
258    async fn storage_get_set() -> TribResult<()> {
259        let storage = MemStorage::new();
260        assert_eq!(
261            true,
262            storage
263                .set(&KeyValue {
264                    key: "test".to_string(),
265                    value: "test-value".to_string()
266                })
267                .await?
268        );
269        assert_eq!(Some("test-value".to_string()), storage.get("test").await?);
270        Ok(())
271    }
272
273    #[tokio::test]
274    async fn storage_get_empty() -> TribResult<()> {
275        let storage = setup_test_storage().await;
276        assert_eq!(None, storage.get("test2").await?);
277        Ok(())
278    }
279
280    #[tokio::test]
281    async fn storage_keys() {
282        let storage = setup_test_storage().await;
283        let p1 = Pattern {
284            prefix: "test".to_string(),
285            suffix: "test".to_string(),
286        };
287        let p2 = Pattern {
288            prefix: "".to_string(),
289            suffix: "test".to_string(),
290        };
291        let p3 = Pattern {
292            prefix: "test".to_string(),
293            suffix: "".to_string(),
294        };
295        let p4 = Pattern {
296            prefix: "wrong".to_string(),
297            suffix: "right".to_string(),
298        };
299        let p5 = Pattern {
300            prefix: "".to_string(),
301            suffix: "".to_string(),
302        };
303        assert_eq!(1, storage.keys(&p1).await.unwrap().0.len());
304        assert_eq!(1, storage.keys(&p2).await.unwrap().0.len());
305        assert_eq!(1, storage.keys(&p3).await.unwrap().0.len());
306        assert_eq!(0, storage.keys(&p4).await.unwrap().0.len());
307        assert_eq!(1, storage.keys(&p5).await.unwrap().0.len());
308    }
309
310    #[tokio::test]
311    async fn storage_keys_unset() {
312        let s = setup_test_storage().await;
313        assert_eq!(1, s.keys(&Pattern::default()).await.unwrap().0.len());
314        let _ = s.set(&KeyValue::new("test", "")).await.unwrap();
315        assert_eq!(0, s.keys(&Pattern::default()).await.unwrap().0.len())
316    }
317
318    #[tokio::test]
319    async fn storage_list_keys_unset() {
320        let s = setup_test_storage().await;
321        assert_eq!(1, s.list_keys(&Pattern::default()).await.unwrap().0.len());
322        let _ = s
323            .list_remove(&KeyValue::new("test", "test-value"))
324            .await
325            .unwrap();
326        assert_eq!(0, s.list_keys(&Pattern::default()).await.unwrap().0.len())
327    }
328
329    #[tokio::test]
330    async fn storage_get_list() {
331        let storage = setup_test_storage().await;
332        assert_eq!("test-value", storage.list_get("test").await.unwrap().0[0]);
333    }
334
335    #[tokio::test]
336    async fn storage_get_list_empty() {
337        let storage = setup_test_storage().await;
338        assert_eq!(0, storage.list_get("test2").await.unwrap().0.len());
339    }
340
341    #[tokio::test]
342    async fn storage_get_list_append() -> TribResult<()> {
343        let storage = setup_test_storage().await;
344        let res = storage
345            .list_append(&KeyValue {
346                key: "test".to_string(),
347                value: "val2".to_string(),
348            })
349            .await?;
350        assert_eq!(true, res);
351        assert_eq!(2, storage.list_get("test").await.unwrap().0.len());
352        Ok(())
353    }
354
355    #[tokio::test]
356    async fn storage_get_list_remove() {
357        let storage = setup_test_storage().await;
358        let kv = KeyValue {
359            key: "test".to_string(),
360            value: "val2".to_string(),
361        };
362        assert_eq!(true, storage.list_append(&kv).await.unwrap());
363        assert_eq!(true, storage.list_append(&kv).await.unwrap());
364        assert_eq!(true, storage.list_append(&kv).await.unwrap());
365        assert_eq!(3, storage.list_remove(&kv).await.unwrap());
366        println!("{:?}", storage.list_get("test").await.unwrap().0);
367        assert_eq!("test-value", storage.list_get("test").await.unwrap().0[0]);
368    }
369
370    #[tokio::test]
371    async fn storage_list_keys() {
372        let storage = setup_test_storage().await;
373        let p1 = Pattern {
374            prefix: "test".to_string(),
375            suffix: "test".to_string(),
376        };
377        let p2 = Pattern {
378            prefix: "".to_string(),
379            suffix: "test".to_string(),
380        };
381        let p3 = Pattern {
382            prefix: "test".to_string(),
383            suffix: "".to_string(),
384        };
385        let p4 = Pattern {
386            prefix: "wrong".to_string(),
387            suffix: "right".to_string(),
388        };
389        let p5 = Pattern {
390            prefix: "".to_string(),
391            suffix: "".to_string(),
392        };
393        assert_eq!(1, storage.list_keys(&p1).await.unwrap().0.len());
394        assert_eq!(1, storage.list_keys(&p2).await.unwrap().0.len());
395        assert_eq!(1, storage.list_keys(&p3).await.unwrap().0.len());
396        assert_eq!(0, storage.list_keys(&p4).await.unwrap().0.len());
397        assert_eq!(1, storage.list_keys(&p5).await.unwrap().0.len());
398    }
399
400    #[tokio::test]
401    async fn clock_at_least() {
402        let storage = setup_test_storage().await;
403        assert_eq!(1234, storage.clock(1234).await.unwrap());
404    }
405
406    #[tokio::test]
407    async fn clock_ge() {
408        let storage = setup_test_storage().await;
409        let c1 = storage.clock(1234).await.unwrap();
410        let c2 = storage.clock(0).await.unwrap();
411        assert_eq!(true, c2 > c1);
412    }
413}
tribbler/storage.rs

tribbler/
storage.rs
