From a477046b1347cd1bfb50dec60e462b3d049b93f4 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Wed, 19 Aug 2026 01:24:58 -0500 Subject: [PATCH] Cache rebalance setting in memory Trigger checks repeatedly fetched this setting from remote storage. --- orange-sdk/src/event.rs | 16 +++++-- orange-sdk/src/lib.rs | 5 +-- orange-sdk/src/rebalancer.rs | 12 ++--- orange-sdk/src/store.rs | 87 ++++++++++++++++++++++++++++-------- 4 files changed, 87 insertions(+), 33 deletions(-) diff --git a/orange-sdk/src/event.rs b/orange-sdk/src/event.rs index 848d86f..c00f053 100644 --- a/orange-sdk/src/event.rs +++ b/orange-sdk/src/event.rs @@ -1,5 +1,5 @@ use crate::logging::Logger; -use crate::store::{self, MppOutcome, PaymentId, TxMetadataStore, TxType}; +use crate::store::{self, MppOutcome, PaymentId, RebalanceEnabledCache, TxMetadataStore, TxType}; use crate::dyn_store::DynStore; use ldk_node::bitcoin::hashes::Hash; @@ -212,6 +212,7 @@ pub struct EventQueue { pending_mpp_events: Arc>>>, waker: Arc>>, kv_store: Arc, + rebalance_enabled: RebalanceEnabledCache, tx_metadata: TxMetadataStore, logger: Arc, } @@ -223,7 +224,16 @@ impl EventQueue { let queue = Arc::new(Mutex::new(VecDeque::new())); let pending_mpp_events = Arc::new(Mutex::new(HashMap::new())); let waker = Arc::new(Mutex::new(None)); - Self { queue, pending_mpp_events, waker, kv_store, tx_metadata, logger } + let rebalance_enabled = RebalanceEnabledCache::new(Arc::clone(&kv_store)); + Self { queue, pending_mpp_events, waker, kv_store, rebalance_enabled, tx_metadata, logger } + } + + pub(crate) async fn get_rebalance_enabled(&self) -> bool { + self.rebalance_enabled.get().await + } + + pub(crate) async fn set_rebalance_enabled(&self, enabled: bool) { + self.rebalance_enabled.set(enabled).await; } /// Starts buffering terminal events for a multi-path payment while its leg metadata is being @@ -641,7 +651,7 @@ impl LdkEventHandler { } => { // We experienced a channel close, we disable rebalancing so we don't automatically // try to reopen the channel. - store::set_rebalance_enabled(self.event_queue.kv_store.as_ref(), false).await; + self.event_queue.set_rebalance_enabled(false).await; if let Err(e) = self .event_queue diff --git a/orange-sdk/src/lib.rs b/orange-sdk/src/lib.rs index db1a457..b51d72d 100644 --- a/orange-sdk/src/lib.rs +++ b/orange-sdk/src/lib.rs @@ -722,7 +722,6 @@ impl Wallet { tunables, tx_metadata.clone(), Arc::clone(&event_queue), - Arc::clone(&store), Arc::clone(&logger), )); @@ -777,12 +776,12 @@ impl Wallet { /// Sets whether the wallet should automatically rebalance from trusted/onchain to lightning. pub async fn set_rebalance_enabled(&self, value: bool) { - store::set_rebalance_enabled(self.inner.store.as_ref(), value).await + self.inner.event_queue.set_rebalance_enabled(value).await } /// Whether the wallet should automatically rebalance from trusted/onchain to lightning. pub async fn get_rebalance_enabled(&self) -> bool { - store::get_rebalance_enabled(self.inner.store.as_ref()).await + self.inner.event_queue.get_rebalance_enabled().await } /// Returns the lightning wallet's node id. diff --git a/orange-sdk/src/rebalancer.rs b/orange-sdk/src/rebalancer.rs index b2a017a..0db8cfc 100644 --- a/orange-sdk/src/rebalancer.rs +++ b/orange-sdk/src/rebalancer.rs @@ -1,12 +1,11 @@ use crate::bitcoin::Txid; use crate::bitcoin::hashes::Hash; use crate::bitcoin::hex::DisplayHex; -use crate::dyn_store::DynStore; use crate::lightning_wallet::LightningWallet; use crate::logging::Logger; use crate::store::{PaymentId, TxMetadata, TxMetadataStore, TxType}; use crate::trusted_wallet::DynTrustedWalletInterface; -use crate::{Event, EventQueue, PaymentType, Tunables, store}; +use crate::{Event, EventQueue, PaymentType, Tunables}; use bitcoin_payment_instructions::amount::Amount; use graduated_rebalancer::{RebalanceTrigger, RebalancerEvent, TriggerParams}; use ldk_node::lightning::util::logger::Logger as _; @@ -29,8 +28,6 @@ pub(crate) struct OrangeTrigger { tx_metadata: TxMetadataStore, /// The event handler for processing wallet events. event_queue: Arc, - /// Key-value store for persistent storage. - store: Arc, /// Time of the last on-chain sync, used to determine when to trigger rebalances. onchain_sync_time: AtomicU64, /// Logger for logging events and errors. @@ -42,7 +39,7 @@ impl OrangeTrigger { pub(crate) fn new( ln_wallet: Arc, trusted: Arc>, tunables: Tunables, tx_metadata: TxMetadataStore, event_queue: Arc, - store: Arc, logger: Arc, + logger: Arc, ) -> Self { let start = ln_wallet.inner.ldk_node.status().latest_onchain_wallet_sync_timestamp.unwrap_or(0); @@ -52,7 +49,6 @@ impl OrangeTrigger { tunables, tx_metadata, event_queue, - store, onchain_sync_time: AtomicU64::new(start), logger, } @@ -62,7 +58,7 @@ impl OrangeTrigger { impl RebalanceTrigger for OrangeTrigger { fn needs_trusted_rebalance(&self) -> impl Future> + Send { async move { - let rebalance_enabled = store::get_rebalance_enabled(self.store.as_ref()).await; + let rebalance_enabled = self.event_queue.get_rebalance_enabled().await; if !rebalance_enabled { return None; } @@ -144,7 +140,7 @@ impl RebalanceTrigger for OrangeTrigger { fn needs_onchain_rebalance(&self) -> impl Future> + Send { async move { - let rebalance_enabled = store::get_rebalance_enabled(self.store.as_ref()).await; + let rebalance_enabled = self.event_queue.get_rebalance_enabled().await; if !rebalance_enabled { return None; } diff --git a/orange-sdk/src/store.rs b/orange-sdk/src/store.rs index 38af996..b478f56 100644 --- a/orange-sdk/src/store.rs +++ b/orange-sdk/src/store.rs @@ -29,6 +29,7 @@ use std::fmt; use std::str::FromStr; use std::sync::{Arc, RwLock, RwLockReadGuard}; use std::time::Duration; +use tokio::sync::Mutex; const STORE_PRIMARY_KEY: &str = "orange_sdk"; const STORE_SECONDARY_KEY: &str = "payment_store"; @@ -733,27 +734,55 @@ impl TxMetadataStore { const REBALANCE_ENABLED_KEY: &str = "rebalance_enabled"; -pub(crate) async fn get_rebalance_enabled(store: &dyn DynStore) -> bool { - match KVStore::read(store, STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY).await { - Ok(bytes) => Readable::read(&mut &bytes[..]).expect("Invalid data in rebalance_enabled"), - Err(e) if e.kind() == io::ErrorKind::NotFound => { - // if rebalance_enabled is not found, default to true - // and write it to the store so we don't have to do this again - let rebalance_enabled = true; - set_rebalance_enabled(store, rebalance_enabled).await; - rebalance_enabled - }, - Err(e) => { - panic!("Failed to read rebalance_enabled: {e}"); - }, - } +pub(crate) struct RebalanceEnabledCache { + store: Arc, + value: Mutex>, } -pub(crate) async fn set_rebalance_enabled(store: &dyn DynStore, enabled: bool) { - let bytes = enabled.encode(); - KVStore::write(store, STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY, bytes) - .await - .expect("Failed to write rebalance_enabled"); +impl RebalanceEnabledCache { + pub(crate) fn new(store: Arc) -> Self { + Self { store, value: Mutex::new(None) } + } + + pub(crate) async fn get(&self) -> bool { + let mut cached = self.value.lock().await; + if let Some(value) = *cached { + return value; + } + + let value = + match KVStore::read(self.store.as_ref(), STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY) + .await + { + Ok(bytes) => { + Readable::read(&mut &bytes[..]).expect("Invalid data in rebalance_enabled") + }, + Err(e) if e.kind() == io::ErrorKind::NotFound => { + // If rebalance_enabled is not found, default to true and persist the value. + let value = true; + Self::write(self.store.as_ref(), value).await; + value + }, + Err(e) => { + panic!("Failed to read rebalance_enabled: {e}"); + }, + }; + *cached = Some(value); + value + } + + pub(crate) async fn set(&self, enabled: bool) { + let mut cached = self.value.lock().await; + Self::write(self.store.as_ref(), enabled).await; + *cached = Some(enabled); + } + + async fn write(store: &dyn DynStore, enabled: bool) { + let bytes = enabled.encode(); + KVStore::write(store, STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY, bytes) + .await + .expect("Failed to write rebalance_enabled"); + } } pub(crate) async fn write_splice_out(store: &dyn DynStore, details: &PaymentDetails) { @@ -813,6 +842,26 @@ mod tests { PaymentId::SelfCustodial(LIGHTNING_LEG) } + #[tokio::test] + async fn rebalance_enabled_set_updates_cache_and_store() { + let (_path, store) = temp_sqlite_store(); + let cache = RebalanceEnabledCache::new(Arc::clone(&store)); + + cache.set(false).await; + assert!(!cache.get().await); + let stored = KVStore::read(store.as_ref(), STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY) + .await + .unwrap(); + assert!(!bool::read(&mut &stored[..]).unwrap()); + + cache.set(true).await; + assert!(cache.get().await); + let stored = KVStore::read(store.as_ref(), STORE_PRIMARY_KEY, "", REBALANCE_ENABLED_KEY) + .await + .unwrap(); + assert!(bool::read(&mut &stored[..]).unwrap()); + } + fn mpp_metadata() -> TxMetadata { TxMetadata { ty: TxType::MppPayment {