WIP

Author: evanlinjin

crates/wallet/src/wallet/persisted.rs

@@ -10,48 +10,103 @@ use chain::Merge;
 
 use crate::{descriptor::DescriptorError, ChangeSet, CreateParams, LoadParams, Wallet};
 
-/// A trait for persisting [`Wallet`].
+/// Trait that persists [`Wallet`].
+///
+/// For an async version, use [`AsyncWalletPersister`].
+///
+/// Associated functions of this trait should not be called directly, and the trait is designed so
+/// that associated functions are hard to find (since they are not methods!). [`WalletPersister`] is
+/// used by [`PersistedWallet`] (a light wrapper around [`Wallet`]) which enforces some level of
+/// safety. Refer to [`PersistedWallet`] for more about the safety checks.
 pub trait WalletPersister {
-    /// Error when initializing the persister.
+    /// Error type of the persister.
     type Error;
 
-    /// Initialize the persister.
-    fn initialize(persister: &mut Self) -> Result<(), Self::Error>;
+    /// Initialize the `persister` and load all data.
+    ///
+    /// This is called by [`PersistedWallet::create`] and [`PersistedWallet::load`] to ensure
+    /// the [`WalletPersister`] is initialized and returns all data in the `persister`.
+    ///
+    /// # Implementation Details
+    ///
+    /// The database schema of the `persister` (if any), should be initialized and migrated here.
+    ///
+    /// The implementation must return all data currently stored in the `persister`. If there is no
+    /// data, return an empty changeset (using [`ChangeSet::default()`]).
+    ///
+    /// Error should only occur on database failure. Multiple calls to `initialize` should not
+    /// error. Calling [`persist`] before calling `initialize` should not error either.
+    ///
+    /// [`persist`]: WalletPersister::persist
+    fn initialize(persister: &mut Self) -> Result<ChangeSet, Self::Error>;
 
-    /// Persist changes.
+    /// Persist the given `changeset` to the `persister`.
+    ///
+    /// This method can fail if the `persister` is not [`initialize`]d.
+    ///
+    /// [`initialize`]: WalletPersister::initialize
     fn persist(persister: &mut Self, changeset: &ChangeSet) -> Result<(), Self::Error>;
-
-    /// Load changes.
-    fn load(persister: &mut Self) -> Result<Option<ChangeSet>, Self::Error>;
 }
 
 type FutureResult<'a, T, E> = Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>;
 
-/// Async version of [`WalletPersister`].
+/// Async trait that persists [`Wallet`].
+///
+/// For a blocking version, use [`WalletPersister`].
+///
+/// Associated functions of this trait should not be called directly, and the trait is designed so
+/// that associated functions are hard to find (since they are not methods!). [`WalletPersister`] is
+/// used by [`PersistedWallet`] (a light wrapper around [`Wallet`]) which enforces some level of
+/// safety. Refer to [`PersistedWallet`] for more about the safety checks.
 pub trait AsyncWalletPersister {
-    /// Error with persistence.
+    /// Error type of the persister.
     type Error;
 
-    /// Initialize the persister.
-    fn initialize<'a>(persister: &'a mut Self) -> FutureResult<'a, (), Self::Error>
+    /// Initialize the `persister` and load all data.
+    ///
+    /// This is called by [`PersistedWallet::create_async`] and [`PersistedWallet::load_async`] to
+    /// ensure the [`WalletPersister`] is initialized and returns all data in the `persister`.
+    ///
+    /// # Implementation Details
+    ///
+    /// The database schema of the `persister` (if any), should be initialized and migrated here.
+    ///
+    /// The implementation must return all data currently stored in the `persister`. If there is no
+    /// data, return an empty changeset (using [`ChangeSet::default()`]).
+    ///
+    /// Error should only occur on database failure. Multiple calls to `initialize` should not
+    /// error. Calling [`persist`] before calling `initialize` should not error either.
+    ///
+    /// [`persist`]: AsyncWalletPersister::persist
+    fn initialize<'a>(persister: &'a mut Self) -> FutureResult<'a, ChangeSet, Self::Error>
     where
         Self: 'a;
 
