[storage] StateDB migration towards supporting get_account_count.

A new storage config (`account_count_migration`) is put in place, upon which the state db starts to write down JellifishMerkleTree internal nodes in a new format in which leaf counts on children will be recorded if available. Once enabled, the DB is incompatible with older versions of the node. Once all accounts get updated at least once, the leaf count should be accessible on the root node, which indicates the completion of the migration.

Closes: aptos-labs#9304

config/management/genesis/src/verify.rs

@@ -214,8 +214,14 @@ fn compute_genesis(
     genesis_path: &Path,
     db_path: &Path,
 ) -> Result<(DbReaderWriter, Waypoint), Error> {
-    let diemdb = DiemDB::open(db_path, false, None, RocksdbConfig::default())
-        .map_err(|e| Error::UnexpectedError(e.to_string()))?;
+    let diemdb = DiemDB::open(
+        db_path,
+        false,
+        None,
+        RocksdbConfig::default(),
+        true, /* account_count_migration */
+    )
+    .map_err(|e| Error::UnexpectedError(e.to_string()))?;
     let db_rw = DbReaderWriter::new(diemdb);
 
     let mut file = File::open(genesis_path)

config/management/genesis/src/waypoint.rs

@@ -40,8 +40,14 @@ impl CreateWaypoint {
 
 pub fn create_genesis_waypoint(genesis: &Transaction) -> Result<Waypoint, Error> {
     let path = TempPath::new();
-    let diemdb = DiemDB::open(&path, false, None, RocksdbConfig::default())
-        .map_err(|e| Error::UnexpectedError(e.to_string()))?;
+    let diemdb = DiemDB::open(
+        &path,
+        false,
+        None,
+        RocksdbConfig::default(),
+        true, /* account_count_migration */
+    )
+    .map_err(|e| Error::UnexpectedError(e.to_string()))?;
     let db_rw = DbReaderWriter::new(diemdb);
 
     db_bootstrapper::generate_waypoint::<DiemVM>(&db_rw, genesis)

config/src/config/storage_config.rs

@@ -48,6 +48,13 @@ pub struct StorageConfig {
     pub timeout_ms: u64,
     /// Rocksdb-specific configurations
     pub rocksdb_config: RocksdbConfig,
+    /// If enabled, leaf counts of the children of a JellyfishMerkelTree internal node are written
+    /// down.
+    /// This is supposed to be enabled after other aspects of the new release has been confirmed
+    /// safe. And once enabled, an older node won't be able to read the state DB, unless the DB is
+    /// wiped and re-synced.
+    #[serde(default, skip_serializing)]
+    pub account_count_migration: bool,
 }
 
 impl Default for StorageConfig {
@@ -68,6 +75,7 @@ impl Default for StorageConfig {
             // Default read/write/connection timeout, in milliseconds
             timeout_ms: 30_000,
             rocksdb_config: RocksdbConfig::default(),
+            account_count_migration: false,
         }
     }
 }

diem-node/src/lib.rs

@@ -281,6 +281,7 @@ pub fn setup_environment(node_config: &NodeConfig, logger: Option<Arc<Logger>>)
             false, /* readonly */
             node_config.storage.prune_window,
             node_config.storage.rocksdb_config,
+            node_config.storage.account_count_migration,
         )
         .expect("DB should open."),
     );

execution/db-bootstrapper/src/bin/db-bootstrapper.rs

@@ -33,6 +33,9 @@ struct Opt {
 
     #[structopt(long, requires("waypoint-to-verify"))]
     commit: bool,
+
+    #[structopt(long)]
+    account_count_migration: bool,
 }
 
 fn main() -> Result<()> {
@@ -53,6 +56,7 @@ fn main() -> Result<()> {
             false,
             None, /* pruner */
             RocksdbConfig::default(),
+            opt.account_count_migration,
         )
     } else {
         // When not committing, we open the DB as secondary so the tool is usable along side a

execution/executor-benchmark/src/db_generator.rs

@@ -46,6 +46,7 @@ pub fn run(
             false,        /* readonly */
             prune_window, /* pruner */
             RocksdbConfig::default(),
+            true, /* account_count_migration */
         )
         .expect("DB should open."),
     );

