[storage] remove leaf count migration code

Taking the chance of a new start.

Closes: aptos-labs#299

aptos-move/aptos-validator-interface/src/storage_interface.rs

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

aptos-node/src/lib.rs

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

config/management/genesis/src/verify.rs

@@ -214,7 +214,6 @@ fn compute_genesis(
         false,
         NO_OP_STORAGE_PRUNER_CONFIG,
         RocksdbConfig::default(),
-        true, /* account_count_migration */
     )
     .map_err(|e| Error::UnexpectedError(e.to_string()))?;
     let db_rw = DbReaderWriter::new(aptosdb);

config/management/genesis/src/waypoint.rs

@@ -45,7 +45,6 @@ pub fn create_genesis_waypoint(genesis: &Transaction) -> Result<Waypoint, Error>
         false,
         NO_OP_STORAGE_PRUNER_CONFIG,
         RocksdbConfig::default(),
-        true, /* account_count_migration */
     )
     .map_err(|e| Error::UnexpectedError(e.to_string()))?;
     let db_rw = DbReaderWriter::new(aptosdb);

config/src/config/storage_config.rs

@@ -46,13 +46,6 @@ 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)]
-    pub account_count_migration: bool,
 }
 
 pub const NO_OP_STORAGE_PRUNER_CONFIG: StoragePrunerConfig = StoragePrunerConfig {
@@ -113,7 +106,6 @@ impl Default for StorageConfig {
             // Default read/write/connection timeout, in milliseconds
             timeout_ms: 30_000,
             rocksdb_config: RocksdbConfig::default(),
-            account_count_migration: false,
         }
     }
 }

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

@@ -33,9 +33,6 @@ struct Opt {
 
     #[structopt(long, requires("waypoint-to-verify"))]
     commit: bool,
-
-    #[structopt(long)]
-    account_count_migration: bool,
 }
 
 fn main() -> Result<()> {
@@ -56,7 +53,6 @@ fn main() -> Result<()> {
             false,
             NO_OP_STORAGE_PRUNER_CONFIG, /* 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

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

execution/executor-benchmark/src/lib.rs

@@ -31,7 +31,6 @@ pub fn init_db_and_executor(config: &NodeConfig) -> (Arc<dyn DbReader>, BlockExe
             false,                       /* readonly */
             NO_OP_STORAGE_PRUNER_CONFIG, /* pruner */
             RocksdbConfig::default(),
-            true, /* account_count_migration */
         )
         .expect("DB should open."),
     );
@@ -60,7 +59,6 @@ pub fn run_benchmark(
         true,                        /* readonly */
         NO_OP_STORAGE_PRUNER_CONFIG, /* pruner */
         RocksdbConfig::default(),
-        true, /* account_count_migration */
     )
     .expect("db open failure.")
     .create_checkpoint(checkpoint_dir.as_ref().join("aptosdb"))

execution/executor/tests/db_bootstrapper_test.rs

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

state-sync/state-sync-v2/state-sync-multiplexer/src/lib.rs

@@ -201,7 +201,6 @@ mod tests {
             false,
             NO_OP_STORAGE_PRUNER_CONFIG,
             RocksdbConfig::default(),
-            true,
         )
         .unwrap();
         let (_, db_rw) = DbReaderWriter::wrap(db);

storage/aptosdb/src/aptossum/mod.rs

@@ -25,7 +25,6 @@ impl Aptossum {
             true,                        /* read only */
             NO_OP_STORAGE_PRUNER_CONFIG, /* no prune_window */
             RocksdbConfig::default(),
-            true, /* account_count_migration, ignored anyway */
         )?;
         Ok(Aptossum { db })
     }

storage/aptosdb/src/backup/restore_handler.rs

@@ -54,13 +54,11 @@ 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,
         )
     }
 

storage/aptosdb/src/lib.rs

@@ -249,11 +249,7 @@ impl AptosDB {
         ]
     }
 