-    /// Persist changes.
+    /// Persist the given `changeset` to the `persister`.
+    ///
+    /// This method can fail if the `persister` is not [`initialize`]d.
+    ///
+    /// [`initialize`]: AsyncWalletPersister::initialize
     fn persist<'a>(
         persister: &'a mut Self,
         changeset: &'a ChangeSet,
     ) -> FutureResult<'a, (), Self::Error>
     where
         Self: 'a;
-
-    /// Load changes.
-    fn load<'a>(persister: &'a mut Self) -> FutureResult<'a, Option<ChangeSet>, Self::Error>
-    where
-        Self: 'a;
 }
 
 /// Represents a persisted wallet.
+///
+/// This is a light wrapper around [`Wallet`] that enforces some level of safety-checking when used
+/// with a [`WalletPersister`] or [`AsyncWalletPersister`] implementation. Safety checks assume that
+/// [`WalletPersister`] and/or [`AsyncWalletPersister`] are implemented correctly.
+///
+/// Checks include:
+///
+/// * Ensure the persister is initialized before data is persisted.
+/// * Ensure there were no previously persisted wallet data before creating a fresh wallet and
+///     persisting it.
+/// * Only clear the staged changes of [`Wallet`] after persisting succeeds.
 #[derive(Debug)]
 pub struct PersistedWallet(pub(crate) Wallet);
 
