loro_internal/
kv_store.rs

1use crate::sync::Mutex;
2use bytes::Bytes;
3pub use loro_kv_store::compress::CompressionType;
4pub use loro_kv_store::MemKvStore;
5use std::sync::Arc;
6use std::{collections::BTreeMap, ops::Bound};
7
8pub trait KvStore: std::fmt::Debug + Send + Sync {
9    fn get(&self, key: &[u8]) -> Option<Bytes>;
10    fn set(&mut self, key: &[u8], value: Bytes);
11    fn compare_and_swap(&mut self, key: &[u8], old: Option<Bytes>, new: Bytes) -> bool;
12    fn remove(&mut self, key: &[u8]) -> Option<Bytes>;
13    fn contains_key(&self, key: &[u8]) -> bool;
14    fn scan(
15        &self,
16        start: Bound<&[u8]>,
17        end: Bound<&[u8]>,
18    ) -> Box<dyn DoubleEndedIterator<Item = (Bytes, Bytes)> + '_>;
19    fn len(&self) -> usize;
20    fn is_empty(&self) -> bool;
21    fn size(&self) -> usize;
22    fn export_all(&mut self) -> Bytes;
23    fn import_all(&mut self, bytes: Bytes) -> Result<(), String>;
24    /// Like [`Self::import_all`], but for callers that already guarantee the
25    /// blob's integrity via an outer checksum, so per-block checksums may be
26    /// skipped. The default falls back to the checked [`Self::import_all`]; only
27    /// override it where the fast path is safe (e.g. SsTable-backed stores).
28    fn import_all_unchecked(&mut self, bytes: Bytes) -> Result<(), String> {
29        self.import_all(bytes)
30    }
31    fn clone_store(&self) -> Arc<Mutex<dyn KvStore>>;
32}
33
34fn get_common_prefix_len_and_strip<'a, T: AsRef<[u8]> + ?Sized>(
35    this: &'a T,
36    last: &T,
37) -> (u8, &'a [u8]) {
38    let mut common_prefix_len = 0;
39    for (i, (a, b)) in this.as_ref().iter().zip(last.as_ref().iter()).enumerate() {
40        if a != b || i == 255 {
41            common_prefix_len = i;
42            break;
43        }
44    }
45
46    let suffix = &this.as_ref()[common_prefix_len..];
47    (common_prefix_len as u8, suffix)
48}
49
50impl KvStore for MemKvStore {
51    fn get(&self, key: &[u8]) -> Option<Bytes> {
52        self.get(key)
53    }
54
55    fn set(&mut self, key: &[u8], value: Bytes) {
56        self.set(key, value)
57    }
58
59    fn compare_and_swap(&mut self, key: &[u8], old: Option<Bytes>, new: Bytes) -> bool {
60        self.compare_and_swap(key, old, new)
61    }
62
63    fn remove(&mut self, key: &[u8]) -> Option<Bytes> {
64        let ans = self.get(key);
65        self.remove(key);
66        ans
67    }
68
69    fn contains_key(&self, key: &[u8]) -> bool {
70        self.contains_key(key)
71    }
72
73    fn scan(
74        &self,
75        start: Bound<&[u8]>,
76        end: Bound<&[u8]>,
77    ) -> Box<dyn DoubleEndedIterator<Item = (Bytes, Bytes)> + '_> {
78        self.scan(start, end)
79    }
80
81    fn len(&self) -> usize {
82        self.len()
83    }
84
85    fn is_empty(&self) -> bool {
86        self.is_empty()
87    }
88
89    fn size(&self) -> usize {
90        self.size()
91    }
92
93    fn export_all(&mut self) -> Bytes {
94        self.export_all()
95    }
96
97    fn import_all(&mut self, bytes: Bytes) -> Result<(), String> {
98        self.import_all(bytes)
99    }
100
101    fn import_all_unchecked(&mut self, bytes: Bytes) -> Result<(), String> {
102        self.import_all_unchecked(bytes)
103    }
104
105    fn clone_store(&self) -> Arc<Mutex<dyn KvStore>> {
106        Arc::new(Mutex::new(self.clone()))
107    }
108}
109
110mod default_binary_format {
111    //! Default binary format for the key-value store.
112    //!
113    //! It will compress the prefix of the keys that are common with the previous key.
114
115    use bytes::Bytes;
116
117    use super::get_common_prefix_len_and_strip;
118
119    pub fn export_by_scan(store: &impl super::KvStore) -> bytes::Bytes {
120        let mut buf = Vec::new();
121        let mut last_key: Option<Bytes> = None;
122        for (k, v) in store.scan(std::ops::Bound::Unbounded, std::ops::Bound::Unbounded) {
123            {
124                // Write the key
125                match last_key.take() {
126                    None => {
127                        leb128::write::unsigned(&mut buf, k.len() as u64).unwrap();
128                        buf.extend_from_slice(&k);
129                    }
130                    Some(last) => {
131                        let (common, suffix) = get_common_prefix_len_and_strip(&k, &last);
132                        buf.push(common);
133                        leb128::write::unsigned(&mut buf, suffix.len() as u64).unwrap();
134                        buf.extend_from_slice(suffix);
135                    }
136                };
137
138                last_key = Some(k);
139            }
140
141            // Write the value
142            leb128::write::unsigned(&mut buf, v.len() as u64).unwrap();
143            buf.extend_from_slice(&v);
144        }
145
146        buf.into()
147    }
148
149    pub fn import(store: &mut impl super::KvStore, bytes: bytes::Bytes) -> Result<(), String> {
150        let mut bytes: &[u8] = &bytes;
151        let mut last_key = Vec::new();
152
153        while !bytes.is_empty() {
154            // Read the key
155            let mut key = Vec::new();
156            if last_key.is_empty() {
157                let key_len = leb128::read::unsigned(&mut bytes).map_err(|e| e.to_string())?;
158                let key_len = key_len as usize;
159                key.extend_from_slice(&bytes[..key_len]);
160                bytes = &bytes[key_len..];
161            } else {
162                let common_prefix_len = bytes[0] as usize;
163                bytes = &bytes[1..];
164                key.extend_from_slice(&last_key[..common_prefix_len]);
165                let suffix_len = leb128::read::unsigned(&mut bytes).map_err(|e| e.to_string())?;
166                let suffix_len = suffix_len as usize;
167                key.extend_from_slice(&bytes[..suffix_len]);
168                bytes = &bytes[suffix_len..];
169            }
170
171            // Read the value
172            let value_len = leb128::read::unsigned(&mut bytes).map_err(|e| e.to_string())?;
173            let value_len = value_len as usize;
174            let value = Bytes::copy_from_slice(&bytes[..value_len]);
175            bytes = &bytes[value_len..];
176
177            // Store the key-value pair
178            store.set(&key, value);
179
180            last_key = key;
181        }
182
183        Ok(())
184    }
185}
186
187impl KvStore for BTreeMap<Bytes, Bytes> {
188    fn get(&self, key: &[u8]) -> Option<Bytes> {
189        self.get(key).cloned()
190    }
191
192    fn set(&mut self, key: &[u8], value: Bytes) {
193        self.insert(Bytes::copy_from_slice(key), value);
194    }
195
196    fn compare_and_swap(&mut self, key: &[u8], old: Option<Bytes>, new: Bytes) -> bool {
197        let key = Bytes::copy_from_slice(key);
198        match self.get_mut(&key) {
199            Some(v) => {
200                if old.as_ref() == Some(v) {
201                    self.insert(key, new);
202                    true
203                } else {
204                    false
205                }
206            }
207            None => {
208                if old.is_none() {
209                    self.insert(key, new);
210                    true
211                } else {
212                    false
213                }
214            }
215        }
216    }
217
218    fn remove(&mut self, key: &[u8]) -> Option<Bytes> {
219        self.remove(key)
220    }
221
222    fn contains_key(&self, key: &[u8]) -> bool {
223        self.contains_key(key)
224    }
225
226    fn scan(
227        &self,
228        start: Bound<&[u8]>,
229        end: Bound<&[u8]>,
230    ) -> Box<dyn DoubleEndedIterator<Item = (Bytes, Bytes)> + '_> {
231        Box::new(
232            self.range::<[u8], _>((start, end))
233                .map(|(k, v)| (k.clone(), v.clone())),
234        )
235    }
236
237    fn len(&self) -> usize {
238        self.len()
239    }
240
241    fn is_empty(&self) -> bool {
242        self.is_empty()
243    }
244
245    fn size(&self) -> usize {
246        self.iter().fold(0, |acc, (k, v)| acc + k.len() + v.len())
247    }
248
249    fn export_all(&mut self) -> Bytes {
250        default_binary_format::export_by_scan(self)
251    }
252
253    fn import_all(&mut self, bytes: Bytes) -> Result<(), String> {
254        default_binary_format::import(self, bytes)
255    }
256
257    fn clone_store(&self) -> Arc<Mutex<dyn KvStore>> {
258        Arc::new(Mutex::new(self.clone()))
259    }
260}
loro_internal/kv_store.rs

loro_internal/
kv_store.rs
