[Storage] Also write usage for the last checkpoint in the chunk. (apt…

…os-labs#10917)
zi0Black · Nov 16, 2023 · d813b32 · d813b32
1 parent 71137ee
commit d813b32
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 5 deletions.
diff --git a/storage/aptosdb/src/lib.rs b/storage/aptosdb/src/lib.rs
@@ -1026,6 +1026,9 @@ impl AptosDB {
             // TODO(grao): remove after APIs migrated off the DB to the indexer.
             self.state_store.state_kv_db.enabled_sharding(),
             skip_index_and_usage,
+            txns_to_commit
+                .iter()
+                .rposition(|txn| txn.is_state_checkpoint()),
         )?;
 
         // Write block index if event index is skipped.

diff --git a/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs b/storage/aptosdb/src/pruner/state_merkle_pruner/test.rs
@@ -64,6 +64,7 @@ fn put_value_set(
             &state_kv_metadata_batch,
             /*put_state_value_indices=*/ false,
             /*skip_usage=*/ false,
+            /*last_checkpoint_index=*/ None,
         )
         .unwrap();
     state_store

diff --git a/storage/aptosdb/src/state_store/mod.rs b/storage/aptosdb/src/state_store/mod.rs
@@ -651,6 +651,7 @@ impl StateStore {
             batch,
             sharded_state_kv_batches,
             /*skip_usage=*/ false,
+            None,
         )?;
 
         self.put_state_values(
@@ -676,6 +677,7 @@ impl StateStore {
         state_kv_metadata_batch: &SchemaBatch,
         put_state_value_indices: bool,
         skip_usage: bool,
+        last_checkpoint_index: Option<usize>,
     ) -> Result<()> {
         let _timer = OTHER_TIMERS_SECONDS
             .with_label_values(&["put_value_sets"])
@@ -689,6 +691,7 @@ impl StateStore {
             ledger_batch,
             sharded_state_kv_batches,
             skip_usage,
+            last_checkpoint_index,
         )?;
 
         let _timer = OTHER_TIMERS_SECONDS
@@ -775,6 +778,7 @@ impl StateStore {
         batch: &SchemaBatch,
         sharded_state_kv_batches: &ShardedStateKvSchemaBatch,
         skip_usage: bool,
+        last_checkpoint_index: Option<usize>,
     ) -> Result<()> {
         let _timer = OTHER_TIMERS_SECONDS
             .with_label_values(&["put_stats_and_indices"])
@@ -907,9 +911,11 @@ impl StateStore {
                 (usage.items() as i64 + items_delta) as usize,
                 (usage.bytes() as i64 + bytes_delta) as usize,
             );
-            if !skip_usage || i == num_versions - 1 {
+            let should_write_index_for_version =
+                (i == num_versions - 1) || Some(i) == last_checkpoint_index;
+            if !skip_usage || should_write_index_for_version {
                 let version = first_version + i as u64;
-                if i == num_versions - 1 {
+                if should_write_index_for_version {
                     info!("Write usage at version {version}, {usage:?}, skip_usage: {skip_usage}.");
                 }
                 batch

diff --git a/storage/aptosdb/src/state_store/state_merkle_batch_committer.rs b/storage/aptosdb/src/state_store/state_merkle_batch_committer.rs
@@ -94,9 +94,8 @@ impl StateMerkleBatchCommitter {
                         .epoch_snapshot_pruner
                         .maybe_set_pruner_target_db_version(current_version);
 
-                    if !self.state_db.skip_usage {
-                        self.check_usage_consistency(&state_delta).unwrap();
-                    }
+                    self.check_usage_consistency(&state_delta).unwrap();
+
                     state_delta.base.log_generation("buffered_state_commit");
                     state_delta
                         .current

diff --git a/storage/aptosdb/src/state_store/state_store_test.rs b/storage/aptosdb/src/state_store/state_store_test.rs
@@ -59,6 +59,7 @@ fn put_value_set(
             &state_kv_metadata_batch,
             /*put_state_value_indices=*/ false,
             /*skip_usage=*/ false,
+            /*last_checkpoint_index=*/ None,
         )
         .unwrap();
     state_store

diff --git a/storage/aptosdb/src/test_helper.rs b/storage/aptosdb/src/test_helper.rs
@@ -78,6 +78,7 @@ pub(crate) fn update_store(
                 &state_kv_metadata_batch,
                 /*put_state_value_indices=*/ false,
                 /*skip_usage=*/ false,
+                /*last_checkpoint_index=*/ None,
             )
             .unwrap();
         store