solana_oasis_node/
state.rs

1use anyhow::Result;
2use rocksdb::{ColumnFamilyDescriptor, Options, DB};
3use serde;
4use std::path::Path;
5
6pub struct StateManager {
7    db: DB,
8    current_root: StateRoot,
9}
10
11#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
12pub struct StateRoot {
13    pub root: [u8; 32],
14    pub height: u64,
15}
16
17impl StateManager {
18    pub fn new(path: &impl AsRef<Path>) -> Result<Self> {
19        let mut opts = Options::default();
20        opts.create_if_missing(true);
21        opts.create_missing_column_families(true);
22
23        let cf_roots = ColumnFamilyDescriptor::new("roots", Options::default());
24        let cf_data = ColumnFamilyDescriptor::new("data", Options::default());
25
26        let db = DB::open_cf_descriptors(&opts, path, vec![cf_roots, cf_data])?;
27
28        let current_root = match db.get_cf(db.cf_handle("roots").unwrap(), "current")? {
29            Some(data) => bincode::deserialize(&data)?,
30            None => StateRoot {
31                root: [0u8; 32],
32                height: 0,
33            },
34        };
35
36        Ok(Self { db, current_root })
37    }
38
39    pub async fn get_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
40        let cf_data = self.db.cf_handle("data").unwrap();
41        Ok(self.db.get_cf(cf_data, key)?)
42    }
43
44    pub async fn set_value(&mut self, key: &[u8], value: Vec<u8>) -> Result<()> {
45        let cf_data = self.db.cf_handle("data").unwrap();
46        self.db.put_cf(cf_data, key, value)?;
47        Ok(())
48    }
49
50    pub async fn update_root(&mut self, root: StateRoot) -> Result<()> {
51        let cf_roots = self.db.cf_handle("roots").unwrap();
52        self.db
53            .put_cf(cf_roots, "current", bincode::serialize(&root)?)?;
54        self.current_root = root;
55        Ok(())
56    }
57
58    pub fn get_current_root(&self) -> &StateRoot {
59        &self.current_root
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use tempfile::tempdir;
67
68    #[tokio::test]
69    async fn test_state_manager() -> Result<()> {
70        let temp_dir = tempdir()?;
71        let mut state = StateManager::new(&temp_dir)?;
72
73        let key = b"test_key";
74        let value = b"test_value".to_vec();
75        state.set_value(key, value.clone()).await?;
76
77        let retrieved = state.get_value(key).await?;
78        assert_eq!(retrieved, Some(value));
79
80        let new_root = StateRoot {
81            root: [1u8; 32],
82            height: 1,
83        };
84        state.update_root(new_root.clone()).await?;
85        assert_eq!(state.get_current_root().root, new_root.root);
86        assert_eq!(state.get_current_root().height, new_root.height);
87
88        Ok(())
89    }
90}