Merge #672: Fix wallet sync not finding coins of addresses which are not cached

5c940c3 Fix wallet sync not finding coins of addresses which are not cached (志宇)

Pull request description:

  Fixes #521
  Fixes #451

  ^ However, only for electrum-based `Blockchain` implementations. For RPC and Compact Block Filters, syncing works differently, and so are the bugs - I have created a separate ticket for this (#677).

  ### Description

  Previously, electrum-based blockchain implementations only synced for `scriptPubKey`s that are already cached in `Database`.

  This PR introduces a feedback mechanism, that uses `stop_gap` and the difference between "current index" and "last active index" to determine whether we need to cache more `scriptPubKeys`.

  The `WalletSync::wallet_setup` trait now may return an `Error::MissingCachedScripts` error which contains the number of extra `scriptPubKey`s to cache, in order to satisfy `stop_gap` for the next call.

  `Wallet::sync` now calls `WalletSync` in a loop, caching in-between subsequent calls (if needed).

  #### Notes to reviewers

  1. The caveat to this solution is that it is not the most efficient. Every call to `WalletSync::wallet_setup` starts polling the Electrum-based server for `scriptPubKey`s starting from index 0.

      However, I feel like this solution is the least "destructive" to the API of `Blockchain`. Also, once the `bdk_core` sync logic is integration, we can select specific ranges of `scriptPubKey`s to sync.

  2. Also note that this PR only fixes the issue for electrum-based `Blockchain` implementations (currently `blockchain::electrum` and `blockchain::esplora` only).

  3. Another thing to note is that, although `Database` assumes 1-2 keychains, the current `WalletSync` "feedback" only returns one number (which is interpreted as the larger "missing count" of the two keychains). This is done for simplicity, and because we are planning to only have one keychain per database in the future.

      https://github.com/bitcoindevkit/bdk/blob/f0c876e7bf38566c0d224cbe421ee312ffc06660/src/blockchain/mod.rs#L157-L161

  4. Please have a read of #672 (comment) for additional context.

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo fmt` and `cargo clippy` before committing

  #### Bugfixes:

  * [x] This pull request breaks the existing API
  * [x] I've added tests to reproduce the issue which are now passing
  * [x] I'm linking the issue being fixed by this PR

ACKs for top commit:
  afilini:
    ACK 5c940c3

Tree-SHA512: aee917ed4821438fc0675241432a7994603a09a77d5a72e96bad863e7cdd55a9bc6fbd931ce096fef1153905cf1b786e1d8d932dc19032d549480bcda7c75d1b

Author: afilini

src/blockchain/esplora/reqwest.rs

@@ -213,7 +213,6 @@ impl WalletSync for EsploraBlockchain {
         };
 
         database.commit_batch(batch_update)?;
-
         Ok(())
     }
 }

src/blockchain/script_sync.rs

@@ -5,6 +5,7 @@ returns associated transactions i.e. electrum.
 #![allow(dead_code)]
 use crate::{
     database::{BatchDatabase, BatchOperations, DatabaseUtils},
+    error::MissingCachedScripts,
     wallet::time::Instant,
     BlockTime, Error, KeychainKind, LocalUtxo, TransactionDetails,
 };
@@ -34,11 +35,12 @@ pub fn start<D: BatchDatabase>(db: &D, stop_gap: usize) -> Result<Request<'_, D>
     let scripts_needed = db
         .iter_script_pubkeys(Some(keychain))?
         .into_iter()
-        .collect();
+        .collect::<VecDeque<_>>();
     let state = State::new(db);
 
     Ok(Request::Script(ScriptReq {
         state,
+        initial_scripts_needed: scripts_needed.len(),
         scripts_needed,
         script_index: 0,
         stop_gap,
@@ -50,6 +52,7 @@ pub fn start<D: BatchDatabase>(db: &D, stop_gap: usize) -> Result<Request<'_, D>
 pub struct ScriptReq<'a, D: BatchDatabase> {
     state: State<'a, D>,
     script_index: usize,
+    initial_scripts_needed: usize, // if this is 1, we assume the descriptor is not derivable
     scripts_needed: VecDeque<Script>,
     stop_gap: usize,
     keychain: KeychainKind,
@@ -113,43 +116,71 @@ impl<'a, D: BatchDatabase> ScriptReq<'a, D> {
             self.script_index += 1;
         }
 
-        for _ in txids {
-            self.scripts_needed.pop_front();
-        }
+        self.scripts_needed.drain(..txids.len());
 
-        let last_active_index = self
+        // last active index: 0 => No last active
+        let last = self
             .state
             .last_active_index
             .get(&self.keychain)
-            .map(|x| x + 1)
-            .unwrap_or(0); // so no addresses active maps to 0
-
-        Ok(
-            if self.script_index > last_active_index + self.stop_gap
-                || self.scripts_needed.is_empty()
-            {
-                debug!(
-                    "finished scanning for transactions for keychain {:?} at index {}",
-                    self.keychain, last_active_index
-                );
-                // we're done here -- check if we need to do the next keychain
-                if let Some(keychain) = self.next_keychains.pop() {
-                    self.keychain = keychain;
-                    self.script_index = 0;
-                    self.scripts_needed = self
-                        .state
-                        .db
-                        .iter_script_pubkeys(Some(keychain))?
-                        .into_iter()
-                        .collect();
-                    Request::Script(self)
-                } else {
-                    Request::Tx(TxReq { state: self.state })
-                }
-            } else {
-                Request::Script(self)
-            },
-        )
+            .map(|&l| l + 1)
+            .unwrap_or(0);
+        // remaining scripts left to check
+        let remaining = self.scripts_needed.len();
+        // difference between current index and last active index
+        let current_gap = self.script_index - last;
+
+        // this is a hack to check whether the scripts are coming from a derivable descriptor
+        // we assume for non-derivable descriptors, the initial script count is always 1
+        let is_derivable = self.initial_scripts_needed > 1;
+
+        debug!(
+            "sync: last={}, remaining={}, diff={}, stop_gap={}",
+            last, remaining, current_gap, self.stop_gap
+        );
+
+        if is_derivable {
+            if remaining > 0 {
+                // we still have scriptPubKeys to do requests for
+                return Ok(Request::Script(self));
+            }
+
+            if last > 0 && current_gap < self.stop_gap {
+                // current gap is not large enough to stop, but we are unable to keep checking since
+                // we have exhausted cached scriptPubKeys, so return error
+                let err = MissingCachedScripts {
+                    last_count: self.script_index,
+                    missing_count: self.stop_gap - current_gap,
+                };
+                return Err(Error::MissingCachedScripts(err));
+            }
+
+            // we have exhausted cached scriptPubKeys and found no txs, continue
+        }
+
+        debug!(
+            "finished scanning for txs of keychain {:?} at index {:?}",
+            self.keychain, last
+        );
+
+        if let Some(keychain) = self.next_keychains.pop() {
+            // we still have another keychain to request txs with
+            let scripts_needed = self
+                .state
+                .db
+                .iter_script_pubkeys(Some(keychain))?
+                .into_iter()
+                .collect::<VecDeque<_>>();
+
+            self.keychain = keychain;
+            self.script_index = 0;
+            self.initial_scripts_needed = scripts_needed.len();
+            self.scripts_needed = scripts_needed;
+            return Ok(Request::Script(self));
+        }
+
+        // We have finished requesting txids, let's get the actual txs.
+        Ok(Request::Tx(TxReq { state: self.state }))
     }
 }
 
@@ -294,6 +325,8 @@ struct State<'a, D> {
     tx_missing_conftime: BTreeMap<Txid, TransactionDetails>,
     /// The start of the sync
     start_time: Instant,
+    /// Missing number of scripts to cache per keychain
+    missing_script_counts: HashMap<KeychainKind, usize>,
 }
 
 impl<'a, D: BatchDatabase> State<'a, D> {
@@ -305,6 +338,7 @@ impl<'a, D: BatchDatabase> State<'a, D> {
             tx_needed: BTreeSet::default(),
             tx_missing_conftime: BTreeMap::default(),
             start_time: Instant::new(),
+            missing_script_counts: HashMap::default(),
         }
     }
     fn into_db_update(self) -> Result<D::Batch, Error> {

src/error.rs

@@ -13,7 +13,7 @@ use std::fmt;
 
 use crate::bitcoin::Network;
 use crate::{descriptor, wallet, wallet::address_validator};
-use bitcoin::OutPoint;
+use bitcoin::{OutPoint, Txid};
 
 /// Errors that can be thrown by the [`Wallet`](crate::wallet::Wallet)
 #[derive(Debug)]
@@ -125,6 +125,10 @@ pub enum Error {
     //DifferentDescriptorStructure,
     //Uncapable(crate::blockchain::Capability),
     //MissingCachedAddresses,
+    /// [`crate::blockchain::WalletSync`] sync attempt failed due to missing scripts in cache which
+    /// are needed to satisfy `stop_gap`.
+    MissingCachedScripts(MissingCachedScripts),
+
     #[cfg(feature = "electrum")]
     /// Electrum client error
     Electrum(electrum_client::Error),
@@ -145,6 +149,16 @@ pub enum Error {
     Rusqlite(rusqlite::Error),
 }
 
+/// Represents the last failed [`crate::blockchain::WalletSync`] sync attempt in which we were short
+/// on cached `scriptPubKey`s.
+#[derive(Debug)]
+pub struct MissingCachedScripts {
+    /// Number of scripts in which txs were requested during last request.
+    pub last_count: usize,
+    /// Minimum number of scripts to cache more of in order to satisfy `stop_gap`.
+    pub missing_count: usize,
+}
+
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{:?}", self)

src/testutils/blockchain_tests.rs

@@ -743,7 +743,6 @@ macro_rules! bdk_blockchain_tests {
 
                 blockchain.broadcast(&tx1).expect("broadcasting first");
                 blockchain.broadcast(&tx2).expect("broadcasting replacement");
-
                 receiver_wallet.sync(&blockchain, SyncOptions::default()).expect("syncing receiver");
                 assert_eq!(receiver_wallet.get_balance().expect("balance"), 49_000, "should have received coins once and only once");
             }

src/testutils/configurable_blockchain_tests.rs

@@ -29,6 +29,8 @@ pub trait ConfigurableBlockchainTester<B: ConfigurableBlockchain>: Sized {
 
         if self.config_with_stop_gap(test_client, 0).is_some() {
             test_wallet_sync_with_stop_gaps(test_client, self);
+            test_wallet_sync_fulfills_missing_script_cache(test_client, self);
+            test_wallet_sync_self_transfer_tx(test_client, self);
         } else {
             println!(
                 "{}: Skipped tests requiring config_with_stop_gap.",
@@ -113,16 +115,21 @@ where
         } else {
             max_balance
         };
+        let details = format!(
+            "test_vector: [stop_gap: {}, actual_gap: {}, addrs_before: {}, addrs_after: {}]",
+            stop_gap, actual_gap, addrs_before, addrs_after,
+        );
+        println!("{}", details);
 
         // perform wallet sync
         wallet.sync(&blockchain, Default::default()).unwrap();
 
         let wallet_balance = wallet.get_balance().unwrap();
-
-        let details = format!(
-            "test_vector: [stop_gap: {}, actual_gap: {}, addrs_before: {}, addrs_after: {}]",
-            stop_gap, actual_gap, addrs_before, addrs_after,
+        println!(
+            "max: {}, min: {}, actual: {}",
+            max_balance, min_balance, wallet_balance
         );
+
         assert!(
             wallet_balance <= max_balance,
             "wallet balance is greater than received amount: {}",
@@ -138,3 +145,113 @@ where
         test_client.generate(1, None);
     }
 }
+
+/// With a `stop_gap` of x and every x addresses having a balance of 1000 (for y addresses),
+/// we expect `Wallet::sync` to correctly self-cache addresses, so that the resulting balance,
+/// after sync, should be y * 1000.
+fn test_wallet_sync_fulfills_missing_script_cache<T, B>(test_client: &mut TestClient, tester: &T)
+where
+    T: ConfigurableBlockchainTester<B>,
+    B: ConfigurableBlockchain,
+{
+    // wallet descriptor
+    let descriptor = "wpkh([c258d2e4/84h/1h/0h]tpubDDYkZojQFQjht8Tm4jsS3iuEmKjTiEGjG6KnuFNKKJb5A6ZUCUZKdvLdSDWofKi4ToRCwb9poe1XdqfUnP4jaJjCB2Zwv11ZLgSbnZSNecE/200/*)";
+
+    // amount in sats per tx
+    const AMOUNT_PER_TX: u64 = 1000;
+
+    // addr constants
+    const ADDR_COUNT: usize = 6;
+    const ADDR_GAP: usize = 60;
+
+    let blockchain =
+        B::from_config(&tester.config_with_stop_gap(test_client, ADDR_GAP).unwrap()).unwrap();
+
+    let wallet = Wallet::new(descriptor, None, Network::Regtest, MemoryDatabase::new()).unwrap();
+
+    let expected_balance = (0..ADDR_COUNT).fold(0_u64, |sum, i| {
+        let addr_i = i * ADDR_GAP;
+        let address = wallet.get_address(AddressIndex::Peek(addr_i as _)).unwrap();
+
+        println!(
+            "tx: {} sats => [{}] {}",
+            AMOUNT_PER_TX,
+            addr_i,
+            address.to_string()
+        );
+
+        test_client.receive(testutils! {
+            @tx ( (@addr address.address) => AMOUNT_PER_TX )
+        });
+        test_client.generate(1, None);
+
+        sum + AMOUNT_PER_TX
+    });
+    println!("expected balance: {}, syncing...", expected_balance);
+
+    // perform sync
+    wallet.sync(&blockchain, Default::default()).unwrap();
+    println!("sync done!");
+
+    let balance = wallet.get_balance().unwrap();
+    assert_eq!(balance, expected_balance);
+}
+
+/// Given a `stop_gap`, a wallet with a 2 transactions, one sending to `scriptPubKey` at derivation
+/// index of `stop_gap`, and the other spending from the same `scriptPubKey` into another
+/// `scriptPubKey` at derivation index of `stop_gap * 2`, we expect `Wallet::sync` to perform
+/// correctly, so that we detect the total balance.
+fn test_wallet_sync_self_transfer_tx<T, B>(test_client: &mut TestClient, tester: &T)
+where
+    T: ConfigurableBlockchainTester<B>,
+    B: ConfigurableBlockchain,
+{
+    const TRANSFER_AMOUNT: u64 = 10_000;
+    const STOP_GAP: usize = 75;
+
+    let descriptor = "wpkh(tprv8i8F4EhYDMquzqiecEX8SKYMXqfmmb1Sm7deoA1Hokxzn281XgTkwsd6gL8aJevLE4aJugfVf9MKMvrcRvPawGMenqMBA3bRRfp4s1V7Eg3/*)";
+
+    let blockchain =
+        B::from_config(&tester.config_with_stop_gap(test_client, STOP_GAP).unwrap()).unwrap();
+
+    let wallet = Wallet::new(descriptor, None, Network::Regtest, MemoryDatabase::new()).unwrap();
+
+    let address1 = wallet
+        .get_address(AddressIndex::Peek(STOP_GAP as _))
+        .unwrap();
+    let address2 = wallet
+        .get_address(AddressIndex::Peek((STOP_GAP * 2) as _))
+        .unwrap();
+
+    test_client.receive(testutils! {
+        @tx ( (@addr address1.address) => TRANSFER_AMOUNT )
+    });
+    test_client.generate(1, None);
+
+    wallet.sync(&blockchain, Default::default()).unwrap();
+
+    let mut builder = wallet.build_tx();
+    builder.add_recipient(address2.script_pubkey(), TRANSFER_AMOUNT / 2);
+    let (mut psbt, details) = builder.finish().unwrap();
+    assert!(wallet.sign(&mut psbt, Default::default()).unwrap());
+    blockchain.broadcast(&psbt.extract_tx()).unwrap();
+
+    test_client.generate(1, None);
+
+    // obtain what is expected
+    let fee = details.fee.unwrap();
+    let expected_balance = TRANSFER_AMOUNT - fee;
+    println!("fee={}, expected_balance={}", fee, expected_balance);
+
+    // actually test the wallet
+    wallet.sync(&blockchain, Default::default()).unwrap();
+    let balance = wallet.get_balance().unwrap();
+    assert_eq!(balance, expected_balance);
+
+    // now try with a fresh wallet
+    let fresh_wallet =
+        Wallet::new(descriptor, None, Network::Regtest, MemoryDatabase::new()).unwrap();
+    fresh_wallet.sync(&blockchain, Default::default()).unwrap();
+    let fresh_balance = fresh_wallet.get_balance().unwrap();
+    assert_eq!(fresh_balance, expected_balance);
+}