From b2078bab813e00e6e90c01ad22daf264ecaf94bd Mon Sep 17 00:00:00 2001 From: daywalker90 Date: Sun, 9 Aug 2026 00:13:06 +0200 Subject: [PATCH] remove budget task, update budget on demand --- src/nwc.rs | 21 +--------- src/nwc_balance.rs | 7 +++- src/nwc_keysend.rs | 19 ++++++--- src/nwc_pay.rs | 38 +++++++++++------- src/parse.rs | 4 +- src/rpc.rs | 24 +++++++----- src/structs.rs | 4 +- src/tasks.rs | 50 ------------------------ src/util.rs | 96 ++++++++++++++++++++++++++++++++++++++++++++++ 9 files changed, 157 insertions(+), 106 deletions(-) diff --git a/src/nwc.rs b/src/nwc.rs index 9677fa6..c761614 100644 --- a/src/nwc.rs +++ b/src/nwc.rs @@ -15,7 +15,7 @@ use nostr_sdk::{ error::Error, relay::RelayStatus, }; -use tokio::{sync::oneshot, time}; +use tokio::time; use crate::{ OPT_NOTIFICATIONS, @@ -32,7 +32,6 @@ use crate::{ nwc_lookups::{list_transactions_response, lookup_invoice_response}, nwc_pay::pay_invoice_response, structs::{ID_MAX_AGE, NwcStore, PluginState, WalletService}, - tasks::budget_task, util::{build_capabilities, build_notifications_vec, is_read_only_nwc, save_event_id}, }; @@ -57,7 +56,7 @@ pub async fn run_nwc( } if nwc_store.interval_config.is_some() { - start_nwc_budget_job(&plugin, label.clone()); + log::debug!("NWC {label} uses an interval budget"); } let nostr_client_clone = nostr_client.clone(); @@ -207,22 +206,6 @@ pub async fn stop_nwc(plugin: Plugin, label: &String) { if let Some(wallet_service) = locked_handles.remove(label) { wallet_service.client.shutdown().await; } - - stop_nwc_budget_job(&plugin, label); -} - -pub fn start_nwc_budget_job(plugin: &Plugin, label: String) { - let (tx, rx) = oneshot::channel::<()>(); - tokio::spawn(budget_task(rx, plugin.clone(), label.clone())); - plugin.state().budget_jobs.lock().insert(label, tx); -} - -pub fn stop_nwc_budget_job(plugin: &Plugin, label: &String) { - let mut budget_jobs = plugin.state().budget_jobs.lock(); - let job = budget_jobs.remove(label); - if let Some(j) = job { - let _ = j.send(()); - } } #[allow(clippy::too_many_lines)] diff --git a/src/nwc_balance.rs b/src/nwc_balance.rs index 1c3c1ea..2fde546 100644 --- a/src/nwc_balance.rs +++ b/src/nwc_balance.rs @@ -2,7 +2,10 @@ use cln_plugin::Plugin; use cln_rpc::{model::requests::ListpeerchannelsRequest, primitives::ChannelState}; use nostr::nips::nip47; -use crate::{structs::PluginState, util::load_nwc_store}; +use crate::{ + structs::PluginState, + util::{get_budget_msat, load_nwc_store}, +}; pub async fn get_balance_response( plugin: Plugin, @@ -40,7 +43,7 @@ async fn get_balance( message: e.to_string(), })?; - let balance = if let Some(bdgt_amt) = nwc_store.budget_msat { + let balance = if let Some(bdgt_amt) = get_budget_msat(&nwc_store) { bdgt_amt } else { let listpeerchannels = rpc diff --git a/src/nwc_keysend.rs b/src/nwc_keysend.rs index dc1ee5e..c53c88d 100644 --- a/src/nwc_keysend.rs +++ b/src/nwc_keysend.rs @@ -10,7 +10,14 @@ use nostr::nips::nip47::{self}; use crate::{ structs::{NwcStore, PluginState}, - util::{at_or_above_version, budget_amount_check, load_nwc_store, update_nwc_store}, + util::{ + at_or_above_version, + budget_amount_check, + get_budget_msat, + load_nwc_store, + update_budget_msat, + update_nwc_store, + }, }; pub async fn pay_keysend_response( @@ -65,7 +72,7 @@ async fn pay_keysend( message: e.to_string(), })?; - budget_amount_check(Some(params.amount), None, nwc_store.budget_msat).map_err(|e| { + budget_amount_check(Some(params.amount), None, get_budget_msat(&nwc_store)).map_err(|e| { nip47::NIP47Error { code: nip47::ErrorCode::QuotaExceeded, message: e.to_string(), @@ -120,8 +127,8 @@ async fn xkeysend( .await { Ok(o) => { - if let Some(ref mut bdg) = nwc_store.budget_msat { - *bdg = bdg.saturating_sub(o.amount_sent_msat.msat()); + if nwc_store.budget_msat.is_some() { + update_budget_msat(&mut nwc_store, o.amount_sent_msat.msat()); update_nwc_store(rpc, label, nwc_store) .await .map_err(|e| nip47::NIP47Error { @@ -200,8 +207,8 @@ async fn keysend( .await { Ok(o) => { - if let Some(ref mut bdg) = nwc_store.budget_msat { - *bdg = bdg.saturating_sub(o.amount_sent_msat.msat()); + if nwc_store.budget_msat.is_some() { + update_budget_msat(&mut nwc_store, o.amount_sent_msat.msat()); update_nwc_store(rpc, label, nwc_store) .await .map_err(|e| nip47::NIP47Error { diff --git a/src/nwc_pay.rs b/src/nwc_pay.rs index 4f4e859..d41210b 100644 --- a/src/nwc_pay.rs +++ b/src/nwc_pay.rs @@ -12,7 +12,14 @@ use nostr::nips::nip47; use crate::{ structs::{NOT_INV_ERR, NwcStore, PluginState}, - util::{at_or_above_version, budget_amount_check, load_nwc_store, update_nwc_store}, + util::{ + at_or_above_version, + budget_amount_check, + get_budget_msat, + load_nwc_store, + update_budget_msat, + update_nwc_store, + }, }; pub async fn pay_invoice_response( @@ -172,17 +179,20 @@ async fn load_nwc_and_check_budget( ) })?; - budget_amount_check(params.amount, Some(invoice_amt_msat), nwc_store.budget_msat).map_err( - |e| { - ( - nip47::NIP47Error { - code: nip47::ErrorCode::QuotaExceeded, - message: e.to_string(), - }, - Some(id.to_owned()), - ) - }, - )?; + budget_amount_check( + params.amount, + Some(invoice_amt_msat), + get_budget_msat(&nwc_store), + ) + .map_err(|e| { + ( + nip47::NIP47Error { + code: nip47::ErrorCode::QuotaExceeded, + message: e.to_string(), + }, + Some(id.to_owned()), + ) + })?; Ok(nwc_store) } @@ -211,8 +221,8 @@ async fn update_budget_and_create_response( preimage: Secret, id: &str, ) -> Result<(nip47::PayInvoiceResponse, Option), (nip47::NIP47Error, Option)> { - if let Some(ref mut bdg) = nwc_store.budget_msat { - *bdg = bdg.saturating_sub(amount_sent_msat); + if nwc_store.budget_msat.is_some() { + update_budget_msat(nwc_store, amount_sent_msat); update_nwc_store(rpc, label, nwc_store.clone()) .await .map_err(|e| { diff --git a/src/parse.rs b/src/parse.rs index b1dffed..efd1e36 100644 --- a/src/parse.rs +++ b/src/parse.rs @@ -2,12 +2,12 @@ use std::path::Path; use anyhow::anyhow; use cln_plugin::ConfiguredPlugin; -use cln_rpc::{model::requests::GetinfoRequest, ClnRpc}; +use cln_rpc::{ClnRpc, model::requests::GetinfoRequest}; use nostr::types::RelayUrl; use crate::{ - structs::{PluginState, TimeUnit}, OPT_RELAYS, + structs::{PluginState, TimeUnit}, }; pub async fn read_startup_options( diff --git a/src/rpc.rs b/src/rpc.rs index fea29c7..1aba835 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -15,10 +15,16 @@ use serde_json::json; use crate::{ PLUGIN_NAME, - nwc::{run_nwc, send_nwc_info_event, start_nwc_budget_job, stop_nwc, stop_nwc_budget_job}, + nwc::{run_nwc, send_nwc_info_event, stop_nwc}, parse::parse_time_period, structs::{BudgetIntervalConfig, NwcStore, PluginState}, - util::{build_capabilities, is_read_only_nwc, load_nwc_store, update_nwc_store}, + util::{ + build_capabilities, + get_budget_msat, + is_read_only_nwc, + load_nwc_store, + update_nwc_store, + }, }; pub async fn nwc_create( @@ -62,6 +68,7 @@ pub async fn nwc_create( interval_secs: interval, reset_budget_msat: bgt_msat, last_reset: Timestamp::now().as_secs(), + spend_since_last_reset: 0, }; result.insert("interval_config".to_owned(), serde_json::to_value(&conf)?); Some(conf) @@ -119,8 +126,6 @@ pub async fn nwc_budget( let (label, budget_msat, interval_secs) = parse_full_args(args)?; - stop_nwc_budget_job(&plugin, &label); - let mut nwc_store = load_nwc_store(&mut rpc, &label).await?; let is_old_nwc_read_only = is_read_only_nwc(&nwc_store); @@ -132,6 +137,7 @@ pub async fn nwc_budget( interval_secs: interval, reset_budget_msat: budget, last_reset: Timestamp::now().as_secs(), + spend_since_last_reset: 0, }; nwc_store.interval_config = Some(interval_config.clone()); } else { @@ -146,10 +152,6 @@ pub async fn nwc_budget( update_nwc_store(&mut rpc, &label, nwc_store.clone()).await?; - if nwc_store.interval_config.is_some() { - start_nwc_budget_job(&plugin, label.clone()); - } - if is_old_nwc_read_only != is_new_nwc_read_only { let wallet_keys = Keys::new(SecretKey::from_hex(&nwc_store.walletkey)?); let (method_capabilities, _) = build_capabilities(is_new_nwc_read_only, &plugin); @@ -181,7 +183,8 @@ pub async fn nwc_list( let mut nwcs = Vec::new(); if let Some(lbl) = label { - let nwc_store = load_nwc_store(&mut rpc, &lbl).await?; + let mut nwc_store = load_nwc_store(&mut rpc, &lbl).await?; + nwc_store.budget_msat = get_budget_msat(&nwc_store); let wallet_key = Keys::new(SecretKey::from_hex(&nwc_store.walletkey)?); let client_key = Keys::new(nwc_store.uri.secret.clone()); let mut nwc_json = json!(nwc_store); @@ -206,7 +209,8 @@ pub async fn nwc_list( for datastore in all_stored_nwcs { let label = datastore.key.last().unwrap().to_owned(); - let nwc_store = load_nwc_store(&mut rpc, &label).await?; + let mut nwc_store = load_nwc_store(&mut rpc, &label).await?; + nwc_store.budget_msat = get_budget_msat(&nwc_store); let wallet_key = Keys::new(SecretKey::from_hex(&nwc_store.walletkey)?); let client_key = Keys::new(nwc_store.uri.secret.clone()); let mut nwc_json = json!(nwc_store); diff --git a/src/structs.rs b/src/structs.rs index a60c49a..6018153 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -9,7 +9,6 @@ use nostr::{ use nostr_sdk::client::Client; use parking_lot::Mutex; use serde::{Deserialize, Serialize}; -use tokio::sync::oneshot; use tonic::transport::Channel; use crate::hold::hold_client::HoldClient; @@ -23,7 +22,6 @@ pub struct PluginState { pub config: Arc>, pub handles: Arc>>, pub rpc_lock: Arc>, - pub budget_jobs: Arc>>>, pub hold_client: Arc>>>, } impl PluginState { @@ -32,7 +30,6 @@ impl PluginState { config: Arc::new(Mutex::new(Config::default())), handles: Arc::new(tokio::sync::Mutex::new(HashMap::new())), rpc_lock: Arc::new(tokio::sync::Mutex::new(ClnRpc::new(path).await?)), - budget_jobs: Arc::new(Mutex::new(HashMap::new())), hold_client: Arc::new(Mutex::new(None)), }) } @@ -86,6 +83,7 @@ pub struct BudgetIntervalConfig { pub interval_secs: u64, pub reset_budget_msat: u64, pub last_reset: u64, + pub spend_since_last_reset: u64, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src/tasks.rs b/src/tasks.rs index 3427834..42c9c61 100644 --- a/src/tasks.rs +++ b/src/tasks.rs @@ -1,67 +1,17 @@ use std::{path::Path, str::FromStr, time::Duration}; -use anyhow::anyhow; use cln_plugin::Plugin; use cln_rpc::{ ClnRpc, model::requests::{DeldatastoreRequest, ListdatastoreRequest}, }; use nostr::types::Timestamp; -use tokio::{sync::oneshot, time}; use crate::{ PLUGIN_NAME, structs::{ID_MAX_AGE, ID_STORE, PluginState}, - util::{load_nwc_store, update_nwc_store}, }; -pub async fn budget_task( - mut rx: oneshot::Receiver<()>, - plugin: Plugin, - label: String, -) -> Result<(), anyhow::Error> { - let mut rpc = ClnRpc::new( - Path::new(&plugin.configuration().lightning_dir).join(&plugin.configuration().rpc_file), - ) - .await?; - loop { - let mut nwc_store = load_nwc_store(&mut rpc, &label).await?; - let interval_config = nwc_store - .interval_config - .as_mut() - .ok_or_else(|| anyhow!("interval_config disappeared!"))?; - let now = Timestamp::now().as_secs(); - log::debug!( - "interval:{} now:{} prev:{}", - interval_config.interval_secs, - now, - interval_config.last_reset - ); - let next_reset = std::cmp::max( - interval_config - .interval_secs - .saturating_sub(now.saturating_sub(interval_config.last_reset)), - 1, - ); - tokio::select! { - _ = &mut rx => { - log::info!("Stopping budget task for {label}"); - break; - } - () = time::sleep(Duration::from_secs(next_reset)) => { - log::info!("Refreshing budget for {label}"); - *nwc_store.budget_msat - .as_mut() - .ok_or_else(||anyhow!("budget_msat missing"))? = interval_config.reset_budget_msat; - interval_config.last_reset = Timestamp::now().as_secs(); - update_nwc_store(&mut rpc, &label, nwc_store).await?; - log::info!("Done refreshing budget for {label}"); - } - } - } - Ok(()) -} - pub async fn cleanup_event_ids(plugin: Plugin) -> Result<(), anyhow::Error> { let mut rpc = ClnRpc::new( Path::new(&plugin.configuration().lightning_dir).join(&plugin.configuration().rpc_file), diff --git a/src/util.rs b/src/util.rs index f6064f7..2012e6c 100644 --- a/src/util.rs +++ b/src/util.rs @@ -17,6 +17,46 @@ use crate::{ structs::{ID_STORE, NwcStore, PluginState}, }; +pub fn get_budget_msat(nwc_store: &NwcStore) -> Option { + match nwc_store.budget_msat { + Some(b) => { + if let Some(conf) = &nwc_store.interval_config { + let now = Timestamp::now().as_secs(); + let spend = if now.saturating_sub(conf.last_reset) >= conf.interval_secs { + 0 + } else { + conf.spend_since_last_reset + }; + Some(conf.reset_budget_msat.saturating_sub(spend)) + } else { + Some(b) + } + } + None => None, + } +} + +pub fn update_budget_msat(nwc_store: &mut NwcStore, amount_spent_msat: u64) { + if let Some(bdg) = nwc_store.budget_msat.as_mut() { + if let Some(conf) = nwc_store.interval_config.as_mut() { + let now = Timestamp::now().as_secs(); + if now.saturating_sub(conf.last_reset) >= conf.interval_secs { + conf.last_reset = now; + conf.spend_since_last_reset = amount_spent_msat; + } else { + conf.spend_since_last_reset = conf + .spend_since_last_reset + .saturating_add(amount_spent_msat); + } + *bdg = conf + .reset_budget_msat + .saturating_sub(conf.spend_since_last_reset); + } else { + *bdg = bdg.saturating_sub(amount_spent_msat); + } + } +} + pub fn budget_amount_check( request_amt_msat: Option, invoice_amt_msat: Option, @@ -224,3 +264,59 @@ fn test_budget_check() { assert!(budget_amount_check(Some(0), Some(0), Some(1)).is_ok()); assert!(budget_amount_check(Some(0), Some(0), Some(0)).is_ok()); } + +#[test] +fn test_budget_interval_helpers() { + use crate::structs::BudgetIntervalConfig; + + let now = Timestamp::now().as_secs(); + let conf = BudgetIntervalConfig { + interval_secs: 10, + reset_budget_msat: 1000, + last_reset: now, + spend_since_last_reset: 0, + }; + let mut store = NwcStore { + uri: nostr::nips::nip47::NostrWalletConnectUri::new( + nostr::key::Keys::generate().public_key(), + vec![], + nostr::key::Keys::generate().secret_key().clone(), + None, + ), + walletkey: "test".to_owned(), + budget_msat: Some(1000), + interval_config: Some(conf), + }; + + assert_eq!(get_budget_msat(&store), Some(1000)); + + update_budget_msat(&mut store, 300); + assert_eq!( + store + .interval_config + .as_ref() + .unwrap() + .spend_since_last_reset, + 300 + ); + assert_eq!(get_budget_msat(&store), Some(700)); + + update_budget_msat(&mut store, 500); + assert_eq!(get_budget_msat(&store), Some(200)); + + store.interval_config.as_mut().unwrap().last_reset = now.saturating_sub(20); + let old_last_reset = store.interval_config.as_ref().unwrap().last_reset; + assert_eq!(get_budget_msat(&store), Some(1000)); + + update_budget_msat(&mut store, 400); + assert!(store.interval_config.as_ref().unwrap().last_reset > old_last_reset); + assert_eq!( + store + .interval_config + .as_ref() + .unwrap() + .spend_since_last_reset, + 400 + ); + assert_eq!(get_budget_msat(&store), Some(600)); +}