Prefactor: Move ChainSource creation before Wallet creation

In the following commits we will use the chain source to poll a best tip
before intializing the listener objects. As a prefactor, we here move
the creation of our onchain wallet after creation of the chain source,
which in turn means we'll need to use the same pattern as for the other
listeners, i.e., not giving the wallet reference to `ChainSource` on
creation but rather handing it in when it's being used at runtime.

Author: tnull

src/builder.rs

@@ -1178,54 +1178,6 @@ fn build_with_store_internal(
 			}
 		},
 	};
-
-	// Initialize the on-chain wallet and chain access
-	let xprv = bitcoin::bip32::Xpriv::new_master(config.network, &seed_bytes).map_err(|e| {
-		log_error!(logger, "Failed to derive master secret: {}", e);
-		BuildError::InvalidSeedBytes
-	})?;
-
-	let descriptor = Bip84(xprv, KeychainKind::External);
-	let change_descriptor = Bip84(xprv, KeychainKind::Internal);
-	let mut wallet_persister =
-		KVStoreWalletPersister::new(Arc::clone(&kv_store), Arc::clone(&logger));
-	let wallet_opt = BdkWallet::load()
-		.descriptor(KeychainKind::External, Some(descriptor.clone()))
-		.descriptor(KeychainKind::Internal, Some(change_descriptor.clone()))
-		.extract_keys()
-		.check_network(config.network)
-		.load_wallet(&mut wallet_persister)
-		.map_err(|e| match e {
-			bdk_wallet::LoadWithPersistError::InvalidChangeSet(
-				bdk_wallet::LoadError::Mismatch(bdk_wallet::LoadMismatch::Network {
-					loaded,
-					expected,
-				}),
-			) => {
-				log_error!(
-					logger,
-					"Failed to setup wallet: Networks do not match. Expected {} but got {}",
-					expected,
-					loaded
-				);
-				BuildError::NetworkMismatch
-			},
-			_ => {
-				log_error!(logger, "Failed to set up wallet: {}", e);
-				BuildError::WalletSetupFailed
-			},
-		})?;
-	let bdk_wallet = match wallet_opt {
-		Some(wallet) => wallet,
-		None => BdkWallet::create(descriptor, change_descriptor)
-			.network(config.network)
-			.create_wallet(&mut wallet_persister)
-			.map_err(|e| {
-				log_error!(logger, "Failed to set up wallet: {}", e);
-				BuildError::WalletSetupFailed
-			})?,
-	};
-
 	let tx_broadcaster = Arc::new(TransactionBroadcaster::new(Arc::clone(&logger)));
 	let fee_estimator = Arc::new(OnchainFeeEstimator::new());
 
@@ -1243,24 +1195,13 @@ fn build_with_store_internal(
 		},
 	};
 
