[Storage][Sharding] Sharded state merkle pruner.

aptos-labs · May 30, 2023 · a438cea · a438cea
1 parent c77bca8
commit a438cea
Show file tree

Hide file tree

Showing 5 changed files with 208 additions and 128 deletions.
diff --git a/storage/aptosdb/src/lib.rs b/storage/aptosdb/src/lib.rs
@@ -57,7 +57,7 @@ use crate::{
         ledger_store::ledger_store_pruner::LedgerPruner, pruner_manager::PrunerManager,
         pruner_utils, state_kv_pruner::StateKvPruner,
         state_kv_pruner_manager::StateKvPrunerManager,
-        state_merkle_pruner_manager::StateMerklePrunerManager, state_store::StateMerklePruner,
+        state_merkle_pruner_manager::StateMerklePrunerManager,
     },
     schema::*,
     stale_node_index::StaleNodeIndexSchema,
@@ -2176,12 +2176,8 @@ impl DbWriter for AptosDB {
                 .pruner()
                 .save_min_readable_version(version, &batch)?;
 
-            let mut state_merkle_batch = SchemaBatch::new();
-            StateMerklePruner::prune_genesis(
-                self.state_merkle_db.clone(),
-                &mut state_merkle_batch,
-            )?;
-
+            // TODO(joshlind): Figure out a way to not write genesis.
+            let state_merkle_batch = SchemaBatch::new();
             self.state_store
                 .state_merkle_pruner
                 .pruner()

diff --git a/storage/aptosdb/src/pruner/state_store/generics.rs b/storage/aptosdb/src/pruner/state_store/generics.rs
@@ -12,13 +12,17 @@ pub trait StaleNodeIndexSchemaTrait: Schema<Key = StaleNodeIndex>
 where
     StaleNodeIndex: KeyCodec<Self>,
 {
-    fn tag() -> DbMetadataKey;
+    fn tag(shard_id: Option<u8>) -> DbMetadataKey;
     fn name() -> &'static str;
 }
 
 impl StaleNodeIndexSchemaTrait for StaleNodeIndexSchema {
-    fn tag() -> DbMetadataKey {
-        DbMetadataKey::StateMerklePrunerProgress
+    fn tag(shard_id: Option<u8>) -> DbMetadataKey {
+        if let Some(shard_id) = shard_id {
+            DbMetadataKey::StateMerkleShardPrunerProgress(shard_id as usize)
+        } else {
+            DbMetadataKey::StateMerklePrunerProgress
+        }
     }
 
     fn name() -> &'static str {
@@ -27,8 +31,12 @@ impl StaleNodeIndexSchemaTrait for StaleNodeIndexSchema {
 }
 
 impl StaleNodeIndexSchemaTrait for StaleNodeIndexCrossEpochSchema {
-    fn tag() -> DbMetadataKey {
-        DbMetadataKey::EpochEndingStateMerklePrunerProgress
+    fn tag(shard_id: Option<u8>) -> DbMetadataKey {
+        if let Some(shard_id) = shard_id {
+            DbMetadataKey::EpochEndingStateMerkleShardPrunerProgress(shard_id as usize)
+        } else {
+            DbMetadataKey::EpochEndingStateMerklePrunerProgress
+        }
     }
 
     fn name() -> &'static str {