@@ -78,7 +133,10 @@ impl PersistedWallet {
     where
         P: WalletPersister,
     {
-        P::initialize(persister).map_err(CreateWithPersistError::Persist)?;
+        let existing = P::initialize(persister).map_err(CreateWithPersistError::Persist)?;
+        if !existing.is_empty() {
+            return Err(CreateWithPersistError::DataAlreadyExists(existing));
+        }
         let mut inner =
             Wallet::create_with_params(params).map_err(CreateWithPersistError::Descriptor)?;
         if let Some(changeset) = inner.take_staged() {
@@ -95,9 +153,12 @@ impl PersistedWallet {
     where
         P: AsyncWalletPersister,
     {
-        P::initialize(persister)
+        let existing = P::initialize(persister)
             .await
             .map_err(CreateWithPersistError::Persist)?;
+        if !existing.is_empty() {
+            return Err(CreateWithPersistError::DataAlreadyExists(existing));
+        }
         let mut inner =
             Wallet::create_with_params(params).map_err(CreateWithPersistError::Descriptor)?;
         if let Some(changeset) = inner.take_staged() {
@@ -116,11 +177,7 @@ impl PersistedWallet {
     where
         P: WalletPersister,
     {
-        P::initialize(persister).map_err(LoadWithPersistError::Persist)?;
-        let changeset = match P::load(persister).map_err(LoadWithPersistError::Persist)? {
-            Some(changeset) => changeset,
-            None => return Ok(None),
-        };
+        let changeset = P::initialize(persister).map_err(LoadWithPersistError::Persist)?;
         Wallet::load_with_params(changeset, params)
             .map(|opt| opt.map(PersistedWallet))
             .map_err(LoadWithPersistError::InvalidChangeSet)
@@ -134,16 +191,9 @@ impl PersistedWallet {
     where
         P: AsyncWalletPersister,
     {
-        P::initialize(persister)
+        let changeset = P::initialize(persister)
             .await
             .map_err(LoadWithPersistError::Persist)?;
-        let changeset = match P::load(persister)
-            .await
-            .map_err(LoadWithPersistError::Persist)?
-        {
-            Some(changeset) => changeset,
-            None => return Ok(None),
-        };
         Wallet::load_with_params(changeset, params)
             .map(|opt| opt.map(PersistedWallet))
             .map_err(LoadWithPersistError::InvalidChangeSet)
@@ -188,41 +238,33 @@ impl PersistedWallet {
 impl<'c> WalletPersister for bdk_chain::rusqlite::Transaction<'c> {
     type Error = bdk_chain::rusqlite::Error;
 
-    fn initialize(persister: &mut Self) -> Result<(), Self::Error> {
-        ChangeSet::init_sqlite_tables(&*persister)
+    fn initialize(persister: &mut Self) -> Result<ChangeSet, Self::Error> {
+        ChangeSet::init_sqlite_tables(&*persister)?;
+        ChangeSet::from_sqlite(persister)
     }
 
     fn persist(persister: &mut Self, changeset: &ChangeSet) -> Result<(), Self::Error> {
         changeset.persist_to_sqlite(persister)
     }
-
-    fn load(persister: &mut Self) -> Result<Option<ChangeSet>, Self::Error> {
-        ChangeSet::from_sqlite(persister).map(Some)
-    }
 }
 
 #[cfg(feature = "rusqlite")]
 impl WalletPersister for bdk_chain::rusqlite::Connection {
     type Error = bdk_chain::rusqlite::Error;
 
-    fn initialize(persister: &mut Self) -> Result<(), Self::Error> {
+    fn initialize(persister: &mut Self) -> Result<ChangeSet, Self::Error> {
         let db_tx = persister.transaction()?;
         ChangeSet::init_sqlite_tables(&db_tx)?;
-        db_tx.commit()
+        let changeset = ChangeSet::from_sqlite(&db_tx)?;
+        db_tx.commit()?;
+        Ok(changeset)
     }
 
     fn persist(persister: &mut Self, changeset: &ChangeSet) -> Result<(), Self::Error> {
         let db_tx = persister.transaction()?;
         changeset.persist_to_sqlite(&db_tx)?;
         db_tx.commit()
     }
-
-    fn load(persister: &mut Self) -> Result<Option<ChangeSet>, Self::Error> {
-        let db_tx = persister.transaction()?;
-        let changeset = ChangeSet::from_sqlite(&db_tx).map(Some)?;
-        db_tx.commit()?;
-        Ok(changeset)
-    }
 }
 
 /// Error for [`bdk_file_store`]'s implementation of [`WalletPersister`].
@@ -253,21 +295,18 @@ impl std::error::Error for FileStoreError {}
 impl WalletPersister for bdk_file_store::Store<ChangeSet> {
     type Error = FileStoreError;
 
-    fn initialize(_persister: &mut Self) -> Result<(), Self::Error> {
-        Ok(())
+    fn initialize(persister: &mut Self) -> Result<ChangeSet, Self::Error> {
+        persister
+            .aggregate_changesets()
+            .map(Option::unwrap_or_default)
+            .map_err(FileStoreError::Load)
     }
 
     fn persist(persister: &mut Self, changeset: &ChangeSet) -> Result<(), Self::Error> {
         persister
             .append_changeset(changeset)
             .map_err(FileStoreError::Write)
     }
-
-    fn load(persister: &mut Self) -> Result<Option<ChangeSet>, Self::Error> {
-        persister
-            .aggregate_changesets()
-            .map_err(FileStoreError::Load)
-    }
 }
 
 /// Error type for [`PersistedWallet::load`].
@@ -296,6 +335,8 @@ impl<E: fmt::Debug + fmt::Display> std::error::Error for LoadWithPersistError<E>
 pub enum CreateWithPersistError<E> {
     /// Error from persistence.
     Persist(E),
+    /// Persister already has wallet data.
+    DataAlreadyExists(ChangeSet),
     /// Occurs when the loaded changeset cannot construct [`Wallet`].
     Descriptor(DescriptorError),
 }
@@ -304,6 +345,11 @@ impl<E: fmt::Display> fmt::Display for CreateWithPersistError<E> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Self::Persist(err) => fmt::Display::fmt(err, f),
+            Self::DataAlreadyExists(changeset) => write!(
+                f,
+                "Cannot create wallet in persister which already contains wallet data: {:?}",
+                changeset
+            ),
             Self::Descriptor(err) => fmt::Display::fmt(&err, f),
         }
     }