fix: wallet validation (#6765)

Description --- Fixes wallet validation to always correctly trigger Motivation and Context --- Currently when a new block comes, validation is triggered, but if validation is already running, nothing happens. We need to stop the current validation as we need to validate with the newest information. How Has This Been Tested? --- Manual
tari-project · Jan 28, 2025 · afee5ad · afee5ad
1 parent 4d68ece
commit afee5ad
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 48 deletions.
diff --git a/base_layer/wallet/src/output_manager_service/service.rs b/base_layer/wallet/src/output_manager_service/service.rs
@@ -567,6 +567,7 @@ where
         }
     }
 
+    #[allow(clippy::too_many_lines)]
     fn validate_outputs(&mut self) -> Result<u64, OutputManagerError> {
         let current_base_node = self
             .resources
@@ -586,6 +587,7 @@ where
         let mut base_node_watch = self.resources.connectivity.get_current_base_node_watcher();
         let event_publisher = self.resources.event_publisher.clone();
         let validation_in_progress = self.validation_in_progress.clone();
+        let mut base_node_service_event_stream = self.base_node_service.get_event_stream();
         tokio::spawn(async move {
             // Note: We do not want the validation task to be queued
             let mut _lock = match validation_in_progress.try_lock() {
@@ -604,60 +606,69 @@ where
                     return;
                 },
             };
-
-            let exec_fut = txo_validation.execute();
-            tokio::pin!(exec_fut);
-            loop {
-                tokio::select! {
-                    result = &mut exec_fut => {
-                        match result {
-                            Ok(id) => {
-                                info!(
-                                    target: LOG_TARGET,
-                                    "UTXO Validation Protocol (Id: {}) completed successfully", id
-                                );
-                                return;
-                            },
-                            Err(OutputManagerProtocolError { id, error }) => {
-                                warn!(
-                                    target: LOG_TARGET,
-                                    "Error completing UTXO Validation Protocol (Id: {}): {}", id, error
-                                );
-                                let event_payload = match error {
-                                    OutputManagerError::InconsistentBaseNodeDataError(_) |
-                                    OutputManagerError::BaseNodeChanged |
-                                    OutputManagerError::Shutdown |
-                                    OutputManagerError::RpcError(_) =>
-                                        OutputManagerEvent::TxoValidationCommunicationFailure(id),
-                                    _ => OutputManagerEvent::TxoValidationInternalFailure(id),
-                                };
-                                if let Err(e) = event_publisher.send(Arc::new(event_payload)) {
+            'outer: loop {
+                let local_run = txo_validation.clone();
+                let exec_fut = local_run.execute();
+                tokio::pin!(exec_fut);
+                loop {
+                    tokio::select! {
+                        result = &mut exec_fut => {
+                            match result {
+                                Ok(id) => {
+                                    info!(
+                                        target: LOG_TARGET,
+                                        "UTXO Validation Protocol (Id: {}) completed successfully", id
+                                    );
+                                    return;
+                                },
+                                Err(OutputManagerProtocolError { id, error }) => {
+                                    warn!(
+                                        target: LOG_TARGET,
+                                        "Error completing UTXO Validation Protocol (Id: {}): {}", id, error
+                                    );
+                                    let event_payload = match error {
+                                        OutputManagerError::InconsistentBaseNodeDataError(_) |
+                                        OutputManagerError::BaseNodeChanged |
+                                        OutputManagerError::Shutdown |
+                                        OutputManagerError::RpcError(_) =>
+                                            OutputManagerEvent::TxoValidationCommunicationFailure(id),
+                                        _ => OutputManagerEvent::TxoValidationInternalFailure(id),
+                                    };
+                                    if let Err(e) = event_publisher.send(Arc::new(event_payload)) {
+                                        debug!(
+                                            target: LOG_TARGET,
+                                            "Error sending event because there are no subscribers: {:?}", e
+                                        );
+                                    }
+
+                                    return;
+                                },
+                            }
+                        },
+                        _ = shutdown.wait() => {
+                            debug!(target: LOG_TARGET, "TXO Validation Protocol (Id: {}) shutting down because the system \
+                                is shutting down", id);
+                            return;
+                        },
+                        event = base_node_service_event_stream.recv() => {
+                            if let Ok(bn_event) = event {
+                                if let BaseNodeEvent::NewBlockDetected(_hash, _height) = *bn_event {
+                                    debug!(target: LOG_TARGET, "TXO Validation Protocol (Id: {}) resetting because base node height changed", id);
+                                    continue 'outer;
+                              }
+                            }
+                        }
+                        _ = base_node_watch.changed() => {
+                            if let Some(peer) = base_node_watch.borrow().as_ref() {
+                                if peer.get_current_peer().node_id != current_base_node {
                                     debug!(
                                         target: LOG_TARGET,
-                                        "Error sending event because there are no subscribers: {:?}", e
+                                        "TXO Validation Protocol (Id: {}) resetting to run again because base node changed", id
                                     );
+                                    continue 'outer;
                                 }
-
-                                return;
-                            },
-                        }
-                    },
-                    _ = shutdown.wait() => {
-                        debug!(target: LOG_TARGET, "TXO Validation Protocol (Id: {}) shutting down because the system \
-                            is shutting down", id);
-                        return;
-                    },
-                    _ = base_node_watch.changed() => {
-                        if let Some(peer) = base_node_watch.borrow().as_ref() {
-                            if peer.get_current_peer().node_id != current_base_node {
-                                debug!(
-                                    target: LOG_TARGET,
-                                    "TXO Validation Protocol (Id: {}) cancelled because base node changed", id
-                                );
-                                return;
                             }
                         }
-
                     }
                 }
             }

diff --git a/base_layer/wallet/src/output_manager_service/tasks/txo_validation_task.rs b/base_layer/wallet/src/output_manager_service/tasks/txo_validation_task.rs
@@ -53,6 +53,7 @@ use crate::{
 
 const LOG_TARGET: &str = "wallet::output_service::txo_validation_task";
 
+#[derive(Clone)]
 pub struct TxoValidationTask<TBackend, TWalletConnectivity> {
     operation_id: u64,
     db: OutputManagerDatabase<TBackend>,