Merge pull request fedimint#6285 from dpc/24-11-04-more-efficient-out…

…put-state-machines refactor(client): handle multiple outputs with a single state machine
AreaLayer · Nov 7, 2024 · e5563bb · e5563bb
2 parents 714542b + 8d156fd
commit e5563bb
Show file tree

Hide file tree

Showing 9 changed files with 297 additions and 43 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -160,6 +160,7 @@ ln-gateway = { package = "fedimint-ln-gateway", path = "./gateway/ln-gateway", v
 miniscript = "12.2.0"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
+rayon = "1.10.0"
 reqwest = { version = "0.12.9", features = [
     "json",
     "rustls-tls",

diff --git a/fedimint-server/Cargo.toml b/fedimint-server/Cargo.toml
@@ -42,7 +42,7 @@ parity-scale-codec = "3.6.12"
 pin-project = "1.1.7"
 rand = { workspace = true }
 rand_chacha = { workspace = true }
-rayon = "1.10.0"
+rayon = { workspace = true }
 rcgen = "=0.13.1"
 serde = { workspace = true }
 serde_json = { workspace = true }

diff --git a/modules/fedimint-mint-client/Cargo.toml b/modules/fedimint-mint-client/Cargo.toml
@@ -43,6 +43,7 @@ fedimint-mint-common = { workspace = true }
 futures = { workspace = true }
 hex = { workspace = true }
 itertools = { workspace = true }
+rayon = { workspace = true }
 serde = { workspace = true }
 serde-big-array = { workspace = true }
 serde_json = { workspace = true }

diff --git a/modules/fedimint-mint-client/src/backup.rs b/modules/fedimint-mint-client/src/backup.rs
@@ -93,13 +93,40 @@ impl MintClientModule {
                 MintClientStateMachines::Output(MintOutputStateMachine {
                     common,
                     state: crate::output::MintOutputStates::Created(created_state),
-                }) => Some((
-                    common.out_point,
+                }) => Some(vec![(
+                    OutPoint {
+                        txid: common.txid,
+                        // MintOutputStates::Created always has one out_idx
+                        out_idx: *common.out_idxs.start(),
+                    },
                     created_state.amount,
                     created_state.issuance_request,
-                )),
+                )]),
+                MintClientStateMachines::Output(MintOutputStateMachine {
+                    common,
+                    state: crate::output::MintOutputStates::CreatedMulti(created_state),
+                }) => Some(
+                    common
+                        .out_idxs
+                        .map(|out_idx| {
+                            let issuance_request = created_state
+                                .issuance_requests
+                                .get(&out_idx)
+                                .expect("Must have corresponding out_idx");
+                            (
+                                OutPoint {
+                                    txid: common.txid,
+                                    out_idx,
+                                },
+                                issuance_request.0,
+                                issuance_request.1,
+                            )
+                        })
+                        .collect(),
+                ),
                 _ => None,
             })
+            .flatten()
             .collect::<Vec<_>>();
 
         let mut idxes = vec![];

diff --git a/modules/fedimint-mint-client/src/backup/recovery.rs b/modules/fedimint-mint-client/src/backup/recovery.rs
@@ -1,6 +1,6 @@
 use std::cmp::max;
 use std::collections::BTreeMap;
-use std::fmt;
+use std::{fmt, ops};
 
 use fedimint_client::module::init::recovery::{RecoveryFromHistory, RecoveryFromHistoryCommon};
 use fedimint_client::module::init::ClientModuleRecoverArgs;
