update save file naming and remove unnecessary merkle_root_dir arg in favor of save_path

tomjohn1028 · tomjohn1028 · commit e7fdd013f9c7 · 2025-02-12T19:55:11.000-05:00
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -178,7 +178,7 @@ jobs:
           name: jito_tip_router_program.so
           path: integration_tests/tests/fixtures/
       - uses: taiki-e/install-action@nextest
-      - run: cargo nextest run --all-features -E 'not test(ledger_utils::tests) and not test(test_meta_merkle_creation_from_ledger)'
+      - run: cargo nextest run --all-features -E 'not test(ledger_utils::tests)'
         env:
           SBF_OUT_DIR: ${{ github.workspace }}/integration_tests/tests/fixtures
 
diff --git a/format.sh b/format.sh
@@ -14,7 +14,7 @@ cargo fmt --all
 
 print_executing "cargo nextest run --all-features"
 cargo build-sbf --sbf-out-dir integration_tests/tests/fixtures
-SBF_OUT_DIR=integration_tests/tests/fixtures cargo nextest run --all-features -E 'not test(ledger_utils::tests::test_get_bank_from_ledger_success) and not test(test_meta_merkle_creation_from_ledger)'
+SBF_OUT_DIR=integration_tests/tests/fixtures cargo nextest run --all-features -E 'not test(ledger_utils::tests::test_get_bank_from_ledger_success)'
 
 # Code coverage only runs with flag
 if [[ "$*" == *"--code-coverage"* ]]; then
