diff --git a/crates/tsa-daemon/src/http_server.rs b/crates/tsa-daemon/src/http_server.rs index f250a60d..dea85680 100644 --- a/crates/tsa-daemon/src/http_server.rs +++ b/crates/tsa-daemon/src/http_server.rs @@ -15,6 +15,7 @@ use kumo_log_types::*; use kumo_server_common::http_server::auth::TrustedIpRequired; use kumo_server_common::http_server::{AppError, RouterAndDocs}; use message::message::QueueNameComponents; +use parking_lot::Mutex; use rfc5321::ForwardPath; use serde::Serialize; use serde_json::Value as JsonValue; @@ -22,8 +23,9 @@ use sha2::{Digest, Sha256}; use sqlite::{Connection, ConnectionThreadSafe}; use std::collections::HashMap; use std::hash::Hash; -use std::sync::{LazyLock, Mutex}; +use std::sync::{Arc, LazyLock}; use tokio::sync::broadcast::{channel, Sender}; +use tokio::task::spawn_blocking; use toml_edit::{value, Value as TomlValue}; use utoipa::OpenApi; @@ -33,7 +35,7 @@ static HISTORY: LazyLock = LazyLock::new(|| open_history_d static SUSPENSION_TX: LazyLock = LazyLock::new(|| SubscriberMgr::new()); pub fn open_history_db() -> anyhow::Result { - let path = DB_PATH.lock().unwrap().clone(); + let path = DB_PATH.lock().clone(); let mut db = Connection::open_thread_safe(&path) .with_context(|| format!("opening TSA database {path}"))?; @@ -123,8 +125,8 @@ enum PreferRollup { No, } -fn create_config( - db: &ConnectionThreadSafe, +async fn create_config( + db: &Arc, rule_hash: &str, rule: &Rule, record: &JsonLogRecord, @@ -133,40 +135,47 @@ fn create_config( source: &str, prefer_rollup: PreferRollup, ) -> anyhow::Result<()> { - let mut upsert = db.prepare( - "INSERT INTO config + let db = db.clone(); + let source = source.to_string(); + let domain = domain.to_string(); + let name = config.name.to_string(); + let value = serde_json::to_string(&config.value)?; + let expires = (record.timestamp + chrono::Duration::from_std(rule.duration)?).to_rfc3339(); + let site = record.site.to_string(); + let rule_hash = rule_hash.to_string(); + let reason = format!("automation rule: {}", regex_list_to_string(&rule.regex)); + let mx_rollup = if prefer_rollup == PreferRollup::Yes && rule.was_rollup { + 1 + } else { + 0 + }; + + spawn_blocking(move || { + let mut upsert = db.prepare( + "INSERT INTO config (rule_hash, site_name, domain, mx_rollup, source, name, value, reason, expires) VALUES ($hash, $site, $domain, $mx_rollup, $source, $name, $value, $reason, $expires) ON CONFLICT (rule_hash, site_name) DO UPDATE SET expires=$expires", - )?; + )?; - let expires = (record.timestamp + chrono::Duration::from_std(rule.duration)?).to_rfc3339(); + upsert.bind(("$hash", rule_hash.as_str()))?; + upsert.bind(("$site", site.as_str()))?; + upsert.bind(("$domain", domain.as_str()))?; + upsert.bind(("$mx_rollup", mx_rollup))?; + upsert.bind(("$source", source.as_str()))?; + upsert.bind(("$name", name.as_str()))?; + upsert.bind(("$value", value.as_str()))?; - upsert.bind(("$hash", rule_hash))?; - upsert.bind(("$site", record.site.as_str()))?; - upsert.bind(("$domain", domain))?; - upsert.bind(( - "$mx_rollup", - if prefer_rollup == PreferRollup::Yes && rule.was_rollup { - 1 - } else { - 0 - }, - ))?; - upsert.bind(("$source", source))?; - upsert.bind(("$name", config.name.as_str()))?; - let value = serde_json::to_string(&config.value)?; - upsert.bind(("$value", value.as_str()))?; + upsert.bind(("$reason", reason.as_str()))?; + upsert.bind(("$expires", expires.as_str()))?; - let reason = format!("automation rule: {}", regex_list_to_string(&rule.regex)); - upsert.bind(("$reason", reason.as_str()))?; - upsert.bind(("$expires", expires.as_str()))?; + upsert.next()?; - upsert.next()?; - - Ok(()) + Ok(()) + }) + .await? } fn regex_list_to_string(list: &[Regex]) -> String { @@ -197,8 +206,8 @@ enum UseTenant { No, } -fn create_bounce( - db: &ConnectionThreadSafe, +async fn create_bounce( + db: &Arc, rule_hash: &str, rule: &Rule, record: &JsonLogRecord, @@ -226,26 +235,6 @@ fn create_bounce( } else { None }; - - let mut upsert = db - .prepare( - "INSERT INTO sched_q_bounces - (rule_hash, campaign, tenant, domain, reason, expires) - VALUES - ($hash, $campaign, $tenant, $domain, $reason, $expires) - ON CONFLICT (rule_hash, campaign, tenant, domain) - DO UPDATE SET expires=$expires", - ) - .context("prepare sched_q_bounces upsert")?; - - let expires = record.timestamp + chrono::Duration::from_std(rule.duration)?; - let expires_str = expires.to_rfc3339(); - - upsert.bind(("$hash", rule_hash))?; - upsert.bind(("$campaign", campaign))?; - upsert.bind(("$tenant", tenant))?; - upsert.bind(("$domain", components.domain))?; - let mut reason = format!( "automation rule: {} domain={}", regex_list_to_string(&rule.regex), @@ -257,10 +246,43 @@ fn create_bounce( if let Some(campaign) = &campaign { reason.push_str(&format!(" campaign={campaign}")); } - upsert.bind(("$reason", reason.as_str()))?; - upsert.bind(("$expires", expires_str.as_str()))?; + let expires = record.timestamp + chrono::Duration::from_std(rule.duration)?; - upsert.next().context("execute sched_q_bounces upsert")?; + { + let db = db.clone(); + let reason = reason.clone(); + let domain = components.domain.to_string(); + let campaign = campaign.as_ref().map(|c| c.to_string()); + let tenant = tenant.as_ref().map(|c| c.to_string()); + let rule_hash = rule_hash.to_string(); + spawn_blocking(move || { + let mut upsert = db + .prepare( + "INSERT INTO sched_q_bounces + (rule_hash, campaign, tenant, domain, reason, expires) + VALUES + ($hash, $campaign, $tenant, $domain, $reason, $expires) + ON CONFLICT (rule_hash, campaign, tenant, domain) + DO UPDATE SET expires=$expires", + ) + .context("prepare sched_q_bounces upsert")?; + + let expires_str = expires.to_rfc3339(); + + upsert.bind(("$hash", rule_hash.as_str()))?; + upsert.bind(("$campaign", campaign.as_deref()))?; + upsert.bind(("$tenant", tenant.as_deref()))?; + upsert.bind(("$domain", domain.as_str()))?; + + upsert.bind(("$reason", reason.as_str()))?; + upsert.bind(("$expires", expires_str.as_str()))?; + + upsert.next().context("execute sched_q_bounces upsert")?; + + Ok::<_, anyhow::Error>(()) + }) + .await??; + } events.push(SubscriptionItem::SchedQBounce(SchedQBounce { rule_hash: rule_hash.to_string(), @@ -274,8 +296,8 @@ fn create_bounce( Ok(()) } -fn create_tenant_suspension( - db: &ConnectionThreadSafe, +async fn create_tenant_suspension( + db: &Arc, rule_hash: &str, rule: &Rule, record: &JsonLogRecord, @@ -297,26 +319,7 @@ fn create_tenant_suspension( } else { None }; - - let mut upsert = db - .prepare( - "INSERT INTO sched_q_suspensions - (rule_hash, campaign, tenant, domain, reason, expires) - VALUES - ($hash, $campaign, $tenant, $domain, $reason, $expires) - ON CONFLICT (rule_hash, campaign, tenant, domain) - DO UPDATE SET expires=$expires", - ) - .context("prepare sched_q_suspensions upsert")?; - let expires = record.timestamp + chrono::Duration::from_std(rule.duration)?; - let expires_str = expires.to_rfc3339(); - - upsert.bind(("$hash", rule_hash))?; - upsert.bind(("$campaign", campaign))?; - upsert.bind(("$tenant", tenant))?; - upsert.bind(("$domain", components.domain))?; - let mut reason = format!( "automation rule: {} tenant={tenant} domain={}", regex_list_to_string(&rule.regex), @@ -325,12 +328,44 @@ fn create_tenant_suspension( if let Some(campaign) = &campaign { reason.push_str(&format!(" campaign={campaign}")); } - upsert.bind(("$reason", reason.as_str()))?; - upsert.bind(("$expires", expires_str.as_str()))?; - upsert - .next() - .context("execute sched_q_suspensions upsert")?; + { + let reason = reason.to_string(); + let rule_hash = rule_hash.to_string(); + let campaign = campaign.as_ref().map(|c| c.to_string()); + let tenant = tenant.to_string(); + let domain = components.domain.to_string(); + + let db = db.clone(); + spawn_blocking(move || { + let mut upsert = db + .prepare( + "INSERT INTO sched_q_suspensions + (rule_hash, campaign, tenant, domain, reason, expires) + VALUES + ($hash, $campaign, $tenant, $domain, $reason, $expires) + ON CONFLICT (rule_hash, campaign, tenant, domain) + DO UPDATE SET expires=$expires", + ) + .context("prepare sched_q_suspensions upsert")?; + + let expires_str = expires.to_rfc3339(); + + upsert.bind(("$hash", rule_hash.as_str()))?; + upsert.bind(("$campaign", campaign.as_deref()))?; + upsert.bind(("$tenant", tenant.as_str()))?; + upsert.bind(("$domain", domain.as_str()))?; + + upsert.bind(("$reason", reason.as_str()))?; + upsert.bind(("$expires", expires_str.as_str()))?; + + upsert + .next() + .context("execute sched_q_suspensions upsert")?; + Ok::<_, anyhow::Error>(()) + }) + .await??; + } events.push(SubscriptionItem::SchedQSuspension(SchedQSuspension { rule_hash: rule_hash.to_string(), @@ -344,35 +379,48 @@ fn create_tenant_suspension( Ok(()) } -fn create_ready_q_suspension( - db: &ConnectionThreadSafe, +async fn create_ready_q_suspension( + db: &Arc, rule_hash: &str, rule: &Rule, record: &JsonLogRecord, source: &str, events: &mut Vec, ) -> anyhow::Result<()> { - let mut upsert = db.prepare( - "INSERT INTO ready_q_suspensions + let expires = record.timestamp + chrono::Duration::from_std(rule.duration)?; + let reason = format!("automation rule: {}", regex_list_to_string(&rule.regex)); + + { + let db = db.clone(); + let reason = reason.to_string(); + let source = source.to_string(); + let site = record.site.to_string(); + let rule_hash = rule_hash.to_string(); + + spawn_blocking(move || { + let mut upsert = db.prepare( + "INSERT INTO ready_q_suspensions (rule_hash, site_name, source, reason, expires) VALUES ($hash, $site, $source, $reason, $expires) ON CONFLICT (rule_hash, site_name) DO UPDATE SET expires=$expires", - )?; + )?; - let expires = record.timestamp + chrono::Duration::from_std(rule.duration)?; - let expires_str = expires.to_rfc3339(); + let expires_str = expires.to_rfc3339(); - upsert.bind(("$hash", rule_hash))?; - upsert.bind(("$site", record.site.as_str()))?; - upsert.bind(("$source", source))?; + upsert.bind(("$hash", rule_hash.as_str()))?; + upsert.bind(("$site", site.as_str()))?; + upsert.bind(("$source", source.as_str()))?; - let reason = format!("automation rule: {}", regex_list_to_string(&rule.regex)); - upsert.bind(("$reason", reason.as_str()))?; - upsert.bind(("$expires", expires_str.as_str()))?; + upsert.bind(("$reason", reason.as_str()))?; + upsert.bind(("$expires", expires_str.as_str()))?; - upsert.next()?; + upsert.next()?; + Ok::<_, anyhow::Error>(()) + }) + .await??; + } events.push(SubscriptionItem::ReadyQSuspension(ReadyQSuspension { rule_hash: rule_hash.to_string(), @@ -385,72 +433,93 @@ fn create_ready_q_suspension( Ok(()) } -fn insert_record( - db: &ConnectionThreadSafe, +async fn insert_record( + db: &Arc, rule_hash: &str, record: &JsonLogRecord, record_hash: &str, ) -> anyhow::Result<()> { let unix: i64 = record.timestamp.format("%s").to_string().parse()?; - let mut insert = - db.prepare("INSERT INTO event_history (rule_hash, record_hash, ts) values (?, ?, ?)")?; - insert.bind((1, rule_hash))?; - insert.bind((2, record_hash))?; - insert.bind((3, unix))?; - insert.next()?; - Ok(()) + let db = db.clone(); + let rule_hash = rule_hash.to_string(); + let record_hash = record_hash.to_string(); + spawn_blocking(move || { + let mut insert = + db.prepare("INSERT INTO event_history (rule_hash, record_hash, ts) values (?, ?, ?)")?; + insert.bind((1, rule_hash.as_str()))?; + insert.bind((2, record_hash.as_str()))?; + insert.bind((3, unix))?; + insert.next()?; + Ok(()) + }) + .await? } -fn prune_old_records( - db: &ConnectionThreadSafe, +async fn prune_old_records( + db: &Arc, rule: &Rule, rule_hash: &str, ) -> anyhow::Result<()> { match rule.trigger { Trigger::Immediate => Ok(()), Trigger::Threshold(spec) => { - let mut query = db.prepare( - "delete from event_history where rule_hash = ? and ts < unixepoch() - ?", - )?; - query.bind((1, rule_hash))?; - // Keep up to 2x the period - query.bind((2, 2 * spec.period as i64))?; - query.next()?; - Ok(()) + let db = db.clone(); + let rule_hash = rule_hash.to_string(); + spawn_blocking(move || { + let mut query = db.prepare( + "delete from event_history where rule_hash = ? and ts < unixepoch() - ?", + )?; + query.bind((1, rule_hash.as_str()))?; + // Keep up to 2x the period + query.bind((2, 2 * spec.period as i64))?; + query.next()?; + Ok(()) + }) + .await? } } } -fn count_matching_records( - db: &ConnectionThreadSafe, +async fn count_matching_records( + db: &Arc, rule: &Rule, rule_hash: &str, ) -> anyhow::Result { match rule.trigger { Trigger::Immediate => Ok(0), Trigger::Threshold(spec) => { - let mut query = db.prepare( - "SELECT COUNT(ts) from event_history where rule_hash = ? and ts >= unixepoch() - ?", - )?; - query.bind((1, rule_hash))?; - query.bind((2, spec.period as i64))?; - query.next()?; + let db = db.clone(); + let rule_hash = rule_hash.to_string(); + spawn_blocking(move || { + let mut query = db.prepare( + "SELECT COUNT(ts) from event_history \ + where rule_hash = ? and ts >= unixepoch() - ?", + )?; + query.bind((1, rule_hash.as_str()))?; + query.bind((2, spec.period as i64))?; + query.next()?; - let count: i64 = query.read(0)?; - Ok(count as u64) + let count: i64 = query.read(0)?; + Ok(count as u64) + }) + .await? } } } pub async fn publish_log_batch( - db: &ConnectionThreadSafe, + db: &Arc, records: &mut Vec, ) -> anyhow::Result<()> { let shaping = get_shaping(); let mut events = vec![]; - db.execute("BEGIN")?; + spawn_blocking({ + let db = db.clone(); + move || db.execute("BEGIN") + }) + .await??; let now = Utc::now(); @@ -460,7 +529,11 @@ pub async fn publish_log_batch( } } - db.execute("COMMIT")?; + spawn_blocking({ + let db = db.clone(); + move || db.execute("COMMIT") + }) + .await??; for event in events { SubscriberMgr::submit(event); @@ -471,7 +544,7 @@ pub async fn publish_log_batch( async fn publish_log_v1_impl( now: &DateTime, - db: &ConnectionThreadSafe, + db: &Arc, shaping: &Shaping, record: JsonLogRecord, events: &mut Vec, @@ -511,10 +584,10 @@ async fn publish_log_v1_impl( let m_hash = match_hash(m); let rule_hash = format!("{store_key}-{m_hash}"); - insert_record(db, &rule_hash, &record, &record_hash)?; - prune_old_records(db, m, &rule_hash)?; + insert_record(db, &rule_hash, &record, &record_hash).await?; + prune_old_records(db, m, &rule_hash).await?; - let count = count_matching_records(db, m, &rule_hash)?; + let count = count_matching_records(db, m, &rule_hash).await?; count >= spec.limit } @@ -536,7 +609,8 @@ async fn publish_log_v1_impl( tracing::debug!("{action:?} for {record:?}"); match action { Action::Suspend => { - create_ready_q_suspension(db, &rule_hash, m, &record, &source, events)?; + create_ready_q_suspension(db, &rule_hash, m, &record, &source, events) + .await?; } Action::SuspendTenant => { create_tenant_suspension( @@ -546,7 +620,8 @@ async fn publish_log_v1_impl( &record, UseCampaign::No, events, - )?; + ) + .await?; } Action::SuspendCampaign => { create_tenant_suspension( @@ -556,7 +631,8 @@ async fn publish_log_v1_impl( &record, UseCampaign::Yes, events, - )?; + ) + .await?; } Action::SetConfig(config) => { create_config( @@ -568,7 +644,8 @@ async fn publish_log_v1_impl( &domain, &source, PreferRollup::Yes, - )?; + ) + .await?; } Action::SetDomainConfig(config) => { create_config( @@ -580,7 +657,8 @@ async fn publish_log_v1_impl( &domain, &source, PreferRollup::No, - )?; + ) + .await?; } Action::Bounce => { create_bounce( @@ -591,7 +669,8 @@ async fn publish_log_v1_impl( UseTenant::No, UseCampaign::No, events, - )?; + ) + .await?; } Action::BounceTenant => { create_bounce( @@ -602,7 +681,8 @@ async fn publish_log_v1_impl( UseTenant::Yes, UseCampaign::No, events, - )?; + ) + .await?; } Action::BounceCampaign => { create_bounce( @@ -613,7 +693,8 @@ async fn publish_log_v1_impl( UseTenant::Yes, UseCampaign::Yes, events, - )?; + ) + .await?; } } } @@ -717,7 +798,7 @@ fn json_to_toml_value(item_value: &JsonValue) -> anyhow::Result { }) } -async fn do_get_config() -> anyhow::Result { +fn do_get_config() -> anyhow::Result { use toml_edit::Item; let mut doc = toml_edit::DocumentMut::new(); @@ -782,11 +863,11 @@ async fn do_get_config() -> anyhow::Result { } async fn get_config_v1(_: TrustedIpRequired) -> Result { - let result = do_get_config().await?; + let result = spawn_blocking(do_get_config).await??; Ok(result) } -async fn do_get_suspension() -> anyhow::Result> { +fn do_get_suspension() -> anyhow::Result> { let mut suspensions = Suspensions::default(); let mut stmt = HISTORY.prepare( @@ -883,7 +964,7 @@ async fn do_get_suspension() -> anyhow::Result> { } async fn get_suspension_v1(_: TrustedIpRequired) -> Result, AppError> { - let result = do_get_suspension().await?; + let result = spawn_blocking(do_get_suspension).await??; Ok(result) } @@ -912,7 +993,7 @@ async fn process_suspension_subscription_inner(mut socket: WebSocket) -> anyhow: // send the current set of suspensions first { - let suspensions = do_get_suspension().await?.0; + let suspensions = spawn_blocking(do_get_suspension).await??.0; for record in suspensions.ready_q { let json = serde_json::to_string(&SuspensionEntry::ReadyQ(record))?; socket.send(Message::Text(json)).await?; @@ -953,7 +1034,7 @@ pub async fn subscribe_suspension_v1( ws.on_upgrade(|socket| process_suspension_subscription(socket)) } -async fn do_get_bounces() -> anyhow::Result>> { +fn do_get_bounces() -> anyhow::Result>> { let mut stmt = HISTORY.prepare( "SELECT * from sched_q_bounces where unixepoch(expires) - unixepoch() > 0 @@ -1001,7 +1082,7 @@ async fn do_get_bounces() -> anyhow::Result>> { } async fn get_bounce_v1(_: TrustedIpRequired) -> Result>, AppError> { - let result = do_get_bounces().await?; + let result = spawn_blocking(do_get_bounces).await??; Ok(result) } @@ -1010,7 +1091,7 @@ async fn process_event_subscription_inner(mut socket: WebSocket) -> anyhow::Resu // send the current set of suspensions first { - let suspensions = do_get_suspension().await?.0; + let suspensions = spawn_blocking(do_get_suspension).await??.0; for record in suspensions.ready_q { let json = serde_json::to_string(&SubscriptionItem::ReadyQSuspension(record))?; socket.send(Message::Text(json)).await?; @@ -1022,7 +1103,7 @@ async fn process_event_subscription_inner(mut socket: WebSocket) -> anyhow::Resu } // and then bounces { - let bounces = do_get_bounces().await?.0; + let bounces = spawn_blocking(do_get_bounces).await??.0; for record in bounces { let json = serde_json::to_string(&SubscriptionItem::SchedQBounce(record))?; socket.send(Message::Text(json)).await?; diff --git a/crates/tsa-daemon/src/mod_auto.rs b/crates/tsa-daemon/src/mod_auto.rs index 03e7f4aa..7e2bcfc8 100644 --- a/crates/tsa-daemon/src/mod_auto.rs +++ b/crates/tsa-daemon/src/mod_auto.rs @@ -21,7 +21,7 @@ pub fn register(lua: &Lua) -> anyhow::Result<()> { tsa_mod.set( "configure_tsa_db_path", lua.create_function(|_lua, file_name: String| { - *crate::http_server::DB_PATH.lock().unwrap() = file_name; + *crate::http_server::DB_PATH.lock() = file_name; Ok(()) })?, )?; diff --git a/crates/tsa-daemon/src/publish.rs b/crates/tsa-daemon/src/publish.rs index 2300f3a9..a4afa1c4 100644 --- a/crates/tsa-daemon/src/publish.rs +++ b/crates/tsa-daemon/src/publish.rs @@ -1,7 +1,7 @@ use crate::http_server::{open_history_db, publish_log_batch}; use kumo_log_types::JsonLogRecord; use parking_lot::Mutex; -use std::sync::LazyLock; +use std::sync::{Arc, LazyLock}; use tokio::sync::Notify; use tokio::task::LocalSet; use tokio::time::{Duration, Instant}; @@ -118,7 +118,7 @@ fn grab_segment() -> Option> { } async fn run_processor() { - let db = open_history_db().unwrap(); + let db = Arc::new(open_history_db().unwrap()); loop { NOTIFY_CONSUMER.notified().await; while let Some(mut batch) = grab_segment() { diff --git a/docs/changelog/main.md b/docs/changelog/main.md index 6b4a81b3..b0ae5757 100644 --- a/docs/changelog/main.md +++ b/docs/changelog/main.md @@ -122,3 +122,4 @@ * Using `prefer_openssl` with a `routing_domain` could result in TLS certificate verification failure due to a trailing `.` on the canonical DNS name passed to the verification routine. +* tsa-daemon: sqlite db lock could block tokio scheduler threads