typed_store/rocks/
memstore.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use std::collections::{BTreeMap, HashMap};
use std::sync::{Arc, RwLock};

type InMemoryStoreInternal = Arc<RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>>;

#[derive(Clone, Debug)]
pub struct InMemoryDB {
    data: InMemoryStoreInternal,
}

#[derive(Clone, Debug)]
enum InMemoryChange {
    Delete((String, Vec<u8>)),
    Put((String, Vec<u8>, Vec<u8>)),
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryBatch {
    data: Vec<InMemoryChange>,
}

impl InMemoryBatch {
    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf_name: &str, key: K) {
        self.data.push(InMemoryChange::Delete((
            cf_name.to_string(),
            key.as_ref().to_vec(),
        )));
    }

    pub fn put_cf<K, V>(&mut self, cf_name: &str, key: K, value: V)
    where
        K: AsRef<[u8]>,
        V: AsRef<[u8]>,
    {
        self.data.push(InMemoryChange::Put((
            cf_name.to_string(),
            key.as_ref().to_vec(),
            value.as_ref().to_vec(),
        )));
    }
}

impl InMemoryDB {
    pub fn get<K: AsRef<[u8]>>(&self, cf_name: &str, key: K) -> Option<Vec<u8>> {
        let data = self.data.read().expect("can't read data");
        match data.get(cf_name) {
            Some(cf) => cf.get(key.as_ref()).cloned(),
            None => None,
        }
    }

    pub fn multi_get<I, K>(&self, cf_name: &str, keys: I) -> Vec<Option<Vec<u8>>>
    where
        I: IntoIterator<Item = K>,
        K: AsRef<[u8]>,
    {
        let data = self.data.read().expect("can't read data");
        match data.get(cf_name) {
            Some(cf) => keys
                .into_iter()
                .map(|k| cf.get(k.as_ref()).cloned())
                .collect(),
            None => vec![],
        }
    }

    pub fn delete(&self, cf_name: &str, key: &[u8]) {
        let mut data = self.data.write().expect("can't write data");
        data.entry(cf_name.to_string()).or_default().remove(key);
    }

    pub fn put(&self, cf_name: &str, key: Vec<u8>, value: Vec<u8>) {
        let mut data = self.data.write().expect("can't write data");
        data.entry(cf_name.to_string())
            .or_default()
            .insert(key, value);
    }

    pub fn write(&self, batch: InMemoryBatch) {
        for change in batch.data {
            match change {
                InMemoryChange::Delete((cf_name, key)) => self.delete(&cf_name, &key),
                InMemoryChange::Put((cf_name, key, value)) => self.put(&cf_name, key, value),
            }
        }
    }

    pub fn drop_cf(&self, name: &str) {
        self.data.write().expect("can't write data").remove(name);
    }
}