Repository: aptos-labs/aptos-core
File: storage/aptosdb/src/db_debugger/state_tree/check_stale_nodes.rs (Rust)

Code before:

        // Step 1: Read pruner progress.
        let p1 = get_progress(
            state_merkle_db.metadata_db(),
            &DbMetadataKey::StateMerklePrunerProgress,
        )?
        .unwrap_or(0);
        let p2 = get_progress(
            state_merkle_db.metadata_db(),
            &DbMetadataKey::EpochEndingStateMerklePrunerProgress,
        )?
        .unwrap_or(0);

        println!("State Merkle Pruner progress: {p1}");
        println!("Epoch Ending State Merkle Pruner progress: {p2}");

        // Step 2: Scan stale node indices across metadata DB and all shard DBs.
        let num_shards = state_merkle_db.hack_num_real_shards();
        let mut dbs: Vec<(&str, &DB)> = vec![("metadata", state_merkle_db.metadata_db())];
        for shard_id in 0..num_shards {
            // Leak the string so we get a &str with the right lifetime — fine for a CLI tool.
            let name: &str = Box::leak(format!("shard_{shard_id}").into_boxed_str());
            dbs.push((name, state_merkle_db.db_shard(shard_id)));
        }

        let mut regular_leaked: Vec<(String, StaleNodeIndex)> = Vec::new();
        let mut cross_epoch_leaked: Vec<(String, StaleNodeIndex)> = Vec::new();

        for &(db_name, db) in &dbs {
            // Scan StaleNodeIndexSchema: entries with stale_since_version <= P1 should have been
            // pruned by the regular pruner.
            {
                let mut iter = db.iter::<StaleNodeIndexSchema>()?;
                iter.seek_to_first();
                for item in iter {
                    let (index, ()) = item?;
                    if index.stale_since_version > p1 {
                        break;
                    }
                    regular_leaked.push((db_name.to_string(), index));
                }
            }

            // Scan StaleNodeIndexCrossEpochSchema: entries with stale_since_version <= P2 should
            // have been pruned by the cross-epoch pruner.
            {
                let mut iter = db.iter::<StaleNodeIndexCrossEpochSchema>()?;
                iter.seek_to_first();
                for item in iter {
                    let (index, ()) = item?;
                    if index.stale_since_version > p2 {

Review comment (performance):
It should be clearer / more efficient to use a counter? Like `map<db_name, num_nodes>` or `map<db_name, vec<stale_node_index>>`? This is repeating the db name for a million times if there are a million leaked nodes. Also not sure if we need to print all the node keys.

Code after:

impl Cmd {
    pub fn run(self) -> Result<()> {
        let state_merkle_db = self.db_dir.open_state_merkle_db()?;

        // Step 1: Read pruner progress.
        let p1 = get_progress(
            state_merkle_db.metadata_db(),
            &DbMetadataKey::StateMerklePrunerProgress,
        )?
        .unwrap_or(0);
        let p2 = get_progress(
            state_merkle_db.metadata_db(),
            &DbMetadataKey::EpochEndingStateMerklePrunerProgress,
        )?
        .unwrap_or(0);

        println!("State Merkle Pruner progress: {p1}");
        println!("Epoch Ending State Merkle Pruner progress: {p2}");

        // Step 2: Collect DB handles.
        let num_shards = state_merkle_db.hack_num_real_shards();
        let mut dbs: Vec<(String, &DB)> =
            vec![("metadata".to_string(), state_merkle_db.metadata_db())];
        for shard_id in 0..num_shards {
            dbs.push((
                format!("shard_{shard_id}"),
                state_merkle_db.db_shard(shard_id),
            ));
        }

        // Step 3: Collect leaked stale node indices per DB.
        let mut regular_leaked: BTreeMap<String, Vec<StaleNodeIndex>> = BTreeMap::new();
        let mut cross_epoch_leaked: BTreeMap<String, Vec<StaleNodeIndex>> = BTreeMap::new();

        for (db_name, db) in &dbs {
            println!("Scanning {db_name} for stale node indices...");

            let regular = collect_stale_indices::<StaleNodeIndexSchema>(db, p1, self.limit_per_db)?;
            let cross_epoch =
                collect_stale_indices::<StaleNodeIndexCrossEpochSchema>(db, p2, self.limit_per_db)?;

            println!(
                "Done scanning {db_name}. regular: {}, cross_epoch: {}",
                regular.len(),
                cross_epoch.len(),
            );

            if !regular.is_empty() {
                regular_leaked.insert(db_name.clone(), regular);
            }