execution/executor-benchmark/src/lib.rs

@@ -34,6 +34,7 @@ pub fn create_storage_service_and_executor(
             false, /* readonly */
             None,  /* pruner */
             RocksdbConfig::default(),
+            true, /* account_count_migration */
         )
         .expect("DB should open."),
     );
@@ -66,6 +67,7 @@ pub fn run_benchmark(
             false, /* readonly */
             None,  /* pruner */
             RocksdbConfig::default(),
+            true, /* account_count_migration */
         )
         .expect("db open failure.")
         .create_checkpoint(checkpoint_dir.as_ref().join("diemdb"))

execution/executor/tests/db_bootstrapper_test.rs

@@ -242,7 +242,9 @@ fn restore_state_to_db(
     version: Version,
 ) {
     let rh = db.get_restore_handler();
-    let mut receiver = rh.get_state_restore_receiver(version, root_hash).unwrap();
+    let mut receiver = rh
+        .get_state_restore_receiver(version, root_hash, true /* account_count_migration */)
+        .unwrap();
     for (chunk, proof) in vec![(accounts, proof)].into_iter() {
         receiver.add_chunk(chunk, proof).unwrap();
     }

language/diem-tools/diem-validator-interface/src/storage_interface.rs

@@ -25,6 +25,7 @@ impl DBDebuggerInterface {
             true,
             None,
             RocksdbConfig::default(),
+            true, /* account_count_migration, ignored anyway */
         )?)))
     }
 }

storage/backup/backup-cli/src/backup_types/epoch_ending/tests.rs

@@ -15,7 +15,7 @@ use crate::{
     },
 };
 use backup_service::start_backup_service;
