[CP][Preview] State kv db pruner, and restore.

config/src/config/storage_config.rs

@@ -117,6 +117,11 @@ pub const NO_OP_STORAGE_PRUNER_CONFIG: PrunerConfig = PrunerConfig {
         prune_window: 0,
         batch_size: 0,
     },
+    state_kv_pruner_config: StateKvPrunerConfig {
+        enable: false,
+        prune_window: 0,
+        batch_size: 0,
+    },
 };
 
 #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
@@ -140,13 +145,12 @@ pub struct LedgerPrunerConfig {
 #[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
 #[serde(default, deny_unknown_fields)]
 pub struct StateMerklePrunerConfig {
-    /// Boolean to enable/disable the state store pruner. The state pruner is responsible for
-    /// pruning state tree nodes.
+    /// Boolean to enable/disable the state merkle pruner. The state merkle pruner is responsible
+    /// for pruning state tree nodes.
     pub enable: bool,
-    /// The size of the window should be calculated based on disk space availability and system TPS.
+    /// Window size in versions.
     pub prune_window: u64,
-    /// Similar to the variable above but for state store pruner. It means the number of stale
-    /// nodes to prune a time.
+    /// Number of stale nodes to prune a time.
     pub batch_size: usize,
 }
 
@@ -161,6 +165,19 @@ pub struct EpochSnapshotPrunerConfig {
     pub batch_size: usize,
 }
 
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct StateKvPrunerConfig {
+    /// Boolean to enable/disable the state kv pruner. The state pruner is responsible for
+    /// pruning state tree nodes.
+    pub enable: bool,
+    /// Window size in versions.
+    pub prune_window: u64,
+    /// Similar to the variable above but for state kv pruner. It means the number of versions to
+    /// prune a time.
+    pub batch_size: usize,
+}
+
 // Config for the epoch ending state pruner is actually in the same format as the state merkle
 // pruner, but it has it's own type hence separate default values. This converts it to the same
 // type, to use the same pruner implementation (but parameterized on the stale node index DB schema).
@@ -180,6 +197,7 @@ pub struct PrunerConfig {
     pub ledger_pruner_config: LedgerPrunerConfig,
     pub state_merkle_pruner_config: StateMerklePrunerConfig,
     pub epoch_snapshot_pruner_config: EpochSnapshotPrunerConfig,
+    pub state_kv_pruner_config: StateKvPrunerConfig,
 }
 
 impl Default for LedgerPrunerConfig {
@@ -230,6 +248,17 @@ impl Default for EpochSnapshotPrunerConfig {
     }
 }
 
+impl Default for StateKvPrunerConfig {
+    fn default() -> Self {
+        Self {
+            // TODO(grao): Keep it the same as ledger pruner config for now, will revisit later.
+            enable: true,
+            prune_window: 150_000_000,
+            batch_size: 500,
+        }
+    }
+}
+
 impl Default for StorageConfig {
     fn default() -> StorageConfig {
         StorageConfig {

execution/executor-benchmark/src/main.rs

@@ -3,7 +3,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use aptos_config::config::{
-    EpochSnapshotPrunerConfig, LedgerPrunerConfig, PrunerConfig, StateMerklePrunerConfig,
+    EpochSnapshotPrunerConfig, LedgerPrunerConfig, PrunerConfig, StateKvPrunerConfig,
+    StateMerklePrunerConfig,
 };
 use aptos_executor::block_executor::TransactionBlockExecutor;
 use aptos_executor_benchmark::{
@@ -29,6 +30,9 @@ struct PrunerOpt {
     #[structopt(long)]
     enable_ledger_pruner: bool,
 
+    #[structopt(long)]
+    enable_state_kv_pruner: bool,
+
     #[structopt(long, default_value = "100000")]
     state_prune_window: u64,
 
@@ -38,6 +42,9 @@ struct PrunerOpt {
     #[structopt(long, default_value = "100000")]
     ledger_prune_window: u64,
 
+    #[structopt(long, default_value = "100000")]
+    state_kv_prune_window: u64,
+
     #[structopt(long, default_value = "500")]
     ledger_pruning_batch_size: usize,
 
@@ -46,6 +53,9 @@ struct PrunerOpt {
 
     #[structopt(long, default_value = "500")]
     epoch_snapshot_pruning_batch_size: usize,
+
+    #[structopt(long, default_value = "500")]
+    state_kv_pruning_batch_size: usize,
 }
 
 impl PrunerOpt {
@@ -67,6 +77,11 @@ impl PrunerOpt {
                 batch_size: self.ledger_pruning_batch_size,
                 user_pruning_window_offset: 0,
             },
+            state_kv_pruner_config: StateKvPrunerConfig {
+                enable: self.enable_state_kv_pruner,
+                prune_window: self.state_kv_prune_window,
+                batch_size: self.state_kv_pruning_batch_size,
+            },
         }
     }
 }

state-sync/storage-service/server/src/lib.rs

@@ -1092,7 +1092,7 @@ impl StorageReader {
     ) -> Result<Option<CompleteDataRange<Version>>, Error> {
         let pruner_enabled = self
             .storage
-            .is_state_pruner_enabled()
+            .is_state_merkle_pruner_enabled()
             .map_err(|error| Error::StorageErrorEncountered(error.to_string()))?;
         if !pruner_enabled {
             return Ok(*transactions_range);

state-sync/storage-service/server/src/tests.rs

@@ -1658,7 +1658,7 @@ async fn test_get_storage_server_summary() {
         .times(1)
         .returning(move || Ok(state_prune_window));
     db_reader
-        .expect_is_state_pruner_enabled()
+        .expect_is_state_merkle_pruner_enabled()
         .returning(move || Ok(true));
 
     // Create the storage client and server
@@ -3133,7 +3133,7 @@ fn create_mock_db_for_subscription(
         .expect_get_epoch_snapshot_prune_window()
         .returning(move || Ok(100));
     db_reader
-        .expect_is_state_pruner_enabled()
+        .expect_is_state_merkle_pruner_enabled()
         .returning(move || Ok(true));
     db_reader
 }
@@ -3775,6 +3775,6 @@ mock! {
 
         fn get_epoch_snapshot_prune_window(&self) -> Result<usize>;
 
-        fn is_state_pruner_enabled(&self) -> Result<bool>;
+        fn is_state_merkle_pruner_enabled(&self) -> Result<bool>;
     }
 }

storage/aptosdb/src/aptosdb_test.rs

@@ -4,15 +4,16 @@
 use crate::{
     get_first_seq_num_and_limit,
     pruner::{
-        ledger_pruner_manager::LedgerPrunerManager, state_pruner_manager::StatePrunerManager,
+        ledger_pruner_manager::LedgerPrunerManager,
+        state_merkle_pruner_manager::StateMerklePrunerManager,
     },
     test_helper,
     test_helper::{arb_blocks_to_commit, put_as_state_root, put_transaction_info},
     AptosDB, PrunerManager, StaleNodeIndexSchema,
 };
 use aptos_config::config::{
     EpochSnapshotPrunerConfig, LedgerPrunerConfig, PrunerConfig, RocksdbConfigs,
-    StateMerklePrunerConfig, BUFFERED_STATE_TARGET_ITEMS,
+    StateKvPrunerConfig, StateMerklePrunerConfig, BUFFERED_STATE_TARGET_ITEMS,
     DEFAULT_MAX_NUM_NODES_PER_LRU_CACHE_SHARD,
 };
 use aptos_crypto::{hash::CryptoHash, HashValue};
@@ -91,27 +92,24 @@ fn test_pruner_config() {
     let tmp_dir = TempPath::new();
     let aptos_db = AptosDB::new_for_test(&tmp_dir);
     for enable in [false, true] {
-        let state_pruner = StatePrunerManager::<StaleNodeIndexSchema>::new(
+        let state_merkle_pruner = StateMerklePrunerManager::<StaleNodeIndexSchema>::new(
             Arc::clone(&aptos_db.state_merkle_db),
             StateMerklePrunerConfig {
                 enable,
                 prune_window: 20,
                 batch_size: 1,
             },
         );
-        assert_eq!(state_pruner.is_pruner_enabled(), enable);
-        assert_eq!(state_pruner.get_prune_window(), 20);
+        assert_eq!(state_merkle_pruner.is_pruner_enabled(), enable);
+        assert_eq!(state_merkle_pruner.get_prune_window(), 20);
 
-        let ledger_pruner = LedgerPrunerManager::new(
-            Arc::clone(&aptos_db.ledger_db),
-            Arc::clone(&aptos_db.state_store),
-            LedgerPrunerConfig {
+        let ledger_pruner =
+            LedgerPrunerManager::new(Arc::clone(&aptos_db.ledger_db), LedgerPrunerConfig {
                 enable,
                 prune_window: 100,
                 batch_size: 1,
                 user_pruning_window_offset: 0,
-            },
-        );
+            });
         assert_eq!(ledger_pruner.is_pruner_enabled(), enable);
         assert_eq!(ledger_pruner.get_prune_window(), 100);
     }
@@ -123,7 +121,7 @@ fn test_error_if_version_pruned() {
     let db = AptosDB::new_for_test(&tmp_dir);
     db.state_store
         .state_db
-        .state_pruner
+        .state_merkle_pruner
         .testonly_update_min_version(5);
     db.ledger_pruner.testonly_update_min_version(10);
     assert_eq!(
@@ -202,6 +200,11 @@ pub fn test_state_merkle_pruning_impl(
                 prune_window: 10,
                 batch_size: 1,
             },
+            state_kv_pruner_config: StateKvPrunerConfig {
+                enable: true,
+                prune_window: 10,
+                batch_size: 1,
+            },
         },
         RocksdbConfigs::default(),
         false, /* enable_indexer */
@@ -252,7 +255,7 @@ pub fn test_state_merkle_pruning_impl(
             .collect();
 
         // Prune till the oldest snapshot readable.
-        let pruner = &db.state_store.state_db.state_pruner;
+        let pruner = &db.state_store.state_db.state_merkle_pruner;
         let epoch_snapshot_pruner = &db.state_store.state_db.epoch_snapshot_pruner;
         pruner
             .pruner_worker

storage/aptosdb/src/fake_aptosdb.rs

@@ -811,8 +811,8 @@ impl DbReader for FakeAptosDB {
             .get_state_value_chunk_with_proof(version, start_idx, chunk_size)
     }
 
-    fn is_state_pruner_enabled(&self) -> Result<bool> {
-        self.inner.is_state_pruner_enabled()
+    fn is_state_merkle_pruner_enabled(&self) -> Result<bool> {
+        self.inner.is_state_merkle_pruner_enabled()
     }
 
     fn get_epoch_snapshot_prune_window(&self) -> Result<usize> {