Cache rebalance setting in memory

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<Mutex<HashMap<PaymentHash, Vec<Event>>>>,
 	waker: Arc<Mutex<Option<Waker>>>,
 	kv_store: Arc<dyn DynStore>,
+	rebalance_enabled: RebalanceEnabledCache,
 	tx_metadata: TxMetadataStore,
 	logger: Arc<Logger>,
 }
@@ -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

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.

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<EventQueue>,
-	/// Key-value store for persistent storage.
-	store: Arc<dyn DynStore>,
 	/// 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<LightningWallet>, trusted: Arc<Box<DynTrustedWalletInterface>>,
 		tunables: Tunables, tx_metadata: TxMetadataStore, event_queue: Arc<EventQueue>,
-		store: Arc<dyn DynStore>, logger: Arc<Logger>,
+		logger: Arc<Logger>,
 	) -> 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<Output = Option<TriggerParams>> + 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<Output = Option<TriggerParams>> + 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;
 			}

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<dyn DynStore>,
+	value: Mutex<Option<bool>>,
 }
 
-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<dyn DynStore>) -> 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 {