@@ -217,7 +217,11 @@ impl RecoveryFromHistory for MintRecovery {
                             MintOutputStateMachine {
                                 common: MintOutputCommon {
                                     operation_id: OperationId::new_random(),
-                                    out_point,
+                                    txid: out_point.txid,
+                                    out_idxs: ops::RangeInclusive::new(
+                                        out_point.out_idx,
+                                        out_point.out_idx,
+                                    ),
                                 },
                                 state: crate::output::MintOutputStates::Created(
                                     MintOutputStatesCreated {

diff --git a/modules/fedimint-mint-client/src/client_db.rs b/modules/fedimint-mint-client/src/client_db.rs
@@ -13,6 +13,7 @@ use strum_macros::EnumIter;
 use crate::backup::recovery::MintRecoveryState;
 use crate::input::{MintInputCommon, MintInputStateMachine, MintInputStateMachineV1};
 use crate::oob::{MintOOBStateMachine, MintOOBStateMachineV1, MintOOBStates, MintOOBStatesV1};
+use crate::output::{MintOutputCommon, MintOutputStateMachine, MintOutputStateMachineV1};
 use crate::{MintClientStateMachines, SpendableNoteUndecoded};
 
 #[repr(u8)]
@@ -126,6 +127,18 @@ pub(crate) fn migrate_state_to_v2(
     let mint_client_state_machine_variant = u16::consensus_decode(cursor, &decoders)?;
 
     let new_mint_state_machine = match mint_client_state_machine_variant {
+        0 => {
+            let old_state = MintOutputStateMachineV1::consensus_decode(cursor, &decoders)?;
+            MintClientStateMachines::Output(MintOutputStateMachine {
+                common: MintOutputCommon {
+                    operation_id: old_state.common.operation_id,
+                    txid: old_state.common.out_point.txid,
+                    out_idxs: old_state.common.out_point.out_idx
+                        ..=old_state.common.out_point.out_idx,
+                },
+                state: old_state.state,
+            })
+        }
         1 => {
             let old_state = MintInputStateMachineV1::consensus_decode(cursor, &decoders)?;
 

diff --git a/modules/fedimint-mint-client/src/lib.rs b/modules/fedimint-mint-client/src/lib.rs
@@ -80,6 +80,7 @@ use futures::{pin_mut, StreamExt};
 use hex::ToHex;
 use input::MintInputStateCreatedBundle;
 use oob::MintOOBStatesCreatedMulti;
+use output::MintOutputStatesCreatedMulti;
 use serde::{Deserialize, Serialize};
 use strum::IntoEnumIterator;
 use tbs::{AggregatePublicKey, Signature};
@@ -94,8 +95,7 @@ use crate::client_db::{
 use crate::input::{MintInputCommon, MintInputStateMachine, MintInputStates};
 use crate::oob::{MintOOBStateMachine, MintOOBStates};
 use crate::output::{
-    MintOutputCommon, MintOutputStateMachine, MintOutputStates, MintOutputStatesCreated,
-    NoteIssuanceRequest,
+    MintOutputCommon, MintOutputStateMachine, MintOutputStates, NoteIssuanceRequest,
 };
 
 const MINT_E_CASH_TYPE_CHILD_ID: ChildId = ChildId(0);
@@ -1064,7 +1064,7 @@ impl MintClientModule {
         );
 
         let mut outputs = Vec::new();
-        let mut output_states = Vec::new();
+        let mut issuance_requests = Vec::new();
 
         for (amount, num) in denominations.iter() {
             for _ in 0..num {
@@ -1080,28 +1080,22 @@ impl MintClientModule {
                     amount,
                 });
 
-                output_states.push(MintOutputStatesCreated {
-                    amount,
-                    issuance_request,
-                });
+                issuance_requests.push((amount, issuance_request));
             }
         }
 
         let state_generator = Arc::new(move |txid, out_idxs: RangeInclusive<u64>| {
-            out_idxs
-                .clone()
-                .flat_map(|out_idx| {
-                    let output_i = (out_idx - out_idxs.clone().start()) as usize;
-                    let output_state = output_states.get(output_i).copied().unwrap();
-                    vec![MintClientStateMachines::Output(MintOutputStateMachine {
-                        common: MintOutputCommon {
-                            operation_id,
-                            out_point: OutPoint { txid, out_idx },
-                        },
-                        state: MintOutputStates::Created(output_state),
-                    })]
-                })
-                .collect()
+            assert_eq!(out_idxs.clone().count(), issuance_requests.len());
+            vec![MintClientStateMachines::Output(MintOutputStateMachine {
+                common: MintOutputCommon {
+                    operation_id,
+                    txid,
+                    out_idxs: out_idxs.clone(),
+                },
+                state: MintOutputStates::CreatedMulti(MintOutputStatesCreatedMulti {
+                    issuance_requests: out_idxs.zip(issuance_requests.clone()).collect(),
+                }),
+            })]
         });
 
         assert!(!outputs.is_empty());
@@ -1145,7 +1139,9 @@ impl MintClientModule {
                     return None;
                 };
 
-                if state.common.out_point != out_point {
+                if state.common.txid != out_point.txid
+                    || !state.common.out_idxs.contains(&out_point.out_idx)
+                {
                     return None;
                 }
 
@@ -1156,7 +1152,7 @@ impl MintClientModule {
                         "Failed to finalize transaction: {}",
                         failed.error
                     ))),
-                    MintOutputStates::Created(_) => None,
+                    MintOutputStates::Created(_) | MintOutputStates::CreatedMulti(_) => None,
                 }
             });
         pin_mut!(stream);