-use diem_config::{config::RocksdbConfig, utils::get_available_port};
+use diem_config::utils::get_available_port;
 use diem_temppath::TempPath;
 use diem_types::{
     ledger_info::LedgerInfoWithSignatures,
@@ -83,6 +83,7 @@ fn end_to_end() {
                 trusted_waypoints: TrustedWaypointOpt::default(),
                 rocksdb_opt: RocksdbOpt::default(),
                 concurernt_downloads: ConcurrentDownloadsOpt::default(),
+                account_count_migration: true,
             }
             .try_into()
             .unwrap(),
@@ -102,13 +103,7 @@ fn end_to_end() {
         .map(|li| li.ledger_info().next_block_epoch())
         .unwrap_or(0);
 
-    let tgt_db = DiemDB::open(
-        &tgt_db_dir,
-        true, /* read_only */
-        None, /* pruner */
-        RocksdbConfig::default(),
-    )
-    .unwrap();
+    let tgt_db = DiemDB::new_for_test(&tgt_db_dir);
     assert_eq!(
         tgt_db
             .get_epoch_ending_ledger_infos(0, target_version_next_block_epoch)
@@ -220,6 +215,7 @@ async fn test_trusted_waypoints_impl(
             trusted_waypoints: TrustedWaypointOpt::default(),
             rocksdb_opt: RocksdbOpt::default(),
             concurernt_downloads: ConcurrentDownloadsOpt::default(),
+            account_count_migration: true,
         }
         .try_into()
         .unwrap(),
@@ -240,6 +236,7 @@ async fn test_trusted_waypoints_impl(
             },
             rocksdb_opt: RocksdbOpt::default(),
             concurernt_downloads: ConcurrentDownloadsOpt::default(),
+            account_count_migration: true,
         }
         .try_into()
         .unwrap(),

storage/backup/backup-cli/src/backup_types/state_snapshot/restore.rs

@@ -50,6 +50,7 @@ pub struct StateSnapshotRestoreController {
     /// nothing will be done, otherwise, this has no effect.
     target_version: Version,
     epoch_history: Option<Arc<EpochHistory>>,
+    account_count_migration: bool,
 }
 
 impl StateSnapshotRestoreController {
@@ -66,6 +67,7 @@ impl StateSnapshotRestoreController {
             manifest_handle: opt.manifest_handle,
             target_version: global_opt.target_version,
             epoch_history,
+            account_count_migration: global_opt.account_count_migration,
         }
     }
 
@@ -110,9 +112,11 @@ impl StateSnapshotRestoreController {
             epoch_history.verify_ledger_info(&li)?;
         }
 
-        let mut receiver = self
-            .run_mode
-            .get_state_restore_receiver(self.version, manifest.root_hash)?;
+        let mut receiver = self.run_mode.get_state_restore_receiver(
+            self.version,
+            manifest.root_hash,
+            self.account_count_migration,
+        )?;
 
         let (ver_gauge, tgt_leaf_idx, leaf_idx) = if self.run_mode.is_verify() {
             (

storage/backup/backup-cli/src/backup_types/state_snapshot/tests.rs

@@ -13,7 +13,6 @@ use crate::{
         ConcurrentDownloadsOpt, GlobalBackupOpt, GlobalRestoreOpt, RocksdbOpt, TrustedWaypointOpt,
     },
 };
-use diem_config::config::RocksdbConfig;
 use diem_temppath::TempPath;
 use diem_types::transaction::PRE_GENESIS_VERSION;
 use diemdb::DiemDB;
@@ -67,6 +66,7 @@ fn end_to_end() {
                 trusted_waypoints: TrustedWaypointOpt::default(),
                 rocksdb_opt: RocksdbOpt::default(),
                 concurernt_downloads: ConcurrentDownloadsOpt::default(),
+                account_count_migration: true,
             }
             .try_into()
             .unwrap(),
@@ -77,13 +77,7 @@ fn end_to_end() {
     )
     .unwrap();
 
-    let tgt_db = DiemDB::open(
-        &tgt_db_dir,
-        true, /* read_only */
-        None, /* pruner */
-        RocksdbConfig::default(),
-    )
-    .unwrap();
+    let tgt_db = DiemDB::new_for_test(&tgt_db_dir);
     assert_eq!(
         tgt_db
             .get_latest_tree_state()

storage/backup/backup-cli/src/backup_types/tests.rs

@@ -19,7 +19,6 @@ use crate::{
         RocksdbOpt, TrustedWaypointOpt,
     },
 };
-use diem_config::config::RocksdbConfig;
 use diem_temppath::TempPath;
 use diem_types::transaction::Version;
 use diemdb::DiemDB;
@@ -114,6 +113,7 @@ fn test_end_to_end_impl(d: TestData) {
         trusted_waypoints: TrustedWaypointOpt::default(),
         rocksdb_opt: RocksdbOpt::default(),
         concurernt_downloads: ConcurrentDownloadsOpt::default(),
+        account_count_migration: true,
     }
     .try_into()
     .unwrap();
@@ -149,13 +149,7 @@ fn test_end_to_end_impl(d: TestData) {
     .unwrap();
 
     // Check
-    let tgt_db = DiemDB::open(
-        &tgt_db_dir,
-        false, /* read_only */
-        None,  /* pruner */
-        RocksdbConfig::default(),
-    )
-    .unwrap();
+    let tgt_db = DiemDB::new_for_test(&tgt_db_dir);
     assert_eq!(
         d.db.get_transactions(
             d.txn_start_ver,

storage/backup/backup-cli/src/backup_types/transaction/tests.rs

@@ -13,7 +13,6 @@ use crate::{
         ConcurrentDownloadsOpt, GlobalBackupOpt, GlobalRestoreOpt, RocksdbOpt, TrustedWaypointOpt,
     },
 };
-use diem_config::config::RocksdbConfig;
 use diem_temppath::TempPath;
 use diem_types::transaction::Version;
 use diemdb::DiemDB;
@@ -85,6 +84,7 @@ fn end_to_end() {
                 trusted_waypoints: TrustedWaypointOpt::default(),
                 rocksdb_opt: RocksdbOpt::default(),
                 concurernt_downloads: ConcurrentDownloadsOpt::default(),
+                account_count_migration: true,
             }
             .try_into()
             .unwrap(),
@@ -98,13 +98,7 @@ fn end_to_end() {
     // We don't write down any ledger infos when recovering transactions. State-sync needs to take
     // care of it before running consensus. The latest transactions are deemed "synced" instead of
     // "committed" most likely.
-    let tgt_db = DiemDB::open(
-        &tgt_db_dir,
-        true, /* read_only */
-        None, /* pruner */
-        RocksdbConfig::default(),
-    )
-    .unwrap();
+    let tgt_db = DiemDB::new_for_test(&tgt_db_dir);
     assert_eq!(
         tgt_db
             .get_latest_transaction_info_option()

storage/backup/backup-cli/src/bin/replay-verify.rs

@@ -58,6 +58,7 @@ async fn main_impl() -> Result<()> {
         false, /* read_only */
         None,  /* pruner */
         opt.rocksdb_opt.into(),
+        true, /* account_count_migration */
     )?)
     .get_restore_handler();
     ReplayVerifyCoordinator::new(

storage/backup/backup-cli/src/coordinators/replay_verify.rs

@@ -96,6 +96,7 @@ impl ReplayVerifyCoordinator {
                 restore_handler: self.restore_handler,
             }),
             concurrent_downloads: self.concurrent_downloads,
+            account_count_migration: true,
         };
 
         if let Some(backup) = state_snapshot {

storage/backup/backup-cli/src/coordinators/verify.rs

@@ -79,6 +79,7 @@ impl VerifyCoordinator {
             trusted_waypoints: Arc::new(self.trusted_waypoints_opt.verify()?),
             run_mode: Arc::new(RestoreRunMode::Verify),
             concurrent_downloads: self.concurrent_downloads,
+            account_count_migration: true,
         };
 
         let epoch_history = Arc::new(

storage/backup/backup-cli/src/utils/mod.rs

@@ -92,6 +92,13 @@ pub struct GlobalRestoreOpt {
 
     #[structopt(flatten)]
     pub concurernt_downloads: ConcurrentDownloadsOpt,
+
+    #[structopt(
+        long,
+        help = "If StateDB is written into, write in the new format supporting account counting, \
+        but incompatible with older node versions."
+    )]
+    pub account_count_migration: bool,
 }
 
 pub enum RestoreRunMode {
@@ -126,15 +133,19 @@ impl RestoreRunMode {
         &self,
         version: Version,
         expected_root_hash: HashValue,
+        account_count_migration: bool,
     ) -> Result<JellyfishMerkleRestore<AccountStateBlob>> {
         match self {
-            Self::Restore { restore_handler } => {
-                restore_handler.get_state_restore_receiver(version, expected_root_hash)
-            }
+            Self::Restore { restore_handler } => restore_handler.get_state_restore_receiver(
+                version,
+                expected_root_hash,
+                account_count_migration,
+            ),
             Self::Verify => JellyfishMerkleRestore::new_overwrite(
                 Arc::new(MockTreeWriter),
                 version,
                 expected_root_hash,
+                account_count_migration,
             ),
         }
     }
@@ -146,6 +157,7 @@ pub struct GlobalRestoreOptions {
     pub trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
     pub run_mode: Arc<RestoreRunMode>,
     pub concurrent_downloads: usize,
+    pub account_count_migration: bool,
 }
 
 impl TryFrom<GlobalRestoreOpt> for GlobalRestoreOptions {
@@ -160,6 +172,7 @@ impl TryFrom<GlobalRestoreOpt> for GlobalRestoreOptions {
                 false, /* read_only */
                 None,  /* pruner */
                 opt.rocksdb_opt.into(),
+                opt.account_count_migration,
             )?)
             .get_restore_handler();
             RestoreRunMode::Restore { restore_handler }
@@ -171,6 +184,7 @@ impl TryFrom<GlobalRestoreOpt> for GlobalRestoreOptions {
             trusted_waypoints: Arc::new(opt.trusted_waypoints.verify()?),
             run_mode: Arc::new(run_mode),
             concurrent_downloads,
+            account_count_migration: opt.account_count_migration,
         })
     }
 }

storage/diemdb/src/backup/restore_handler.rs

@@ -54,11 +54,13 @@ impl RestoreHandler {
         &self,
         version: Version,
         expected_root_hash: HashValue,
+        account_count_migration: bool,
     ) -> Result<JellyfishMerkleRestore<AccountStateBlob>> {
         JellyfishMerkleRestore::new_overwrite(
             Arc::clone(&self.state_store),
             version,
             expected_root_hash,
+            account_count_migration,
         )
     }