feat: remove spawn blocking calls from wallet db (key manager service) (

#4564)

Description
---
Removed spawn blocking calls for db operations from the wallet in the key manager service. (This is another PR in a couple of PRs required to implement this fully throughout the wallet code.)

Motivation and Context
---
As per #3982 and #4555

How Has This Been Tested?
---
Unit tests
Cucumber tests

base_layer/wallet/src/key_manager_service/error.rs

@@ -67,8 +67,6 @@ pub enum KeyManagerStorageError {
     DieselConnectionError(#[from] diesel::ConnectionError),
     #[error("Database migration error: `{0}`")]
     DatabaseMigrationError(String),
-    #[error("Blocking task spawn error: `{0}`")]
-    BlockingTaskSpawnError(String),
     #[error("Wallet db is already encrypted and cannot be encrypted until the previous encryption is removed")]
     AlreadyEncrypted,
     #[error("Wallet db is currently encrypted, decrypt before use")]

base_layer/wallet/src/key_manager_service/handle.rs

@@ -67,15 +67,14 @@ where TBackend: KeyManagerBackend + 'static
             .write()
             .await
             .add_key_manager_branch(branch.into())
-            .await
     }
 
     async fn apply_encryption(&self, cipher: XChaCha20Poly1305) -> Result<(), KeyManagerServiceError> {
-        (*self.key_manager_inner).write().await.apply_encryption(cipher).await
+        (*self.key_manager_inner).write().await.apply_encryption(cipher)
     }
 
     async fn remove_encryption(&self) -> Result<(), KeyManagerServiceError> {
-        (*self.key_manager_inner).write().await.remove_encryption().await
+        (*self.key_manager_inner).write().await.remove_encryption()
     }
 
     async fn get_next_key<T: Into<String> + Send>(&self, branch: T) -> Result<NextKeyResult, KeyManagerServiceError> {

base_layer/wallet/src/key_manager_service/service.rs

@@ -56,19 +56,19 @@ where TBackend: KeyManagerBackend + 'static
         }
     }
 
-    pub async fn add_key_manager_branch(&mut self, branch: String) -> Result<AddResult, KeyManagerServiceError> {
+    pub fn add_key_manager_branch(&mut self, branch: String) -> Result<AddResult, KeyManagerServiceError> {
         let result = if self.key_managers.contains_key(&branch) {
             AddResult::AlreadyExists
         } else {
             AddResult::NewEntry
         };
-        let state = match self.db.get_key_manager_state(branch.clone()).await? {
+        let state = match self.db.get_key_manager_state(branch.clone())? {
             None => {
                 let starting_state = KeyManagerState {
                     branch_seed: branch.to_string(),
                     primary_key_index: 0,
                 };
-                self.db.set_key_manager_state(starting_state.clone()).await?;
+                self.db.set_key_manager_state(starting_state.clone())?;
                 starting_state
             },
             Some(km) => km,
@@ -92,7 +92,7 @@ where TBackend: KeyManagerBackend + 'static
             .lock()
             .await;
         let key = km.next_key()?;
-        self.db.increment_key_index(branch).await?;
+        self.db.increment_key_index(branch)?;
         Ok(NextKeyResult {
             key: key.k,
             index: km.key_index(),
@@ -110,13 +110,13 @@ where TBackend: KeyManagerBackend + 'static
         Ok(key.k)
     }
 
-    pub async fn apply_encryption(&self, cipher: XChaCha20Poly1305) -> Result<(), KeyManagerServiceError> {
-        self.db.apply_encryption(cipher).await?;
+    pub fn apply_encryption(&self, cipher: XChaCha20Poly1305) -> Result<(), KeyManagerServiceError> {
+        self.db.apply_encryption(cipher)?;
         Ok(())
     }
 
-    pub async fn remove_encryption(&self) -> Result<(), KeyManagerServiceError> {
-        self.db.remove_encryption().await?;
+    pub fn remove_encryption(&self) -> Result<(), KeyManagerServiceError> {
+        self.db.remove_encryption()?;
         Ok(())
     }
 
@@ -156,7 +156,7 @@ where TBackend: KeyManagerBackend + 'static
         let current_index = km.key_index();
         if index > current_index {
             km.update_key_index(index);
-            self.db.set_key_index(branch, index).await?;
+            self.db.set_key_index(branch, index)?;
             trace!(target: LOG_TARGET, "Updated UTXO Key Index to {}", index);
         }
         Ok(())

base_layer/wallet/src/key_manager_service/storage/database/mod.rs

@@ -52,63 +52,35 @@ where T: KeyManagerBackend + 'static
 
     /// Retrieves the key manager state of the provided branch
     /// Returns None if the request branch does not exist.
-    pub async fn get_key_manager_state(
-        &self,
-        branch: String,
-    ) -> Result<Option<KeyManagerState>, KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.get_key_manager(branch))
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))
-            .and_then(|inner_result| inner_result)
+    pub fn get_key_manager_state(&self, branch: String) -> Result<Option<KeyManagerState>, KeyManagerStorageError> {
+        self.db.get_key_manager(branch)
     }
 
     /// Saves the specified key manager state to the backend database.
-    pub async fn set_key_manager_state(&self, state: KeyManagerState) -> Result<(), KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.add_key_manager(state))
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))??;
-
-        Ok(())
+    pub fn set_key_manager_state(&self, state: KeyManagerState) -> Result<(), KeyManagerStorageError> {
+        self.db.add_key_manager(state)
     }
 
     /// Increment the key index of the provided branch of the key manager.
     /// Will error if the branch does not exist.
-    pub async fn increment_key_index(&self, branch: String) -> Result<(), KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.increment_key_index(branch))
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))??;
-        Ok(())
+    pub fn increment_key_index(&self, branch: String) -> Result<(), KeyManagerStorageError> {
+        self.db.increment_key_index(branch)
     }
 
     /// Sets the key index of the provided branch of the key manager.
     /// Will error if the branch does not exist.
-    pub async fn set_key_index(&self, branch: String, index: u64) -> Result<(), KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.set_key_index(branch, index))
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))??;
-        Ok(())
+    pub fn set_key_index(&self, branch: String, index: u64) -> Result<(), KeyManagerStorageError> {
+        self.db.set_key_index(branch, index)
     }
 
     /// Encrypts the entire key manager with all branches.
     /// This will only encrypt the index used, as the master seed phrase is not directly stored with the key manager.
-    pub async fn apply_encryption(&self, cipher: XChaCha20Poly1305) -> Result<(), KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.apply_encryption(cipher))
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))
-            .and_then(|inner_result| inner_result)
+    pub fn apply_encryption(&self, cipher: XChaCha20Poly1305) -> Result<(), KeyManagerStorageError> {
+        self.db.apply_encryption(cipher)
     }
 
     /// Decrypts the entire key manager.
-    pub async fn remove_encryption(&self) -> Result<(), KeyManagerStorageError> {
-        let db_clone = self.db.clone();
-        tokio::task::spawn_blocking(move || db_clone.remove_encryption())
-            .await
-            .map_err(|err| KeyManagerStorageError::BlockingTaskSpawnError(err.to_string()))
-            .and_then(|inner_result| inner_result)
+    pub fn remove_encryption(&self) -> Result<(), KeyManagerStorageError> {
+        self.db.remove_encryption()
     }
 }