Fix race condition in rebalances

Before we would return a successful rebalance on when we received the
final payment, but we did not wait for the trusted wallet to get the
succcess back. This could cause a race where if we called
list_transactions() after the rebalance but before the trusted wallet
got the success event, list_transactions would not correctly reflect the
rebalance and would hit a debug_assert.

This fixes by waiting for the payment success event on the trusted side.

Author: benthecarman

graduated-rebalancer/src/lib.rs

@@ -79,9 +79,10 @@ pub trait TrustedWallet: Send + Sync {
 		&self, method: PaymentMethod, amount: Amount,
 	) -> Pin<Box<dyn Future<Output = Result<[u8; 32], Self::Error>> + Send + '_>>;
 
-	/// Get transaction fee paid
-	fn get_tx_fee(&self, id: [u8; 32])
-		-> Pin<Box<dyn Future<Output = Option<Amount>> + Send + '_>>;
+	/// Wait for a payment success notification
+	fn await_payment_success(
+		&self, payment_hash: [u8; 32],
+	) -> Pin<Box<dyn Future<Output = Option<ReceivedLightningPayment>> + Send + '_>>;
 }
 
 /// Trait representing a lightning wallet backend
@@ -265,10 +266,6 @@ where
 						amount_msat: transfer_amt.milli_sats(),
 					});
 
-					// get the fee of the rebalance transaction
-					let rebalance_tx_fee =
-						self.trusted.get_tx_fee(rebalance_id).await.unwrap_or(Amount::ZERO);
-
 					let ln_payment = match self
 						.ln_wallet
 						.await_payment_receipt(expected_hash.to_byte_array())
@@ -281,6 +278,18 @@ where
 						},
 					};
 
