Incomplete State DB pruning

bin/citrea/tests/mock/pruning.rs

@@ -2,20 +2,149 @@ use std::collections::BTreeMap;
 use std::panic::AssertUnwindSafe;
 use std::str::FromStr;
 
-use alloy_primitives::Address;
+use alloy_primitives::{Address, U256};
 /// Testing if the sequencer and full node can handle system transactions correctly (the full node should have the same system transactions as the sequencer)
 use citrea_storage_ops::pruning::PruningConfig;
 use futures::FutureExt;
-use reth_primitives::BlockNumberOrTag;
+use reth_primitives::{BlockId, BlockNumberOrTag};
 use sov_mock_da::{MockAddress, MockDaService};
 
 use super::{initialize_test, TestConfig};
 use crate::common::helpers::{tempdir_with_children, wait_for_l1_block, wait_for_l2_block};
 
+/// Trigger pruning state DB data.
+#[tokio::test(flavor = "multi_thread")]
+#[ignore]
+async fn test_state_db_pruning() -> Result<(), anyhow::Error> {
+    citrea::initialize_logging(tracing::Level::INFO);
+    let storage_dir = tempdir_with_children(&["DA", "sequencer", "full-node"]);
+    let da_db_dir = storage_dir.path().join("DA").to_path_buf();
+    let sequencer_db_dir = storage_dir.path().join("sequencer").to_path_buf();
+    let fullnode_db_dir = storage_dir.path().join("full-node").to_path_buf();
+
+    let da_service = MockDaService::new(MockAddress::default(), &da_db_dir.clone());
+
+    // start rollup on da block 3
+    for _ in 0..3 {
+        da_service.publish_test_block().await.unwrap();
+    }
+    wait_for_l1_block(&da_service, 3, None).await;
+
+    let addr = Address::from_str("0xd39Fd6e51aad88F6F4ce6aB8827279cffFb92266").unwrap();
+
+    let (seq_test_client, full_node_test_client, seq_task, full_node_task, _) =
+        initialize_test(TestConfig {
+            da_path: da_db_dir,
+            sequencer_path: sequencer_db_dir,
+            fullnode_path: fullnode_db_dir,
+            pruning_config: Some(PruningConfig { distance: 20 }),
+            ..Default::default()
+        })
+        .await;
+
+    for i in 1..=50 {
+        // send one ether to some address
+        let _pending = seq_test_client
+            .send_eth(addr, None, None, None, 1e18 as u128)
+            .await
+            .unwrap();
+
+        seq_test_client.send_publish_batch_request().await;
+
+        if i % 5 == 0 {
+            wait_for_l2_block(&seq_test_client, i, None).await;
+            da_service.publish_test_block().await.unwrap();
+            wait_for_l1_block(&da_service, 3 + (i / 5), None).await;
+        }
+    }
+
+    wait_for_l2_block(&full_node_test_client, 50, None).await;
+
+    tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
+
+    // Old blocks balance information should have been pruned
+    let get_balance_result = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(2))))
+        .await;
+    assert!(get_balance_result.is_ok());
+    assert_eq!(get_balance_result.unwrap(), U256::from(0));
+
+    let get_balance_result = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(19))))
+        .await;
+    assert!(get_balance_result.is_ok());
+    assert_eq!(get_balance_result.unwrap(), U256::from(0));
+
+    // Non pruned block balances should be available
+    let balance = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(20))))
+        .await
+        .unwrap();
+    assert_eq!(balance, U256::from(20000000000000000000u128));
+
+    let balance = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(50))))
+        .await
+        .unwrap();
+    assert_eq!(balance, U256::from(50000000000000000000u128));
+
+    // Continnue block production
+    for i in 51..=101 {
+        // send one ether to some address
+        let _pending = seq_test_client
+            .send_eth(addr, None, None, None, 1e18 as u128)
+            .await
+            .unwrap();
+
+        seq_test_client.send_publish_batch_request().await;
+
+        if i % 5 == 0 {
+            wait_for_l2_block(&seq_test_client, i, None).await;
+            da_service.publish_test_block().await.unwrap();
+            wait_for_l1_block(&da_service, 3 + (i / 5), None).await;
+        }
+    }
+    wait_for_l2_block(&full_node_test_client, 101, None).await;
+
+    tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
+
+    // Old blocks balance information should have been pruned
+    let get_balance_result = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(42))))
+        .await;
+    assert!(get_balance_result.is_ok());
+    assert_eq!(get_balance_result.unwrap(), U256::from(0));
+
+    let get_balance_result = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(59))))
+        .await;
+    assert!(get_balance_result.is_ok());
+    assert_eq!(get_balance_result.unwrap(), U256::from(0));
+
+    // Non pruned block balances should be available
+    let balance = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(80))))
+        .await
+        .unwrap();
+    assert_eq!(balance, U256::from(80000000000000000000u128));
+
+    let balance = full_node_test_client
+        .eth_get_balance(addr, Some(BlockId::Number(BlockNumberOrTag::Number(100))))
+        .await
+        .unwrap();
+    assert_eq!(balance, U256::from(100000000000000000000u128));
+
+    seq_task.abort();
+    full_node_task.abort();
+
+    Ok(())
+}
+
 /// Trigger pruning native DB data.
 #[tokio::test(flavor = "multi_thread")]
 async fn test_native_db_pruning() -> Result<(), anyhow::Error> {
-    citrea::initialize_logging(tracing::Level::DEBUG);
+    // citrea::initialize_logging(tracing::Level::DEBUG);
+
     let storage_dir = tempdir_with_children(&["DA", "sequencer", "full-node"]);
     let da_db_dir = storage_dir.path().join("DA").to_path_buf();
     let sequencer_db_dir = storage_dir.path().join("sequencer").to_path_buf();

bin/cli/src/commands/prune.rs

@@ -22,17 +22,17 @@ pub(crate) async fn prune(db_path: PathBuf, distance: u64) -> anyhow::Result<()>
     let native_db = NativeDB::<SnapshotManager>::setup_schema_db(&rocksdb_config)?;
     let state_db = StateDB::<SnapshotManager>::setup_schema_db(&rocksdb_config)?;
 
-    let Some((soft_confirmation_number, _)) = ledger_db.get_head_soft_confirmation()? else {
+    let Some(soft_confirmation_number) = ledger_db.get_head_soft_confirmation_height()? else {
         return Ok(());
     };
 
     debug!(
         "Pruning up to latest soft confirmation number: {}, taking into consideration the configured distance of {}",
-        soft_confirmation_number.0, distance
+        soft_confirmation_number, distance
     );
 
     let pruner = Pruner::new(config, ledger_db, Arc::new(state_db), Arc::new(native_db));
-    pruner.prune(soft_confirmation_number.0).await;
+    pruner.prune(soft_confirmation_number).await;
 
     Ok(())
 }

crates/sovereign-sdk/full-node/db/sov-db/src/schema/tables.rs

@@ -10,7 +10,7 @@
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
-use jmt::storage::{NibblePath, Node, NodeKey};
+use jmt::storage::{NibblePath, Node, NodeKey, StaleNodeIndex};
 use jmt::Version;
 use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::mmr::{MMRChunk, MMRNodeHash, Wtxid};
@@ -39,7 +39,10 @@ pub const MMR_TABLES: &[&str] = &[
 pub const STATE_TABLES: &[&str] = &[
     KeyHashToKey::table_name(),
     JmtValues::table_name(),
+    // when iterating we get bigger versions first
     JmtNodes::table_name(),
+    // when iterating we get smaller stale since versions first
+    StaleNodes::table_name(),
 ];
 
 /// A list of all tables used by Sequencer LedgerDB
@@ -374,6 +377,11 @@ define_table_without_codec!(
     (JmtNodes) NodeKey => Node
 );
 
+define_table_with_default_codec!(
+    /// The list of stale nodes in JMT
+    (StaleNodes) StaleNodeIndex => ()
+);
+
 define_table_with_default_codec!(
     /// Light client proof data by l1 height
     (LightClientProofBySlotNumber) SlotNumber => StoredLightClientProof

crates/sovereign-sdk/full-node/db/sov-db/src/state_db.rs

@@ -1,12 +1,13 @@
+use std::collections::BTreeSet;
 use std::sync::{Arc, Mutex};
 
-use jmt::storage::{HasPreimage, TreeReader, TreeWriter};
+use jmt::storage::{HasPreimage, StaleNodeIndex, TreeReader, TreeWriter};
 use jmt::{KeyHash, Version};
 use sov_schema_db::snapshot::{DbSnapshot, QueryManager, ReadOnlyDbSnapshot};
 use sov_schema_db::SchemaBatch;
 
 use crate::rocks_db_config::RocksdbConfig;
-use crate::schema::tables::{JmtNodes, JmtValues, KeyHashToKey, STATE_TABLES};
+use crate::schema::tables::{JmtNodes, JmtValues, KeyHashToKey, StaleNodes, STATE_TABLES};
 use crate::schema::types::StateKey;
 
 /// A typed wrapper around the db for storing rollup state. Internally,
@@ -103,6 +104,16 @@ impl<Q: QueryManager> StateDB<Q> {
         }
     }
 
+    /// Record stale nodes in state
+    pub fn set_stale_nodes(&self, stale_nodes: &BTreeSet<StaleNodeIndex>) -> anyhow::Result<()> {
+        let mut batch = SchemaBatch::new();
+        for index in stale_nodes {
+            batch.put::<StaleNodes>(index, &())?;
+        }
+        self.db.write_many(batch)?;
+        Ok(())
+    }
+
     /// Increment the `next_version` counter by 1.
     pub fn inc_next_version(&self) {
         let mut version = self.next_version.lock().unwrap();

crates/sovereign-sdk/module-system/sov-state/src/prover_storage.rs

@@ -1,6 +1,7 @@
+use std::collections::BTreeSet;
 use std::sync::Arc;
 
-use jmt::storage::{NodeBatch, TreeWriter};
+use jmt::storage::{NodeBatch, StaleNodeIndex, TreeWriter};
 use jmt::{JellyfishMerkleTree, KeyHash, Version};
 use sov_db::native_db::NativeDB;
 use sov_db::schema::{QueryManager, ReadOnlyDbSnapshot};
@@ -68,6 +69,7 @@ where
 pub struct ProverStateUpdate {
     pub(crate) node_batch: NodeBatch,
     pub key_preimages: Vec<(KeyHash, CacheKey)>,
+    pub stale_state: BTreeSet<StaleNodeIndex>,
 }
 
 impl<Q> Storage for ProverStorage<Q>
@@ -184,6 +186,7 @@ where
         let state_update = ProverStateUpdate {
             node_batch: tree_update.node_batch,
             key_preimages,
+            stale_state: tree_update.stale_node_index_batch,
         };
 
         // We need the state diff to be calculated only inside zk context.
@@ -242,6 +245,11 @@ where
             .write_node_batch(&state_update.node_batch)
             .expect("db write must succeed");
 
+        // Write the stale state nodes which will be used by the pruner at a later stage for cleanup.
+        self.db
+            .set_stale_nodes(&state_update.stale_state)
+            .expect("db set stale nodes must succeed");
+
         // Finally, update our in-memory view of the current item numbers
         self.db.inc_next_version();
     }

crates/storage-ops/Cargo.toml

@@ -12,12 +12,14 @@ repository.workspace = true
 # Citrea dependencies
 
 # Sov SDK deps
+jmt = { workspace = true }
 sov-db = { path = "../sovereign-sdk/full-node/db/sov-db" }
 sov-schema-db = { path = "../sovereign-sdk/full-node/db/sov-schema-db" }
 
 # 3rd-party dependencies
 anyhow = { workspace = true }
 futures = { workspace = true }
+hex = { workspace = true }
 serde = { workspace = true }
 tokio = { workspace = true }
 tokio-util = { workspace = true }

crates/storage-ops/src/pruning/components/mod.rs

@@ -4,4 +4,5 @@ mod state_db;
 
 pub(crate) use ledger_db::*;
 pub(crate) use native_db::*;
+#[allow(unused_imports)]
 pub(crate) use state_db::*;

crates/storage-ops/src/pruning/components/native_db.rs

@@ -13,7 +13,7 @@ pub(crate) fn prune_native_db(native_db: Arc<sov_schema_db::DB>, up_to_block: u6
 
     iter.seek_to_first();
 
-    let mut counter = 1u32;
+    let mut counter = 0u32;
     let mut keys_to_delete = vec![];
     while let Some(Ok(entry)) = iter.next() {
         let version = entry.key.1;