Skip to main content

actix_cloud/memorydb/
interface.rs

1use std::time::Duration;
2
3use async_trait::async_trait;
4
5use crate::Result;
6
7/// Object-safe async key-value database trait.
8///
9/// All values are strings and all TTLs are millisecond-precision [`Duration`]s.
10/// This is the abstraction behind session storage and CSRF token storage; you can
11/// implement it to plug in your own backend (see [`DefaultBackend`](super::default::DefaultBackend)
12/// and [`RedisBackend`](super::redis::RedisBackend) for reference implementations).
13///
14/// The trait is implemented for `Arc<dyn MemoryDB>`, so backends are shared as
15/// `Arc<dyn MemoryDB>` throughout the framework.
16#[async_trait]
17pub trait MemoryDB: Send + Sync {
18    /// Set `key` to `value` without expiration.
19    async fn set(&self, key: &str, value: &str) -> Result<()>;
20
21    /// Get the value of `key`, or `None` if it does not exist or has expired.
22    async fn get(&self, key: &str) -> Result<Option<String>>;
23
24    /// Get the value of `key` and delete it atomically.
25    async fn get_del(&self, key: &str) -> Result<Option<String>>;
26
27    /// Get the value of `key` and reset its TTL to `ttl`.
28    async fn get_ex(&self, key: &str, ttl: &Duration) -> Result<Option<String>>;
29
30    /// Set `key` to `value` with an expiration of `ttl`.
31    async fn set_ex(&self, key: &str, value: &str, ttl: &Duration) -> Result<()>;
32
33    /// Delete `key`. Returns whether the key existed.
34    async fn del(&self, key: &str) -> Result<bool>;
35
36    /// Set the TTL of `key` to `ttl`. Returns whether the key exists.
37    /// A zero `ttl` deletes the key.
38    async fn expire(&self, key: &str, ttl: &Duration) -> Result<bool>;
39
40    /// Remove all keys from the database.
41    ///
42    /// # Warning
43    /// On shared backends such as Redis this flushes the whole database (`FLUSHDB`).
44    async fn flush(&self) -> Result<()>;
45
46    /// Get all keys matching the glob `key` pattern (e.g. `prefix*`).
47    async fn keys(&self, key: &str) -> Result<Vec<String>>;
48
49    /// Delete multiple keys. Returns the number of keys deleted.
50    async fn dels(&self, keys: &[String]) -> Result<u64>;
51
52    /// Get the remaining TTL of `key`, or `None` if the key does not exist or has no TTL.
53    async fn ttl(&self, key: &str) -> Result<Option<Duration>>;
54}
55
56#[cfg(test)]
57mod tests {
58    use std::time::Duration;
59    use tokio::time::sleep;
60
61    use super::*;
62    use crate::memorydb::default::DefaultBackend;
63
64    #[cfg(feature = "redis")]
65    async fn setup_redis() -> impl MemoryDB {
66        crate::memorydb::redis::RedisBackend::new("redis://127.0.0.1:6379/0")
67            .await
68            .unwrap()
69    }
70
71    fn setup_default() -> impl MemoryDB {
72        DefaultBackend::new(None)
73    }
74
75    #[tokio::test]
76    async fn test_normal() {
77        test_normal_fn("default", setup_default()).await;
78        #[cfg(feature = "redis")]
79        test_normal_fn("redis", setup_redis().await).await;
80    }
81
82    async fn test_normal_fn(name: &str, r: impl MemoryDB) {
83        let key = "_actix_cloud_key1";
84        let value1 = "value1";
85        let value2 = "value2";
86
87        println!("Backend: {}", name);
88
89        let _ = r.del(key).await;
90
91        assert_eq!(r.get(key).await.unwrap(), None);
92
93        r.set(key, value1).await.unwrap();
94        assert_eq!(r.get(key).await.unwrap().unwrap(), value1);
95        r.set(key, value2).await.unwrap();
96        assert_eq!(r.get(key).await.unwrap().unwrap(), value2);
97
98        assert_eq!(r.del(key).await.unwrap(), true);
99        assert_eq!(r.del(key).await.unwrap(), false);
100        assert_eq!(r.get(key).await.unwrap(), None);
101    }
102
103    #[tokio::test]
104    async fn test_ex() {
105        test_ex_fn("default", setup_default()).await;
106        #[cfg(feature = "redis")]
107        test_ex_fn("redis", setup_redis().await).await;
108    }
109
110    async fn test_ex_fn(name: &str, r: impl MemoryDB) {
111        let key = "_actix_cloud_key2";
112        let value = "value";
113
114        println!("Backend: {}", name);
115
116        let _ = r.del(key).await;
117
118        r.set(key, value).await.unwrap();
119        assert_eq!(r.ttl(key).await.unwrap(), None);
120        assert_eq!(r.get_del(key).await.unwrap().unwrap(), value);
121        assert_eq!(r.get(key).await.unwrap(), None);
122
123        r.set_ex(key, value, &Duration::from_secs(2)).await.unwrap();
124        assert_eq!(r.get(key).await.unwrap().unwrap(), value);
125        assert!(r.ttl(key).await.unwrap().unwrap().as_millis() > 1500);
126        sleep(Duration::from_secs(1)).await;
127        assert!(r.ttl(key).await.unwrap().unwrap().as_millis() > 500);
128        assert_eq!(
129            r.get_ex(key, &Duration::from_secs(2))
130                .await
131                .unwrap()
132                .unwrap(),
133            value
134        );
135        assert!(r.ttl(key).await.unwrap().unwrap().as_millis() > 1500);
136        sleep(Duration::from_secs(1)).await;
137        assert!(r.ttl(key).await.unwrap().unwrap().as_millis() > 500);
138        assert_eq!(r.get(key).await.unwrap().unwrap(), value);
139        sleep(Duration::from_secs(2)).await;
140        assert_eq!(r.ttl(key).await.unwrap(), None);
141        assert_eq!(r.get(key).await.unwrap(), None);
142    }
143
144    #[tokio::test]
145    async fn test_expire() {
146        test_expire_fn("default", setup_default()).await;
147        #[cfg(feature = "redis")]
148        test_expire_fn("redis", setup_redis().await).await;
149    }
150
151    async fn test_expire_fn(name: &str, r: impl MemoryDB) {
152        let key = "_actix_cloud_key3";
153        let value = "value";
154
155        println!("Backend: {}", name);
156
157        let _ = r.del(key).await;
158
159        r.set(key, value).await.unwrap();
160        assert_eq!(r.get(key).await.unwrap().unwrap(), value);
161        assert_eq!(r.expire(key, &Duration::from_secs(1)).await.unwrap(), true);
162        sleep(Duration::from_secs(2)).await;
163        assert_eq!(r.get(key).await.unwrap(), None);
164        assert_eq!(r.expire(key, &Duration::from_secs(1)).await.unwrap(), false);
165
166        r.set_ex(key, value, &Duration::from_secs(1)).await.unwrap();
167        assert_eq!(r.expire(key, &Duration::from_secs(3)).await.unwrap(), true);
168        sleep(Duration::from_secs(2)).await;
169        assert_eq!(r.get(key).await.unwrap().unwrap(), value);
170        assert_eq!(r.expire(key, &Duration::default()).await.unwrap(), true);
171        assert_eq!(r.get(key).await.unwrap(), None);
172        assert_eq!(r.expire(key, &Duration::default()).await.unwrap(), false);
173    }
174
175    #[tokio::test]
176    async fn test_batch() {
177        test_batch_fn("default", setup_default()).await;
178        #[cfg(feature = "redis")]
179        test_batch_fn("redis", setup_redis().await).await;
180    }
181
182    async fn test_batch_fn(name: &str, r: impl MemoryDB) {
183        let key1 = "_actix_cloud_bkey1";
184        let key2 = "_actix_cloud_bkey2";
185        let value = "value";
186
187        println!("Backend: {}", name);
188
189        let _ = r.del(key1).await;
190        let _ = r.del(key2).await;
191
192        r.set(key1, value).await.unwrap();
193        r.set(key2, value).await.unwrap();
194        let ret = r.keys("_actix_cl?ud_bkey*").await.unwrap();
195        assert_eq!(ret.len(), 2);
196        assert_eq!(
197            (ret[0] == key1 && ret[1] == key2) || (ret[1] == key1 && ret[0] == key2),
198            true
199        );
200        assert_eq!(
201            r.dels(&vec![key1.to_owned(), key2.to_owned()])
202                .await
203                .unwrap(),
204            2
205        );
206        assert_eq!(r.keys("_actix_cl?ud_bkey*").await.unwrap().len(), 0);
207    }
208
209    #[tokio::test]
210    async fn test_default_capacity() {
211        let key1 = "_actix_cloud_ckey1";
212        let key2 = "_actix_cloud_ckey2";
213        let key3 = "_actix_cloud_ckey3";
214        let value = "value";
215
216        let r = DefaultBackend::new(Some(2));
217        let _ = r.del(key1).await;
218        let _ = r.del(key2).await;
219        let _ = r.del(key3).await;
220
221        assert_eq!(r.set(key1, value).await.is_ok(), true);
222        assert_eq!(r.set(key2, value).await.is_ok(), true);
223        assert_eq!(r.set(key3, value).await.is_ok(), false);
224        assert_eq!(r.set(key1, value).await.is_ok(), true);
225        assert_eq!(r.set(key2, value).await.is_ok(), true);
226
227        assert_eq!(
228            r.set_ex(key2, value, &Duration::from_secs(1)).await.is_ok(),
229            true
230        );
231        sleep(Duration::from_secs(2)).await;
232        assert_eq!(r.set(key3, value).await.is_ok(), true);
233        assert_eq!(r.get(key1).await.unwrap().unwrap(), value);
234        assert_eq!(r.get(key2).await.unwrap(), None);
235
236        assert_eq!(
237            r.set_ex(key3, value, &Duration::from_secs(3)).await.is_ok(),
238            true
239        );
240        sleep(Duration::from_secs(1)).await;
241        assert_eq!(r.set(key2, value).await.is_ok(), true);
242        assert_eq!(r.get(key1).await.unwrap().unwrap(), value);
243        assert_eq!(r.get(key3).await.unwrap(), None);
244
245        assert_eq!(r.del(key1).await.unwrap(), true);
246        assert_eq!(r.del(key2).await.unwrap(), true);
247        assert_eq!(
248            r.set_ex(key1, value, &Duration::from_secs(3)).await.is_ok(),
249            true
250        );
251        assert_eq!(
252            r.set_ex(key2, value, &Duration::from_secs(2)).await.is_ok(),
253            true
254        );
255        assert_eq!(
256            r.set_ex(key3, value, &Duration::from_secs(2)).await.is_ok(),
257            true
258        );
259        assert_eq!(r.get(key1).await.unwrap().unwrap(), value);
260        assert_eq!(r.get(key2).await.unwrap(), None);
261        assert_eq!(r.get(key3).await.unwrap().unwrap(), value);
262
263        assert_eq!(r.del(key1).await.unwrap(), true);
264        assert_eq!(r.del(key3).await.unwrap(), true);
265    }
266}