diff --git a/tip-router-operator-cli/src/claim.rs b/tip-router-operator-cli/src/claim.rs
@@ -32,6 +32,7 @@ use solana_sdk::{
 use thiserror::Error;
 
 use crate::{
+    merkle_tree_collection_file_name,
     rpc_utils::{get_batched_accounts, send_until_blockhash_expires},
     Cli,
 };
@@ -77,10 +78,9 @@ pub async fn claim_mev_tips_with_emit(
     let keypair = read_keypair_file(cli.keypair_path.clone())
         .map_err(|e| anyhow::anyhow!("Failed to read keypair file: {:?}", e))?;
     let keypair = Arc::new(keypair);
-    let meta_merkle_tree_dir = cli.meta_merkle_tree_dir.clone();
+    let meta_merkle_tree_dir = cli.save_path.clone();
     let rpc_url = cli.rpc_url.clone();
-    let merkle_tree_coll_path =
-        meta_merkle_tree_dir.join(format!("generated_merkle_tree_{}.json", epoch));
+    let merkle_tree_coll_path = meta_merkle_tree_dir.join(merkle_tree_collection_file_name(epoch));
     let merkle_tree_coll = GeneratedMerkleTreeCollection::new_from_file(&merkle_tree_coll_path)
         .map_err(|e| anyhow::anyhow!(e))?;
     let start = Instant::now();
diff --git a/tip-router-operator-cli/src/cli.rs b/tip-router-operator-cli/src/cli.rs
@@ -30,9 +30,6 @@ pub struct Cli {
     #[arg(short, long, env)]
     pub snapshot_output_dir: PathBuf,
 
-    #[arg(short, long, env)]
-    pub meta_merkle_tree_dir: PathBuf,
-
     #[arg(long, env, default_value = "false")]
     pub submit_as_memo: bool,
 
diff --git a/tip-router-operator-cli/src/lib.rs b/tip-router-operator-cli/src/lib.rs
@@ -12,8 +12,7 @@ pub mod process_epoch;
 pub mod rpc_utils;
 pub mod submit;
 
-use std::fs::{self, File};
-use std::io::{BufWriter, Write};
+use std::fs;
 use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::sync::Arc;
@@ -27,15 +26,14 @@ use jito_tip_payment_sdk::{
     TIP_ACCOUNT_SEED_7,
 };
 use ledger_utils::get_bank_from_ledger;
-use log::{error, info};
-use meta_merkle_tree::generated_merkle_tree::MerkleRootGeneratorError;
+use log::info;
 use meta_merkle_tree::generated_merkle_tree::StakeMetaCollection;
 use meta_merkle_tree::{
     generated_merkle_tree::GeneratedMerkleTreeCollection, meta_merkle_tree::MetaMerkleTree,
 };
 use solana_metrics::{datapoint_error, datapoint_info};
 use solana_runtime::bank::Bank;
-use solana_sdk::{account::AccountSharedData, pubkey::Pubkey, slot_history::Slot};
+use solana_sdk::{account::AccountSharedData, pubkey::Pubkey};
 use stake_meta_generator::generate_stake_meta_collection;
 
 // TODO: Should this be loaded from somewhere?
@@ -52,6 +50,18 @@ pub enum OperatorState {
     WaitForNextEpoch,
 }
 
+pub fn stake_meta_file_name(epoch: u64) -> String {
+    format!("{}_stake_meta_collection.json", epoch)
+}
+
+pub fn merkle_tree_collection_file_name(epoch: u64) -> String {
+    format!("{}_merkle_tree_collection.json", epoch)
+}
+
+pub fn meta_merkle_tree_file_name(epoch: u64) -> String {
+    format!("{}_meta_merkle_tree.json", epoch)
+}
+
 // STAGE 1 LoadBankFromSnapshot
 pub fn load_bank_from_snapshot(cli: Cli, slot: u64, store_snapshot: bool) -> Arc<Bank> {
     let ledger_path = cli.ledger_path.clone();
@@ -118,7 +128,7 @@ pub fn create_stake_meta(
     if save {
         // Note: We have the epoch come before the file name so ordering is neat on a machine
         //  with multiple epochs saved.
-        let file = save_path.join(format!("{}_stake_meta_collection.json", epoch));
+        let file = save_path.join(stake_meta_file_name(epoch));
         stake_meta_coll.write_to_file(&file);
     }
 
@@ -189,7 +199,7 @@ pub fn create_merkle_tree_collection(
     if save {
         // Note: We have the epoch come before the file name so ordering is neat on a machine
         //  with multiple epochs saved.
-        let file = save_path.join(format!("{}_merkle_tree_collection.json", epoch));
+        let file = save_path.join(merkle_tree_collection_file_name(epoch));
         match merkle_tree_coll.write_to_file(&file) {
             Ok(_) => {}
             Err(e) => {
@@ -253,7 +263,7 @@ pub fn create_meta_merkle_tree(
     if save {
         // Note: We have the epoch come before the file name so ordering is neat on a machine
         //  with multiple epochs saved.
-        let file = save_path.join(format!("{}_meta_merkle_tree.json", epoch));
+        let file = save_path.join(meta_merkle_tree_file_name(epoch));
         match meta_merkle_tree.write_to_file(&file) {
             Ok(_) => {}
             Err(e) => {
@@ -323,235 +333,13 @@ fn derive_tip_payment_pubkeys(program_id: &Pubkey) -> TipPaymentPubkeys {
     }
 }
 
-fn write_to_json_file(
-    merkle_tree_coll: &GeneratedMerkleTreeCollection,
-    file_path: &PathBuf,
-) -> std::result::Result<(), MerkleRootGeneratorError> {
-    let file = File::create(file_path)?;
-    let mut writer = BufWriter::new(file);
-    let json = serde_json::to_string_pretty(&merkle_tree_coll).unwrap();
-    writer.write_all(json.as_bytes())?;
-    writer.flush()?;
-
-    Ok(())
-}
-
 /// Convenience wrapper around [TipDistributionAccount]
 pub struct TipDistributionAccountWrapper {
     pub tip_distribution_account: TipDistributionAccount,
     pub account_data: AccountSharedData,
     pub tip_distribution_pubkey: Pubkey,
 }
 
-#[allow(clippy::too_many_arguments)]
-pub fn get_meta_merkle_root(
-    ledger_path: &Path,
-    account_paths: Vec<PathBuf>,
-    full_snapshots_path: PathBuf,
-    incremental_snapshots_path: PathBuf,
-    desired_slot: &Slot,
-    tip_distribution_program_id: &Pubkey,
-    out_path: &str,
-    tip_payment_program_id: &Pubkey,
-    tip_router_program_id: &Pubkey,
-    ncn_address: &Pubkey,
-    operator_address: &Pubkey,
-    epoch: u64,
-    protocol_fee_bps: u64,
-    snapshots_enabled: bool,
-    meta_merkle_tree_dir: &Path,
-) -> std::result::Result<MetaMerkleTree, MerkleRootError> {
-    let start = Instant::now();
-
-    datapoint_info!(
-        "tip_router_cli.get_meta_merkle_root",
-        ("operator_address", operator_address.to_string(), String),
-        ("state", "stake_meta_generation", String),
-        ("step", 1, i64),
-        ("epoch", epoch, i64),
-        ("duration_ms", start.elapsed().as_millis() as i64, i64)
-    );
-
-    // cleanup tmp files - update with path where stake meta is written
-    match cleanup_tmp_files(&incremental_snapshots_path) {
-        Ok(_) => {}
-        Err(e) => {
-            datapoint_info!(
-                "tip_router_cli.get_meta_merkle_root",
-                ("operator_address", operator_address.to_string(), String),
-                ("state", "cleanup_tmp_files", String),
-                ("error", format!("{:?}", e), String),
-                ("epoch", epoch, i64),
-                ("duration_ms", start.elapsed().as_millis() as i64, i64)
-            );
-        }
-    }
-
-    // Get stake meta collection
-    let stake_meta_collection = stake_meta_generator::generate_stake_meta(
-        operator_address,
-        ledger_path,
-        account_paths,
-        full_snapshots_path,
-        incremental_snapshots_path.clone(),
-        desired_slot,
-        tip_distribution_program_id,
-        out_path,
-        tip_payment_program_id,
-        snapshots_enabled,
-    )
-    .map_err(|e| {
-        MerkleRootError::StakeMetaGeneratorError(format!("Failed to generate stake meta: {:?}", e))
-    })?;
-
-    info!(
-        "Created StakeMetaCollection:\n - epoch: {:?}\n - slot: {:?}\n - num stake metas: {:?}\n - bank_hash: {:?}",
-        stake_meta_collection.epoch,
-        stake_meta_collection.slot,
-        stake_meta_collection.stake_metas.len(),
-        stake_meta_collection.bank_hash
-    );
-
-    datapoint_info!(
-        "tip_router_cli.get_meta_merkle_root",
-        ("operator_address", operator_address.to_string(), String),
-        ("state", "generated_merkle_tree_collection", String),
-        ("step", 2, i64),
-        ("epoch", epoch, i64),
-        ("duration_ms", start.elapsed().as_millis() as i64, i64)
-    );
-
-    // Cleanup tmp files
-    match cleanup_tmp_files(&incremental_snapshots_path) {
-        Ok(_) => {}
-        Err(e) => {
-            datapoint_info!(
-                "tip_router_cli.get_meta_merkle_root",
-                ("operator_address", operator_address.to_string(), String),
-                ("state", "cleanup_tmp_files", String),
-                ("error", format!("{:?}", e), String),
-                ("epoch", epoch, i64),
-                ("duration_ms", start.elapsed().as_millis() as i64, i64)
-            );
-        }
-    }
-
-    // Generate merkle tree collection
-    let merkle_tree_coll = GeneratedMerkleTreeCollection::new_from_stake_meta_collection(
-        stake_meta_collection,
-        ncn_address,
-        epoch,
-        protocol_fee_bps,
-        tip_router_program_id,
-    )
-    .map_err(|_| {
-        MerkleRootError::MerkleRootGeneratorError(
-            "Failed to generate merkle tree collection".to_string(),
-        )
-    })?;
-
-    info!(
-        "Created GeneratedMerkleTreeCollection:\n - epoch: {:?}\n - slot: {:?}\n - num generated merkle trees: {:?}\n - bank_hash: {:?}",
-        merkle_tree_coll.epoch,
-        merkle_tree_coll.slot,
-        merkle_tree_coll.generated_merkle_trees.len(),
-        merkle_tree_coll.bank_hash
-    );
-
-    // Write GeneratedMerkleTreeCollection to file for debugging/verification
-    let generated_merkle_tree_path = incremental_snapshots_path.join(format!(
-        "generated_merkle_tree_{}.json",
-        merkle_tree_coll.epoch
-    ));
-    match write_to_json_file(&merkle_tree_coll, &generated_merkle_tree_path) {
-        Ok(_) => {
-            info!(
-                "Wrote GeneratedMerkleTreeCollection to {}",
-                generated_merkle_tree_path.display()
-            );
-        }
-        Err(e) => {
-            error!(
-                "Failed to write GeneratedMerkleTreeCollection to file {}: {:?}",
-                generated_merkle_tree_path.display(),
-                e
-            );
-        }
-    }
-
-    datapoint_info!(
-        "tip_router_cli.get_meta_merkle_root",
-        ("operator_address", operator_address.to_string(), String),
-        ("state", "meta_merkle_tree_creation", String),
-        ("step", 3, i64),
-        ("epoch", epoch, i64),
-        ("duration_ms", start.elapsed().as_millis() as i64, i64)
-    );
-
-    // TODO: Hide this behind a flag when the process gets split up into the various stages and
-    //  checkpoints.
-
-    // Write GeneratedMerkleTreeCollection to disk. Required for Claiming
-    let merkle_tree_coll_path =
-        meta_merkle_tree_dir.join(format!("generated_merkle_tree_{}.json", epoch));
-    let generated_merkle_tree_col_json = match serde_json::to_string(&merkle_tree_coll) {
-        Ok(json) => json,
-        Err(e) => {
-            datapoint_error!(
-                "tip_router_cli.process_epoch",
-                ("operator_address", operator_address.to_string(), String),
-                ("epoch", epoch, i64),
-                ("status", "error", String),
-                ("error", format!("{:?}", e), String),
-                ("state", "merkle_root_serialization", String),
-                ("duration_ms", start.elapsed().as_millis() as i64, i64)
-            );
-            return Err(MerkleRootError::MerkleRootGeneratorError(
-                "Failed to serialize merkle tree collection".to_string(),
-            ));
-        }
-    };
-
-    if let Err(e) = std::fs::write(merkle_tree_coll_path, generated_merkle_tree_col_json) {
-        datapoint_error!(
-            "tip_router_cli.process_epoch",
-            ("operator_address", operator_address.to_string(), String),
-            ("epoch", epoch, i64),
-            ("status", "error", String),
-            ("error", format!("{:?}", e), String),
-            ("state", "merkle_root_file_write", String),
-            ("duration_ms", start.elapsed().as_millis() as i64, i64)
-        );
-        return Err(MerkleRootError::MerkleRootGeneratorError(
-            "Failed to write meta merkle tree to file".to_string(),
-        ));
-    }
-
-    // Convert to MetaMerkleTree
-    let meta_merkle_tree = MetaMerkleTree::new_from_generated_merkle_tree_collection(
-        merkle_tree_coll,
-    )
-    .map_err(|e| {
-        MerkleRootError::MerkleTreeError(format!("Failed to create meta merkle tree: {:?}", e))
-    })?;
-
-    info!(
-        "Created MetaMerkleTree:\n - num nodes: {:?}\n - merkle root: {:?}",
-        meta_merkle_tree.num_nodes, meta_merkle_tree.merkle_root
-    );
-
-    datapoint_info!(
-        "tip_router_cli.get_meta_merkle_root",
-        ("operator_address", operator_address.to_string(), String),
-        ("state", "meta_merkle_tree_creation", String),
-        ("step", 4, i64),
-        ("epoch", epoch, i64),
-        ("duration_ms", start.elapsed().as_millis() as i64, i64)
-    );
-
-    Ok(meta_merkle_tree)
-}
-
 fn get_validator_cmdline() -> Result<String> {
     let output = Command::new("pgrep").arg("solana-validator").output()?;
 
diff --git a/tip-router-operator-cli/src/main.rs b/tip-router-operator-cli/src/main.rs
@@ -182,8 +182,8 @@ async fn main() -> Result<()> {
             set_merkle_roots,
         } => {
             let meta_merkle_tree_path = PathBuf::from(format!(
-                "{}/meta_merkle_tree_{}.json",
-                cli.meta_merkle_tree_dir.display(),
+                "{}/{}_meta_merkle_tree.json",
+                cli.save_path.display(),
                 epoch
             ));
             info!(
diff --git a/tip-router-operator-cli/src/process_epoch.rs b/tip-router-operator-cli/src/process_epoch.rs
diff --git a/tip-router-operator-cli/src/submit.rs b/tip-router-operator-cli/src/submit.rs
diff --git a/tip-router-operator-cli/tests/integration_tests.rs b/tip-router-operator-cli/tests/integration_tests.rs