-    fn new_with_db(
-        db: DB,
-        storage_pruner_config: StoragePrunerConfig,
-        account_count_migration: bool,
-    ) -> Self {
+    fn new_with_db(db: DB, storage_pruner_config: StoragePrunerConfig) -> Self {
         let db = Arc::new(db);
         let transaction_store = Arc::new(TransactionStore::new(Arc::clone(&db)));
         let event_store = Arc::new(EventStore::new(Arc::clone(&db)));
@@ -264,7 +260,7 @@ impl AptosDB {
             db: Arc::clone(&db),
             event_store: Arc::clone(&event_store),
             ledger_store: Arc::clone(&ledger_store),
-            state_store: Arc::new(StateStore::new(Arc::clone(&db), account_count_migration)),
+            state_store: Arc::new(StateStore::new(Arc::clone(&db))),
             transaction_store: Arc::clone(&transaction_store),
             system_store: Arc::clone(&system_store),
             rocksdb_property_reporter: RocksdbPropertyReporter::new(Arc::clone(&db)),
@@ -286,7 +282,6 @@ impl AptosDB {
         readonly: bool,
         storage_pruner_config: StoragePrunerConfig,
         rocksdb_config: RocksdbConfig,
-        account_count_migration: bool, // ignored when opening readonly
     ) -> Result<Self> {
         ensure!(
             storage_pruner_config.eq(&NO_OP_STORAGE_PRUNER_CONFIG) || !readonly,
@@ -298,31 +293,25 @@ impl AptosDB {
 
         let mut rocksdb_opts = gen_rocksdb_options(&rocksdb_config);
 
-        let (db, account_count_migration) = if readonly {
-            (
-                DB::open_readonly(
-                    path.clone(),
-                    "aptosdb_ro",
-                    Self::column_families(),
-                    &rocksdb_opts,
-                )?,
-                true,
-            )
+        let db = if readonly {
+            DB::open_readonly(
+                path.clone(),
+                "aptosdb_ro",
+                Self::column_families(),
+                &rocksdb_opts,
+            )?
         } else {
             rocksdb_opts.create_if_missing(true);
             rocksdb_opts.create_missing_column_families(true);
-            (
-                DB::open(
-                    path.clone(),
-                    "aptosdb",
-                    Self::column_families(),
-                    &rocksdb_opts,
-                )?,
-                account_count_migration,
-            )
+            DB::open(
+                path.clone(),
+                "aptosdb",
+                Self::column_families(),
+                &rocksdb_opts,
+            )?
         };
 
-        let ret = Self::new_with_db(db, storage_pruner_config, account_count_migration);
+        let ret = Self::new_with_db(db, storage_pruner_config);
         info!(
             path = path,
             time_ms = %instant.elapsed().as_millis(),
@@ -351,7 +340,6 @@ impl AptosDB {
                 &rocksdb_opts,
             )?,
             NO_OP_STORAGE_PRUNER_CONFIG,
-            true, // account_count_migration
         ))
     }
 
@@ -363,7 +351,6 @@ impl AptosDB {
             false,                       /* readonly */
             NO_OP_STORAGE_PRUNER_CONFIG, /* pruner */
             RocksdbConfig::default(),
-            true, /* account_count_migration */
         )
         .expect("Unable to open AptosDB")
     }
@@ -1198,14 +1185,7 @@ impl DbReader for AptosDB {
 
     fn get_account_count(&self, version: Version) -> Result<usize> {
         gauged_api("get_account_count", || {
-            self.state_store
-                .get_account_count(version)
-                .and_then(|count| {
-                    count.ok_or_else(|| {
-                        AptosDbError::NotFound(format!("Account count at version {}", version))
-                            .into()
-                    })
-                })
+            self.state_store.get_account_count(version)
         })
     }
 
@@ -1343,8 +1323,7 @@ impl DbWriter for AptosDB {
                 DIEM_STORAGE_LATEST_ACCOUNT_COUNT.set(
                     self.state_store
                         .get_account_count(last_version)
-                        .map(|opt| opt.map_or(-1, |n| n as i64))
-                        .unwrap_or(-2),
+                        .map_or(-1, |c| c as i64),
                 );
 
                 self.wake_pruner(last_version);

storage/aptosdb/src/pruner/state_store/test.rs

@@ -49,7 +49,7 @@ fn test_state_store_pruner() {
     let tmp_dir = TempPath::new();
     let aptos_db = AptosDB::new_for_test(&tmp_dir);
     let db = aptos_db.db;
-    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);
+    let state_store = &StateStore::new(Arc::clone(&db));
     let transaction_store = &aptos_db.transaction_store;
     let pruner = Pruner::new(
         Arc::clone(&db),
@@ -137,7 +137,7 @@ fn test_worker_quit_eagerly() {
     let tmp_dir = TempPath::new();
     let aptos_db = AptosDB::new_for_test(&tmp_dir);
     let db = aptos_db.db;
-    let state_store = &StateStore::new(Arc::clone(&db), true /* account_count_migration */);
+    let state_store = &StateStore::new(Arc::clone(&db));
 
     let _root0 = put_account_state_set(
         &db,

storage/aptosdb/src/state_store/mod.rs

@@ -39,15 +39,11 @@ type NodeBatch = aptos_jellyfish_merkle::NodeBatch<AccountStateBlob>;
 #[derive(Debug)]
 pub(crate) struct StateStore {
     db: Arc<DB>,
-    pub account_count_migration: bool,
 }
 
 impl StateStore {
-    pub fn new(db: Arc<DB>, account_count_migration: bool) -> Self {
-        Self {
-            db,
-            account_count_migration,
-        }
+    pub fn new(db: Arc<DB>) -> Self {
+        Self { db }
     }
 
     /// Get the account state blob given account address and root hash of state Merkle tree
@@ -59,8 +55,7 @@ impl StateStore {
         Option<AccountStateBlob>,
         SparseMerkleProof<AccountStateBlob>,
     )> {
-        JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-            .get_with_proof(address.hash(), version)
+        JellyfishMerkleTree::new(self).get_with_proof(address.hash(), version)
     }
 
     /// Gets the proof that proves a range of accounts.
@@ -69,8 +64,7 @@ impl StateStore {
         rightmost_key: HashValue,
         version: Version,
     ) -> Result<SparseMerkleRangeProof> {
-        JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-            .get_range_proof(rightmost_key, version)
+        JellyfishMerkleTree::new(self).get_range_proof(rightmost_key, version)
     }
 
     /// Put the results generated by `account_state_sets` to `batch` and return the result root
@@ -92,9 +86,8 @@ impl StateStore {
             })
             .collect::<Vec<_>>();
 
-        let (new_root_hash_vec, tree_update_batch) =
-            JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-                .batch_put_value_sets(blob_sets, node_hashes, first_version)?;
+        let (new_root_hash_vec, tree_update_batch) = JellyfishMerkleTree::new(self)
+            .batch_put_value_sets(blob_sets, node_hashes, first_version)?;
 
         let num_versions = new_root_hash_vec.len();
         assert_eq!(num_versions, tree_update_batch.node_stats.len());
@@ -122,13 +115,11 @@ impl StateStore {
     }
 
     pub fn get_root_hash(&self, version: Version) -> Result<HashValue> {
-        JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-            .get_root_hash(version)
+        JellyfishMerkleTree::new(self).get_root_hash(version)
     }
 
     pub fn get_root_hash_option(&self, version: Version) -> Result<Option<HashValue>> {
-        JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-            .get_root_hash_option(version)
+        JellyfishMerkleTree::new(self).get_root_hash_option(version)
     }
 
     /// Finds the rightmost leaf by scanning the entire DB.
@@ -157,9 +148,8 @@ impl StateStore {
         Ok(ret)
     }
 
-    pub fn get_account_count(&self, version: Version) -> Result<Option<usize>> {
-        JellyfishMerkleTree::new_migration(self, self.account_count_migration)
-            .get_leaf_count(version)
+    pub fn get_account_count(&self, version: Version) -> Result<usize> {
+        JellyfishMerkleTree::new(self).get_leaf_count(version)
     }
 
     pub fn get_account_chunk_with_proof(
@@ -203,7 +193,6 @@ impl StateStore {
             Arc::clone(self),
             version,
             expected_root_hash,
-            self.account_count_migration,
         )?))
     }
 }