Move on-chain syncing to ChainSource

tnull · tnull · commit 94cc44887068 · 2024-10-07T16:26:12.000+02:00
diff --git a/src/builder.rs b/src/builder.rs
@@ -5,7 +5,7 @@
 // http://opensource.org/licenses/MIT>, at your option. You may not use this file except in
 // accordance with one or both of these licenses.
 
-use crate::chain::{DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS, DEFAULT_ESPLORA_SERVER_URL};
+use crate::chain::{ChainSource, DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS, DEFAULT_ESPLORA_SERVER_URL};
 use crate::config::{default_user_config, Config, WALLET_KEYS_SEED_LEN};
 
 use crate::connection::ConnectionManager;
@@ -585,58 +585,78 @@ fn build_with_store_internal(
 			})?,
 	};
 
-	let (esplora_client, tx_sync, tx_broadcaster, fee_estimator) = match chain_data_source_config {
-		Some(ChainDataSourceConfig::Esplora(server_url)) => {
-			let mut client_builder = esplora_client::Builder::new(&server_url.clone());
-			client_builder = client_builder.timeout(DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS);
-			let esplora_client = client_builder.build_async().unwrap();
-			let tx_sync = Arc::new(EsploraSyncClient::from_client(
-				esplora_client.clone(),
-				Arc::clone(&logger),
-			));
-			let tx_broadcaster = Arc::new(TransactionBroadcaster::new(
-				tx_sync.client().clone(),
-				Arc::clone(&logger),
-			));
-			let fee_estimator = Arc::new(OnchainFeeEstimator::new(
-				tx_sync.client().clone(),
-				Arc::clone(&config),
-				Arc::clone(&logger),
-			));
-			(esplora_client, tx_sync, tx_broadcaster, fee_estimator)
-		},
-		None => {
-			// Default to Esplora client.
-			let server_url = DEFAULT_ESPLORA_SERVER_URL.to_string();
-			let mut client_builder = esplora_client::Builder::new(&server_url);
-			client_builder = client_builder.timeout(DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS);
-			let esplora_client = client_builder.build_async().unwrap();
-			let tx_sync = Arc::new(EsploraSyncClient::from_client(
-				esplora_client.clone(),
-				Arc::clone(&logger),
-			));
-			let tx_broadcaster = Arc::new(TransactionBroadcaster::new(
-				tx_sync.client().clone(),
-				Arc::clone(&logger),
-			));
-			let fee_estimator = Arc::new(OnchainFeeEstimator::new(
-				tx_sync.client().clone(),
-				Arc::clone(&config),
-				Arc::clone(&logger),
-			));
-			(esplora_client, tx_sync, tx_broadcaster, fee_estimator)
-		},
-	};
+	let (wallet, chain_source, tx_sync, tx_broadcaster, fee_estimator) =
+		match chain_data_source_config {
+			Some(ChainDataSourceConfig::Esplora(server_url)) => {
+				let mut client_builder = esplora_client::Builder::new(&server_url.clone());
+				client_builder = client_builder.timeout(DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS);
+				let esplora_client = client_builder.build_async().unwrap();
+				let tx_sync =
+					Arc::new(EsploraSyncClient::from_client(esplora_client, Arc::clone(&logger)));
+				let tx_broadcaster = Arc::new(TransactionBroadcaster::new(
+					tx_sync.client().clone(),
+					Arc::clone(&logger),
+				));
+				let fee_estimator = Arc::new(OnchainFeeEstimator::new(
+					tx_sync.client().clone(),
+					Arc::clone(&config),
+					Arc::clone(&logger),
+				));
+
+				let wallet = Arc::new(Wallet::new(
+					bdk_wallet,
+					wallet_persister,
+					Arc::clone(&tx_broadcaster),
+					Arc::clone(&fee_estimator),
+					Arc::clone(&logger),
+				));
+
+				let chain_source = Arc::new(ChainSource::new_esplora(
+					server_url.clone(),
+					Arc::clone(&wallet),
+					Arc::clone(&logger),
+				));
+				(wallet, chain_source, tx_sync, tx_broadcaster, fee_estimator)
+			},
+			None => {
+				// Default to Esplora client.
+				let server_url = DEFAULT_ESPLORA_SERVER_URL.to_string();
+				let mut client_builder = esplora_client::Builder::new(&server_url);
+				client_builder = client_builder.timeout(DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS);
+				let esplora_client = client_builder.build_async().unwrap();
+				let tx_sync = Arc::new(EsploraSyncClient::from_client(
+					esplora_client.clone(),
+					Arc::clone(&logger),
+				));
+				let tx_broadcaster = Arc::new(TransactionBroadcaster::new(
+					tx_sync.client().clone(),
+					Arc::clone(&logger),
+				));
+				let fee_estimator = Arc::new(OnchainFeeEstimator::new(
+					tx_sync.client().clone(),
+					Arc::clone(&config),
+					Arc::clone(&logger),
+				));
+
+				let wallet = Arc::new(Wallet::new(
+					bdk_wallet,
+					wallet_persister,
+					Arc::clone(&tx_broadcaster),
+					Arc::clone(&fee_estimator),
+					Arc::clone(&logger),
+				));
+
+				let chain_source = Arc::new(ChainSource::new_esplora(
+					server_url.clone(),
+					Arc::clone(&wallet),
+					Arc::clone(&logger),
+				));
+
+				(wallet, chain_source, tx_sync, tx_broadcaster, fee_estimator)
+			},
+		};
 
 	let runtime = Arc::new(RwLock::new(None));
