fix(sync): add cooperative hard-cancel support

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bdk"
-version = "0.30.2"
+version = "0.30.4"
 authors = ["Alekos Filini <alekos.filini@gmail.com>", "Riccardo Casatta <riccardo@casatta.it>"]
 homepage = "https://bitcoindevkit.org"
 repository = "https://github.com/bitcoindevkit/bdk"

src/blockchain/any.rs

@@ -129,6 +129,21 @@ impl GetBlockHash for AnyBlockchain {
 
 #[maybe_async]
 impl WalletSync for AnyBlockchain {
+    fn wallet_setup_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(impl_inner_method!(
+            self,
+            wallet_setup_with_control,
+            database,
+            progress_update,
+            control
+        ))
+    }
+
     fn wallet_sync<D: BatchDatabase>(
         &self,
         database: &RefCell<D>,
@@ -154,6 +169,21 @@ impl WalletSync for AnyBlockchain {
             progress_update
         ))
     }
+
+    fn wallet_sync_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(impl_inner_method!(
+            self,
+            wallet_sync_with_control,
+            database,
+            progress_update,
+            control
+        ))
+    }
 }
 
 impl_from!(boxed electrum::ElectrumBlockchain, AnyBlockchain, Electrum, #[cfg(feature = "electrum")]);

src/blockchain/electrum.rs

@@ -114,12 +114,14 @@ impl GetBlockHash for ElectrumBlockchain {
     }
 }
 
