feat: add gen limit to dag extension for regular updates

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();

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

sn_cli/src/bin/subcommands/wallet/audit.rs

@@ -21,14 +21,14 @@ async fn gather_spend_dag(client: &Client, root_dir: &Path) -> Result<SpendDag>
     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?
         }
     };
 

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

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<SpendDag> {
+    pub async fn spend_dag_build_from(
+        &self,
+        spend_addr: SpendAddress,
+        max_depth: Option<u32>,
+    ) -> WalletResult<SpendDag> {
         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<u32>,
+    ) -> 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(())
     }
 }