-	let wallet = Arc::new(Wallet::new(
-		bdk_wallet,
-		wallet_persister,
-		esplora_client,
-		Arc::clone(&tx_broadcaster),
-		Arc::clone(&fee_estimator),
-		Arc::clone(&logger),
-	));
 
 	// Initialize the ChainMonitor
 	let chain_monitor: Arc<ChainMonitor> = Arc::new(chainmonitor::ChainMonitor::new(
@@ -1012,6 +1032,7 @@ fn build_with_store_internal(
 		event_handling_stopped_sender,
 		config,
 		wallet,
+		chain_source,
 		tx_sync,
 		tx_broadcaster,
 		fee_estimator,
diff --git a/src/chain/mod.rs b/src/chain/mod.rs
@@ -5,33 +5,165 @@
 // http://opensource.org/licenses/MIT>, at your option. You may not use this file except in
 // accordance with one or both of these licenses.
 
-use crate::logger::FilesystemLogger;
+use crate::config::{BDK_CLIENT_CONCURRENCY, BDK_CLIENT_STOP_GAP, BDK_WALLET_SYNC_TIMEOUT_SECS};
+use crate::logger::{log_error, log_info, FilesystemLogger, Logger};
+use crate::types::Wallet;
+use crate::Error;
+
+use bdk_esplora::EsploraAsyncExt;
 
 use esplora_client::AsyncClient as EsploraAsyncClient;
 
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
 
 // The default Esplora server we're using.
 pub(crate) const DEFAULT_ESPLORA_SERVER_URL: &str = "https://blockstream.info/api";
 
 // The default Esplora client timeout we're using.
 pub(crate) const DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS: u64 = 10;
 
+pub(crate) enum WalletSyncStatus {
+	Completed,
+	InProgress { subscribers: tokio::sync::broadcast::Sender<Result<(), Error>> },
+}
+
+impl WalletSyncStatus {
+	fn register_or_subscribe_pending_sync(
+		&mut self,
+	) -> Option<tokio::sync::broadcast::Receiver<Result<(), Error>>> {
+		match self {
+			WalletSyncStatus::Completed => {
+				// We're first to register for a sync.
+				let (tx, _) = tokio::sync::broadcast::channel(1);
+				*self = WalletSyncStatus::InProgress { subscribers: tx };
+				None
+			},
+			WalletSyncStatus::InProgress { subscribers } => {
+				// A sync is in-progress, we subscribe.
+				let rx = subscribers.subscribe();
+				Some(rx)
+			},
+		}
+	}
+
+	fn propagate_result_to_subscribers(&mut self, res: Result<(), Error>) {
+		// Send the notification to any other tasks that might be waiting on it by now.
+		{
+			match self {
+				WalletSyncStatus::Completed => {
+					// No sync in-progress, do nothing.
+					return;
+				},
+				WalletSyncStatus::InProgress { subscribers } => {
+					// A sync is in-progress, we notify subscribers.
+					if subscribers.receiver_count() > 0 {
+						match subscribers.send(res) {
+							Ok(_) => (),
+							Err(e) => {
+								debug_assert!(
+									false,
+									"Failed to send wallet sync result to subscribers: {:?}",
+									e
+								);
+							},
+						}
+					}
+					*self = WalletSyncStatus::Completed;
+				},
+			}
+		}
+	}
+}
+
 pub(crate) enum ChainSource {
 	Esplora {
 		esplora_client: EsploraAsyncClient,
+		onchain_wallet: Arc<Wallet>,
+		onchain_wallet_sync_status: Mutex<WalletSyncStatus>,
 		logger: Arc<FilesystemLogger>,
 	},
 }
 
 impl ChainSource {
 	pub(crate) fn new_esplora(
-		server_url: String,
-		logger: Arc<FilesystemLogger>,
+		server_url: String, onchain_wallet: Arc<Wallet>, logger: Arc<FilesystemLogger>,
 	) -> Self {
 		let mut client_builder = esplora_client::Builder::new(&server_url.clone());
 		client_builder = client_builder.timeout(DEFAULT_ESPLORA_CLIENT_TIMEOUT_SECS);
 		let esplora_client = client_builder.build_async().unwrap();
-		Self::Esplora { esplora_client, logger }
+		let onchain_wallet_sync_status = Mutex::new(WalletSyncStatus::Completed);
+		Self::Esplora {
+			esplora_client,
+			onchain_wallet,
+			onchain_wallet_sync_status,
+			logger,
+		}
+	}
+
+	pub(crate) async fn sync_onchain_wallet(&self) -> Result<(), Error> {
+		match self {
+			Self::Esplora {
+				esplora_client,
+				onchain_wallet,
+				onchain_wallet_sync_status,
+				logger,
+				..
+			} => {
+				let receiver_res = {
+					let mut status_lock = onchain_wallet_sync_status.lock().unwrap();
+					status_lock.register_or_subscribe_pending_sync()
+				};
+				if let Some(mut sync_receiver) = receiver_res {
+					log_info!(logger, "Sync in progress, skipping.");
+					return sync_receiver.recv().await.map_err(|e| {
+						debug_assert!(false, "Failed to receive wallet sync result: {:?}", e);
+						log_error!(logger, "Failed to receive wallet sync result: {:?}", e);
+						Error::WalletOperationFailed
+					})?;
+				}
+
+				let res = {
+					let full_scan_request = onchain_wallet.get_full_scan_request();
+
+					let wallet_sync_timeout_fut = tokio::time::timeout(
+						Duration::from_secs(BDK_WALLET_SYNC_TIMEOUT_SECS),
+						esplora_client.full_scan(
+							full_scan_request,
+							BDK_CLIENT_STOP_GAP,
+							BDK_CLIENT_CONCURRENCY,
+						),
+					);
+
+					match wallet_sync_timeout_fut.await {
+						Ok(res) => match res {
+							Ok(update) => onchain_wallet.apply_update(update),
+							Err(e) => match *e {
+								esplora_client::Error::Reqwest(he) => {
+									log_error!(
+										logger,
+										"Sync failed due to HTTP connection error: {}",
+										he
+									);
+									Err(Error::WalletOperationFailed)
+								},
+								_ => {
+									log_error!(logger, "Sync failed due to Esplora error: {}", e);
+									Err(Error::WalletOperationFailed)
+								},
+							},
+						},
+						Err(e) => {
+							log_error!(logger, "On-chain wallet sync timed out: {}", e);
+							Err(Error::WalletOperationTimeout)
+						},
+					}
+				};
+
+				onchain_wallet_sync_status.lock().unwrap().propagate_result_to_subscribers(res);
+
+				res
+			},
+		}
 	}
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -127,6 +127,7 @@ use config::{
 	WALLET_SYNC_INTERVAL_MINIMUM_SECS,
 };
 use connection::ConnectionManager;
+use chain::ChainSource;
 use event::{EventHandler, EventQueue};
 use gossip::GossipSource;
 use graph::NetworkGraph;
@@ -180,6 +181,7 @@ pub struct Node {
 	event_handling_stopped_sender: tokio::sync::watch::Sender<()>,
 	config: Arc<Config>,
 	wallet: Arc<Wallet>,
+	chain_source: Arc<ChainSource>,
 	tx_sync: Arc<EsploraSyncClient<Arc<FilesystemLogger>>>,
 	tx_broadcaster: Arc<Broadcaster>,
 	fee_estimator: Arc<FeeEstimator>,
@@ -274,7 +276,7 @@ impl Node {
 		})?;
 
 		// Setup wallet sync
-		let wallet = Arc::clone(&self.wallet);
+		let chain_source = Arc::clone(&self.chain_source);
 		let sync_logger = Arc::clone(&self.logger);
 		let sync_onchain_wallet_timestamp = Arc::clone(&self.latest_onchain_wallet_sync_timestamp);
 		let mut stop_sync = self.stop_sender.subscribe();
@@ -298,7 +300,7 @@ impl Node {
 					}
 					_ = onchain_wallet_sync_interval.tick() => {
 						let now = Instant::now();
-						match wallet.sync().await {
+						match chain_source.sync_onchain_wallet().await {
 							Ok(()) => {
 								log_trace!(
 									sync_logger,
@@ -1368,7 +1370,7 @@ impl Node {
 			return Err(Error::NotRunning);
 		}
 
-		let wallet = Arc::clone(&self.wallet);
+		let chain_source = Arc::clone(&self.chain_source);
 		let tx_sync = Arc::clone(&self.tx_sync);
 		let sync_cman = Arc::clone(&self.channel_manager);
 		let archive_cman = Arc::clone(&self.channel_manager);
@@ -1394,7 +1396,7 @@ impl Node {
 					let now = Instant::now();
 					// We don't add an additional timeout here, as `Wallet::sync` already returns
 					// after a timeout.
-					match wallet.sync().await {
+					match chain_source.sync_onchain_wallet().await {
 						Ok(()) => {
 							log_info!(
 								sync_logger,
diff --git a/src/wallet/mod.rs b/src/wallet/mod.rs
