Aleks/wallet atomic

crates/apps/src/lib/bench_utils.rs

@@ -953,18 +953,24 @@ impl Default for BenchShieldedCtx {
         let mut chain_ctx = ctx.take_chain_or_exit();
 
         // Generate spending key for Albert and Bertha
-        chain_ctx.wallet.gen_store_spending_key(
-            ALBERT_SPENDING_KEY.to_string(),
-            None,
-            true,
-            &mut OsRng,
-        );
-        chain_ctx.wallet.gen_store_spending_key(
-            BERTHA_SPENDING_KEY.to_string(),
-            None,
-            true,
-            &mut OsRng,
-        );
+        chain_ctx
+            .wallet
+            .gen_store_spending_key_atomic(
+                ALBERT_SPENDING_KEY.to_string(),
+                None,
+                true,
+                &mut OsRng,
+            )
+            .expect("Failed to update the wallet storage.");
+        chain_ctx
+            .wallet
+            .gen_store_spending_key_atomic(
+                BERTHA_SPENDING_KEY.to_string(),
+                None,
+                true,
+                &mut OsRng,
+            )
+            .expect("Failed to update the wallet storage.");
         crate::wallet::save(&chain_ctx.wallet).unwrap();
 
         // Generate payment addresses for both Albert and Bertha
@@ -977,7 +983,8 @@ impl Default for BenchShieldedCtx {
             let viewing_key: FromContext<ExtendedViewingKey> = FromContext::new(
                 chain_ctx
                     .wallet
-                    .find_viewing_key(viewing_alias)
+                    .find_viewing_key_atomic(viewing_alias)
+                    .expect("Failed to read from the wallet storage.")
                     .unwrap()
                     .to_string(),
             );
@@ -991,11 +998,12 @@ impl Default for BenchShieldedCtx {
             let payment_addr = viewing_key.to_payment_address(div).unwrap();
             let _ = chain_ctx
                 .wallet
-                .insert_payment_addr(
+                .insert_payment_addr_atomic(
                     alias,
                     PaymentAddress::from(payment_addr).pinned(false),
                     true,
                 )
+                .expect("Failed to update the wallet storage.")
                 .unwrap();
         }
 
@@ -1020,7 +1028,8 @@ impl BenchShieldedCtx {
         let async_runtime = tokio::runtime::Runtime::new().unwrap();
         let spending_key = self
             .wallet
-            .find_spending_key(ALBERT_SPENDING_KEY, None)
+            .find_spending_key_atomic(ALBERT_SPENDING_KEY, None)
+            .expect("Failed to read from the wallet storage.")
             .unwrap();
         self.shielded = async_runtime
             .block_on(crate::client::masp::syncing(

crates/apps/src/lib/cli.rs

@@ -4295,7 +4295,7 @@ pub mod args {
     impl CliToSdk<TxInitAccount<SdkTypes>> for TxInitAccount<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> TxInitAccount<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             TxInitAccount::<SdkTypes> {
                 tx,
                 vp_code_path: self.vp_code_path,
@@ -4350,7 +4350,7 @@ pub mod args {
     impl CliToSdk<TxBecomeValidator<SdkTypes>> for TxBecomeValidator<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> TxBecomeValidator<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             TxBecomeValidator::<SdkTypes> {
                 tx,
                 scheme: self.scheme,
@@ -4470,7 +4470,7 @@ pub mod args {
     impl CliToSdk<TxInitValidator<SdkTypes>> for TxInitValidator<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> TxInitValidator<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             TxInitValidator::<SdkTypes> {
                 tx,
                 scheme: self.scheme,
@@ -4619,7 +4619,7 @@ pub mod args {
     impl CliToSdk<TxUpdateAccount<SdkTypes>> for TxUpdateAccount<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> TxUpdateAccount<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             TxUpdateAccount::<SdkTypes> {
                 tx,
                 vp_code_path: self.vp_code_path,
@@ -5021,7 +5021,7 @@ pub mod args {
     impl CliToSdk<RevealPk<SdkTypes>> for RevealPk<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> RevealPk<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             RevealPk::<SdkTypes> {
                 tx,
                 public_key: chain_ctx.get(&self.public_key),
@@ -5588,7 +5588,7 @@ pub mod args {
     impl CliToSdk<ConsensusKeyChange<SdkTypes>> for ConsensusKeyChange<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> ConsensusKeyChange<SdkTypes> {
             let tx = self.tx.to_sdk(ctx);
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             ConsensusKeyChange::<SdkTypes> {
                 tx,
                 validator: chain_ctx.get(&self.validator),
@@ -6412,7 +6412,7 @@ pub mod args {
 
     impl CliToSdk<Query<SdkTypes>> for Query<CliTypes> {
         fn to_sdk(self, ctx: &mut Context) -> Query<SdkTypes> {
-            let chain_ctx = ctx.borrow_mut_chain_or_exit();
+            let chain_ctx = ctx.borrow_chain_or_exit();
             Query::<SdkTypes> {
                 ledger_address: chain_ctx.get(&self.ledger_address),
             }
@@ -6467,9 +6467,9 @@ pub mod args {
 
             use crate::wallet::CliWalletUtils;
 
-            let find_viewing_key = |w: &mut Wallet<CliWalletUtils>| {
-                w.find_viewing_key(&self.viewing_key.raw)
-                    .copied()
+            let find_viewing_key = |w: &Wallet<CliWalletUtils>| {
+                w.find_viewing_key_atomic(&self.viewing_key.raw)
+                    .expect("Failed to read from the wallet storage.")
                     .unwrap_or_else(|_| {
                         eprintln!(
                             "Unknown viewing key {}",
@@ -6481,10 +6481,10 @@ pub mod args {
             let viewing_key = if ctx.global_args.is_pre_genesis {
                 let wallet_path =
                     ctx.global_args.base_dir.join(PRE_GENESIS_DIR);
-                let mut wallet = crate::wallet::load_or_new(&wallet_path);
-                find_viewing_key(&mut wallet)
+                let wallet = crate::wallet::load_or_new(&wallet_path);
+                find_viewing_key(&wallet)
             } else {
-                find_viewing_key(&mut ctx.borrow_mut_chain_or_exit().wallet)
+                find_viewing_key(&ctx.borrow_chain_or_exit().wallet)
             };
             PayAddressGen::<SdkTypes> {
                 alias: self.alias,