From dd63267db7b7464adc809d8e4cd71cf045f7064a Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Mon, 17 Mar 2025 17:08:29 -0700 Subject: [PATCH] tsa-daemon: fix sqlite locking blocking tokio scheduler sqlite doesn't have a native async interface, and instead will use traditional OS-level mutexes to ensure thread safety. Using those when under contention in a tokio scheduler thread can lead to blocking of the tokio scheduler threads, which can prevent timely delivery of data via websockets, or timely processing of incoming log records. This commit fixes up the sqlite access points to use tokio's spawn_blocking function to move that style of mutex acquisition to a more suitable context. --- crates/tsa-daemon/src/http_server.rs | 375 ++++++++++++++++----------- crates/tsa-daemon/src/mod_auto.rs | 2 +- crates/tsa-daemon/src/publish.rs | 4 +- docs/changelog/main.md | 1 + 4 files changed, 232 insertions(+), 150 deletions(-) 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