From 892b97d4416d9d38ebf3d7e0c52636f110042302 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= Date: Fri, 19 Jul 2024 07:05:38 +0000 Subject: [PATCH] feat(chain,wallet)!: Change persist-traits to be "safer" Previously, `Persist{Async}With::persist` can be directly called as a method on the type (i.e. `Wallet`). However, the `db: Db` (which is an input) may not be initialized. We want a design which makes it harder for the caller to make this mistake. We change `Persist{Async}With::persist` to be an "associated function" which takes two inputs: `db: &mut Db` and `changeset`. However, the implementer cannot take directly from `Self` (as it's no longer an input). So we introduce a new trait `Staged` which defines the staged changeset type and a method that gives us a `&mut` of the staged changes. --- crates/chain/src/persist.rs | 57 +++++++++++++++++++++------ crates/wallet/src/wallet/mod.rs | 9 +++++ crates/wallet/src/wallet/persisted.rs | 33 ++++++---------- 3 files changed, 66 insertions(+), 33 deletions(-) diff --git a/crates/chain/src/persist.rs b/crates/chain/src/persist.rs index 5f7b37e4..6bcdb6bd 100644 --- a/crates/chain/src/persist.rs +++ b/crates/chain/src/persist.rs @@ -6,10 +6,21 @@ use core::{ use alloc::boxed::Box; +use crate::Merge; + +/// Represents a type that contains staged changes. +pub trait Staged { + /// Type for staged changes. + type ChangeSet: Merge; + + /// Get mutable reference of staged changes. + fn staged(&mut self) -> &mut Self::ChangeSet; +} + /// Trait that persists the type with `Db`. /// /// Methods of this trait should not be called directly. -pub trait PersistWith: Sized { +pub trait PersistWith: Staged + Sized { /// Parameters for [`PersistWith::create`]. type CreateParams; /// Parameters for [`PersistWith::load`]. @@ -21,20 +32,23 @@ pub trait PersistWith: Sized { /// Error type of [`PersistWith::persist`]. type PersistError; - /// Create the type and initialize the `Db`. + /// Initialize the `Db` and create `Self`. fn create(db: &mut Db, params: Self::CreateParams) -> Result; - /// Load the type from the `Db`. + /// Initialize the `Db` and load a previously-persisted `Self`. fn load(db: &mut Db, params: Self::LoadParams) -> Result, Self::LoadError>; - /// Persist staged changes into `Db`. - fn persist(&mut self, db: &mut Db) -> Result; + /// Persist changes to the `Db`. + fn persist( + db: &mut Db, + changeset: &::ChangeSet, + ) -> Result<(), Self::PersistError>; } type FutureResult<'a, T, E> = Pin> + Send + 'a>>; /// Trait that persists the type with an async `Db`. -pub trait PersistAsyncWith: Sized { +pub trait PersistAsyncWith: Staged + Sized { /// Parameters for [`PersistAsyncWith::create`]. type CreateParams; /// Parameters for [`PersistAsyncWith::load`]. @@ -46,14 +60,17 @@ pub trait PersistAsyncWith: Sized { /// Error type of [`PersistAsyncWith::persist`]. type PersistError; - /// Create the type and initialize the `Db`. + /// Initialize the `Db` and create `Self`. fn create(db: &mut Db, params: Self::CreateParams) -> FutureResult; - /// Load the type from `Db`. + /// Initialize the `Db` and load a previously-persisted `Self`. fn load(db: &mut Db, params: Self::LoadParams) -> FutureResult, Self::LoadError>; - /// Persist staged changes into `Db`. - fn persist<'a>(&'a mut self, db: &'a mut Db) -> FutureResult<'a, bool, Self::PersistError>; + /// Persist changes to the `Db`. + fn persist<'a>( + db: &'a mut Db, + changeset: &'a ::ChangeSet, + ) -> FutureResult<'a, (), Self::PersistError>; } /// Represents a persisted `T`. @@ -102,14 +119,24 @@ impl Persisted { } /// Persist staged changes of `T` into `Db`. + /// + /// If the database errors, the staged changes will not be cleared. pub fn persist(&mut self, db: &mut Db) -> Result where T: PersistWith, { - self.inner.persist(db) + let stage = T::staged(&mut self.inner); + if stage.is_empty() { + return Ok(false); + } + T::persist(db, &*stage)?; + stage.take(); + Ok(true) } /// Persist staged changes of `T` into an async `Db`. + /// + /// If the database errors, the staged changes will not be cleared. pub async fn persist_async<'a, Db>( &'a mut self, db: &'a mut Db, @@ -117,7 +144,13 @@ impl Persisted { where T: PersistAsyncWith, { - self.inner.persist(db).await + let stage = T::staged(&mut self.inner); + if stage.is_empty() { + return Ok(false); + } + T::persist(db, &*stage).await?; + stage.take(); + Ok(true) } } diff --git a/crates/wallet/src/wallet/mod.rs b/crates/wallet/src/wallet/mod.rs index d2d7c0ad..84fe5289 100644 --- a/crates/wallet/src/wallet/mod.rs +++ b/crates/wallet/src/wallet/mod.rs @@ -42,6 +42,7 @@ use bitcoin::{ use bitcoin::{consensus::encode::serialize, transaction, BlockHash, Psbt}; use bitcoin::{constants::genesis_block, Amount}; use bitcoin::{secp256k1::Secp256k1, Weight}; +use chain::Staged; use core::fmt; use core::mem; use core::ops::Deref; @@ -119,6 +120,14 @@ pub struct Wallet { secp: SecpCtx, } +impl Staged for Wallet { + type ChangeSet = ChangeSet; + + fn staged(&mut self) -> &mut Self::ChangeSet { + &mut self.stage + } +} + /// An update to [`Wallet`]. /// /// It updates [`KeychainTxOutIndex`], [`bdk_chain::TxGraph`] and [`local_chain::LocalChain`] atomically. diff --git a/crates/wallet/src/wallet/persisted.rs b/crates/wallet/src/wallet/persisted.rs index 70d9f27c..5dfea3f2 100644 --- a/crates/wallet/src/wallet/persisted.rs +++ b/crates/wallet/src/wallet/persisted.rs @@ -41,14 +41,10 @@ impl<'c> chain::PersistWith> for Wallet { } fn persist( - &mut self, - conn: &mut bdk_chain::sqlite::Transaction, - ) -> Result { - if let Some(changeset) = self.take_staged() { - changeset.persist_to_sqlite(conn)?; - return Ok(true); - } - Ok(false) + db: &mut bdk_chain::sqlite::Transaction<'c>, + changeset: &::ChangeSet, + ) -> Result<(), Self::PersistError> { + changeset.persist_to_sqlite(db) } } @@ -82,13 +78,12 @@ impl chain::PersistWith for Wallet { } fn persist( - &mut self, db: &mut bdk_chain::sqlite::Connection, - ) -> Result { - let mut db_tx = db.transaction()?; - let has_changes = chain::PersistWith::persist(self, &mut db_tx)?; - db_tx.commit()?; - Ok(has_changes) + changeset: &::ChangeSet, + ) -> Result<(), Self::PersistError> { + let db_tx = db.transaction()?; + changeset.persist_to_sqlite(&db_tx)?; + db_tx.commit() } } @@ -126,14 +121,10 @@ impl chain::PersistWith> for Wallet { } fn persist( - &mut self, db: &mut bdk_file_store::Store, - ) -> Result { - if let Some(changeset) = self.take_staged() { - db.append_changeset(&changeset)?; - return Ok(true); - } - Ok(false) + changeset: &::ChangeSet, + ) -> Result<(), Self::PersistError> { + db.append_changeset(changeset) } }