oasis_runtime_sdk/storage/
overlay.rs

1use std::{
2    collections::{btree_map, BTreeMap, HashSet},
3    iter::{Iterator, Peekable},
4};
5
6use oasis_core_runtime::storage::mkvs;
7
8use super::{NestedStore, Prefix, Store};
9
10/// An overlay store which keeps values locally until explicitly committed.
11pub struct OverlayStore<S: Store> {
12    parent: S,
13    overlay: BTreeMap<Vec<u8>, Vec<u8>>,
14    dirty: HashSet<Vec<u8>>,
15}
16
17impl<S: Store> OverlayStore<S> {
18    /// Create a new overlay store.
19    pub fn new(parent: S) -> Self {
20        Self {
21            parent,
22            overlay: BTreeMap::new(),
23            dirty: HashSet::new(),
24        }
25    }
26}
27
28impl<S: Store> NestedStore for OverlayStore<S> {
29    type Inner = S;
30
31    fn commit(mut self) -> Self::Inner {
32        // Insert all items present in the overlay.
33        for (key, value) in self.overlay {
34            self.dirty.remove(&key);
35            self.parent.insert(&key, &value);
36        }
37
38        // Any remaining dirty items must have been removed.
39        for key in self.dirty {
40            self.parent.remove(&key);
41        }
42
43        self.parent
44    }
45
46    fn rollback(self) -> Self::Inner {
47        self.parent
48    }
49
50    fn has_pending_updates(&self) -> bool {
51        !self.dirty.is_empty()
52    }
53
54    fn pending_update_byte_size(&self) -> usize {
55        let updated_size = self
56            .overlay
57            .iter()
58            .map(|(key, value)| key.len() + value.len())
59            .reduce(|acc, bytes| acc + bytes)
60            .unwrap_or_default();
61
62        let removed_size = self
63            .dirty
64            .iter()
65            .filter(|key| !self.overlay.contains_key(key.as_slice()))
66            .map(|key| key.len())
67            .reduce(|acc, bytes| acc + bytes)
68            .unwrap_or_default();
69
70        updated_size.saturating_add(removed_size)
71    }
72}
73
74impl<S: Store> Store for OverlayStore<S> {
75    fn get(&self, key: &[u8]) -> Option<Vec<u8>> {
76        // For dirty values, check the overlay.
77        if self.dirty.contains(key) {
78            return self.overlay.get(key).cloned();
79        }
80
81        // Otherwise fetch from parent store.
82        self.parent.get(key)
83    }
84
85    fn insert(&mut self, key: &[u8], value: &[u8]) {
86        self.overlay.insert(key.to_owned(), value.to_owned());
87        self.dirty.insert(key.to_owned());
88    }
89
90    fn remove(&mut self, key: &[u8]) {
91        // For dirty values, remove from the overlay.
92        if self.dirty.contains(key) {
93            self.overlay.remove(key);
94            return;
95        }
96
97        // Since we don't care about the previous value, we can just record an update.
98        self.dirty.insert(key.to_owned());
99    }
100
101    fn iter(&self) -> Box<dyn mkvs::Iterator + '_> {
102        Box::new(OverlayStoreIterator::new(self))
103    }
104
105    fn prefetch_prefixes(&mut self, prefixes: Vec<Prefix>, limit: u16) {
106        self.parent.prefetch_prefixes(prefixes, limit);
107    }
108}
109
110/// An iterator over the `OverlayStore`.
111pub(crate) struct OverlayStoreIterator<'store, S: Store> {
112    store: &'store OverlayStore<S>,
113
114    parent: Box<dyn mkvs::Iterator + 'store>,
115
116    overlay: Peekable<btree_map::Range<'store, Vec<u8>, Vec<u8>>>,
117    overlay_valid: bool,
118
119    key: Option<Vec<u8>>,
120    value: Option<Vec<u8>>,
121}
122
123impl<'store, S: Store> OverlayStoreIterator<'store, S> {
124    fn new(store: &'store OverlayStore<S>) -> Self {
125        Self {
126            store,
127            parent: store.parent.iter(),
128            overlay: store.overlay.range(vec![]..).peekable(),
129            overlay_valid: true,
130            key: None,
131            value: None,
132        }
133    }
134
135    fn update_iterator_position(&mut self) {
136        // Skip over any dirty entries from the parent iterator.
137        loop {
138            if !self.parent.is_valid()
139                || !self
140                    .store
141                    .dirty
142                    .contains(self.parent.get_key().as_ref().expect("parent.is_valid"))
143            {
144                break;
145            }
146            self.parent.next();
147        }
148
149        let i_key = self.parent.get_key();
150        let o_item = self.overlay.peek();
151        self.overlay_valid = o_item.is_some();
152
153        if self.parent.is_valid()
154            && (!self.overlay_valid
155                || i_key.as_ref().expect("parent.is_valid") < o_item.expect("overlay_valid").0)
156        {
157            // Key of parent iterator is smaller than the key of the overlay iterator.
158            self.key = i_key.clone();
159            self.value = self.parent.get_value().clone();
160        } else if self.overlay_valid {
161            // Key of overlay iterator is smaller than or equal to the key of the parent iterator.
162            let (o_key, o_value) = o_item.expect("overlay_valid");
163            self.key = Some(o_key.to_vec());
164            self.value = Some(o_value.to_vec());
165        } else {
166            // Both iterators are invalid.
167            self.key = None;
168            self.value = None;
169        }
170    }
171
172    fn next(&mut self) {
173        if !self.overlay_valid
174            || (self.parent.is_valid()
175                && self.parent.get_key().as_ref().expect("parent.is_valid")
176                    <= self.overlay.peek().expect("overlay_valid").0)
177        {
178            // Key of parent iterator is smaller or equal than the key of the overlay iterator.
179            self.parent.next();
180        } else {
181            // Key of parent iterator is greater than the key of the overlay iterator.
182            self.overlay.next();
183        }
184
185        self.update_iterator_position();
186    }
187}
188
189impl<S: Store> Iterator for OverlayStoreIterator<'_, S> {
190    type Item = (Vec<u8>, Vec<u8>);
191
192    fn next(&mut self) -> Option<Self::Item> {
193        use mkvs::Iterator;
194
195        if !self.is_valid() {
196            return None;
197        }
198
199        let key = self.key.as_ref().expect("iterator is valid").clone();
200        let value = self.value.as_ref().expect("iterator is valid").clone();
201        OverlayStoreIterator::next(self);
202
203        Some((key, value))
204    }
205}
206
207impl<S: Store> mkvs::Iterator for OverlayStoreIterator<'_, S> {
208    fn set_prefetch(&mut self, prefetch: usize) {
209        self.parent.set_prefetch(prefetch)
210    }
211
212    fn is_valid(&self) -> bool {
213        // If either iterator is valid, the merged iterator is valid.
214        self.parent.is_valid() || self.overlay_valid
215    }
216
217    fn error(&self) -> &Option<anyhow::Error> {
218        self.parent.error()
219    }
220
221    fn rewind(&mut self) {
222        self.seek(&[]);
223    }
224
225    fn seek(&mut self, key: &[u8]) {
226        self.parent.seek(key);
227        self.overlay = self.store.overlay.range(key.to_vec()..).peekable();
228
229        self.update_iterator_position();
230    }
231
232    fn get_key(&self) -> &Option<mkvs::Key> {
233        &self.key
234    }
235
236    fn get_value(&self) -> &Option<Vec<u8>> {
237        &self.value
238    }
239
240    fn next(&mut self) {
241        OverlayStoreIterator::next(self)
242    }
243}
oasis_runtime_sdk/storage/overlay.rs

oasis_runtime_sdk/storage/
overlay.rs
