feat(wallet): add pending-note reservation (#3080)

nodes/node/binary/src/config/wallet/mod.rs

@@ -24,6 +24,7 @@ impl ServiceConfig {
             known_keys: self.user.known_keys,
             voucher_master_key_id: self.user.voucher_master_key_id,
             recovery_path,
+            pending_note_expiry_blocks: self.user.pending_note_expiry_blocks,
         }
     }
 }

nodes/node/binary/src/config/wallet/serde.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use lb_key_management_system_service::{backend::preload::KeyId, keys::ZkPublicKey};
+use lb_wallet_service::default_pending_note_expiry_blocks;
 use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -8,6 +9,8 @@ pub struct Config {
     #[serde(default)]
     pub known_keys: HashMap<KeyId, ZkPublicKey>,
     pub voucher_master_key_id: KeyId,
+    #[serde(default = "default_pending_note_expiry_blocks")]
+    pub pending_note_expiry_blocks: u64,
 }
 
 pub struct RequiredValues {
@@ -24,6 +27,7 @@ impl Config {
         Self {
             known_keys: HashMap::new(),
             voucher_master_key_id,
+            pending_note_expiry_blocks: default_pending_note_expiry_blocks(),
         }
     }
 }

services/wallet/src/lib.rs

@@ -202,6 +202,7 @@ pub struct UtxoWithKeyId {
 struct LeaderClaimTx {
     signed_tx: SignedMantleTx,
     voucher_nullifier: VoucherNullifier,
+    funded_notes: Vec<NoteId>,
 }
 
 struct LeaderClaimTxRequest {
@@ -244,6 +245,17 @@ pub struct WalletServiceSettings {
     pub known_keys: HashMap<KeyId, ZkPublicKey>,
     pub voucher_master_key_id: KeyId,
     pub recovery_path: PathBuf,
+    /// How much LIB progress a pending note reservation survives before being
+    /// evicted. Notes funded into in-flight transactions are excluded from
+    /// funding until they are observed spent in a block or this many immutable
+    /// blocks have passed since the reservation.
+    #[serde(default = "default_pending_note_expiry_blocks")]
+    pub pending_note_expiry_blocks: u64,
+}
+
+#[must_use]
+pub const fn default_pending_note_expiry_blocks() -> u64 {
+    10
 }
 
 impl FileBackendSettings for WalletServiceSettings {
@@ -412,7 +424,7 @@ where
         loop {
             tokio::select! {
                 Some(msg) = service_resources_handle.inbound_relay.recv() => {
-                    Self::handle_wallet_message(msg, &mut state, &voucher_master_key_id, &storage_adapter, &cryptarchia_api, &kms, &epoch_config).await;
+                    Box::pin(Self::handle_wallet_message(msg, &mut state, &voucher_master_key_id, &storage_adapter, &cryptarchia_api, &kms, &epoch_config)).await;
                 }
                 Ok(event) = new_block_receiver.recv() => {
                     Self::handle_new_block(event.block_id, &mut state, &storage_adapter, &cryptarchia_api, &epoch_config).await;
@@ -517,7 +529,7 @@ where
                     }
                 };
 
-                let funded = match state.wallet().fund_tx::<MainnetGasConstants>(
+                let funded = match state.fund_tx::<MainnetGasConstants>(
                     tip,
                     &tx_builder,
                     change_pk,
@@ -531,6 +543,9 @@ where
                     }
                 };
 
+                let funded_notes: Vec<NoteId> = funded.consumed_or_locked_notes().collect();
+                state.reserve_pending_notes(funded_notes.iter().copied());
+
                 if resp_tx
                     .send(Ok(TipResponse {
                         tip,
@@ -538,6 +553,7 @@ where
                     }))
                     .is_err()
                 {
+                    state.release_pending_notes(funded_notes);
                     debug!(target: LOG_TARGET, "Failed to respond to FundTx");
                 }
             }
@@ -568,6 +584,7 @@ where
                 match response {
                     Ok(built_tx) => {
                         let voucher_nullifier = built_tx.voucher_nullifier;
+                        let funded_notes = built_tx.funded_notes;
                         if resp_tx
                             .send(Ok(TipResponse {
                                 tip,
@@ -576,6 +593,7 @@ where
                             .is_err()
                         {
                             state.release_claim_reservation(voucher_nullifier);
+                            state.release_pending_notes(funded_notes);
                             debug!(target: LOG_TARGET, "Failed to respond to BuildLeaderClaimTx");
                         }
                     }
@@ -607,6 +625,8 @@ where
                     }
                 };
 
+                let funded_notes: Vec<NoteId> = tx_builder.consumed_or_locked_notes().collect();
+
                 let resp = Self::sign_tx(tx_builder, tip, ledger, kms, state.wallet())
                     .await
                     .map(|signed_tx| TipResponse {
@@ -614,7 +634,14 @@ where
                         response: signed_tx,
                     });
 
-                if resp_tx.send(resp).is_err() {
+                let signing_failed = resp.is_err();
+                let response_delivered = resp_tx.send(resp).is_ok();
+
+                if signing_failed || !response_delivered {
+                    state.release_pending_notes(funded_notes);
+                }
+
+                if !response_delivered {
                     debug!(target: LOG_TARGET, "Failed to respond to SignTx");
                 }
             }
@@ -1120,9 +1147,10 @@ where
             state.release_claim_reservation(voucher_nullifier);
         }
 
-        result.map(|signed_tx| LeaderClaimTx {
+        result.map(|(signed_tx, funded_notes)| LeaderClaimTx {
             signed_tx,
             voucher_nullifier,
+            funded_notes,
         })
     }
 
@@ -1193,9 +1221,9 @@ where
         request: LeaderClaimTxRequest,
         voucher_nullifier: VoucherNullifier,
         ledger: LedgerState,
-        state: &ServiceState<'_>,
+        state: &mut ServiceState<'_>,
         kms: &KmsServiceApi<Kms, RuntimeServiceId>,
-    ) -> Result<SignedMantleTx, WalletServiceError> {
+    ) -> Result<(SignedMantleTx, Vec<NoteId>), WalletServiceError> {
         let context = ledger.tx_context();
         let tx_builder = MantleTxBuilder::new().push_op(Op::LeaderClaim(LeaderClaimOp {
             rewards_root: request.rewards_root,
@@ -1203,7 +1231,7 @@ where
             pk: request.funding_pk,
         }))?;
 
-        let funded_tx_builder = state.wallet().fund_tx::<MainnetGasConstants>(
+        let funded_tx_builder = state.fund_tx::<MainnetGasConstants>(
             request.tip,
             &tx_builder,
             request.funding_pk,
@@ -1211,6 +1239,28 @@ where
             &context,
         )?;
 
+        let funded_notes: Vec<NoteId> = funded_tx_builder.consumed_or_locked_notes().collect();
+        state.reserve_pending_notes(funded_notes.iter().copied());
+
+        match Self::sign_funded_leader_claim_tx(request, funded_tx_builder, ledger, state, kms)
+            .await
+        {
+            Ok(signed_tx) => Ok((signed_tx, funded_notes)),
+            Err(err) => {
+                state.release_pending_notes(funded_notes);
+                Err(err)
+            }
+        }
+    }
+
+    async fn sign_funded_leader_claim_tx(
+        request: LeaderClaimTxRequest,
+        funded_tx_builder: MantleTxBuilder,
+        ledger: LedgerState,
+        state: &ServiceState<'_>,
+        kms: &KmsServiceApi<Kms, RuntimeServiceId>,
+    ) -> Result<SignedMantleTx, WalletServiceError> {
+        let context = ledger.tx_context();
         let net_balance = funded_tx_builder.net_balance();
         let gas_cost = funded_tx_builder.gas_cost::<MainnetGasConstants>(&context)?;
         debug!(

services/wallet/src/states.rs

@@ -1,9 +1,17 @@
-use std::collections::HashMap;
+use std::{
+    borrow::Borrow,
+    collections::{HashMap, HashSet},
+};
 
 use lb_core::{
     header::HeaderId,
-    mantle::ops::leader_claim::{VoucherCm, VoucherNullifier},
+    mantle::{
+        GasConstants, NoteId,
+        ops::leader_claim::{VoucherCm, VoucherNullifier},
+        transactions::{MantleTxBuilder, MantleTxContext},
+    },
 };
+use lb_key_management_system_service::keys::ZkPublicKey;
 use lb_ledger::LedgerState;
 use lb_log_targets::wallet;
 use lb_wallet::{Voucher, Vouchers, WalletBlock, WalletError, WalletState};
@@ -108,6 +116,90 @@ impl PendingClaims {
     }
 }
 
+/// Notes handed out as funding for transactions that have been built but not
+/// yet observed on chain. Ephemeral: rebuilt empty on restart; entries are
+/// dropped once the note is observed spent in a block, or expire after enough
+/// LIB progress (mirroring [`PendingClaims`]).
+///
+/// Unlike claim reservations, which must hold until the claim is finalised,
+/// funded notes are released as soon as they are observed spent in a tip
+/// block; eviction only covers transactions that never land. Releasing too
+/// early merely reintroduces a conflicting tx, the same failure mode as
+/// having no reservation at all.
+#[derive(Debug, Default)]
+struct PendingNotes {
+    notes: HashMap<NoteId, u64>,
+}
+
+impl PendingNotes {
+    fn reserve(&mut self, note_ids: impl IntoIterator<Item = NoteId>) {
+        for note_id in note_ids {
+            debug!(
+                target: wallet::SERVICE,
+                ?note_id,
+                "Reserved pending note"
+            );
+            self.notes.insert(note_id, 0);
+        }
+    }
+
+    fn release(&mut self, note_ids: impl IntoIterator<Item = NoteId>) {
+        for note_id in note_ids {
+            if self.notes.remove(&note_id).is_some() {
+                debug!(
+                    target: wallet::SERVICE,
+                    ?note_id,
+                    "Released pending note reservation"
+                );
+            }
+        }
+    }
+
+    fn note_ids(&self) -> HashSet<NoteId> {
+        self.notes.keys().copied().collect()
+    }
+
+    fn remove_spent(&mut self, spent: &HashSet<NoteId>) {
+        self.notes.retain(|note_id, _| !spent.contains(note_id));
+    }
+
+    /// Evict reservations whose LIB-progress age reached the configured limit.
+    ///
+    /// Each LIB update adds `new_immutable_blocks_count` to every reservation's
+    /// counter. A reservation expires once that counter reaches
+    /// `max_immutable_blocks_since_reservation`.
+    fn evict_expired(
+        &mut self,
+        new_immutable_blocks_count: u64,
+        max_immutable_blocks_since_reservation: u64,
+    ) {
+        if new_immutable_blocks_count == 0 {
+            return;
+        }
+
+        self.notes
+            .retain(|note_id, immutable_blocks_since_reservation| {
+                *immutable_blocks_since_reservation =
+                    immutable_blocks_since_reservation.saturating_add(new_immutable_blocks_count);
+
+                let expired =
+                    *immutable_blocks_since_reservation >= max_immutable_blocks_since_reservation;
+
+                if expired {
+                    debug!(
+                        target: wallet::SERVICE,
+                        ?note_id,
+                        immutable_blocks_since_reservation,
+                        max_immutable_blocks_since_reservation,
+                        "Removing pending note reservation after LIB progress"
+                    );
+                }
+
+                !expired
+            });
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct RecoveryState {
     next_new_voucher_index: VoucherIndex,
@@ -142,6 +234,8 @@ pub struct ServiceState<'u> {
     lib: HeaderId,
     updater: &'u StateUpdater<Option<RecoveryState>>,
     pending_claims: PendingClaims,
+    pending_notes: PendingNotes,
+    pending_note_expiry_blocks: u64,
     security_param: u64,
 }
 
@@ -184,6 +278,8 @@ impl<'u> ServiceState<'u> {
             lib: wallet_lib,
             updater,
             pending_claims,
+            pending_notes: PendingNotes::default(),
+            pending_note_expiry_blocks: settings.pending_note_expiry_blocks,
             security_param,
         }
     }
@@ -205,6 +301,7 @@ impl<'u> ServiceState<'u> {
 
     pub fn apply_block(&mut self, block: &WalletBlock) -> Result<(), WalletError> {
         self.wallet.apply_block(block)?;
+        self.pending_notes.remove_spent(&block.spent_note_ids());
         self.update_state();
         Ok(())
     }
@@ -221,6 +318,8 @@ impl<'u> ServiceState<'u> {
         self.wallet.prune_vouchers(pruned_nullifiers);
         self.pending_claims
             .evict_expired(new_immutable_blocks_count, self.security_param);
+        self.pending_notes
+            .evict_expired(new_immutable_blocks_count, self.pending_note_expiry_blocks);
         self.update_state();
     }
 
@@ -261,6 +360,34 @@ impl<'u> ServiceState<'u> {
         self.pending_claims.release(nullifier);
     }
 
+    /// Fund `tx_builder` from the wallet's UTXOs at `tip`, excluding notes
+    /// already reserved for in-flight transactions.
+    pub fn fund_tx<G: GasConstants>(
+        &self,
+        tip: HeaderId,
+        tx_builder: &MantleTxBuilder,
+        change_pk: ZkPublicKey,
+        funding_pks: impl IntoIterator<Item = impl Borrow<ZkPublicKey>>,
+        context: &MantleTxContext,
+    ) -> Result<MantleTxBuilder, WalletError> {
+        self.wallet.fund_tx::<G>(
+            tip,
+            tx_builder,
+            change_pk,
+            funding_pks,
+            context,
+            &self.pending_notes.note_ids(),
+        )
+    }
+
+    pub fn reserve_pending_notes(&mut self, note_ids: impl IntoIterator<Item = NoteId>) {
+        self.pending_notes.reserve(note_ids);
+    }
+
+    pub fn release_pending_notes(&mut self, note_ids: impl IntoIterator<Item = NoteId>) {
+        self.pending_notes.release(note_ids);
+    }
+
     fn update_state(&self) {
         let lib_wallet_state = self
             .wallet()
@@ -278,6 +405,8 @@ impl<'u> ServiceState<'u> {
 
 #[cfg(test)]
 mod tests {
+    use lb_groth16::{AdditiveGroup as _, Field as _, Fr};
+
     use super::*;
 
     const EXPIRY_BLOCKS: u64 = 4;
@@ -305,4 +434,53 @@ mod tests {
         pending_claims.evict_expired(1, EXPIRY_BLOCKS);
         assert!(!pending_claims.is_reserved(&nullifier));
     }
+
+    #[test]
+    fn pending_note_does_not_expire_without_lib_progress() {
+        let note_id = NoteId::from(Fr::ONE);
+        let mut pending_notes = PendingNotes::default();
+
+        pending_notes.reserve([note_id]);
+        pending_notes.evict_expired(0, EXPIRY_BLOCKS);
+
+        assert!(pending_notes.note_ids().contains(&note_id));
+    }
+
+    #[test]
+    fn pending_note_expires_after_enough_lib_progress() {
+        let note_id = NoteId::from(Fr::ONE);
+        let mut pending_notes = PendingNotes::default();
+
+        pending_notes.reserve([note_id]);
+        pending_notes.evict_expired(EXPIRY_BLOCKS - 1, EXPIRY_BLOCKS);
+        assert!(pending_notes.note_ids().contains(&note_id));
+
+        pending_notes.evict_expired(1, EXPIRY_BLOCKS);
+        assert!(!pending_notes.note_ids().contains(&note_id));
+    }
+
+    #[test]
+    fn pending_note_removed_when_observed_spent() {
+        let spent = NoteId::from(Fr::ZERO);
+        let kept = NoteId::from(Fr::ONE);
+        let mut pending_notes = PendingNotes::default();
+
+        pending_notes.reserve([spent, kept]);
+        pending_notes.remove_spent(&HashSet::from([spent]));
+
+        let note_ids = pending_notes.note_ids();
+        assert!(!note_ids.contains(&spent));
+        assert!(note_ids.contains(&kept));
+    }
+
+    #[test]
+    fn released_pending_note_becomes_available() {
+        let note_id = NoteId::from(Fr::ONE);
+        let mut pending_notes = PendingNotes::default();
+
+        pending_notes.reserve([note_id]);
+        pending_notes.release([note_id]);
+
+        assert!(pending_notes.note_ids().is_empty());
+    }
 }

tests/src/common/wallet/transaction/builder_funding.rs

@@ -1,6 +1,6 @@
 //! Applies wallet funding decisions to Mantle transaction builders.
 
-use std::cmp::Ordering;
+use std::{cmp::Ordering, collections::HashSet};
 
 use lb_core::mantle::{
     Note, Op, Utxo,
@@ -30,6 +30,7 @@ pub fn fund_builder_from_wallet_source(
         source.public_key(),
         [source.public_key()],
         context,
+        &HashSet::new(),
     )
 }
 

tests/src/tests/mantle/leader.rs

@@ -121,7 +121,6 @@ async fn claim_leader_rewards(node: &NodeHttpClient, duration: Duration) -> TxHa
 }
 
 #[tokio::test]
-#[ignore = "Blocked on zero-fee transaction funding fix in https://github.com/logos-blockchain/logos-blockchain/pull/2970"]
 async fn concurrent_leader_claims() {
     let (_base, nodes, leader_funding_pk) = setup_test_nodes("concurrent_leader_claims").await;
     let node = &nodes[0];

tests/testing_framework/src/framework/local/provisioning.rs

@@ -601,9 +601,11 @@ fn build_run_config(config: Config, genesis_block: &GenesisBlock) -> RunConfig {
 
             wallet::serde::Config {
                 known_keys,
-                voucher_master_key_id: key_id_for_preload_backend(&Key::Zk(
-                    config.consensus_config.known_key.clone(),
-                )),
+                ..wallet::serde::Config::with_required_values(wallet::serde::RequiredValues {
+                    voucher_master_key_id: key_id_for_preload_backend(&Key::Zk(
+                        config.consensus_config.known_key.clone(),
+                    )),
+                })
             }
         },
         kms: config::kms::serde::Config {

wallet/src/lib.rs

@@ -113,6 +113,21 @@ impl WalletBlock {
             txs: transform_txs(block.transactions(), tx_events).collect(),
         }
     }
+
+    /// Note IDs this block spends or locks.
+    #[must_use]
+    pub fn spent_note_ids(&self) -> HashSet<NoteId> {
+        self.txs
+            .iter()
+            .flat_map(|tx| tx.ops.iter())
+            .flat_map(|op| match op {
+                WalletOp::Transfer(transfer) => transfer.inputs.iter().copied().collect::<Vec<_>>(),
+                WalletOp::ChannelDeposit(inputs) => inputs.clone(),
+                WalletOp::Lock(note_id) => vec![*note_id],
+                WalletOp::LeaderClaim(_) | WalletOp::ChannelWithdraw(_) => Vec::new(),
+            })
+            .collect()
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -188,13 +203,16 @@ impl WalletState {
         change_pk: ZkPublicKey,
         pks: impl IntoIterator<Item = impl Borrow<ZkPublicKey>>,
         context: &MantleTxContext,
+        excluded_notes: &HashSet<NoteId>,
     ) -> Result<MantleTxBuilder, WalletError> {
         // Get all UTXOs owned by the provided PKs, excluding the following notes:
         // - Notes that are being consumed/locked by the tx
         // - Notes that are already locked in Ledger
+        // - Notes excluded by the caller (e.g. reserved for in-flight txs)
         let consumed_or_locked = tx_builder
             .consumed_or_locked_notes()
             .chain(self.locked_notes.iter().copied())
+            .chain(excluded_notes.iter().copied())
             .collect::<HashSet<_>>();
         let mut utxos = self
             .utxos_owned_by_pks(pks)
@@ -692,9 +710,15 @@ where
         change_pk: ZkPublicKey,
         funding_pks: impl IntoIterator<Item = impl Borrow<ZkPublicKey>>,
         context: &MantleTxContext,
+        excluded_notes: &HashSet<NoteId>,
     ) -> Result<MantleTxBuilder, WalletError> {
-        self.wallet_state_at(tip)?
-            .fund_tx::<G>(tx_builder, change_pk, funding_pks, context)
+        self.wallet_state_at(tip)?.fund_tx::<G>(
+            tx_builder,
+            change_pk,
+            funding_pks,
+            context,
+            excluded_notes,
+        )
     }
 
     pub fn wallet_state_at(&self, tip: HeaderId) -> Result<WalletState, WalletError> {
@@ -1224,7 +1248,7 @@ mod tests {
 
         // Fund the transaction
         let funded_tx_builder = wallet_state
-            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context)
+            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new())
             .unwrap();
 
         assert_eq!(
@@ -1288,7 +1312,7 @@ mod tests {
         assert_eq!(0, tx_builder.funding_delta::<Gas>(&context).unwrap());
 
         let funded_tx_builder = wallet_state
-            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context)
+            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new())
             .unwrap();
 
         // No input was pulled in (an added input would push the net balance to
@@ -1343,7 +1367,8 @@ mod tests {
         tx_builder = tx_builder.push_op(inscription).unwrap();
 
         // Fund the transaction
-        let fund_attempt = wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context);
+        let fund_attempt =
+            wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new());
 
         assert_eq!(
             fund_attempt.unwrap_err(),
@@ -1371,7 +1396,8 @@ mod tests {
         let tx_builder = MantleTxBuilder::new();
 
         // Fund the transaction
-        let fund_attempt = wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context);
+        let fund_attempt =
+            wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new());
 
         assert_eq!(
             fund_attempt.unwrap_err(),
@@ -1402,7 +1428,8 @@ mod tests {
         let tx_builder = MantleTxBuilder::new();
 
         // Fund the transaction
-        let fund_attempt = wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context);
+        let fund_attempt =
+            wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new());
 
         assert_eq!(
             fund_attempt.unwrap_err(),
@@ -1434,7 +1461,8 @@ mod tests {
         let tx_builder = MantleTxBuilder::new();
 
         // Attempt to fund the transaction with Alice's notes.
-        let fund_attempt = wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context);
+        let fund_attempt =
+            wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new());
 
         assert_eq!(
             fund_attempt.unwrap_err(),
@@ -1443,7 +1471,7 @@ mod tests {
 
         // Fund the transaction with Bob's notes.
         wallet_state
-            .fund_tx::<Gas>(&tx_builder, bob, [bob], &context)
+            .fund_tx::<Gas>(&tx_builder, bob, [bob], &context, &HashSet::new())
             .unwrap(); // succesfully funded;
     }
 
@@ -1483,7 +1511,7 @@ mod tests {
         );
 
         let funded_tx_wo_change = wallet_state
-            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context)
+            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new())
             .unwrap()
             .build()
             .unwrap(); // successfully funded the tx
@@ -1523,7 +1551,8 @@ mod tests {
                 ),
             );
 
-            let fund_attempt = wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context);
+            let fund_attempt =
+                wallet_state.fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new());
 
             assert_eq!(
                 fund_attempt.unwrap_err(),
@@ -1541,7 +1570,7 @@ mod tests {
         );
 
         let funded_tx_wo_change = wallet_state
-            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context)
+            .fund_tx::<Gas>(&tx_builder, alice, [alice], &context, &HashSet::new())
             .unwrap()
             .build()
             .unwrap(); // successfully funded the tx