feat(wallet): maintain derived-account rolling lookahead

Author: ben-kaufman

crates/bdk-wallet-aggregate/src/lib.rs

@@ -746,6 +746,39 @@ where
 		Ok(wallet.start_sync_with_revealed_spks().build())
 	}
 
+	/// Build an incremental sync request including external lookahead scripts.
+	///
+	/// The wallet's configured lookahead determines how many unrevealed scripts are included.
+	pub fn wallet_incremental_sync_request_with_lookahead(
+		&self, key: &K,
+	) -> Result<SyncRequest<(KeychainKind, u32)>, Error> {
+		let wallet = self.wallets.get(key).ok_or(Error::WalletNotFound)?;
+		let mut request = wallet.start_sync_with_revealed_spks();
+		let external_lookahead = wallet.spk_index().lookahead();
+		if external_lookahead == 0 {
+			return Ok(request.build());
+		}
+
+		let start_index = match wallet.derivation_index(KeychainKind::External) {
+			Some(BIP32_MAX_NORMAL_INDEX) => return Ok(request.build()),
+			Some(index) => index + 1,
+			None => 0,
+		};
+		let end_index =
+			start_index.saturating_add(external_lookahead - 1).min(BIP32_MAX_NORMAL_INDEX);
+		let lookahead_spks = (start_index..=end_index)
+			.map(|index| {
+				wallet
+					.spk_index()
+					.spk_at_index(KeychainKind::External, index)
+					.map(|spk| ((KeychainKind::External, index), spk))
+					.ok_or(Error::WalletOperationFailed)
+			})
+			.collect::<Result<Vec<_>, _>>()?;
+		request = request.spks_with_indexes(lookahead_spks);
+		Ok(request.build())
+	}
+
 	/// Apply a chain update to the primary wallet.
 	///
 	/// Returns the wallet events and a list of all transaction IDs in the
