cw_storage_plus/snapshot/
mod.rs

1#![cfg(feature = "iterator")]
2mod item;
3mod map;
4
5pub use item::SnapshotItem;
6pub use map::SnapshotMap;
7
8use crate::bound::Bound;
9use crate::de::KeyDeserialize;
10use crate::namespace::Namespace;
11use crate::{Map, Prefixer, PrimaryKey};
12use cosmwasm_std::{Order, StdError, StdResult, Storage};
13use serde::de::DeserializeOwned;
14use serde::{Deserialize, Serialize};
15
16/// Structure holding a map of checkpoints composited from
17/// height (as u64) and counter of how many times it has
18/// been checkpointed (as u32).
19/// Stores all changes in changelog.
20#[derive(Debug, Clone)]
21pub(crate) struct Snapshot<K, T> {
22    checkpoints: Map<u64, u32>,
23
24    // this stores all changes (key, height). Must differentiate between no data written,
25    // and explicit None (just inserted)
26    pub changelog: Map<(K, u64), ChangeSet<T>>,
27
28    // How aggressive we are about checkpointing all data
29    strategy: Strategy,
30}
31
32impl<K, T> Snapshot<K, T> {
33    /// Creates a new [`Snapshot`] with the given storage keys and strategy.
34    /// This is a const fn only suitable when all the storage keys provided are
35    /// static strings.
36    pub const fn new(
37        checkpoints: &'static str,
38        changelog: &'static str,
39        strategy: Strategy,
40    ) -> Snapshot<K, T> {
41        Snapshot {
42            checkpoints: Map::new(checkpoints),
43            changelog: Map::new(changelog),
44            strategy,
45        }
46    }
47
48    /// Creates a new [`Snapshot`] with the given storage keys and strategy.
49    /// Use this if you might need to handle dynamic strings. Otherwise, you might
50    /// prefer [`Snapshot::new`].
51    pub fn new_dyn(
52        checkpoints: impl Into<Namespace>,
53        changelog: impl Into<Namespace>,
54        strategy: Strategy,
55    ) -> Snapshot<K, T> {
56        Snapshot {
57            checkpoints: Map::new_dyn(checkpoints),
58            changelog: Map::new_dyn(changelog),
59            strategy,
60        }
61    }
62
63    pub fn add_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
64        self.checkpoints
65            .update::<_, StdError>(store, height, |count| Ok(count.unwrap_or_default() + 1))?;
66        Ok(())
67    }
68
69    pub fn remove_checkpoint(&self, store: &mut dyn Storage, height: u64) -> StdResult<()> {
70        let count = self
71            .checkpoints
72            .may_load(store, height)?
73            .unwrap_or_default();
74        if count <= 1 {
75            self.checkpoints.remove(store, height);
76            Ok(())
77        } else {
78            self.checkpoints.save(store, height, &(count - 1))
79        }
80    }
81}
82
83impl<'a, K, T> Snapshot<K, T>
84where
85    T: Serialize + DeserializeOwned + Clone,
86    K: PrimaryKey<'a> + Prefixer<'a> + KeyDeserialize,
87{
88    /// should_checkpoint looks at the strategy and determines if we want to checkpoint
89    pub fn should_checkpoint(&self, store: &dyn Storage, k: &K) -> StdResult<bool> {
90        match self.strategy {
91            Strategy::EveryBlock => Ok(true),
92            Strategy::Never => Ok(false),
93            Strategy::Selected => self.should_checkpoint_selected(store, k),
94        }
95    }
96
97    /// this is just pulled out from above for the selected block
98    fn should_checkpoint_selected(&self, store: &dyn Storage, k: &K) -> StdResult<bool> {
99        // most recent checkpoint
100        let checkpoint = self
101            .checkpoints
102            .range(store, None, None, Order::Descending)
103            .next()
104            .transpose()?;
105        if let Some((height, _)) = checkpoint {
106            // any changelog for the given key since then?
107            let start = Bound::inclusive(height);
108            let first = self
109                .changelog
110                .prefix(k.clone())
111                .range_raw(store, Some(start), None, Order::Ascending)
112                .next()
113                .transpose()?;
114            if first.is_none() {
115                // there must be at least one open checkpoint and no changelog for the given height since then
116                return Ok(true);
117            }
118        }
119        // otherwise, we don't save this
120        Ok(false)
121    }
122
123    // If there is no checkpoint for that height, then we return StdError::NotFound
124    pub fn assert_checkpointed(&self, store: &dyn Storage, height: u64) -> StdResult<()> {
125        let has = match self.strategy {
126            Strategy::EveryBlock => true,
127            Strategy::Never => false,
128            Strategy::Selected => self.checkpoints.may_load(store, height)?.is_some(),
129        };
130        match has {
131            true => Ok(()),
132            false => Err(StdError::not_found("checkpoint")),
133        }
134    }
135
136    pub fn has_changelog(&self, store: &mut dyn Storage, key: K, height: u64) -> StdResult<bool> {
137        Ok(self.changelog.may_load(store, (key, height))?.is_some())
138    }
139
140    pub fn write_changelog(
141        &self,
142        store: &mut dyn Storage,
143        key: K,
144        height: u64,
145        old: Option<T>,
146    ) -> StdResult<()> {
147        self.changelog
148            .save(store, (key, height), &ChangeSet { old })
149    }
150
151    // may_load_at_height reads historical data from given checkpoints.
152    // Returns StdError::NotFound if we have no checkpoint, and can give no data.
153    // Returns Ok(None) if there is a checkpoint, but no cached data (no changes since the
154    // checkpoint. Caller should query current state).
155    // Return Ok(Some(x)) if there is a checkpoint and data written to changelog, returning the state at that time
156    pub fn may_load_at_height(
157        &self,
158        store: &dyn Storage,
159        key: K,
160        height: u64,
161    ) -> StdResult<Option<Option<T>>> {
162        self.assert_checkpointed(store, height)?;
163
164        // this will look for the first snapshot of height >= given height
165        // If None, there is no snapshot since that time.
166        let start = Bound::inclusive(height);
167        let first = self
168            .changelog
169            .prefix(key)
170            .range_raw(store, Some(start), None, Order::Ascending)
171            .next();
172
173        if let Some(r) = first {
174            // if we found a match, return this last one
175            r.map(|(_, v)| Some(v.old))
176        } else {
177            Ok(None)
178        }
179    }
180}
181
182#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
183pub enum Strategy {
184    EveryBlock,
185    Never,
186    /// Only writes for linked blocks - does a few more reads to save some writes.
187    /// Probably uses more gas, but less total disk usage.
188    ///
189    /// Note that you need a trusted source (eg. own contract) to set/remove checkpoints.
190    /// Useful when the checkpoint setting happens in the same contract as the snapshotting.
191    Selected,
192}
193
194#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
195pub struct ChangeSet<T> {
196    pub old: Option<T>,
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use cosmwasm_std::testing::MockStorage;
203
204    type TestSnapshot = Snapshot<&'static str, u64>;
205
206    const NEVER: TestSnapshot = Snapshot::new("never__check", "never__change", Strategy::Never);
207    const EVERY: TestSnapshot =
208        Snapshot::new("every__check", "every__change", Strategy::EveryBlock);
209    const SELECT: TestSnapshot =
210        Snapshot::new("select__check", "select__change", Strategy::Selected);
211
212    const DUMMY_KEY: &str = "dummy";
213
214    #[test]
215    fn should_checkpoint() {
216        let storage = MockStorage::new();
217
218        assert_eq!(NEVER.should_checkpoint(&storage, &DUMMY_KEY), Ok(false));
219        assert_eq!(EVERY.should_checkpoint(&storage, &DUMMY_KEY), Ok(true));
220        assert_eq!(SELECT.should_checkpoint(&storage, &DUMMY_KEY), Ok(false));
221    }
222
223    #[test]
224    fn assert_checkpointed() {
225        let mut storage = MockStorage::new();
226
227        assert_eq!(
228            NEVER.assert_checkpointed(&storage, 1),
229            Err(StdError::not_found("checkpoint"))
230        );
231        assert_eq!(EVERY.assert_checkpointed(&storage, 1), Ok(()));
232        assert_eq!(
233            SELECT.assert_checkpointed(&storage, 1),
234            Err(StdError::not_found("checkpoint"))
235        );
236
237        // Add a checkpoint at 1
238        NEVER.add_checkpoint(&mut storage, 1).unwrap();
239        EVERY.add_checkpoint(&mut storage, 1).unwrap();
240        SELECT.add_checkpoint(&mut storage, 1).unwrap();
241
242        assert_eq!(
243            NEVER.assert_checkpointed(&storage, 1),
244            Err(StdError::not_found("checkpoint"))
245        );
246        assert_eq!(EVERY.assert_checkpointed(&storage, 1), Ok(()));
247        assert_eq!(SELECT.assert_checkpointed(&storage, 1), Ok(()));
248
249        // Remove checkpoint
250        NEVER.remove_checkpoint(&mut storage, 1).unwrap();
251        EVERY.remove_checkpoint(&mut storage, 1).unwrap();
252        SELECT.remove_checkpoint(&mut storage, 1).unwrap();
253
254        assert_eq!(
255            NEVER.assert_checkpointed(&storage, 1),
256            Err(StdError::not_found("checkpoint"))
257        );
258        assert_eq!(EVERY.assert_checkpointed(&storage, 1), Ok(()));
259        assert_eq!(
260            SELECT.assert_checkpointed(&storage, 1),
261            Err(StdError::not_found("checkpoint"))
262        );
263    }
264
265    #[test]
266    fn has_changelog() {
267        let mut storage = MockStorage::new();
268
269        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
270        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
271        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
272
273        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(false));
274        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(false));
275        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(false));
276
277        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
278        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
279        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
280
281        // Write a changelog at 2
282        NEVER
283            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(3))
284            .unwrap();
285        EVERY
286            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(4))
287            .unwrap();
288        SELECT
289            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(5))
290            .unwrap();
291
292        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
293        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
294        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 1), Ok(false));
295
296        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(true));
297        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(true));
298        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 2), Ok(true));
299
300        assert_eq!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
301        assert_eq!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
302        assert_eq!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 3), Ok(false));
303    }
304
305    #[test]
306    fn may_load_at_height() {
307        let mut storage = MockStorage::new();
308
309        assert_eq!(
310            NEVER.may_load_at_height(&storage, DUMMY_KEY, 3),
311            Err(StdError::not_found("checkpoint"))
312        );
313        assert_eq!(EVERY.may_load_at_height(&storage, DUMMY_KEY, 3), Ok(None));
314        assert_eq!(
315            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3),
316            Err(StdError::not_found("checkpoint"))
317        );
318
319        // Add a checkpoint at 3
320        NEVER.add_checkpoint(&mut storage, 3).unwrap();
321        EVERY.add_checkpoint(&mut storage, 3).unwrap();
322        SELECT.add_checkpoint(&mut storage, 3).unwrap();
323
324        assert_eq!(
325            NEVER.may_load_at_height(&storage, DUMMY_KEY, 3),
326            Err(StdError::not_found("checkpoint"))
327        );
328        assert_eq!(EVERY.may_load_at_height(&storage, DUMMY_KEY, 3), Ok(None));
329        assert_eq!(SELECT.may_load_at_height(&storage, DUMMY_KEY, 3), Ok(None));
330
331        // Write a changelog at 3
332        NEVER
333            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(100))
334            .unwrap();
335        EVERY
336            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(101))
337            .unwrap();
338        SELECT
339            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(102))
340            .unwrap();
341
342        assert_eq!(
343            NEVER.may_load_at_height(&storage, DUMMY_KEY, 3),
344            Err(StdError::not_found("checkpoint"))
345        );
346        assert_eq!(
347            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3),
348            Ok(Some(Some(101)))
349        );
350        assert_eq!(
351            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3),
352            Ok(Some(Some(102)))
353        );
354        // Check that may_load_at_height at a previous value will return the first change after that.
355        // (Only with EVERY).
356        assert_eq!(
357            NEVER.may_load_at_height(&storage, DUMMY_KEY, 2),
358            Err(StdError::not_found("checkpoint"))
359        );
360        assert_eq!(
361            EVERY.may_load_at_height(&storage, DUMMY_KEY, 2),
362            Ok(Some(Some(101)))
363        );
364        assert_eq!(
365            SELECT.may_load_at_height(&storage, DUMMY_KEY, 2),
366            Err(StdError::not_found("checkpoint"))
367        );
368
369        // Write a changelog at 4, removing the value
370        NEVER
371            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
372            .unwrap();
373        EVERY
374            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
375            .unwrap();
376        SELECT
377            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
378            .unwrap();
379        // And add a checkpoint at 4
380        NEVER.add_checkpoint(&mut storage, 4).unwrap();
381        EVERY.add_checkpoint(&mut storage, 4).unwrap();
382        SELECT.add_checkpoint(&mut storage, 4).unwrap();
383
384        assert_eq!(
385            NEVER.may_load_at_height(&storage, DUMMY_KEY, 4),
386            Err(StdError::not_found("checkpoint"))
387        );
388        assert_eq!(
389            EVERY.may_load_at_height(&storage, DUMMY_KEY, 4),
390            Ok(Some(None))
391        );
392        assert_eq!(
393            SELECT.may_load_at_height(&storage, DUMMY_KEY, 4),
394            Ok(Some(None))
395        );
396
397        // Confirm old value at 3
398        assert_eq!(
399            NEVER.may_load_at_height(&storage, DUMMY_KEY, 3),
400            Err(StdError::not_found("checkpoint"))
401        );
402        assert_eq!(
403            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3),
404            Ok(Some(Some(101)))
405        );
406        assert_eq!(
407            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3),
408            Ok(Some(Some(102)))
409        );
410    }
411}
cw_storage_plus/snapshot/mod.rs

cw_storage_plus/snapshot/
mod.rs
