Implement on-chain wallet sync via compact block filters

Author: febyeji

src/chain/cbf.rs

@@ -5,18 +5,23 @@
 // 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 std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::net::SocketAddr;
 use std::sync::{Arc, Mutex, RwLock};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
-use bip157::{BlockHash, Builder, Client, Event, Info, Requester, TrustedPeer, Warning};
-use bitcoin::{Script, Transaction, Txid};
+use bdk_chain::{BlockId, ConfirmationBlockTime, TxUpdate};
+use bdk_wallet::Update;
+use bip157::{
+	BlockHash, Builder, Client, Event, Info, Requester, SyncUpdate, TrustedPeer, Warning,
+};
+use bitcoin::{Script, ScriptBuf, Transaction, Txid};
 use lightning::chain::WatchedOutput;
 use lightning::util::ser::Writeable;
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, oneshot};
 
-use crate::config::{CbfSyncConfig, Config};
+use super::WalletSyncStatus;
+use crate::config::{CbfSyncConfig, Config, BDK_CLIENT_STOP_GAP};
 use crate::fee_estimator::{
 	apply_post_estimation_adjustments, get_all_conf_targets, OnchainFeeEstimator,
 };
@@ -35,6 +40,14 @@ pub(super) struct CbfChainSource {
 	cbf_runtime_status: Mutex<CbfRuntimeStatus>,
 	/// Latest chain tip hash, updated by the background event processing task.
 	latest_tip: Arc<Mutex<Option<BlockHash>>>,
+	/// Scripts to match against compact block filters during a scan.
+	watched_scripts: Arc<RwLock<Vec<ScriptBuf>>>,
+	/// Block (height, hash) pairs where filters matched watched scripts.
+	matched_block_hashes: Arc<Mutex<Vec<(u32, BlockHash)>>>,
+	/// One-shot channel sender to signal filter scan completion.
+	sync_completion_tx: Arc<Mutex<Option<oneshot::Sender<SyncUpdate>>>>,
+	/// Deduplicates concurrent on-chain wallet sync requests.
+	onchain_wallet_sync_status: Mutex<WalletSyncStatus>,
 	/// Shared fee rate estimator, updated by this chain source.
 	fee_estimator: Arc<OnchainFeeEstimator>,
 	/// Persistent key-value store for node metrics.
@@ -52,6 +65,14 @@ enum CbfRuntimeStatus {
 	Stopped,
 }
 
+/// Shared state passed to the background event processing task.
+struct CbfEventState {
+	latest_tip: Arc<Mutex<Option<BlockHash>>>,
+	watched_scripts: Arc<RwLock<Vec<ScriptBuf>>>,
+	matched_block_hashes: Arc<Mutex<Vec<(u32, BlockHash)>>>,
+	sync_completion_tx: Arc<Mutex<Option<oneshot::Sender<SyncUpdate>>>>,
+}
+
 impl CbfChainSource {
 	pub(crate) fn new(
 		peers: Vec<String>, sync_config: CbfSyncConfig, fee_estimator: Arc<OnchainFeeEstimator>,
@@ -60,11 +81,19 @@ impl CbfChainSource {
 	) -> Self {
 		let cbf_runtime_status = Mutex::new(CbfRuntimeStatus::Stopped);
 		let latest_tip = Arc::new(Mutex::new(None));
+		let watched_scripts = Arc::new(RwLock::new(Vec::new()));
+		let matched_block_hashes = Arc::new(Mutex::new(Vec::new()));
+		let sync_completion_tx = Arc::new(Mutex::new(None));
+		let onchain_wallet_sync_status = Mutex::new(WalletSyncStatus::Completed);
 		Self {
 			peers,
 			sync_config,
 			cbf_runtime_status,
 			latest_tip,
+			watched_scripts,
+			matched_block_hashes,
+			sync_completion_tx,
+			onchain_wallet_sync_status,
 			fee_estimator,
 			kv_store,
 			config,
@@ -129,11 +158,16 @@ impl CbfChainSource {
 			.spawn_cancellable_background_task(Self::process_warn_messages(warn_rx, warn_logger));
 
 		// Spawn a task to process events.
+		let event_state = CbfEventState {
+			latest_tip: Arc::clone(&self.latest_tip),
+			watched_scripts: Arc::clone(&self.watched_scripts),
+			matched_block_hashes: Arc::clone(&self.matched_block_hashes),
+			sync_completion_tx: Arc::clone(&self.sync_completion_tx),
+		};
 		let event_logger = Arc::clone(&self.logger);
-		let event_tip = Arc::clone(&self.latest_tip);
 		runtime.spawn_cancellable_background_task(Self::process_events(
 			event_rx,
-			event_tip,
+			event_state,
 			event_logger,
 		));
 
@@ -170,20 +204,22 @@ impl CbfChainSource {
 	}
 
 	async fn process_events(
-		mut event_rx: mpsc::UnboundedReceiver<Event>, latest_tip: Arc<Mutex<Option<BlockHash>>>,
-		logger: Arc<Logger>,
+		mut event_rx: mpsc::UnboundedReceiver<Event>, state: CbfEventState, logger: Arc<Logger>,
 	) {
 		while let Some(event) = event_rx.recv().await {
 			match event {
 				Event::FiltersSynced(sync_update) => {
 					let tip = sync_update.tip();
-					*latest_tip.lock().unwrap() = Some(tip.hash);
+					*state.latest_tip.lock().unwrap() = Some(tip.hash);
 					log_info!(
 						logger,
 						"CBF filters synced to tip: height={}, hash={}",
 						tip.height,
 						tip.hash,
 					);
+					if let Some(tx) = state.sync_completion_tx.lock().unwrap().take() {
+						let _ = tx.send(sync_update);
+					}
 				},
 				Event::Block(indexed_block) => {
 					log_trace!(logger, "CBF received block at height {}", indexed_block.height,);
@@ -192,6 +228,14 @@ impl CbfChainSource {
 					log_debug!(logger, "CBF chain update: {:?}", header_changes);
 				},
 				Event::IndexedFilter(indexed_filter) => {
+					let scripts = state.watched_scripts.read().unwrap();
+					if !scripts.is_empty() && indexed_filter.contains_any(scripts.iter()) {
+						state
+							.matched_block_hashes
+							.lock()
+							.unwrap()
+							.push((indexed_filter.height(), indexed_filter.block_hash()));
+					}
 					log_trace!(logger, "CBF received filter at height {}", indexed_filter.height(),);
 				},
 			}
@@ -212,12 +256,144 @@ impl CbfChainSource {
 		}
 	}
 
+	/// Run a CBF filter scan: set watched scripts, trigger a rescan, wait for
+	/// completion, and return the sync update along with matched block hashes.
+	async fn run_filter_scan(
+		&self, scripts: Vec<ScriptBuf>,
+	) -> Result<(SyncUpdate, Vec<(u32, BlockHash)>), Error> {
+		let requester = self.requester()?;
+
+		self.matched_block_hashes.lock().unwrap().clear();
+		*self.watched_scripts.write().unwrap() = scripts;
+
+		let (tx, rx) = oneshot::channel();
+		*self.sync_completion_tx.lock().unwrap() = Some(tx);
+
+		requester.rescan().map_err(|e| {
+			log_error!(self.logger, "Failed to trigger CBF rescan: {:?}", e);
+			Error::WalletOperationFailed
+		})?;
+
+		let sync_update = rx.await.map_err(|e| {
+			log_error!(self.logger, "CBF sync completion channel dropped: {:?}", e);
+			Error::WalletOperationFailed
+		})?;
+
+		self.watched_scripts.write().unwrap().clear();
+		let matched = std::mem::take(&mut *self.matched_block_hashes.lock().unwrap());
+
+		Ok((sync_update, matched))
+	}
+
 	/// Sync the on-chain wallet by scanning compact block filters for relevant transactions.
 	pub(crate) async fn sync_onchain_wallet(
-		&self, _onchain_wallet: Arc<Wallet>,
+		&self, onchain_wallet: Arc<Wallet>,
 	) -> Result<(), Error> {
-		log_error!(self.logger, "On-chain wallet sync via CBF is not yet implemented.");
-		Err(Error::WalletOperationFailed)
+		let receiver_res = {
+			let mut status_lock = self.onchain_wallet_sync_status.lock().unwrap();
+			status_lock.register_or_subscribe_pending_sync()
+		};
+		if let Some(mut sync_receiver) = receiver_res {
+			log_debug!(self.logger, "On-chain wallet sync already in progress, waiting.");
+			return sync_receiver.recv().await.map_err(|e| {
+				debug_assert!(false, "Failed to receive wallet sync result: {:?}", e);
+				log_error!(self.logger, "Failed to receive wallet sync result: {:?}", e);
+				Error::WalletOperationFailed
+			})?;
+		}
+
+		let requester = self.requester()?;
+		let now = Instant::now();
+
+		let scripts = onchain_wallet.get_spks_for_cbf_sync(BDK_CLIENT_STOP_GAP);
+		if scripts.is_empty() {
+			log_debug!(self.logger, "No wallet scripts to sync via CBF.");
+			return Ok(());
+		}
+
+		let timeout_fut = tokio::time::timeout(
+			Duration::from_secs(self.sync_config.timeouts_config.onchain_wallet_sync_timeout_secs),
+			self.sync_onchain_wallet_op(requester, scripts),
+		);
+
+		let (tx_update, sync_update) = match timeout_fut.await {
+			Ok(res) => res?,
+			Err(e) => {
+				log_error!(self.logger, "Sync of on-chain wallet timed out: {}", e);
+				return Err(Error::WalletOperationTimeout);
+			},
+		};
+
+		// Build chain checkpoint extending from the wallet's current tip.
+		let mut cp = onchain_wallet.latest_checkpoint();
+		for (height, header) in sync_update.recent_history() {
+			if *height > cp.height() {
+				let block_id = BlockId { height: *height, hash: header.block_hash() };
+				cp = cp.push(block_id).unwrap_or_else(|old| old);
+			}
+		}
+		let tip = sync_update.tip();
+		if tip.height > cp.height() {
+			let tip_block_id = BlockId { height: tip.height, hash: tip.hash };
+			cp = cp.push(tip_block_id).unwrap_or_else(|old| old);
+		}
+
+		let update = Update { last_active_indices: BTreeMap::new(), tx_update, chain: Some(cp) };
+
+		// Apply update to wallet.
+		onchain_wallet.apply_update(update)?;
+
+		log_debug!(
+			self.logger,
+			"Sync of on-chain wallet via CBF finished in {}ms.",
+			now.elapsed().as_millis()
+		);
+
+		update_node_metrics_timestamp(
+			&self.node_metrics,
+			&*self.kv_store,
+			&*self.logger,
+			|m, t| {
+				m.latest_onchain_wallet_sync_timestamp = t;
+			},
+		)?;
+
+		Ok(())
+	}
+
+	async fn sync_onchain_wallet_op(
+		&self, requester: Requester, scripts: Vec<ScriptBuf>,
+	) -> Result<(TxUpdate<ConfirmationBlockTime>, SyncUpdate), Error> {
+		let (sync_update, matched) = self.run_filter_scan(scripts).await?;
+
+		log_debug!(
+			self.logger,
+			"CBF on-chain filter scan complete: {} matching blocks found.",
+			matched.len()
+		);
+
+		// Fetch matching blocks and include all their transactions.
+		// The compact block filter already matched our scripts (covering both
+		// created outputs and spent inputs), so we include every transaction
+		// from matched blocks and let BDK determine relevance.
+		let mut tx_update = TxUpdate::default();
+		for (height, block_hash) in &matched {
+			let indexed_block = requester.get_block(*block_hash).await.map_err(|e| {
+				log_error!(self.logger, "Failed to fetch block {}: {:?}", block_hash, e);
+				Error::WalletOperationFailed
+			})?;
+			let block = indexed_block.block;
+			let block_id = BlockId { height: *height, hash: block.header.block_hash() };
+			let conf_time =
+				ConfirmationBlockTime { block_id, confirmation_time: block.header.time as u64 };
+			for tx in &block.txdata {
+				let txid = tx.compute_txid();
+				tx_update.txs.push(Arc::new(tx.clone()));
+				tx_update.anchors.insert((conf_time, txid));
+			}
+		}
+
+		Ok((tx_update, sync_update))
 	}
 
 	/// Sync the Lightning wallet by confirming channel transactions via compact block filters.

src/wallet/mod.rs

@@ -122,6 +122,28 @@ impl Wallet {
 		self.inner.lock().unwrap().start_sync_with_revealed_spks().build()
 	}
 
+	pub(crate) fn get_spks_for_cbf_sync(&self, stop_gap: usize) -> Vec<ScriptBuf> {
+		let wallet = self.inner.lock().unwrap();
+		let mut scripts: Vec<ScriptBuf> =
+			wallet.spk_index().revealed_spks(..).map(|((_, _), spk)| spk).collect();
+
+		// For first sync when no scripts have been revealed yet, generate
+		// lookahead scripts up to the stop gap for both keychains.
+		if scripts.is_empty() {
+			for keychain in [KeychainKind::External, KeychainKind::Internal] {
+				for idx in 0..stop_gap as u32 {
+					scripts.push(wallet.peek_address(keychain, idx).address.script_pubkey());
+				}
+			}
+		}
+
+		scripts
+	}
+
+	pub(crate) fn latest_checkpoint(&self) -> bdk_chain::CheckPoint {
+		self.inner.lock().unwrap().latest_checkpoint()
+	}
+
 	pub(crate) fn get_cached_txs(&self) -> Vec<Arc<Transaction>> {
 		self.inner.lock().unwrap().tx_graph().full_txs().map(|tx_node| tx_node.tx).collect()
 	}

tests/common/mod.rs

@@ -603,6 +603,24 @@ pub(crate) async fn wait_for_outpoint_spend<E: ElectrumApi>(electrs: &E, outpoin
 	.await;
 }
 
+pub(crate) async fn wait_for_cbf_sync(node: &TestNode) {
+	let before = node.status().latest_onchain_wallet_sync_timestamp;
+	let mut delay = Duration::from_millis(200);
+	for _ in 0..30 {
+		if node.sync_wallets().is_ok() {
+			let after = node.status().latest_onchain_wallet_sync_timestamp;
+			if after > before {
+				return;
+			}
+		}
+		tokio::time::sleep(delay).await;
+		if delay < Duration::from_secs(2) {
+			delay = delay.mul_f32(1.5);
+		}
+	}
+	panic!("wait_for_cbf_sync: timed out waiting for CBF sync to complete");
+}
+
 pub(crate) async fn exponential_backoff_poll<T, F>(mut poll: F) -> T
 where
 	F: FnMut() -> Option<T>,