-	let wallet = Arc::new(Wallet::new(
-		bdk_wallet,
-		wallet_persister,
-		Arc::clone(&tx_broadcaster),
-		Arc::clone(&fee_estimator),
-		Arc::clone(&payment_store),
-		Arc::clone(&config),
-		Arc::clone(&logger),
-	));
-
 	let chain_source = match chain_data_source_config {
 		Some(ChainDataSourceConfig::Esplora { server_url, headers, sync_config }) => {
 			let sync_config = sync_config.unwrap_or(EsploraSyncConfig::default());
 			Arc::new(ChainSource::new_esplora(
 				server_url.clone(),
 				headers.clone(),
 				sync_config,
-				Arc::clone(&wallet),
 				Arc::clone(&fee_estimator),
 				Arc::clone(&tx_broadcaster),
 				Arc::clone(&kv_store),
@@ -1274,7 +1215,6 @@ fn build_with_store_internal(
 			Arc::new(ChainSource::new_electrum(
 				server_url.clone(),
 				sync_config,
-				Arc::clone(&wallet),
 				Arc::clone(&fee_estimator),
 				Arc::clone(&tx_broadcaster),
 				Arc::clone(&kv_store),
@@ -1295,7 +1235,6 @@ fn build_with_store_internal(
 				*rpc_port,
 				rpc_user.clone(),
 				rpc_password.clone(),
-				Arc::clone(&wallet),
 				Arc::clone(&fee_estimator),
 				Arc::clone(&tx_broadcaster),
 				Arc::clone(&kv_store),
@@ -1309,7 +1248,6 @@ fn build_with_store_internal(
 				*rpc_port,
 				rpc_user.clone(),
 				rpc_password.clone(),
-				Arc::clone(&wallet),
 				Arc::clone(&fee_estimator),
 				Arc::clone(&tx_broadcaster),
 				Arc::clone(&kv_store),
@@ -1327,7 +1265,6 @@ fn build_with_store_internal(
 				server_url.clone(),
 				HashMap::new(),
 				sync_config,
-				Arc::clone(&wallet),
 				Arc::clone(&fee_estimator),
 				Arc::clone(&tx_broadcaster),
 				Arc::clone(&kv_store),
@@ -1338,6 +1275,63 @@ fn build_with_store_internal(
 		},
 	};
 
+	// Initialize the on-chain wallet and chain access
+	let xprv = bitcoin::bip32::Xpriv::new_master(config.network, &seed_bytes).map_err(|e| {
+		log_error!(logger, "Failed to derive master secret: {}", e);
+		BuildError::InvalidSeedBytes
+	})?;
+
+	let descriptor = Bip84(xprv, KeychainKind::External);
+	let change_descriptor = Bip84(xprv, KeychainKind::Internal);
+	let mut wallet_persister =
+		KVStoreWalletPersister::new(Arc::clone(&kv_store), Arc::clone(&logger));
+	let wallet_opt = BdkWallet::load()
+		.descriptor(KeychainKind::External, Some(descriptor.clone()))
+		.descriptor(KeychainKind::Internal, Some(change_descriptor.clone()))
+		.extract_keys()
+		.check_network(config.network)
+		.load_wallet(&mut wallet_persister)
+		.map_err(|e| match e {
+			bdk_wallet::LoadWithPersistError::InvalidChangeSet(
+				bdk_wallet::LoadError::Mismatch(bdk_wallet::LoadMismatch::Network {
+					loaded,
+					expected,
+				}),
+			) => {
+				log_error!(
+					logger,
+					"Failed to setup wallet: Networks do not match. Expected {} but got {}",
+					expected,
+					loaded
+				);
+				BuildError::NetworkMismatch
+			},
+			_ => {
+				log_error!(logger, "Failed to set up wallet: {}", e);
+				BuildError::WalletSetupFailed
+			},
+		})?;
+	let bdk_wallet = match wallet_opt {
+		Some(wallet) => wallet,
+		None => BdkWallet::create(descriptor, change_descriptor)
+			.network(config.network)
+			.create_wallet(&mut wallet_persister)
+			.map_err(|e| {
+				log_error!(logger, "Failed to set up wallet: {}", e);
+				BuildError::WalletSetupFailed
+			})?,
+	};
+
+	let wallet = Arc::new(Wallet::new(
+		bdk_wallet,
+		wallet_persister,
+		Arc::clone(&tx_broadcaster),
+		Arc::clone(&fee_estimator),
+		Arc::clone(&payment_store),
+		Arc::clone(&config),
+		Arc::clone(&logger),
+	));
+
 	// Initialize the KeysManager
 	let cur_time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).map_err(|e| {
 		log_error!(logger, "Failed to get current time: {}", e);

src/chain/bitcoind.rs

@@ -47,7 +47,6 @@ pub(super) struct BitcoindChainSource {
 	api_client: Arc<BitcoindClient>,
 	header_cache: tokio::sync::Mutex<BoundedHeaderCache>,
 	latest_chain_tip: RwLock<Option<ValidatedBlockHeader>>,
-	onchain_wallet: Arc<Wallet>,
 	wallet_polling_status: Mutex<WalletSyncStatus>,
 	fee_estimator: Arc<OnchainFeeEstimator>,
 	kv_store: Arc<DynStore>,
@@ -59,9 +58,8 @@ pub(super) struct BitcoindChainSource {
 impl BitcoindChainSource {
 	pub(crate) fn new_rpc(
 		rpc_host: String, rpc_port: u16, rpc_user: String, rpc_password: String,
-		onchain_wallet: Arc<Wallet>, fee_estimator: Arc<OnchainFeeEstimator>,
-		kv_store: Arc<DynStore>, config: Arc<Config>, logger: Arc<Logger>,
-		node_metrics: Arc<RwLock<NodeMetrics>>,
+		fee_estimator: Arc<OnchainFeeEstimator>, kv_store: Arc<DynStore>, config: Arc<Config>,
+		logger: Arc<Logger>, node_metrics: Arc<RwLock<NodeMetrics>>,
 	) -> Self {
 		let api_client = Arc::new(BitcoindClient::new_rpc(
 			rpc_host.clone(),
@@ -77,7 +75,6 @@ impl BitcoindChainSource {
 			api_client,
 			header_cache,
 			latest_chain_tip,
-			onchain_wallet,
 			wallet_polling_status,
 			fee_estimator,
 			kv_store,
@@ -89,9 +86,9 @@ impl BitcoindChainSource {
 
 	pub(crate) fn new_rest(
 		rpc_host: String, rpc_port: u16, rpc_user: String, rpc_password: String,
-		onchain_wallet: Arc<Wallet>, fee_estimator: Arc<OnchainFeeEstimator>,
-		kv_store: Arc<DynStore>, config: Arc<Config>, rest_client_config: BitcoindRestClientConfig,
-		logger: Arc<Logger>, node_metrics: Arc<RwLock<NodeMetrics>>,
+		fee_estimator: Arc<OnchainFeeEstimator>, kv_store: Arc<DynStore>, config: Arc<Config>,
+		rest_client_config: BitcoindRestClientConfig, logger: Arc<Logger>,
+		node_metrics: Arc<RwLock<NodeMetrics>>,
 	) -> Self {
 		let api_client = Arc::new(BitcoindClient::new_rest(
 			rest_client_config.rest_host,
@@ -111,7 +108,6 @@ impl BitcoindChainSource {
 			header_cache,
 			latest_chain_tip,
 			wallet_polling_status,
-			onchain_wallet,
 			fee_estimator,
 			kv_store,
 			config,
@@ -126,8 +122,8 @@ impl BitcoindChainSource {
 
 	pub(super) async fn continuously_sync_wallets(
 		&self, mut stop_sync_receiver: tokio::sync::watch::Receiver<()>,
-		channel_manager: Arc<ChannelManager>, chain_monitor: Arc<ChainMonitor>,
-		output_sweeper: Arc<Sweeper>,
+		onchain_wallet: Arc<Wallet>, channel_manager: Arc<ChannelManager>,
+		chain_monitor: Arc<ChainMonitor>, output_sweeper: Arc<Sweeper>,
 	) {
 		// First register for the wallet polling status to make sure `Node::sync_wallets` calls
 		// wait on the result before proceeding.
@@ -155,14 +151,10 @@ impl BitcoindChainSource {
 
 			let channel_manager_best_block_hash = channel_manager.current_best_block().block_hash;
 			let sweeper_best_block_hash = output_sweeper.current_best_block().block_hash;
-			let onchain_wallet_best_block_hash =
-				self.onchain_wallet.current_best_block().block_hash;
+			let onchain_wallet_best_block_hash = onchain_wallet.current_best_block().block_hash;
 
 			let mut chain_listeners = vec![
-				(
-					onchain_wallet_best_block_hash,
-					&*self.onchain_wallet as &(dyn Listen + Send + Sync),
-				),
+				(onchain_wallet_best_block_hash, &*onchain_wallet as &(dyn Listen + Send + Sync)),
 				(channel_manager_best_block_hash, &*channel_manager as &(dyn Listen + Send + Sync)),
 				(sweeper_best_block_hash, &*output_sweeper as &(dyn Listen + Send + Sync)),
 			];
@@ -307,6 +299,7 @@ impl BitcoindChainSource {
 							return;
 						}
 						_ = self.poll_and_update_listeners(
+							Arc::clone(&onchain_wallet),
 							Arc::clone(&channel_manager),
 							Arc::clone(&chain_monitor),
 							Arc::clone(&output_sweeper)
@@ -337,8 +330,8 @@ impl BitcoindChainSource {
 	}
 
 	pub(super) async fn poll_and_update_listeners(
-		&self, channel_manager: Arc<ChannelManager>, chain_monitor: Arc<ChainMonitor>,
-		output_sweeper: Arc<Sweeper>,
+		&self, onchain_wallet: Arc<Wallet>, channel_manager: Arc<ChannelManager>,
+		chain_monitor: Arc<ChainMonitor>, output_sweeper: Arc<Sweeper>,
 	) -> Result<(), Error> {
 		let receiver_res = {
 			let mut status_lock = self.wallet_polling_status.lock().unwrap();
@@ -355,7 +348,12 @@ impl BitcoindChainSource {
 		}
 
 		let res = self
-			.poll_and_update_listeners_inner(channel_manager, chain_monitor, output_sweeper)
+			.poll_and_update_listeners_inner(
+				onchain_wallet,
+				channel_manager,
+				chain_monitor,
+				output_sweeper,
+			)
 			.await;
 
 		self.wallet_polling_status.lock().unwrap().propagate_result_to_subscribers(res);
@@ -364,8 +362,8 @@ impl BitcoindChainSource {
 	}
 
 	async fn poll_and_update_listeners_inner(
-		&self, channel_manager: Arc<ChannelManager>, chain_monitor: Arc<ChainMonitor>,
-		output_sweeper: Arc<Sweeper>,
+		&self, onchain_wallet: Arc<Wallet>, channel_manager: Arc<ChannelManager>,
+		chain_monitor: Arc<ChainMonitor>, output_sweeper: Arc<Sweeper>,
 	) -> Result<(), Error> {
 		let latest_chain_tip_opt = self.latest_chain_tip.read().unwrap().clone();
 		let chain_tip = if let Some(tip) = latest_chain_tip_opt {
@@ -386,7 +384,7 @@ impl BitcoindChainSource {
 		let mut locked_header_cache = self.header_cache.lock().await;
 		let chain_poller = ChainPoller::new(Arc::clone(&self.api_client), self.config.network);
 		let chain_listener = ChainListener {
-			onchain_wallet: Arc::clone(&self.onchain_wallet),
+			onchain_wallet: Arc::clone(&onchain_wallet),
 			channel_manager: Arc::clone(&channel_manager),
 			chain_monitor: Arc::clone(&chain_monitor),
 			output_sweeper,
@@ -422,7 +420,7 @@ impl BitcoindChainSource {
 		let cur_height = channel_manager.current_best_block().height;
 
 		let now = SystemTime::now();
-		let bdk_unconfirmed_txids = self.onchain_wallet.get_unconfirmed_txids();
+		let bdk_unconfirmed_txids = onchain_wallet.get_unconfirmed_txids();
 		match self
 			.api_client
 			.get_updated_mempool_transactions(cur_height, bdk_unconfirmed_txids)
@@ -436,11 +434,11 @@ impl BitcoindChainSource {
 					evicted_txids.len(),
 					now.elapsed().unwrap().as_millis()
 				);
-				self.onchain_wallet
-					.apply_mempool_txs(unconfirmed_txs, evicted_txids)
-					.unwrap_or_else(|e| {
+				onchain_wallet.apply_mempool_txs(unconfirmed_txs, evicted_txids).unwrap_or_else(
+					|e| {
 						log_error!(self.logger, "Failed to apply mempool transactions: {:?}", e);
-					});
+					},
+				);
 			},
 			Err(e) => {
 				log_error!(self.logger, "Failed to poll for mempool transactions: {:?}", e);