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::msg("checkpoint not found")),
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 (e.g. 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        assert!(!NEVER.should_checkpoint(&storage, &DUMMY_KEY).unwrap());
218        assert!(EVERY.should_checkpoint(&storage, &DUMMY_KEY).unwrap());
219        assert!(!SELECT.should_checkpoint(&storage, &DUMMY_KEY).unwrap());
220    }
221
222    #[test]
223    fn assert_checkpointed() {
224        let mut storage = MockStorage::new();
225
226        assert_eq!(
227            "kind: Other, error: checkpoint not found",
228            NEVER
229                .assert_checkpointed(&storage, 1)
230                .unwrap_err()
231                .to_string()
232        );
233        assert!(EVERY.assert_checkpointed(&storage, 1).is_ok());
234        assert_eq!(
235            "kind: Other, error: checkpoint not found",
236            SELECT
237                .assert_checkpointed(&storage, 1)
238                .unwrap_err()
239                .to_string()
240        );
241
242        // Add a checkpoint at 1
243        NEVER.add_checkpoint(&mut storage, 1).unwrap();
244        EVERY.add_checkpoint(&mut storage, 1).unwrap();
245        SELECT.add_checkpoint(&mut storage, 1).unwrap();
246
247        assert_eq!(
248            "kind: Other, error: checkpoint not found",
249            NEVER
250                .assert_checkpointed(&storage, 1)
251                .unwrap_err()
252                .to_string()
253        );
254        assert!(EVERY.assert_checkpointed(&storage, 1).is_ok());
255        assert!(SELECT.assert_checkpointed(&storage, 1).is_ok());
256
257        // Remove checkpoint
258        NEVER.remove_checkpoint(&mut storage, 1).unwrap();
259        EVERY.remove_checkpoint(&mut storage, 1).unwrap();
260        SELECT.remove_checkpoint(&mut storage, 1).unwrap();
261
262        assert_eq!(
263            "kind: Other, error: checkpoint not found",
264            NEVER
265                .assert_checkpointed(&storage, 1)
266                .unwrap_err()
267                .to_string()
268        );
269        assert!(EVERY.assert_checkpointed(&storage, 1).is_ok());
270        assert_eq!(
271            "kind: Other, error: checkpoint not found",
272            SELECT
273                .assert_checkpointed(&storage, 1)
274                .unwrap_err()
275                .to_string()
276        );
277    }
278
279    #[test]
280    fn has_changelog() {
281        let mut storage = MockStorage::new();
282
283        assert!(!NEVER.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
284        assert!(!EVERY.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
285        assert!(!SELECT.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
286        assert!(!NEVER.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
287        assert!(!EVERY.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
288        assert!(!SELECT.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
289        assert!(!NEVER.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
290        assert!(!EVERY.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
291        assert!(!SELECT.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
292
293        // Write a changelog at 2
294        NEVER
295            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(3))
296            .unwrap();
297        EVERY
298            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(4))
299            .unwrap();
300        SELECT
301            .write_changelog(&mut storage, DUMMY_KEY, 2, Some(5))
302            .unwrap();
303
304        assert!(!NEVER.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
305        assert!(!EVERY.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
306        assert!(!SELECT.has_changelog(&mut storage, DUMMY_KEY, 1).unwrap());
307        assert!(NEVER.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
308        assert!(EVERY.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
309        assert!(SELECT.has_changelog(&mut storage, DUMMY_KEY, 2).unwrap());
310        assert!(!NEVER.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
311        assert!(!EVERY.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
312        assert!(!SELECT.has_changelog(&mut storage, DUMMY_KEY, 3).unwrap());
313    }
314
315    #[test]
316    fn may_load_at_height() {
317        let mut storage = MockStorage::new();
318
319        assert_eq!(
320            "kind: Other, error: checkpoint not found",
321            NEVER
322                .may_load_at_height(&storage, DUMMY_KEY, 3)
323                .unwrap_err()
324                .to_string()
325        );
326        assert_eq!(
327            None,
328            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
329        );
330        assert_eq!(
331            "kind: Other, error: checkpoint not found",
332            SELECT
333                .may_load_at_height(&storage, DUMMY_KEY, 3)
334                .unwrap_err()
335                .to_string()
336        );
337
338        // Add a checkpoint at 3
339        NEVER.add_checkpoint(&mut storage, 3).unwrap();
340        EVERY.add_checkpoint(&mut storage, 3).unwrap();
341        SELECT.add_checkpoint(&mut storage, 3).unwrap();
342
343        assert_eq!(
344            "kind: Other, error: checkpoint not found",
345            NEVER
346                .may_load_at_height(&storage, DUMMY_KEY, 3)
347                .unwrap_err()
348                .to_string()
349        );
350        assert_eq!(
351            None,
352            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
353        );
354        assert_eq!(
355            None,
356            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
357        );
358
359        // Write a changelog at 3
360        NEVER
361            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(100))
362            .unwrap();
363        EVERY
364            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(101))
365            .unwrap();
366        SELECT
367            .write_changelog(&mut storage, DUMMY_KEY, 3, Some(102))
368            .unwrap();
369
370        assert_eq!(
371            "kind: Other, error: checkpoint not found",
372            NEVER
373                .may_load_at_height(&storage, DUMMY_KEY, 3)
374                .unwrap_err()
375                .to_string()
376        );
377        assert_eq!(
378            Some(Some(101)),
379            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
380        );
381        assert_eq!(
382            Some(Some(102)),
383            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
384        );
385        // Check that may_load_at_height at a previous value will return the first change after that.
386        // (Only with EVERY).
387        assert_eq!(
388            "kind: Other, error: checkpoint not found",
389            NEVER
390                .may_load_at_height(&storage, DUMMY_KEY, 2)
391                .unwrap_err()
392                .to_string()
393        );
394        assert_eq!(
395            Some(Some(101)),
396            EVERY.may_load_at_height(&storage, DUMMY_KEY, 2).unwrap()
397        );
398        assert_eq!(
399            "kind: Other, error: checkpoint not found",
400            SELECT
401                .may_load_at_height(&storage, DUMMY_KEY, 2)
402                .unwrap_err()
403                .to_string()
404        );
405
406        // Write a changelog at 4, removing the value
407        NEVER
408            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
409            .unwrap();
410        EVERY
411            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
412            .unwrap();
413        SELECT
414            .write_changelog(&mut storage, DUMMY_KEY, 4, None)
415            .unwrap();
416        // And add a checkpoint at 4
417        NEVER.add_checkpoint(&mut storage, 4).unwrap();
418        EVERY.add_checkpoint(&mut storage, 4).unwrap();
419        SELECT.add_checkpoint(&mut storage, 4).unwrap();
420
421        assert_eq!(
422            "kind: Other, error: checkpoint not found",
423            NEVER
424                .may_load_at_height(&storage, DUMMY_KEY, 4)
425                .unwrap_err()
426                .to_string()
427        );
428        assert_eq!(
429            Some(None),
430            EVERY.may_load_at_height(&storage, DUMMY_KEY, 4).unwrap()
431        );
432        assert_eq!(
433            Some(None),
434            SELECT.may_load_at_height(&storage, DUMMY_KEY, 4).unwrap()
435        );
436
437        // Confirm old value at 3
438        assert_eq!(
439            "kind: Other, error: checkpoint not found",
440            NEVER
441                .may_load_at_height(&storage, DUMMY_KEY, 3)
442                .unwrap_err()
443                .to_string()
444        );
445        assert_eq!(
446            Some(Some(101)),
447            EVERY.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
448        );
449        assert_eq!(
450            Some(Some(102)),
451            SELECT.may_load_at_height(&storage, DUMMY_KEY, 3).unwrap()
452        );
453    }
454}
cw_storage_plus/snapshot/mod.rs

cw_storage_plus/snapshot/
mod.rs
