sov_db/
state_db.rs

1use std::path::Path;
2use std::sync::{Arc, Mutex};
3
4use jmt::storage::{TreeReader, TreeWriter};
5use jmt::{KeyHash, Version};
6use sov_schema_db::DB;
7
8use crate::rocks_db_config::gen_rocksdb_options;
9use crate::schema::tables::{JmtNodes, JmtValues, KeyHashToKey, STATE_TABLES};
10use crate::schema::types::StateKey;
11
12/// A typed wrapper around the db for storing rollup state. Internally,
13/// this is roughly just an `Arc<SchemaDB>`.
14///
15/// StateDB implements several convenience functions for state storage -
16/// notably the `TreeReader` and `TreeWriter` traits.
17#[derive(Clone, Debug)]
18pub struct StateDB {
19    /// The underlying database instance, wrapped in an [`Arc`] for convenience and [`SchemaDB`] for type safety
20    db: Arc<DB>,
21    /// The [`Version`] that will be used for the next batch of writes to the DB.
22    next_version: Arc<Mutex<Version>>,
23}
24
25const STATE_DB_PATH_SUFFIX: &str = "state";
26
27impl StateDB {
28    /// Open a [`StateDB`] (backed by RocksDB) at the specified path.
29    /// The returned instance will be at the path `{path}/state-db`.
30    pub fn with_path(path: impl AsRef<Path>) -> Result<Self, anyhow::Error> {
31        let path = path.as_ref().join(STATE_DB_PATH_SUFFIX);
32        let inner = DB::open(
33            path,
34            "state-db",
35            STATE_TABLES.iter().copied(),
36            &gen_rocksdb_options(&Default::default(), false),
37        )?;
38
39        let next_version = Self::last_version_written(&inner)?.unwrap_or_default() + 1;
40
41        Ok(Self {
42            db: Arc::new(inner),
43            next_version: Arc::new(Mutex::new(next_version)),
44        })
45    }
46
47    /// Put the preimage of a hashed key into the database. Note that the preimage is not checked for correctness,
48    /// since the DB is unaware of the hash function used by the JMT.
49    pub fn put_preimage(&self, key_hash: KeyHash, key: &Vec<u8>) -> Result<(), anyhow::Error> {
50        self.db.put::<KeyHashToKey>(&key_hash.0, key)
51    }
52
53    /// Get an optional value from the database, given a version and a key hash.
54    pub fn get_value_option_by_key(
55        &self,
56        version: Version,
57        key: &StateKey,
58    ) -> anyhow::Result<Option<jmt::OwnedValue>> {
59        let mut iter = self.db.iter::<JmtValues>()?;
60        // find the latest instance of the key whose version <= target
61        iter.seek_for_prev(&(&key, version))?;
62        let found = iter.next();
63        match found {
64            Some(result) => {
65                let ((found_key, found_version), value) = result?;
66                if &found_key == key {
67                    anyhow::ensure!(found_version <= version, "Bug! iterator isn't returning expected values. expected a version <= {version:} but found {found_version:}");
68                    Ok(value)
69                } else {
70                    Ok(None)
71                }
72            }
73            None => Ok(None),
74        }
75    }
76
77    /// Increment the `next_version` counter by 1.
78    pub fn inc_next_version(&self) {
79        let mut version = self.next_version.lock().unwrap();
80        *version += 1;
81    }
82
83    /// Get the current value of the `next_version` counter
84    pub fn get_next_version(&self) -> Version {
85        let version = self.next_version.lock().unwrap();
86        *version
87    }
88
89    fn last_version_written(db: &DB) -> anyhow::Result<Option<Version>> {
90        let mut iter = db.iter::<JmtValues>()?;
91        iter.seek_to_last();
92
93        let version = match iter.next() {
94            Some(Ok(((_, version), _))) => Some(version),
95            _ => None,
96        };
97        Ok(version)
98    }
99}
100
101impl TreeReader for StateDB {
102    fn get_node_option(
103        &self,
104        node_key: &jmt::storage::NodeKey,
105    ) -> anyhow::Result<Option<jmt::storage::Node>> {
106        self.db.get::<JmtNodes>(node_key)
107    }
108
109    fn get_value_option(
110        &self,
111        version: Version,
112        key_hash: KeyHash,
113    ) -> anyhow::Result<Option<jmt::OwnedValue>> {
114        if let Some(key) = self.db.get::<KeyHashToKey>(&key_hash.0)? {
115            self.get_value_option_by_key(version, &key)
116        } else {
117            Ok(None)
118        }
119    }
120
121    fn get_rightmost_leaf(
122        &self,
123    ) -> anyhow::Result<Option<(jmt::storage::NodeKey, jmt::storage::LeafNode)>> {
124        todo!("StateDB does not support [`TreeReader::get_rightmost_leaf`] yet")
125    }
126}
127
128impl TreeWriter for StateDB {
129    fn write_node_batch(&self, node_batch: &jmt::storage::NodeBatch) -> anyhow::Result<()> {
130        for (node_key, node) in node_batch.nodes() {
131            self.db.put::<JmtNodes>(node_key, node)?;
132        }
133
134        for ((version, key_hash), value) in node_batch.values() {
135            let key_preimage =
136                self.db
137                    .get::<KeyHashToKey>(&key_hash.0)?
138                    .ok_or(anyhow::format_err!(
139                        "Could not find preimage for key hash {key_hash:?}. Has `StateDB::put_preimage` been called for this key?"
140                    ))?;
141            self.db.put::<JmtValues>(&(key_preimage, *version), value)?;
142        }
143        Ok(())
144    }
145}
146
147#[cfg(feature = "arbitrary")]
148pub mod arbitrary {
149    //! Arbitrary definitions for the [`StateDB`].
150
151    use core::ops::{Deref, DerefMut};
152
153    use proptest::strategy::LazyJust;
154    use tempfile::TempDir;
155
156    use super::*;
157
158    /// Arbitrary instance of [`StateDB`].
159    ///
160    /// This is a db wrapper bound to its temporary path that will be deleted once the type is
161    /// dropped.
162    #[derive(Debug)]
163    pub struct ArbitraryDB {
164        /// The underlying RocksDB instance.
165        pub db: StateDB,
166        /// The temporary directory used to create the [`StateDB`].
167        pub path: TempDir,
168    }
169
170    /// A fallible, arbitrary instance of [`StateDB`].
171    ///
172    /// This type is suitable for operations that are expected to be infallible. The internal
173    /// implementation of the db requires I/O to create the temporary dir, making it fallible.
174    #[derive(Debug)]
175    pub struct FallibleArbitraryStateDB {
176        /// The result of the new db instance.
177        pub result: anyhow::Result<ArbitraryDB>,
178    }
179
180    impl Deref for ArbitraryDB {
181        type Target = StateDB;
182
183        fn deref(&self) -> &Self::Target {
184            &self.db
185        }
186    }
187
188    impl DerefMut for ArbitraryDB {
189        fn deref_mut(&mut self) -> &mut Self::Target {
190            &mut self.db
191        }
192    }
193
194    impl<'a> ::arbitrary::Arbitrary<'a> for ArbitraryDB {
195        fn arbitrary(_u: &mut ::arbitrary::Unstructured<'a>) -> ::arbitrary::Result<Self> {
196            let path = TempDir::new().map_err(|_| ::arbitrary::Error::NotEnoughData)?;
197            let db = StateDB::with_path(&path).map_err(|_| ::arbitrary::Error::IncorrectFormat)?;
198            Ok(Self { db, path })
199        }
200    }
201
202    impl proptest::arbitrary::Arbitrary for FallibleArbitraryStateDB {
203        type Parameters = ();
204        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
205            fn gen() -> FallibleArbitraryStateDB {
206                FallibleArbitraryStateDB {
207                    result: TempDir::new()
208                        .map_err(|e| {
209                            anyhow::anyhow!(format!("failed to generate path for StateDB: {e}"))
210                        })
211                        .and_then(|path| {
212                            let db = StateDB::with_path(&path)?;
213                            Ok(ArbitraryDB { db, path })
214                        }),
215                }
216            }
217            LazyJust::new(gen)
218        }
219
220        type Strategy = LazyJust<Self, fn() -> FallibleArbitraryStateDB>;
221    }
222}
223
224#[cfg(test)]
225mod state_db_tests {
226    use jmt::storage::{NodeBatch, TreeReader, TreeWriter};
227    use jmt::KeyHash;
228
229    use super::StateDB;
230
231    #[test]
232    fn test_simple() {
233        let tmpdir = tempfile::tempdir().unwrap();
234        let db = StateDB::with_path(tmpdir.path()).unwrap();
235        let key_hash = KeyHash([1u8; 32]);
236        let key = vec![2u8; 100];
237        let value = [8u8; 150];
238
239        db.put_preimage(key_hash, &key).unwrap();
240        let mut batch = NodeBatch::default();
241        batch.extend(vec![], vec![((0, key_hash), Some(value.to_vec()))]);
242        db.write_node_batch(&batch).unwrap();
243
244        let found = db.get_value(0, key_hash).unwrap();
245        assert_eq!(found, value);
246
247        let found = db.get_value_option_by_key(0, &key).unwrap().unwrap();
248        assert_eq!(found, value);
249    }
250}
sov_db/state_db.rs

sov_db/
state_db.rs
