diff --git a/sn_auditor/src/dag_db.rs b/sn_auditor/src/dag_db.rs index e8469c99bc..0f7f435238 100644 --- a/sn_auditor/src/dag_db.rs +++ b/sn_auditor/src/dag_db.rs @@ -144,7 +144,9 @@ impl SpendDagDb { }; // update that copy - self.client.spend_dag_continue_from_utxos(&mut dag).await?; + self.client + .spend_dag_continue_from_utxos(&mut dag, Some(10)) + .await?; // write update to DAG let dag_ref = self.dag.clone(); diff --git a/sn_auditor/src/main.rs b/sn_auditor/src/main.rs index 10cb9a7a6d..b9cda187b8 100644 --- a/sn_auditor/src/main.rs +++ b/sn_auditor/src/main.rs @@ -147,7 +147,7 @@ async fn initialize_background_spend_dag_collection( .map_err(|e| eyre!("Could not create new DAG from genesis: {e}"))?; tokio::spawn(async move { let _ = client - .spend_dag_continue_from_utxos(&mut genesis_dag) + .spend_dag_continue_from_utxos(&mut genesis_dag, None) .await .map_err(|e| eprintln!("Could not update DAG from genesis: {e}")); let _ = d diff --git a/sn_cli/src/bin/subcommands/wallet/audit.rs b/sn_cli/src/bin/subcommands/wallet/audit.rs index 9edc8655b3..53a60041c6 100644 --- a/sn_cli/src/bin/subcommands/wallet/audit.rs +++ b/sn_cli/src/bin/subcommands/wallet/audit.rs @@ -21,14 +21,14 @@ async fn gather_spend_dag(client: &Client, root_dir: &Path) -> Result let dag = match SpendDag::load_from_file(&dag_path) { Ok(mut dag) => { println!("Starting from the loaded spend dag on disk..."); - client.spend_dag_continue_from_utxos(&mut dag).await?; + client.spend_dag_continue_from_utxos(&mut dag, None).await?; dag } Err(err) => { println!("Starting from Genesis as found no local spend dag on disk..."); info!("Starting from Genesis as failed to load spend dag from disk: {err}"); let genesis_addr = SpendAddress::from_unique_pubkey(&GENESIS_CASHNOTE.unique_pubkey()); - client.spend_dag_build_from(genesis_addr).await? + client.spend_dag_build_from(genesis_addr, None).await? } }; diff --git a/sn_client/src/audit/spend_dag.rs b/sn_client/src/audit/spend_dag.rs index da76f88a10..197f1d73a8 100644 --- a/sn_client/src/audit/spend_dag.rs +++ b/sn_client/src/audit/spend_dag.rs @@ -227,6 +227,11 @@ impl SpendDag { /// Merges the given dag into ours pub fn merge(&mut self, sub_dag: SpendDag) { + info!( + "Merging sub DAG starting at {:?} into our DAG with source {:?}", + sub_dag.source(), + self.source() + ); for (addr, spends) in sub_dag.spends { // only add spends to the dag, ignoring utxos // utxos will be added automatically as their ancestors are added diff --git a/sn_client/src/audit/spend_dag_building.rs b/sn_client/src/audit/spend_dag_building.rs index 402b799bed..c81c53e0eb 100644 --- a/sn_client/src/audit/spend_dag_building.rs +++ b/sn_client/src/audit/spend_dag_building.rs @@ -18,7 +18,11 @@ impl Client { /// Builds a SpendDag from a given SpendAddress recursively following descendants all the way to UTxOs /// Started from Genesis this gives the entire SpendDag of the Network at a certain point in time /// Once the DAG collected, verifies and records errors in the DAG - pub async fn spend_dag_build_from(&self, spend_addr: SpendAddress) -> WalletResult { + pub async fn spend_dag_build_from( + &self, + spend_addr: SpendAddress, + max_depth: Option, + ) -> WalletResult { info!("Building spend DAG from {spend_addr:?}"); let mut dag = SpendDag::new(spend_addr); @@ -46,7 +50,7 @@ impl Client { // use iteration instead of recursion to avoid stack overflow let mut txs_to_follow = BTreeSet::from_iter([first_spend.spend.spent_tx]); let mut known_tx = BTreeSet::new(); - let mut gen = 0; + let mut gen: u32 = 0; let start = std::time::Instant::now(); while !txs_to_follow.is_empty() { @@ -95,12 +99,18 @@ impl Client { } // only follow tx we haven't already gathered - gen += 1; known_tx.extend(txs_to_follow.iter().map(|tx| tx.hash())); txs_to_follow = next_gen_tx .into_iter() .filter(|tx| !known_tx.contains(&tx.hash())) .collect(); + + // go on to next gen + gen += 1; + if gen >= max_depth.unwrap_or(u32::MAX) { + info!("Reached generation {gen}, stopping DAG collection from {spend_addr:?}"); + break; + } } let elapsed = start.elapsed(); @@ -243,14 +253,19 @@ impl Client { /// Extends an existing SpendDag starting from the utxos in this DAG /// Covers the entirety of currently existing Spends if the DAG was built from Genesis /// Records errors in the new DAG branches if any - pub async fn spend_dag_continue_from_utxos(&self, dag: &mut SpendDag) -> WalletResult<()> { - trace!("Gathering spend DAG from utxos..."); + /// Stops gathering after max_depth generations + pub async fn spend_dag_continue_from_utxos( + &self, + dag: &mut SpendDag, + max_depth: Option, + ) -> WalletResult<()> { + info!("Gathering spend DAG from utxos..."); let utxos = dag.get_utxos(); let mut stream = futures::stream::iter(utxos.into_iter()) .map(|utxo| async move { debug!("Queuing task to gather DAG from utxo: {:?}", utxo); - (self.spend_dag_build_from(utxo).await, utxo) + (self.spend_dag_build_from(utxo, max_depth).await, utxo) }) .buffer_unordered(crate::MAX_CONCURRENT_TASKS); @@ -264,7 +279,7 @@ impl Client { dag.record_faults(&dag.source()) .map_err(|e| WalletError::Dag(e.to_string()))?; - trace!("Done gathering spend DAG from utxos"); + info!("Done gathering spend DAG from utxos"); Ok(()) } }