@@ -1698,31 +1731,52 @@ mod tests {
 		P: WalletPersister,
 		P::Error: std::fmt::Debug,
 	{
-		PersistedWallet::create(
-			persister,
-			Wallet::create(
-				Bip84(xprv, KeychainKind::External),
-				Bip84(xprv, KeychainKind::Internal),
-			)
-			.network(Network::Regtest),
+		create_empty_wallet_from_xprv_with_lookahead(persister, xprv, None)
+	}
+
+	fn create_empty_wallet_from_xprv_with_lookahead<P>(
+		persister: &mut P, xprv: Xpriv, lookahead: Option<u32>,
+	) -> PersistedWallet<P>
+	where
+		P: WalletPersister,
+		P::Error: std::fmt::Debug,
+	{
+		let mut params = Wallet::create(
+			Bip84(xprv, KeychainKind::External),
+			Bip84(xprv, KeychainKind::Internal),
 		)
-		.unwrap()
+		.network(Network::Regtest);
+		if let Some(lookahead) = lookahead {
+			params = params.lookahead(lookahead);
+		}
+		PersistedWallet::create(persister, params).unwrap()
 	}
 
 	fn load_empty_wallet<P>(persister: &mut P) -> PersistedWallet<P>
+	where
+		P: WalletPersister,
+		P::Error: std::fmt::Debug,
+	{
+		load_empty_wallet_with_lookahead(persister, None)
+	}
+
+	fn load_empty_wallet_with_lookahead<P>(
+		persister: &mut P, lookahead: Option<u32>,
+	) -> PersistedWallet<P>
 	where
 		P: WalletPersister,
 		P::Error: std::fmt::Debug,
 	{
 		let xprv = test_xprv();
-		Wallet::load()
+		let mut params = Wallet::load()
 			.descriptor(KeychainKind::External, Some(Bip84(xprv, KeychainKind::External)))
 			.descriptor(KeychainKind::Internal, Some(Bip84(xprv, KeychainKind::Internal)))
 			.extract_keys()
-			.check_network(Network::Regtest)
-			.load_wallet(persister)
-			.unwrap()
-			.expect("wallet should have been persisted")
+			.check_network(Network::Regtest);
+		if let Some(lookahead) = lookahead {
+			params = params.lookahead(lookahead);
+		}
+		params.load_wallet(persister).unwrap().expect("wallet should have been persisted")
 	}
 
 	#[test]
@@ -2283,6 +2337,80 @@ mod tests {
 		assert_ne!(next.address, peeked.address);
 	}
 
+	#[test]
+	fn incremental_lookahead_does_not_advance_receive_index() {
+		let mut persister = NoopPersister;
+		let wallet =
+			create_empty_wallet_from_xprv_with_lookahead(&mut persister, test_xprv(), Some(3));
+		let mut aggregate = AggregateWallet::<u8, _>::new(wallet, persister, 0, vec![]);
+
+		let request = aggregate.wallet_incremental_sync_request_with_lookahead(&0).unwrap();
+		assert_eq!(request.progress().spks_remaining, 3);
+		assert_eq!(aggregate.new_address_info_for(&0).unwrap().index, 0);
+	}
+
+	#[test]
+	fn incremental_lookahead_rolls_forward_after_activity() {
+		let mut persister = NoopPersister;
+		let wallet =
+			create_empty_wallet_from_xprv_with_lookahead(&mut persister, test_xprv(), Some(3));
+		let mut aggregate = AggregateWallet::<u8, _>::new(wallet, persister, 0, vec![]);
+		let address = aggregate.address_info_for(&0, KeychainKind::External, 2).unwrap().address;
+		let transaction = Transaction {
+			version: bitcoin::transaction::Version::TWO,
+			lock_time: LockTime::ZERO,
+			input: vec![TxIn {
+				previous_output: OutPoint { txid: Txid::from_byte_array([0x18; 32]), vout: 0 },
+				..Default::default()
+			}],
+			output: vec![TxOut {
+				value: Amount::from_sat(50_000),
+				script_pubkey: address.script_pubkey(),
+			}],
+		};
+
+		aggregate.apply_mempool_txs(vec![(transaction, 0)], vec![]).unwrap();
+
+		let request = aggregate.wallet_incremental_sync_request_with_lookahead(&0).unwrap();
+		assert_eq!(request.progress().spks_remaining, 6);
+		assert_eq!(aggregate.new_address_info_for(&0).unwrap().index, 3);
+	}
+
+	#[test]
+	fn incremental_lookahead_activity_persists_the_active_index() {
+		let mut create_persister = MemoryPersister::default();
+		let wallet = create_empty_wallet_from_xprv_with_lookahead(
+			&mut create_persister,
+			test_xprv(),
+			Some(3),
+		);
+		let aggregate_persister = create_persister.clone();
+		{
+			let mut aggregate =
+				AggregateWallet::<u8, _>::new(wallet, aggregate_persister, 0, vec![]);
+			let address =
+				aggregate.address_info_for(&0, KeychainKind::External, 2).unwrap().address;
+			let transaction = Transaction {
+				version: bitcoin::transaction::Version::TWO,
+				lock_time: LockTime::ZERO,
+				input: vec![TxIn {
+					previous_output: OutPoint { txid: Txid::from_byte_array([0x19; 32]), vout: 0 },
+					..Default::default()
+				}],
+				output: vec![TxOut {
+					value: Amount::from_sat(50_000),
+					script_pubkey: address.script_pubkey(),
+				}],
+			};
+			aggregate.apply_mempool_txs(vec![(transaction, 0)], vec![]).unwrap();
+		}
+
+		let mut reload_persister = create_persister.clone();
+		let wallet = load_empty_wallet_with_lookahead(&mut reload_persister, Some(3));
+		let mut aggregate = AggregateWallet::<u8, _>::new(wallet, reload_persister, 0, vec![]);
+		assert_eq!(aggregate.new_address_info_for(&0).unwrap().index, 3);
+	}
+
 	#[test]
 	fn address_info_for_index_supports_internal_keychain_without_advancing_receive_index() {
 		let mut persister = NoopPersister;

src/builder.rs

@@ -52,8 +52,8 @@ use crate::config::{
 	default_user_config, may_announce_channel, AddressType, AddressTypeRuntimeConfig,
 	AnnounceError, AsyncPaymentsRole, BitcoindRestClientConfig, Config, ElectrumSyncConfig,
 	EsploraSyncConfig, OnchainWalletAccount, RuntimeSyncIntervals, ScoringDecayParameters,
-	ScoringFeeParameters, DEFAULT_ESPLORA_SERVER_URL, DEFAULT_LOG_FILENAME, DEFAULT_LOG_LEVEL,
-	MAX_ONCHAIN_WALLET_ACCOUNT_INDEX, WALLET_KEYS_SEED_LEN,
+	ScoringFeeParameters, BDK_CLIENT_STOP_GAP, DEFAULT_ESPLORA_SERVER_URL, DEFAULT_LOG_FILENAME,
+	DEFAULT_LOG_LEVEL, MAX_ONCHAIN_WALLET_ACCOUNT_INDEX, WALLET_KEYS_SEED_LEN,
 };
 use crate::connection::ConnectionManager;
 use crate::event::EventQueue;
@@ -115,6 +115,18 @@ enum ChainDataSourceConfig {
 	},
 }
 
+fn derived_account_lookahead(chain_data_source_config: Option<&ChainDataSourceConfig>) -> u32 {
+	match chain_data_source_config {
+		Some(ChainDataSourceConfig::Electrum { sync_config, .. }) => {
+			sync_config.unwrap_or_default().additional_wallet_full_scan_stop_gap.max(1)
+		},
+		Some(ChainDataSourceConfig::Bitcoind { .. }) => {
+			bdk_chain::keychain_txout::DEFAULT_LOOKAHEAD
+		},
+		Some(ChainDataSourceConfig::Esplora { .. }) | None => BDK_CLIENT_STOP_GAP as u32,
+	}
+}
+
 #[derive(Debug, Clone)]
 enum EntropySourceConfig {
 	SeedFile(String),
@@ -1350,42 +1362,49 @@ fn account_keychain_derivation_path(
 /// persisted wallet was restored.
 pub(crate) fn get_or_create_wallet_for_account(
 	wallet_account: OnchainWalletAccount, master_xprv: Xpriv, network: Network,
-	persister: &mut KVStoreWalletPersister,
+	persister: &mut KVStoreWalletPersister, lookahead: Option<u32>,
 ) -> Result<(PersistedWallet<KVStoreWalletPersister>, bool), BuildError> {
 	macro_rules! load_or_create {
 		($ext:expr, $int:expr) => {{
-			let wallet_opt = BdkWallet::load()
+			let mut load_params = BdkWallet::load()
 				.descriptor(KeychainKind::External, Some($ext))
 				.descriptor(KeychainKind::Internal, Some($int))
 				.extract_keys()
-				.check_network(network)
-				.load_wallet(persister)
-				.map_err(|e| match e {
-					bdk_wallet::LoadWithPersistError::InvalidChangeSet(
-						bdk_wallet::LoadError::Mismatch(bdk_wallet::LoadMismatch::Network {
-							loaded,
-							expected,
-						}),
-					) => {
-						log::error!(
-							"Failed to setup wallet: Networks do not match. Expected {} but got {}",
-							expected,
-							loaded
-						);
-						BuildError::NetworkMismatch
-					},
-					_ => {
-						log::error!("Failed to set up wallet: {}", e);
-						BuildError::WalletSetupFailed
-					},
-				})?;
+				.check_network(network);
+			if let Some(lookahead) = lookahead {
+				load_params = load_params.lookahead(lookahead);
+			}
+			let wallet_opt = load_params.load_wallet(persister).map_err(|e| match e {
+				bdk_wallet::LoadWithPersistError::InvalidChangeSet(
+					bdk_wallet::LoadError::Mismatch(bdk_wallet::LoadMismatch::Network {
+						loaded,
+						expected,
+					}),
+				) => {
+					log::error!(
+						"Failed to setup wallet: Networks do not match. Expected {} but got {}",
+						expected,
+						loaded
+					);
+					BuildError::NetworkMismatch
+				},
+				_ => {
+					log::error!("Failed to set up wallet: {}", e);
+					BuildError::WalletSetupFailed
+				},
+			})?;
 			match wallet_opt {
 				Some(w) => Ok((w, true)),
-				None => BdkWallet::create($ext, $int)
-					.network(network)
-					.create_wallet(persister)
-					.map(|w| (w, false))
-					.map_err(|_| BuildError::WalletSetupFailed),
+				None => {
+					let mut create_params = BdkWallet::create($ext, $int).network(network);
+					if let Some(lookahead) = lookahead {
+						create_params = create_params.lookahead(lookahead);
+					}
+					create_params
+						.create_wallet(persister)
+						.map(|w| (w, false))
+						.map_err(|_| BuildError::WalletSetupFailed)
+				},
 			}
 		}};
 	}
@@ -1484,20 +1503,25 @@ fn validate_configured_onchain_wallet_accounts(
 
 fn build_wallet_for_account(
 	config: &Config, wallet_account: OnchainWalletAccount, xprv: Xpriv, kv_store: Arc<DynStore>,
-	logger: Arc<Logger>, chain_tip_opt: Option<&BestBlock>,
+	logger: Arc<Logger>, chain_tip_opt: Option<&BestBlock>, lookahead: Option<u32>,
 ) -> Result<
 	(OnchainWalletAccount, PersistedWallet<KVStoreWalletPersister>, KVStoreWalletPersister),
 	BuildError,
 > {
 	let mut persister =
 		KVStoreWalletPersister::new(Arc::clone(&kv_store), Arc::clone(&logger), wallet_account);
 
-	let (mut wallet, loaded_from_store) =
-		get_or_create_wallet_for_account(wallet_account, xprv, config.network, &mut persister)
-			.map_err(|e| {
-				log_error!(logger, "Failed to setup wallet for {:?}: {}", wallet_account, e);
-				e
-			})?;
+	let (mut wallet, loaded_from_store) = get_or_create_wallet_for_account(
+		wallet_account,
+		xprv,
+		config.network,
+		&mut persister,
+		lookahead,
+	)
+	.map_err(|e| {
+		log_error!(logger, "Failed to setup wallet for {:?}: {}", wallet_account, e);
+		e
+	})?;
 
 	// Re-loaded wallets keep their persisted tip so chain backends can catch up. New wallets start
 	// at the known tip; Esplora and Electrum recover history independently with a full scan, while
@@ -1522,6 +1546,7 @@ fn build_wallet_for_account(
 fn build_additional_wallets(
 	config: &Config, configured_accounts: &[OnchainWalletAccount], xprv: Xpriv,
 	kv_store: Arc<DynStore>, logger: Arc<Logger>, chain_tip_opt: Option<&BestBlock>,
+	derived_account_lookahead: u32,
 ) -> Result<
 	Vec<(OnchainWalletAccount, PersistedWallet<KVStoreWalletPersister>, KVStoreWalletPersister)>,
 	BuildError,
@@ -1537,6 +1562,7 @@ fn build_additional_wallets(
 			Arc::clone(&kv_store),
 			Arc::clone(&logger),
 			chain_tip_opt,
+			None,
 		) {
 			Ok(tuple) => {
 				log_info!(logger, "Created additional wallet for {:?}", tuple.0);
@@ -1561,6 +1587,7 @@ fn build_additional_wallets(
 			Arc::clone(&kv_store),
 			Arc::clone(&logger),
 			chain_tip_opt,
+			Some(derived_account_lookahead),
 		)?;
 		log_info!(logger, "Loaded configured wallet for {:?}", wallet_account);
 		additional_wallets.push(wallet);
@@ -1755,6 +1782,7 @@ fn build_with_store_internal(
 		BuildError::InvalidSeedBytes
 	})?;
 	let configured_accounts = validate_configured_onchain_wallet_accounts(&config, xprv)?;
+	let derived_account_lookahead = derived_account_lookahead(chain_data_source_config);
 
 	let tx_broadcaster = Arc::new(TransactionBroadcaster::new(Arc::clone(&logger)));
 	let fee_estimator = Arc::new(OnchainFeeEstimator::new());
@@ -1785,6 +1813,7 @@ fn build_with_store_internal(
 						xprv,
 						network,
 						&mut persister,
+						None,
 					)
 					.map(|(wallet, _loaded)| wallet);
 					(result, persister)
@@ -1951,6 +1980,7 @@ fn build_with_store_internal(
 		Arc::clone(&kv_store),
 		Arc::clone(&logger),
 		chain_tip_opt.as_ref(),
+		derived_account_lookahead,
 	)?;
 
 	let wallet = Arc::new(Wallet::new(
@@ -1966,6 +1996,7 @@ fn build_with_store_internal(
 		Arc::clone(&address_type_runtime_config),
 		Arc::clone(&node_metrics),
 		Arc::clone(&logger),
+		derived_account_lookahead,
 	));
 
 	// Initialize the KeysManager
@@ -3291,7 +3322,8 @@ mod tests {
 		let mut persister =
 			KVStoreWalletPersister::new(Arc::clone(&store), Arc::clone(&logger), account);
 		let (wallet, loaded) =
-			get_or_create_wallet_for_account(account, master, network, &mut persister).unwrap();
+			get_or_create_wallet_for_account(account, master, network, &mut persister, None)
+				.unwrap();
 		assert!(!loaded);
 
 		let descriptor = wallet.public_descriptor(KeychainKind::External).to_string();