Skip to main content

kv_storage/
lib.rs

1//http://sled.rs/
2
3#![allow(dead_code)]
4mod iface;
5mod sled_config;
6mod sled_storage;
7mod test;
8mod test_kv;
9mod test_list;
10mod test_map;
11
12use async_trait::async_trait;
13use core::fmt;
14use iface::*;
15pub use iface::{List, Map};
16use serde::Serialize;
17use serde::de::DeserializeOwned;
18pub use sled_config::Config;
19use sled_storage::{SledStorageDB, SledStorageList, SledStorageMap};
20
21type TimestampMillis = i64;
22type Result<T> = anyhow::Result<T>;
23#[inline]
24fn timestamp_millis() -> TimestampMillis {
25    use std::time::{SystemTime, UNIX_EPOCH};
26    SystemTime::now()
27        .duration_since(UNIX_EPOCH)
28        .map(|dur| dur.as_millis() as TimestampMillis)
29        .unwrap_or_else(|_| {
30            let now = chrono::Local::now();
31            now.timestamp_millis() as TimestampMillis
32        })
33}
34
35#[allow(unused)]
36const SEPARATOR: &[u8] = b"@";
37#[allow(unused)]
38const KEY_PREFIX: &[u8] = b"__sled@";
39#[allow(unused)]
40const KEY_PREFIX_LEN: &[u8] = b"__sled_len@";
41#[allow(unused)]
42const MAP_NAME_PREFIX: &[u8] = b"__sled_map@";
43#[allow(unused)]
44const LIST_NAME_PREFIX: &[u8] = b"__sled_list@";
45
46/// Type alias for storage keys
47type Key = Vec<u8>;
48/// Result type for iteration items (key-value pair)
49type IterItem<V> = Result<(Key, V)>;
50
51pub async fn init_db(cfg: &Config) -> Result<StorageDB> {
52    let db = SledStorageDB::new(cfg.clone()).await?;
53    let db = StorageDB::Sled(db);
54    Ok(db)
55}
56
57#[derive(Clone)]
58pub enum StorageDB {
59    Sled(SledStorageDB),
60}
61
62impl StorageDB {
63    /// Accesses a named map
64    #[inline]
65    pub async fn map<V: AsRef<[u8]> + Sync + Send>(
66        &self,
67        name: V,
68        expire: Option<TimestampMillis>,
69    ) -> Result<StorageMap> {
70        Ok(match self {
71            StorageDB::Sled(db) => StorageMap::Sled(db.map(name, expire).await?),
72        })
73    }
74
75    /// Removes a named map
76    #[inline]
77    pub async fn map_remove<K>(&self, name: K) -> Result<()>
78    where
79        K: AsRef<[u8]> + Sync + Send,
80    {
81        match self {
82            StorageDB::Sled(db) => db.map_remove(name).await,
83        }
84    }
85
86    /// Checks if map exists
87    #[inline]
88    pub async fn map_contains_key<K: AsRef<[u8]> + Sync + Send>(&self, key: K) -> Result<bool> {
89        match self {
90            StorageDB::Sled(db) => db.map_contains_key(key).await,
91        }
92    }
93
94    /// Accesses a named list
95    #[inline]
96    pub async fn list<V: AsRef<[u8]> + Sync + Send>(
97        &self,
98        name: V,
99        expire: Option<TimestampMillis>,
100    ) -> Result<StorageList> {
101        Ok(match self {
102            StorageDB::Sled(db) => StorageList::Sled(db.list(name, expire).await?),
103        })
104    }
105
106    /// Removes a named list
107    #[inline]
108    pub async fn list_remove<K>(&self, name: K) -> Result<()>
109    where
110        K: AsRef<[u8]> + Sync + Send,
111    {
112        match self {
113            StorageDB::Sled(db) => db.list_remove(name).await,
114        }
115    }
116
117    /// Checks if list exists
118    #[inline]
119    pub async fn list_contains_key<K: AsRef<[u8]> + Sync + Send>(&self, key: K) -> Result<bool> {
120        match self {
121            StorageDB::Sled(db) => db.list_contains_key(key).await,
122        }
123    }
124
125    /// Inserts a key-value pair
126    #[inline]
127    pub async fn insert<K, V>(&self, key: K, val: &V) -> Result<()>
128    where
129        K: AsRef<[u8]> + Sync + Send,
130        V: Serialize + Sync + Send,
131    {
132        match self {
133            StorageDB::Sled(db) => db.insert(key, val).await,
134        }
135    }
136
137    /// Retrieves a value by key
138    #[inline]
139    pub async fn get<K, V>(&self, key: K) -> Result<Option<V>>
140    where
141        K: AsRef<[u8]> + Sync + Send,
142        V: DeserializeOwned + Sync + Send,
143    {
144        match self {
145            StorageDB::Sled(db) => db.get(key).await,
146        }
147    }
148
149    /// Removes a key-value pair
150    #[inline]
151    pub async fn remove<K>(&self, key: K) -> Result<()>
152    where
153        K: AsRef<[u8]> + Sync + Send,
154    {
155        match self {
156            StorageDB::Sled(db) => db.remove(key).await,
157        }
158    }
159
160    /// Batch insert of key-value pairs
161    #[inline]
162    pub async fn batch_insert<V>(&self, key_vals: Vec<(Key, V)>) -> Result<()>
163    where
164        V: serde::ser::Serialize + Sync + Send,
165    {
166        match self {
167            StorageDB::Sled(db) => db.batch_insert(key_vals).await,
168        }
169    }
170
171    /// Batch removal of keys
172    #[inline]
173    pub async fn batch_remove(&self, keys: Vec<Key>) -> Result<()> {
174        match self {
175            StorageDB::Sled(db) => db.batch_remove(keys).await,
176        }
177    }
178
179    /// Increments a counter
180    #[inline]
181    pub async fn counter_incr<K>(&self, key: K, increment: isize) -> Result<()>
182    where
183        K: AsRef<[u8]> + Sync + Send,
184    {
185        match self {
186            StorageDB::Sled(db) => db.counter_incr(key, increment).await,
187        }
188    }
189
190    /// Decrements a counter
191    #[inline]
192    pub async fn counter_decr<K>(&self, key: K, decrement: isize) -> Result<()>
193    where
194        K: AsRef<[u8]> + Sync + Send,
195    {
196        match self {
197            StorageDB::Sled(db) => db.counter_decr(key, decrement).await,
198        }
199    }
200
201    /// Gets counter value
202    #[inline]
203    pub async fn counter_get<K>(&self, key: K) -> Result<Option<isize>>
204    where
205        K: AsRef<[u8]> + Sync + Send,
206    {
207        match self {
208            StorageDB::Sled(db) => db.counter_get(key).await,
209        }
210    }
211
212    /// Sets counter value
213    #[inline]
214    pub async fn counter_set<K>(&self, key: K, val: isize) -> Result<()>
215    where
216        K: AsRef<[u8]> + Sync + Send,
217    {
218        match self {
219            StorageDB::Sled(db) => db.counter_set(key, val).await,
220        }
221    }
222
223    /// Gets number of items (requires "len" feature)
224    #[inline]
225    #[cfg(feature = "len")]
226    pub async fn len(&self) -> Result<usize> {
227        match self {
228            StorageDB::Sled(db) => db.len().await,
229        }
230    }
231
232    /// Gets total storage size in bytes
233    #[inline]
234    pub async fn db_size(&self) -> Result<usize> {
235        match self {
236            StorageDB::Sled(db) => db.db_size().await,
237        }
238    }
239
240    /// Checks if key exists
241    #[inline]
242    pub async fn contains_key<K: AsRef<[u8]> + Sync + Send>(&self, key: K) -> Result<bool> {
243        match self {
244            StorageDB::Sled(db) => db.contains_key(key).await,
245        }
246    }
247
248    /// Sets expiration timestamp (requires "ttl" feature)
249    #[inline]
250    #[cfg(feature = "ttl")]
251    pub async fn expire_at<K>(&self, key: K, at: TimestampMillis) -> Result<bool>
252    where
253        K: AsRef<[u8]> + Sync + Send,
254    {
255        match self {
256            StorageDB::Sled(db) => db.expire_at(key, at).await,
257        }
258    }
259
260    /// Sets expiration duration (requires "ttl" feature)
261    #[inline]
262    #[cfg(feature = "ttl")]
263    pub async fn expire<K>(&self, key: K, dur: TimestampMillis) -> Result<bool>
264    where
265        K: AsRef<[u8]> + Sync + Send,
266    {
267        match self {
268            StorageDB::Sled(db) => db.expire(key, dur).await,
269        }
270    }
271
272    /// Gets time-to-live (requires "ttl" feature)
273    #[inline]
274    #[cfg(feature = "ttl")]
275    pub async fn ttl<K>(&self, key: K) -> Result<Option<TimestampMillis>>
276    where
277        K: AsRef<[u8]> + Sync + Send,
278    {
279        match self {
280            StorageDB::Sled(db) => db.ttl(key).await,
281        }
282    }
283
284    /// Iterates over maps
285    #[inline]
286    pub async fn map_iter<'a>(
287        &'a mut self,
288    ) -> Result<Box<dyn AsyncIterator<Item = Result<StorageMap>> + Send + 'a>> {
289        match self {
290            StorageDB::Sled(db) => db.map_iter().await,
291        }
292    }
293
294    /// Iterates over lists
295    #[inline]
296    pub async fn list_iter<'a>(
297        &'a mut self,
298    ) -> Result<Box<dyn AsyncIterator<Item = Result<StorageList>> + Send + 'a>> {
299        match self {
300            StorageDB::Sled(db) => db.list_iter().await,
301        }
302    }
303
304    /// Scans keys matching pattern
305    #[inline]
306    pub async fn scan<'a, P>(
307        &'a mut self,
308        pattern: P,
309    ) -> Result<Box<dyn AsyncIterator<Item = Result<Key>> + Send + 'a>>
310    where
311        P: AsRef<[u8]> + Send + Sync,
312    {
313        match self {
314            StorageDB::Sled(db) => db.scan(pattern).await,
315        }
316    }
317
318    /// Gets storage information
319    #[inline]
320    pub async fn info(&self) -> Result<serde_json::Value> {
321        match self {
322            StorageDB::Sled(db) => db.info().await,
323        }
324    }
325}
326
327#[derive(Clone)]
328pub enum StorageMap {
329    /// Sled map implementation
330    Sled(SledStorageMap),
331}
332
333#[async_trait]
334impl Map for StorageMap {
335    fn name(&self) -> &[u8] {
336        match self {
337            StorageMap::Sled(m) => m.name(),
338        }
339    }
340
341    async fn insert<K, V>(&self, key: K, val: &V) -> Result<()>
342    where
343        K: AsRef<[u8]> + Sync + Send,
344        V: Serialize + Sync + Send + ?Sized,
345    {
346        match self {
347            StorageMap::Sled(m) => m.insert(key, val).await,
348        }
349    }
350
351    async fn get<K, V>(&self, key: K) -> Result<Option<V>>
352    where
353        K: AsRef<[u8]> + Sync + Send,
354        V: DeserializeOwned + Sync + Send,
355    {
356        match self {
357            StorageMap::Sled(m) => m.get(key).await,
358        }
359    }
360
361    async fn remove<K>(&self, key: K) -> Result<()>
362    where
363        K: AsRef<[u8]> + Sync + Send,
364    {
365        match self {
366            StorageMap::Sled(m) => m.remove(key).await,
367        }
368    }
369
370    async fn contains_key<K: AsRef<[u8]> + Sync + Send>(&self, key: K) -> Result<bool> {
371        match self {
372            StorageMap::Sled(m) => m.contains_key(key).await,
373        }
374    }
375
376    #[cfg(feature = "map_len")]
377    async fn len(&self) -> Result<usize> {
378        match self {
379            StorageMap::Sled(m) => m.len().await,
380        }
381    }
382
383    async fn is_empty(&self) -> Result<bool> {
384        match self {
385            StorageMap::Sled(m) => m.is_empty().await,
386        }
387    }
388
389    async fn clear(&self) -> Result<()> {
390        match self {
391            StorageMap::Sled(m) => m.clear().await,
392        }
393    }
394
395    async fn remove_and_fetch<K, V>(&self, key: K) -> Result<Option<V>>
396    where
397        K: AsRef<[u8]> + Sync + Send,
398        V: DeserializeOwned + Sync + Send,
399    {
400        match self {
401            StorageMap::Sled(m) => m.remove_and_fetch(key).await,
402        }
403    }
404
405    async fn remove_with_prefix<K>(&self, prefix: K) -> Result<()>
406    where
407        K: AsRef<[u8]> + Sync + Send,
408    {
409        match self {
410            StorageMap::Sled(m) => m.remove_with_prefix(prefix).await,
411        }
412    }
413
414    async fn batch_insert<V>(&self, key_vals: Vec<(Key, V)>) -> Result<()>
415    where
416        V: Serialize + Sync + Send,
417    {
418        match self {
419            StorageMap::Sled(m) => m.batch_insert(key_vals).await,
420        }
421    }
422
423    async fn batch_remove(&self, keys: Vec<Key>) -> Result<()> {
424        match self {
425            StorageMap::Sled(m) => m.batch_remove(keys).await,
426        }
427    }
428
429    async fn iter<'a, V>(
430        &'a mut self,
431    ) -> Result<Box<dyn AsyncIterator<Item = IterItem<V>> + Send + 'a>>
432    where
433        V: DeserializeOwned + Sync + Send + 'a + 'static,
434    {
435        match self {
436            StorageMap::Sled(m) => m.iter().await,
437        }
438    }
439
440    async fn key_iter<'a>(
441        &'a mut self,
442    ) -> Result<Box<dyn AsyncIterator<Item = Result<Key>> + Send + 'a>> {
443        match self {
444            StorageMap::Sled(m) => m.key_iter().await,
445        }
446    }
447
448    async fn prefix_iter<'a, P, V>(
449        &'a mut self,
450        prefix: P,
451    ) -> Result<Box<dyn AsyncIterator<Item = IterItem<V>> + Send + 'a>>
452    where
453        P: AsRef<[u8]> + Send + Sync,
454        V: DeserializeOwned + Sync + Send + 'a + 'static,
455    {
456        match self {
457            StorageMap::Sled(m) => m.prefix_iter(prefix).await,
458        }
459    }
460
461    #[cfg(feature = "ttl")]
462    async fn expire_at(&self, at: TimestampMillis) -> Result<bool> {
463        match self {
464            StorageMap::Sled(m) => m.expire_at(at).await,
465        }
466    }
467
468    #[cfg(feature = "ttl")]
469    async fn expire(&self, dur: TimestampMillis) -> Result<bool> {
470        match self {
471            StorageMap::Sled(m) => m.expire(dur).await,
472        }
473    }
474
475    #[cfg(feature = "ttl")]
476    async fn ttl(&self) -> Result<Option<TimestampMillis>> {
477        match self {
478            StorageMap::Sled(m) => m.ttl().await,
479        }
480    }
481}
482
483#[derive(Clone)]
484pub enum StorageList {
485    /// Sled list implementation
486    Sled(SledStorageList),
487}
488
489impl fmt::Debug for StorageList {
490    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
491        let name = match self {
492            StorageList::Sled(list) => list.name(),
493        };
494        f.debug_tuple(&format!("StorageList({:?})", String::from_utf8_lossy(name)))
495            .finish()
496    }
497}
498
499#[async_trait]
500impl List for StorageList {
501    fn name(&self) -> &[u8] {
502        match self {
503            StorageList::Sled(m) => m.name(),
504        }
505    }
506
507    async fn push<V>(&self, val: &V) -> Result<()>
508    where
509        V: Serialize + Sync + Send,
510    {
511        match self {
512            StorageList::Sled(list) => list.push(val).await,
513        }
514    }
515
516    async fn pushs<V>(&self, vals: Vec<V>) -> Result<()>
517    where
518        V: serde::ser::Serialize + Sync + Send,
519    {
520        match self {
521            StorageList::Sled(list) => list.pushs(vals).await,
522        }
523    }
524
525    async fn push_limit<V>(
526        &self,
527        val: &V,
528        limit: usize,
529        pop_front_if_limited: bool,
530    ) -> Result<Option<V>>
531    where
532        V: Serialize + Sync + Send,
533        V: DeserializeOwned,
534    {
535        match self {
536            StorageList::Sled(list) => list.push_limit(val, limit, pop_front_if_limited).await,
537        }
538    }
539
540    async fn pop<V>(&self) -> Result<Option<V>>
541    where
542        V: DeserializeOwned + Sync + Send,
543    {
544        match self {
545            StorageList::Sled(list) => list.pop().await,
546        }
547    }
548
549    async fn all<V>(&self) -> Result<Vec<V>>
550    where
551        V: DeserializeOwned + Sync + Send,
552    {
553        match self {
554            StorageList::Sled(list) => list.all().await,
555        }
556    }
557
558    async fn get_index<V>(&self, idx: usize) -> Result<Option<V>>
559    where
560        V: DeserializeOwned + Sync + Send,
561    {
562        match self {
563            StorageList::Sled(list) => list.get_index(idx).await,
564        }
565    }
566
567    async fn len(&self) -> Result<usize> {
568        match self {
569            StorageList::Sled(list) => list.len().await,
570        }
571    }
572
573    async fn is_empty(&self) -> Result<bool> {
574        match self {
575            StorageList::Sled(list) => list.is_empty().await,
576        }
577    }
578
579    async fn clear(&self) -> Result<()> {
580        match self {
581            StorageList::Sled(list) => list.clear().await,
582        }
583    }
584
585    async fn iter<'a, V>(
586        &'a mut self,
587    ) -> Result<Box<dyn AsyncIterator<Item = Result<V>> + Send + 'a>>
588    where
589        V: DeserializeOwned + Sync + Send + 'a + 'static,
590    {
591        match self {
592            StorageList::Sled(list) => list.iter().await,
593        }
594    }
595
596    #[cfg(feature = "ttl")]
597    async fn expire_at(&self, at: TimestampMillis) -> Result<bool> {
598        match self {
599            StorageList::Sled(l) => l.expire_at(at).await,
600        }
601    }
602
603    #[cfg(feature = "ttl")]
604    async fn expire(&self, dur: TimestampMillis) -> Result<bool> {
605        match self {
606            StorageList::Sled(l) => l.expire(dur).await,
607        }
608    }
609
610    #[cfg(feature = "ttl")]
611    async fn ttl(&self) -> Result<Option<TimestampMillis>> {
612        match self {
613            StorageList::Sled(l) => l.ttl().await,
614        }
615    }
616}