+					let trusted_payment = match self
+						.trusted
+						.await_payment_success(expected_hash.to_byte_array())
+						.await
+					{
+						Some(success) => success,
+						None => {
+							log_error!(self.logger, "Failed to send rebalance payment!");
+							return;
+						},
+					};
+
 					log_info!(
 						self.logger,
 						"Rebalance succeeded. Sent trusted tx {} to lightning tx {}",
@@ -294,7 +303,7 @@ where
 						ln_rebalance_payment_id: ln_payment.id,
 						amount_msat: transfer_amt.milli_sats(),
 						fee_msat: ln_payment.fee_paid_msat.unwrap_or_default()
-							+ rebalance_tx_fee.milli_sats(),
+							+ trusted_payment.fee_paid_msat.unwrap_or_default(),
 					});
 				},
 				Err(e) => {

orange-sdk/Cargo.toml

@@ -26,15 +26,15 @@ reqwest = { version = "0.12.23", default-features = false, features = ["rustls-t
 breez-sdk-spark = { git = "https://github.com/breez/spark-sdk.git", rev = "41212dfcfe36e22a55ac224c791b326f259f90d6", default-features = false, features = ["rustls-tls"], optional = true }
 tokio = { version = "1.0", default-features = false, features = ["rt-multi-thread", "sync"] }
 uuid = { version = "1.0", default-features = false, optional = true }
-cdk = { git = "https://github.com/benthecarman/cdk.git", rev = "7d25e9ae5ed7f47f9ae7e87d8a9ee16797fee8cd", default-features = false, features = ["wallet"], optional = true }
+cdk = { git = "https://github.com/benthecarman/cdk.git", rev = "39c1206a4a1dda2adc1f3e23628136ef645f6c6b", default-features = false, features = ["wallet"], optional = true }
 serde_json = { version = "1.0", optional = true }
 async-trait = "0.1"
 log = "0.4.28"
 
 corepc-node = { version = "0.8.0", features = ["29_0", "download"], optional = true }
-cdk-ldk-node = { git = "https://github.com/benthecarman/cdk.git", rev = "7d25e9ae5ed7f47f9ae7e87d8a9ee16797fee8cd", optional = true }
-cdk-sqlite = { git = "https://github.com/benthecarman/cdk.git", rev = "7d25e9ae5ed7f47f9ae7e87d8a9ee16797fee8cd", optional = true }
-cdk-axum = { git = "https://github.com/benthecarman/cdk.git", rev = "7d25e9ae5ed7f47f9ae7e87d8a9ee16797fee8cd", optional = true }
+cdk-ldk-node = { git = "https://github.com/benthecarman/cdk.git", rev = "39c1206a4a1dda2adc1f3e23628136ef645f6c6b", optional = true }
+cdk-sqlite = { git = "https://github.com/benthecarman/cdk.git", rev = "39c1206a4a1dda2adc1f3e23628136ef645f6c6b", optional = true }
+cdk-axum = { git = "https://github.com/benthecarman/cdk.git", rev = "39c1206a4a1dda2adc1f3e23628136ef645f6c6b", optional = true }
 axum = { version = "0.8.1", optional = true }
 
 uniffi = { version = "0.29", features = ["cli", "tokio"], optional = true }

orange-sdk/src/lib.rs

@@ -887,10 +887,10 @@ impl Wallet {
 			}
 		}
 
-		for (_, tx_info) in internal_transfers {
+		for (id, tx_info) in internal_transfers {
 			debug_assert!(
 				tx_info.send_fee.is_some(),
-				"Internal transfers must have a send fee, got {tx_info:?}",
+				"Internal transfers must have a send fee, got {id}: {tx_info:?}",
 			);
 			debug_assert!(tx_info.transaction.is_some());
 

orange-sdk/src/trusted_wallet/cashu/mod.rs

@@ -1,5 +1,6 @@
 //! An implementation of `TrustedWalletInterface` using the Cashu (CDK) SDK.
 
+use crate::bitcoin::hex::DisplayHex;
 use crate::logging::Logger;
 use crate::store::{PaymentId, TxMetadataStore, TxStatus};
 use crate::trusted_wallet::{Payment, TrustedError, TrustedWalletInterface};
@@ -26,8 +27,11 @@ use cdk::wallet::Wallet;
 use cdk::wallet::types::{Transaction, TransactionDirection};
 use cdk::{Amount as CdkAmount, StreamExt};
 
+use graduated_rebalancer::ReceivedLightningPayment;
+
 use tokio::sync::{mpsc, watch};
 
+use std::collections::HashMap;
 use std::future::Future;
 use std::pin::Pin;
 use std::str::FromStr;
@@ -55,6 +59,8 @@ pub struct Cashu {
 	cashu_wallet: Arc<Wallet>,
 	unit: CurrencyUnit,
 	shutdown_sender: watch::Sender<()>,
+	payment_success_sender: watch::Sender<()>,
+	payment_success_flag: watch::Receiver<()>,
 	logger: Arc<Logger>,
 	supports_bolt12: bool,
 	mint_quote_sender: mpsc::Sender<MintQuote>,
@@ -291,8 +297,14 @@ impl TrustedWalletInterface for Cashu {
 			let event_queue = Arc::clone(&self.event_queue);
 			let tx_metadata = self.tx_metadata.clone();
 			let quote_id = quote.id.clone();
+			let payment_success_sender = self.payment_success_sender.clone();
 			self.runtime.spawn(async move {
-				match cashu_wallet.melt(&quote_id).await {
+				let mut metadata = HashMap::new();
+				if let Some(hash) = &payment_hash {
+					metadata.insert(PAYMENT_HASH_METADATA_KEY.to_string(), hash.to_string());
+				}
+
+				match cashu_wallet.melt_with_metadata(&quote_id, metadata).await {
 					Ok(res) => {
 						match res.state {
 							MeltQuoteState::Paid => {
@@ -304,6 +316,8 @@ impl TrustedWalletInterface for Cashu {
 									map.get(&payment_id).is_some_and(|m| m.ty.is_rebalance())
 								};
 								if is_rebalance {
+									// make sure we still send payment success
+									payment_success_sender.send(()).unwrap();
 									return;
 								}
 
@@ -369,6 +383,8 @@ impl TrustedWalletInterface for Cashu {
 									payment_preimage,
 									fee_paid_msat: Some(fee_paid_sat * 1_000), // convert to msats
 								});
+
+								payment_success_sender.send(()).unwrap();
 							},
 							MeltQuoteState::Failed => {
 								log_error!(logger, "Melt failed for quote: {quote_id}");
@@ -418,6 +434,35 @@ impl TrustedWalletInterface for Cashu {
 		})
 	}
 
+	fn await_payment_success(
+		&self, payment_hash: [u8; 32],
+	) -> Pin<Box<dyn Future<Output = Option<ReceivedLightningPayment>> + Send + '_>> {
+		Box::pin(async move {
+			loop {
+				let txs = self
+					.cashu_wallet
+					.list_transactions(Some(TransactionDirection::Outgoing))
+					.await
+					.ok()?;
+
+				let hex = payment_hash.to_lower_hex_string();
+				let tx = txs.iter().find(|tx| {
+					tx.metadata.get(PAYMENT_HASH_METADATA_KEY).is_some_and(|h| h == &hex)
+				});
+
+				if let Some(tx) = tx {
+					let payment_id = Self::id_to_32_byte_array(tx.quote_id.as_ref().expect("safe"));
+					return Some(ReceivedLightningPayment {
+						id: payment_id,
+						fee_paid_msat: Some(convert_amount(tx.fee, &self.unit).ok()?.milli_sats()),
+					});
+				}
+
+				self.await_payment_success().await;
+			}
+		})
+	}
+
 	fn stop(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
 		Box::pin(async move {
 			log_info!(self.logger, "Stopping Cashu wallet");
@@ -426,6 +471,8 @@ impl TrustedWalletInterface for Cashu {
 	}
 }
 
+const PAYMENT_HASH_METADATA_KEY: &str = "payment_hash";
+
 impl Cashu {
 	pub(crate) async fn init(
 		config: &WalletConfig, cashu_config: CashuConfig, store: Arc<dyn KVStore + Sync + Send>,
@@ -485,6 +532,7 @@ impl Cashu {
 			.unwrap_or(false);
 
 		let (shutdown_sender, mut shutdown_receiver) = watch::channel::<()>(());
+		let (payment_success_sender, payment_success_flag) = watch::channel(());
 
 		// Create channel for mint quote monitoring with bounded capacity
 		let (mint_quote_sender, mut mint_quote_receiver) = mpsc::channel::<MintQuote>(32);
@@ -554,6 +602,8 @@ impl Cashu {
 			cashu_wallet,
 			unit: cashu_config.unit,
 			shutdown_sender,
+			payment_success_sender,
+			payment_success_flag,
 			logger,
 			supports_bolt12,
 			mint_quote_sender,
@@ -644,6 +694,12 @@ impl Cashu {
 		}
 		Ok(())
 	}
+
+	pub(crate) async fn await_payment_success(&self) {
+		let mut flag = self.payment_success_flag.clone();
+		flag.mark_unchanged();
+		let _ = flag.changed().await;
+	}
 }
 
 fn convert_amount(cdk_amount: CdkAmount, unit: &CurrencyUnit) -> Result<Amount, TrustedError> {

orange-sdk/src/trusted_wallet/dummy.rs

@@ -8,8 +8,11 @@ use bitcoin_payment_instructions::PaymentMethod;
 use bitcoin_payment_instructions::amount::Amount;
 use corepc_node::client::bitcoin::Network;
 use corepc_node::{Node as Bitcoind, get_available_port};
+use graduated_rebalancer::ReceivedLightningPayment;
+use ldk_node::lightning::ln::channelmanager;
 use ldk_node::lightning::ln::msgs::SocketAddress;
 use ldk_node::lightning_invoice::{Bolt11Invoice, Bolt11InvoiceDescription, Description};
+use ldk_node::payment::{PaymentKind, PaymentStatus};
 use ldk_node::{Event, Node};
 use rand::RngCore;
 use std::env::temp_dir;
@@ -18,6 +21,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::{Arc, RwLock};
 use std::time::Duration;
 use tokio::runtime::Runtime;
+use tokio::sync::watch;
 use uuid::Uuid;
 
 /// A dummy implementation of `TrustedWalletInterface` for testing purposes.
@@ -28,6 +32,7 @@ pub(crate) struct DummyTrustedWallet {
 	current_bal_msats: Arc<AtomicU64>,
 	payments: Arc<RwLock<Vec<Payment>>>,
 	ldk_node: Arc<Node>,
+	payment_success_flag: watch::Receiver<()>,
 }
 
 #[derive(Clone)]
@@ -78,6 +83,8 @@ impl DummyTrustedWallet {
 		let current_bal_msats = Arc::new(AtomicU64::new(0));
 		let payments: Arc<RwLock<Vec<Payment>>> = Arc::new(RwLock::new(vec![]));
 
+		let (payment_success_sender, payment_success_flag) = watch::channel(());
+
 		let events_ref = Arc::clone(&ldk_node);
 		let bal = Arc::clone(&current_bal_msats);
 		let pays = Arc::clone(&payments);
@@ -126,6 +133,8 @@ impl DummyTrustedWallet {
 								})
 								.unwrap();
 						}
+
+						payment_success_sender.send(()).unwrap();
 					},
 					Event::PaymentFailed { payment_id, payment_hash, reason } => {
 						// convert id
@@ -241,7 +250,13 @@ impl DummyTrustedWallet {
 			panic!("No usable channels found {channels:?}");
 		}
 
-		DummyTrustedWallet { current_bal_msats, payments, ldk_node }
+		DummyTrustedWallet { current_bal_msats, payments, ldk_node, payment_success_flag }
+	}
+
+	pub(crate) async fn await_payment_success(&self) {
+		let mut flag = self.payment_success_flag.clone();
+		flag.mark_unchanged();
+		let _ = flag.changed().await;
 	}
 }
 
@@ -359,6 +374,40 @@ impl TrustedWalletInterface for DummyTrustedWallet {
 		})
 	}
 
+	fn await_payment_success(
+		&self, payment_hash: [u8; 32],
+	) -> Pin<Box<dyn Future<Output = Option<ReceivedLightningPayment>> + Send + '_>> {
+		Box::pin(async move {
+			let id = channelmanager::PaymentId(payment_hash);
+			loop {
+				if let Some(payment) = self.ldk_node.payment(&id) {
+					let counterparty_skimmed_fee_msat = match payment.kind {
+						PaymentKind::Bolt11 { hash, .. } => {
+							debug_assert!(hash.0 == payment_hash, "Payment Hash mismatch");
+							None
+						},
+						PaymentKind::Bolt11Jit { hash, counterparty_skimmed_fee_msat, .. } => {
+							debug_assert!(hash.0 == payment_hash, "Payment Hash mismatch");
+							counterparty_skimmed_fee_msat
+						},
+						_ => return None, // Ignore other payment kinds, we only care about the one we just sent.
+					};
+					match payment.status {
+						PaymentStatus::Succeeded => {
+							return Some(ReceivedLightningPayment {
+								id: payment.id.0,
+								fee_paid_msat: counterparty_skimmed_fee_msat,
+							});
+						},
+						PaymentStatus::Pending => {},
+						PaymentStatus::Failed => return None,
+					}
+				}
+				self.await_payment_success().await;
+			}
+		})
+	}
+
 	fn stop(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>> {
 		Box::pin(async move {
 			let _ = self.ldk_node.stop();

orange-sdk/src/trusted_wallet/mod.rs

@@ -7,6 +7,8 @@ use ldk_node::lightning_invoice::Bolt11Invoice;
 use bitcoin_payment_instructions::PaymentMethod;
 use bitcoin_payment_instructions::amount::Amount;
 
+use graduated_rebalancer::ReceivedLightningPayment;
+
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -82,6 +84,12 @@ pub trait TrustedWalletInterface: Send + Sync + private::Sealed {
 		&self, method: PaymentMethod, amount: Amount,
 	) -> Pin<Box<dyn Future<Output = Result<[u8; 32], TrustedError>> + Send + '_>>;
 
+	/// Waits for a payment with the given payment hash to succeed.
+	/// Returns the `ReceivedLightningPayment` if successful, or `None` if it fails or times out.
+	fn await_payment_success(
+		&self, payment_hash: [u8; 32],
+	) -> Pin<Box<dyn Future<Output = Option<ReceivedLightningPayment>> + Send + '_>>;
+
 	/// Stops the wallet, cleaning up any resources.
 	/// This is typically used to gracefully shut down the wallet.
 	fn stop(&self) -> Pin<Box<dyn Future<Output = ()> + Send + '_>>;
@@ -110,13 +118,10 @@ impl<T: ?Sized + TrustedWalletInterface> graduated_rebalancer::TrustedWallet for
 		Box::pin(async move { self.0.pay(method, amount).await })
 	}
 
-	fn get_tx_fee(
-		&self, id: [u8; 32],
-	) -> Pin<Box<dyn Future<Output = Option<Amount>> + Send + '_>> {
-		Box::pin(async move {
-			let trusted_txs = self.0.list_payments().await.unwrap_or_default();
-			trusted_txs.iter().find(|p| p.id == id).map(|p| p.fee)
-		})
+	fn await_payment_success(
+		&self, payment_hash: [u8; 32],
+	) -> Pin<Box<dyn Future<Output = Option<ReceivedLightningPayment>> + Send + '_>> {
+		Box::pin(async move { self.0.await_payment_success(payment_hash).await })
 	}
 }