typed_store/
memstore.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use bincode::Options;
5use serde::de::DeserializeOwned;
6use std::collections::{BTreeMap, Bound, HashMap};
7use std::sync::{Arc, RwLock};
8use typed_store_error::TypedStoreError;
9
10type InMemoryStoreInternal = Arc<RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>>;
11
12#[derive(Clone, Debug)]
13pub struct InMemoryDB {
14    data: InMemoryStoreInternal,
15}
16
17#[derive(Clone, Debug)]
18enum InMemoryChange {
19    Delete((String, Vec<u8>)),
20    Put((String, Vec<u8>, Vec<u8>)),
21}
22
23#[derive(Clone, Debug, Default)]
24pub struct InMemoryBatch {
25    data: Vec<InMemoryChange>,
26}
27
28impl InMemoryBatch {
29    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf_name: &str, key: K) {
30        self.data.push(InMemoryChange::Delete((
31            cf_name.to_string(),
32            key.as_ref().to_vec(),
33        )));
34    }
35
36    pub fn put_cf<K, V>(&mut self, cf_name: &str, key: K, value: V)
37    where
38        K: AsRef<[u8]>,
39        V: AsRef<[u8]>,
40    {
41        self.data.push(InMemoryChange::Put((
42            cf_name.to_string(),
43            key.as_ref().to_vec(),
44            value.as_ref().to_vec(),
45        )));
46    }
47}
48
49impl InMemoryDB {
50    pub fn get<K: AsRef<[u8]>>(&self, cf_name: &str, key: K) -> Option<Vec<u8>> {
51        let data = self.data.read().expect("can't read data");
52        match data.get(cf_name) {
53            Some(cf) => cf.get(key.as_ref()).cloned(),
54            None => None,
55        }
56    }
57
58    pub fn multi_get<I, K>(&self, cf_name: &str, keys: I) -> Vec<Option<Vec<u8>>>
59    where
60        I: IntoIterator<Item = K>,
61        K: AsRef<[u8]>,
62    {
63        let data = self.data.read().expect("can't read data");
64        match data.get(cf_name) {
65            Some(cf) => keys
66                .into_iter()
67                .map(|k| cf.get(k.as_ref()).cloned())
68                .collect(),
69            None => vec![],
70        }
71    }
72
73    pub fn delete(&self, cf_name: &str, key: &[u8]) {
74        let mut data = self.data.write().expect("can't write data");
75        data.entry(cf_name.to_string()).or_default().remove(key);
76    }
77
78    pub fn put(&self, cf_name: &str, key: Vec<u8>, value: Vec<u8>) {
79        let mut data = self.data.write().expect("can't write data");
80        data.entry(cf_name.to_string())
81            .or_default()
82            .insert(key, value);
83    }
84
85    pub fn write(&self, batch: InMemoryBatch) {
86        for change in batch.data {
87            match change {
88                InMemoryChange::Delete((cf_name, key)) => self.delete(&cf_name, &key),
89                InMemoryChange::Put((cf_name, key, value)) => self.put(&cf_name, key, value),
90            }
91        }
92    }
93
94    pub fn drop_cf(&self, name: &str) {
95        self.data.write().expect("can't write data").remove(name);
96    }
97
98    pub fn iterator<K, V>(
99        &self,
100        cf_name: &str,
101        lower_bound: Option<Vec<u8>>,
102        upper_bound: Option<Vec<u8>>,
103        reverse: bool,
104    ) -> Box<dyn Iterator<Item = Result<(K, V), TypedStoreError>> + '_>
105    where
106        K: DeserializeOwned,
107        V: DeserializeOwned,
108    {
109        let config = bincode::DefaultOptions::new()
110            .with_big_endian()
111            .with_fixint_encoding();
112        let lower_bound = lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);
113        let upper_bound = upper_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);
114
115        let data = self.data.read().expect("can't read data");
116        let mut section: Vec<_> = data
117            .get(cf_name)
118            .unwrap_or(&BTreeMap::new())
119            .range((lower_bound, upper_bound))
120            .map(|(k, v)| (k.clone(), v.clone()))
121            .collect();
122        if reverse {
123            section.reverse();
124        }
125        Box::new(section.into_iter().map(move |(raw_key, raw_value)| {
126            let key = config.deserialize(&raw_key).unwrap();
127            let value = bcs::from_bytes(&raw_value).unwrap();
128            Ok((key, value))
129        }))
130    }
131}
typed_store/memstore.rs

typed_store/
memstore.rs
