diff --git a/wallet/src/account/mod.rs b/wallet/src/account/mod.rs index c733f41b2..c9c7ab8c8 100644 --- a/wallet/src/account/mod.rs +++ b/wallet/src/account/mod.rs @@ -2401,6 +2401,11 @@ impl Account { ) } + /// Reset all transactions that are currently in-mempool to inactive state + pub fn reset_inmempool_txs_to_inactive(&mut self) { + self.output_cache.reset_inmempool_txs_to_inactive(); + } + pub fn update_best_block( &mut self, db_tx: &mut impl WalletStorageWriteLocked, diff --git a/wallet/src/account/output_cache/mod.rs b/wallet/src/account/output_cache/mod.rs index bc4bd4437..d059ba853 100644 --- a/wallet/src/account/output_cache/mod.rs +++ b/wallet/src/account/output_cache/mod.rs @@ -1431,6 +1431,11 @@ impl OutputCache { Ok(()) } + /// Reset all transactions that are currently in-mempool to inactive state. + pub fn reset_inmempool_txs_to_inactive(&mut self) { + self.txs.values_mut().for_each(|tx| tx.reset_inmempool_to_inactive()); + } + fn is_consumed(&self, utxo_states: UtxoStates, outpoint: &UtxoOutPoint) -> bool { self.consumed .get(outpoint) diff --git a/wallet/src/account/output_cache/tests.rs b/wallet/src/account/output_cache/tests.rs index 979d67ee5..f7a74328f 100644 --- a/wallet/src/account/output_cache/tests.rs +++ b/wallet/src/account/output_cache/tests.rs @@ -1608,10 +1608,47 @@ fn update_conflicting_txs_order_v1_freeze(#[case] seed: Seed) { assert!(!output_cache.orders[&order_id].is_concluded); } +#[rstest] +#[trace] +#[case(Seed::from_entropy())] +fn reset_inmempool_txs_to_inactive(#[case] seed: Seed) { + let mut rng = make_seedable_rng(seed); + let chain_config = create_unit_test_config(); + let mut output_cache = OutputCache::empty(); + + // add 10 random txs + for _ in 0..10 { + add_random_transfer_tx_with_state( + &mut output_cache, + &chain_config, + TxStateTag::InMempool, + &mut rng, + ); + } + + // Reset + output_cache.reset_inmempool_txs_to_inactive(); + + // Check none of the txs are InMempool state + for tx in output_cache.txs.values() { + assert!(!matches!(tx.state(), TxState::InMempool(_))); + } +} + fn add_random_transfer_tx( output_cache: &mut OutputCache, chain_config: &ChainConfig, mut rng: impl Rng, +) { + let state_tag = TxStateTag::iter().choose(&mut rng).unwrap(); + add_random_transfer_tx_with_state(output_cache, chain_config, state_tag, rng); +} + +fn add_random_transfer_tx_with_state( + output_cache: &mut OutputCache, + chain_config: &ChainConfig, + state_tag: TxStateTag, + mut rng: impl Rng, ) { let random_tx_id = Id::::random_using(&mut rng); let random_block_id = Id::::random_using(&mut rng); @@ -1627,7 +1664,7 @@ fn add_random_transfer_tx( .build(); let tx_id = tx.transaction().get_id(); - let tx_state = match TxStateTag::iter().choose(&mut rng).unwrap() { + let tx_state = match state_tag { TxStateTag::Confirmed => TxState::Confirmed( BlockHeight::new(rng.random_range(0..100)), BlockTimestamp::from_int_seconds(rng.random_range(0..100)), diff --git a/wallet/src/wallet/mod.rs b/wallet/src/wallet/mod.rs index 2414b8253..d1a9fad65 100644 --- a/wallet/src/wallet/mod.rs +++ b/wallet/src/wallet/mod.rs @@ -812,6 +812,13 @@ where Ok(()) } + /// Resets all transactions in all accounts from in-mempool state to inactive. + pub fn reset_inmempool_txs_to_inactive(&mut self) { + self.accounts + .values_mut() + .for_each(|account| account.reset_inmempool_txs_to_inactive()); + } + fn reset_wallet_transactions( chain_config: Arc, db_tx: &mut impl WalletStorageWriteLocked, @@ -973,6 +980,8 @@ where .into_iter() .map(|account| (account.account_index(), account)) .collect(); + // reset all inmempool transactions to inactive so we can set them properly from the current node again + accounts.values_mut().for_each(|a| a.reset_inmempool_txs_to_inactive()); let latest_median_time = db_tx.get_median_time()?.unwrap_or(chain_config.genesis_block().timestamp()); diff --git a/wallet/types/src/wallet_tx.rs b/wallet/types/src/wallet_tx.rs index 938ae7448..b9cbe6830 100644 --- a/wallet/types/src/wallet_tx.rs +++ b/wallet/types/src/wallet_tx.rs @@ -164,6 +164,17 @@ impl WalletTx { } } + pub fn reset_inmempool_to_inactive(&mut self) { + match self { + WalletTx::Block(_) => {} // Blocks are never InMempool + WalletTx::Tx(tx) => { + if let TxState::InMempool(counter) = tx.state { + tx.state = TxState::Inactive(counter); + } + } + } + } + pub fn inputs(&self) -> &[TxInput] { match self { WalletTx::Block(block) => block.kernel_inputs(), diff --git a/wallet/wallet-controller/src/lib.rs b/wallet/wallet-controller/src/lib.rs index ddf536225..236012f88 100644 --- a/wallet/wallet-controller/src/lib.rs +++ b/wallet/wallet-controller/src/lib.rs @@ -54,7 +54,6 @@ use types::{ SeedWithPassPhrase, SignatureStats, TransactionToInspect, ValidatedSignatures, WalletInfo, WalletTypeArgsComputed, }; -use utils::set_flag::SetFlag; use wallet_storage::DefaultBackend; use read::ReadOnlyController; @@ -238,7 +237,7 @@ pub struct Controller { wallet_events: W, mempool_events: MempoolEvents, - finished_initial_sync: SetFlag, + should_rescan_mempool_txs: bool, } impl std::fmt::Debug for Controller { @@ -261,8 +260,21 @@ where wallet: RuntimeWallet, wallet_events: W, ) -> Result> { - let mut controller = - Self::new_unsynced(chain_config, rpc_client, wallet, wallet_events).await?; + let wallet_mode = rpc_client.is_cold_wallet_node().await; + let mempool_events = rpc_client + .mempool_subscribe_to_events() + .await + .map_err(ControllerError::NodeCallError)?; + let mut controller = Self { + chain_config, + rpc_client, + wallet, + wallet_mode, + staking_started: BTreeSet::new(), + wallet_events, + mempool_events, + should_rescan_mempool_txs: true, + }; // In the cold mode, try_sync_once is a no-op, so it doesn't matter whether we call it. // We omit the call to avoid printing the "Syncing the wallet" log line, which looks @@ -299,7 +311,7 @@ where staking_started: BTreeSet::new(), wallet_events, mempool_events, - finished_initial_sync: SetFlag::new(), + should_rescan_mempool_txs: true, }) } @@ -1569,8 +1581,8 @@ where match self.wallet_mode { WalletControllerMode::Hot => { - // after the first successful sync to the tip fetch all mempool transactions - if !self.finished_initial_sync.test() { + // fetch all mempool transactions after a broken connection or after the initial sync + if self.should_rescan_mempool_txs { let txs = self.rpc_client.mempool_get_transactions().await; match txs { @@ -1580,7 +1592,7 @@ where { log::error!("Error adding mempool transactions: {err}"); } else { - self.finished_initial_sync.set(); + self.should_rescan_mempool_txs = false } } Err(err) => { @@ -1609,8 +1621,11 @@ where Some(e) => e, None => { log::error!("Mempool notifications channel is closed"); - tokio::time::sleep(ERROR_DELAY).await; + // reset in-mempool transactions to inactive so we can rescan them when we connect again + self.wallet.reset_inmempool_txs_to_inactive(); + self.should_rescan_mempool_txs = true; + tokio::time::sleep(ERROR_DELAY).await; match self.rpc_client .mempool_subscribe_to_events() .await { diff --git a/wallet/wallet-controller/src/runtime_wallet.rs b/wallet/wallet-controller/src/runtime_wallet.rs index b693edfec..5b02831af 100644 --- a/wallet/wallet-controller/src/runtime_wallet.rs +++ b/wallet/wallet-controller/src/runtime_wallet.rs @@ -102,6 +102,16 @@ where } } + pub fn reset_inmempool_txs_to_inactive(&mut self) { + match self { + RuntimeWallet::Software(w) => w.reset_inmempool_txs_to_inactive(), + #[cfg(feature = "trezor")] + RuntimeWallet::Trezor(w) => w.reset_inmempool_txs_to_inactive(), + #[cfg(feature = "ledger")] + RuntimeWallet::Ledger(w) => w.reset_inmempool_txs_to_inactive(), + } + } + pub fn find_account_destination(&self, acc_outpoint: &AccountOutPoint) -> Option { match self { RuntimeWallet::Software(w) => w.find_account_destination(acc_outpoint),