-impl WalletSync for ElectrumBlockchain {
-    fn wallet_setup<D: BatchDatabase>(
+impl ElectrumBlockchain {
+    fn wallet_setup_impl<D: BatchDatabase>(
         &self,
         database: &RefCell<D>,
-        _progress_update: Box<dyn Progress>,
+        control: &SyncControl,
     ) -> Result<(), Error> {
+        control.check_cancelled()?;
+
         let mut database = database.borrow_mut();
         let database = database.deref_mut();
         let mut request = script_sync::start(database, self.stop_gap)?;
@@ -138,6 +140,8 @@ impl WalletSync for ElectrumBlockchain {
         let batch_update = loop {
             request = match request {
                 Request::Script(script_req) => {
+                    control.check_cancelled()?;
+
                     let scripts = script_req.request().take(chunk_size);
                     let txids_per_script: Vec<Vec<_>> = self
                         .client
@@ -164,6 +168,8 @@ impl WalletSync for ElectrumBlockchain {
                 }
 
                 Request::Conftime(conftime_req) => {
+                    control.check_cancelled()?;
+
                     // collect up to chunk_size heights to fetch from electrum
                     let needs_block_height = conftime_req
                         .request()
@@ -202,8 +208,10 @@ impl WalletSync for ElectrumBlockchain {
                     conftime_req.satisfy(conftimes)?
                 }
                 Request::Tx(tx_req) => {
+                    control.check_cancelled()?;
+
                     let needs_full = tx_req.request().take(chunk_size);
-                    tx_cache.save_txs(needs_full.clone())?;
+                    tx_cache.save_txs(needs_full.clone(), control)?;
                     let full_transactions = needs_full
                         .map(|txid| tx_cache.get(*txid).ok_or_else(electrum_goof))
                         .collect::<Result<Vec<_>, _>>()?;
@@ -213,7 +221,7 @@ impl WalletSync for ElectrumBlockchain {
                             .filter(|input| !input.previous_output.is_null())
                             .map(|input| &input.previous_output.txid)
                     });
-                    tx_cache.save_txs(input_txs)?;
+                    tx_cache.save_txs(input_txs, control)?;
 
                     let full_details = full_transactions
                         .into_iter()
@@ -247,11 +255,40 @@ impl WalletSync for ElectrumBlockchain {
             }
         };
 
+        control.check_cancelled()?;
         database.commit_batch(batch_update)?;
         Ok(())
     }
 }
 
+impl WalletSync for ElectrumBlockchain {
+    fn wallet_setup<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        _progress_update: Box<dyn Progress>,
+    ) -> Result<(), Error> {
+        self.wallet_setup_impl(database, &SyncControl::default())
+    }
+
+    fn wallet_setup_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        _progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        self.wallet_setup_impl(database, &control)
+    }
+
+    fn wallet_sync_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        _progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        self.wallet_setup_impl(database, &control)
+    }
+}
+
 struct TxCache<'a, 'b, D> {
     db: &'a D,
     client: &'b Client,
@@ -266,9 +303,14 @@ impl<'a, 'b, D: Database> TxCache<'a, 'b, D> {
             cache: HashMap::default(),
         }
     }
-    fn save_txs<'c>(&mut self, txids: impl Iterator<Item = &'c Txid>) -> Result<(), Error> {
+    fn save_txs<'c>(
+        &mut self,
+        txids: impl Iterator<Item = &'c Txid>,
+        control: &SyncControl,
+    ) -> Result<(), Error> {
         let mut need_fetch = vec![];
         for txid in txids {
+            control.check_cancelled()?;
             if self.cache.contains_key(txid) {
                 continue;
             } else if let Some(transaction) = self.db.get_raw_tx(txid)? {
@@ -282,6 +324,7 @@ impl<'a, 'b, D: Database> TxCache<'a, 'b, D> {
         // of transactions at once, which creates enormous memory pressure. By chunking the batch
         // into more reasonably sized sub-queries, we allow time for memory to be freed.
         for chunk in need_fetch.chunks(1000) {
+            control.check_cancelled()?;
             let txs = self
                 .client
                 .batch_transaction_get(chunk)

src/blockchain/mod.rs

@@ -19,8 +19,10 @@
 use std::cell::RefCell;
 use std::collections::HashSet;
 use std::ops::Deref;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc::{channel, Receiver, Sender};
 use std::sync::Arc;
+use std::time::Instant;
 
 use bitcoin::{BlockHash, Transaction, Txid};
 
@@ -119,6 +121,35 @@ pub trait GetBlockHash {
 }
 
 /// Trait for blockchains that can sync by updating the database directly.
+#[derive(Debug, Clone, Default)]
+pub struct SyncControl {
+    /// Shared cancellation flag observed by sync checkpoints.
+    pub cancel: Option<Arc<AtomicBool>>,
+    /// Deadline after which sync should abort at the next checkpoint.
+    pub deadline: Option<Instant>,
+}
+
+impl SyncControl {
+    /// Return true if cancellation should occur.
+    pub fn is_cancelled(&self) -> bool {
+        self.cancel
+            .as_ref()
+            .map_or(false, |cancel| cancel.load(Ordering::Relaxed))
+            || self
+                .deadline
+                .map_or(false, |deadline| Instant::now() >= deadline)
+    }
+
+    /// Return an error when cancellation is requested.
+    pub fn check_cancelled(&self) -> Result<(), Error> {
+        if self.is_cancelled() {
+            return Err(Error::SyncCancelled);
+        }
+
+        Ok(())
+    }
+}
+
 #[maybe_async]
 pub trait WalletSync {
     /// Setup the backend and populate the internal database for the first time
@@ -138,6 +169,18 @@ pub trait WalletSync {
         progress_update: Box<dyn Progress>,
     ) -> Result<(), Error>;
 
+    /// Like [`Self::wallet_setup`], but with cooperative cancellation support.
+    ///
+    /// Default implementation preserves current behavior and ignores `control`.
+    fn wallet_setup_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        _control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(self.wallet_setup(database, progress_update))
+    }
+
     /// If not overridden, it defaults to calling [`Self::wallet_setup`] internally.
     ///
     /// This method should implement the logic required to iterate over the list of the wallet's
@@ -162,6 +205,18 @@ pub trait WalletSync {
     ) -> Result<(), Error> {
         maybe_await!(self.wallet_setup(database, progress_update))
     }
+
+    /// Like [`Self::wallet_sync`], but with cooperative cancellation support.
+    ///
+    /// Default implementation preserves current behavior and ignores `control`.
+    fn wallet_sync_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        _control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(self.wallet_sync(database, progress_update))
+    }
 }
 
 /// Trait for [`Blockchain`] types that can be created given a configuration
@@ -381,11 +436,33 @@ impl<T: WalletSync> WalletSync for Arc<T> {
         maybe_await!(self.deref().wallet_setup(database, progress_update))
     }
 
+    fn wallet_setup_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(self
+            .deref()
+            .wallet_setup_with_control(database, progress_update, control))
+    }
+
     fn wallet_sync<D: BatchDatabase>(
         &self,
         database: &RefCell<D>,
         progress_update: Box<dyn Progress>,
     ) -> Result<(), Error> {
         maybe_await!(self.deref().wallet_sync(database, progress_update))
     }
+
+    fn wallet_sync_with_control<D: BatchDatabase>(
+        &self,
+        database: &RefCell<D>,
+        progress_update: Box<dyn Progress>,
+        control: SyncControl,
+    ) -> Result<(), Error> {
+        maybe_await!(self
+            .deref()
+            .wallet_sync_with_control(database, progress_update, control))
+    }
 }

src/error.rs

@@ -130,6 +130,8 @@ pub enum Error {
     /// [`crate::blockchain::WalletSync`] sync attempt failed due to missing scripts in cache which
     /// are needed to satisfy `stop_gap`.
     MissingCachedScripts(MissingCachedScripts),
+    /// [`crate::blockchain::WalletSync`] sync was cancelled by cooperative control.
+    SyncCancelled,
 
     #[cfg(feature = "electrum")]
     /// Electrum client error
@@ -258,6 +260,7 @@ impl fmt::Display for Error {
             Self::MissingCachedScripts(missing_cached_scripts) => {
                 write!(f, "Missing cached scripts: {:?}", missing_cached_scripts)
             }
+            Self::SyncCancelled => write!(f, "Wallet sync cancelled"),
             #[cfg(feature = "electrum")]
             Self::Electrum(err) => write!(f, "Electrum client error: {}", err),
             #[cfg(feature = "esplora")]