remove budget task, update budget on demand

This commit is contained in:
daywalker90 2026-08-09 00:13:06 +02:00
parent 651b11b229
commit 2214e07d81
No known key found for this signature in database
GPG key ID: 7AB4802ED5A639F3
9 changed files with 157 additions and 106 deletions

View file

@ -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<PluginState>, 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<PluginState>, 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<PluginState>, 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)]

View file

@ -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<PluginState>,
@ -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

View file

@ -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 {

View file

@ -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<String>), (nip47::NIP47Error, Option<String>)> {
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| {

View file

@ -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(

View file

@ -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);

View file

@ -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<Mutex<Config>>,
pub handles: Arc<tokio::sync::Mutex<HashMap<String, WalletService>>>,
pub rpc_lock: Arc<tokio::sync::Mutex<ClnRpc>>,
pub budget_jobs: Arc<Mutex<HashMap<String, oneshot::Sender<()>>>>,
pub hold_client: Arc<Mutex<Option<HoldClient<Channel>>>>,
}
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)]

View file

@ -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<PluginState>,
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<PluginState>) -> Result<(), anyhow::Error> {
let mut rpc = ClnRpc::new(
Path::new(&plugin.configuration().lightning_dir).join(&plugin.configuration().rpc_file),

View file

@ -17,6 +17,46 @@ use crate::{
structs::{ID_STORE, NwcStore, PluginState},
};
pub fn get_budget_msat(nwc_store: &NwcStore) -> Option<u64> {
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<u64>,
invoice_amt_msat: Option<u64>,
@ -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));
}