From 344ee0ee1d1f69d543913a42e36d3d3b996ffd5d Mon Sep 17 00:00:00 2001 From: Wez Furlong Date: Fri, 11 Apr 2025 13:36:34 -0700 Subject: [PATCH] tsa: migrate config overrides from sqlite to new state This moves another set of data from the sqlite db and into our own data structures. The sqlite data will be copied into the new structure on startup if we haven't already populated that data. --- crates/kumo-api-types/src/shaping.rs | 2 +- crates/tsa-daemon/src/http_server.rs | 262 ++++++++++++--------------- crates/tsa-daemon/src/state.rs | 216 +++++++++++++++++++++- 3 files changed, 324 insertions(+), 156 deletions(-) diff --git a/crates/kumo-api-types/src/shaping.rs b/crates/kumo-api-types/src/shaping.rs index 22d9f0ad..d86a0bf7 100644 --- a/crates/kumo-api-types/src/shaping.rs +++ b/crates/kumo-api-types/src/shaping.rs @@ -158,7 +158,7 @@ pub struct EgressPathConfigValue { /// The validation is performed by the TryFrom impl that is used to convert to the /// checked form below. #[derive(Deserialize, Serialize, Debug, Clone)] -struct EgressPathConfigValueUnchecked { +pub struct EgressPathConfigValueUnchecked { pub name: String, pub value: toml::Value, } diff --git a/crates/tsa-daemon/src/http_server.rs b/crates/tsa-daemon/src/http_server.rs index 6320449d..28ba77fe 100644 --- a/crates/tsa-daemon/src/http_server.rs +++ b/crates/tsa-daemon/src/http_server.rs @@ -1,14 +1,16 @@ use crate::database::Database; use crate::publish::submit_record; use crate::shaping_config::get_shaping; -use crate::state::{MatchingScope, TSA_STATE}; +use crate::state::{ActionHash, ConfigurationOverride, MatchingScope, TsaState, TSA_STATE}; use anyhow::{anyhow, Context}; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; use axum::response::IntoResponse; use axum::routing::{get, post}; use axum::{Json, Router}; use chrono::{DateTime, Utc}; -use kumo_api_types::shaping::{Action, EgressPathConfigValue, Regex, Rule, Shaping, Trigger}; +use kumo_api_types::shaping::{ + Action, EgressPathConfigValueUnchecked, Regex, Rule, Shaping, Trigger, +}; use kumo_api_types::tsa::{ ReadyQSuspension, SchedQBounce, SchedQSuspension, SubscriptionItem, SuspensionEntry, Suspensions, @@ -27,7 +29,6 @@ use std::hash::Hash; use std::sync::{Arc, LazyLock}; use std::time::Instant; use tokio::sync::broadcast::{channel, Sender}; -use toml_edit::{value, Value as TomlValue}; use utoipa::OpenApi; pub static DB_PATH: LazyLock> = @@ -58,64 +59,12 @@ pub fn make_router() -> RouterAndDocs { } #[derive(PartialEq, Clone, Copy)] -enum PreferRollup { +pub enum PreferRollup { Yes, No, } -async fn create_config( - db: &Arc, - rule_hash: &str, - rule: &Rule, - record: &JsonLogRecord, - config: &EgressPathConfigValue, - domain: &str, - source: &str, - prefer_rollup: PreferRollup, -) -> anyhow::Result<()> { - 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 - }; - - db.perform("create_config", move |db| { - 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", - )?; - - 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(("$reason", reason.as_str()))?; - upsert.bind(("$expires", expires.as_str()))?; - - upsert.next()?; - - Ok(()) - }) - .await -} - -fn regex_list_to_string(list: &[Regex]) -> String { +pub fn regex_list_to_string(list: &[Regex]) -> String { if list.len() == 1 { list[0].to_string() } else { @@ -439,13 +388,15 @@ async fn publish_log_v1_impl( continue; } + let matching_scope = MatchingScope::from_rule_and_record(m, &record); + let triggered = match m.trigger { Trigger::Immediate => true, Trigger::Threshold(spec) => { let count = TSA_STATE .get() .expect("state not initialized") - .record_event(&MatchingScope::from_rule_and_record(m, &record), m, &record); + .record_event(&matching_scope, m, &record); count >= spec.limit } @@ -463,6 +414,7 @@ async fn publish_log_v1_impl( // action that we are iterating let a_hash = action_hash(m, action); let rule_hash = format!("{store_key}-{a_hash}"); + let action_hash = ActionHash::from_rule_and_record(m, action, &record); tracing::debug!("{action:?} for {record:?}"); match action { @@ -493,30 +445,32 @@ async fn publish_log_v1_impl( .await?; } Action::SetConfig(config) => { - create_config( - db, - &rule_hash, - m, - &record, - config, - &domain, - source, - PreferRollup::Yes, - ) - .await?; + TSA_STATE + .get() + .expect("tsa_state missing") + .create_config_override( + &action_hash, + m, + &record, + config, + &domain, + source, + PreferRollup::Yes, + ); } Action::SetDomainConfig(config) => { - create_config( - db, - &rule_hash, - m, - &record, - config, - &domain, - source, - PreferRollup::No, - ) - .await?; + TSA_STATE + .get() + .expect("tsa_state missing") + .create_config_override( + &action_hash, + m, + &record, + config, + &domain, + source, + PreferRollup::No, + ); } Action::Bounce => { create_bounce( @@ -616,30 +570,29 @@ async fn publish_log_v1( }) } -fn json_to_toml_value(item_value: &JsonValue) -> anyhow::Result { - use toml_edit::Formatted; +fn json_to_toml_value(item_value: &JsonValue) -> anyhow::Result { Ok(match item_value { - JsonValue::Bool(b) => TomlValue::Boolean(Formatted::new(*b)), - JsonValue::String(s) => TomlValue::String(Formatted::new(s.to_string())), + JsonValue::Bool(b) => toml::Value::Boolean(*b), + JsonValue::String(s) => toml::Value::String(s.to_string()), JsonValue::Array(a) => { - let mut res = toml_edit::Array::new(); + let mut res = toml::value::Array::new(); for item in a { res.push(json_to_toml_value(item)?); } - TomlValue::Array(res) + toml::Value::Array(res) } JsonValue::Object(o) => { - let mut tbl = toml_edit::InlineTable::new(); + let mut tbl = toml::Table::new(); for (k, v) in o.iter() { - tbl.insert(k, json_to_toml_value(v)?); + tbl.insert(k.to_string(), json_to_toml_value(v)?); } - TomlValue::InlineTable(tbl) + toml::Value::Table(tbl) } JsonValue::Number(n) => { if let Some(i) = n.as_i64() { - TomlValue::Integer(Formatted::new(i)) + toml::Value::Integer(i) } else if let Some(f) = n.as_f64() { - TomlValue::Float(Formatted::new(f)) + toml::Value::Float(f) } else { anyhow::bail!("impossible number value {n:?}"); } @@ -648,72 +601,79 @@ fn json_to_toml_value(item_value: &JsonValue) -> anyhow::Result { }) } -fn do_get_config(db: &ConnectionThreadSafe) -> anyhow::Result { - use toml_edit::Item; - let mut doc = toml_edit::DocumentMut::new(); - - let mut stmt = db.prepare( - "SELECT * from config where - unixepoch(expires) - unixepoch() > 0 - order by expires, domain, source, name", - )?; - let mut num_entries = 0; - while let Ok(sqlite::State::Row) = stmt.next() { - num_entries += 1; - let reason: String = stmt.read("reason")?; - let domain: String = stmt.read("domain")?; - let mx_rollup: i64 = stmt.read("mx_rollup")?; - let source: String = stmt.read("source")?; - let name: String = stmt.read("name")?; - let config_value: String = stmt.read("value")?; - let expires: String = stmt.read("expires")?; - - let config_value = serde_json::from_str(&config_value)?; - let config_value = json_to_toml_value(&config_value)?; - - let domain_entry = doc - .entry(&domain) - .or_insert_with(|| { - let mut tbl = toml_edit::Table::new(); - tbl["mx_rollup"] = value(mx_rollup != 0); - Item::Table(tbl) - }) - .as_table_mut() - .unwrap(); - let sources = domain_entry - .entry("sources") - .or_insert_with(|| { - let tbl = toml_edit::Table::new(); - Item::Table(tbl) - }) - .as_table_mut() - .unwrap(); - let source_entry = sources - .entry(&source) - .or_insert_with(|| { - let tbl = toml_edit::Table::new(); - Item::Table(tbl) - }) - .as_table_mut() - .unwrap(); - - let item = Item::Value(config_value); - source_entry.insert(&name, item); - - if let Some(mut key) = source_entry.key_mut(&name) { - key.leaf_decor_mut() - .set_prefix(format!("# reason: {reason}\n# expires: {expires}\n")); +pub fn toml_to_toml_edit_value(v: toml::Value) -> toml_edit::Value { + use toml_edit::Formatted; + match v { + toml::Value::String(s) => toml_edit::Value::String(Formatted::new(s)), + toml::Value::Integer(s) => toml_edit::Value::Integer(Formatted::new(s)), + toml::Value::Float(s) => toml_edit::Value::Float(Formatted::new(s)), + toml::Value::Boolean(s) => toml_edit::Value::Boolean(Formatted::new(s)), + toml::Value::Datetime(s) => toml_edit::Value::Datetime(Formatted::new(s)), + toml::Value::Array(s) => { + let mut array = toml_edit::Array::new(); + for item in s.into_iter().map(toml_to_toml_edit_value) { + array.push(item); + } + toml_edit::Value::Array(array) + } + toml::Value::Table(t) => { + let mut tbl = toml_edit::InlineTable::new(); + for (k, v) in t { + tbl.insert(k, toml_to_toml_edit_value(v)); + } + toml_edit::Value::InlineTable(tbl) } } +} - Ok(format!( - "# Generated by tsa-daemon\n# Number of entries: {num_entries}\n\n{}", - doc - )) +pub async fn import_configs_from_sqlite(state: Arc) -> anyhow::Result<()> { + HISTORY + .perform("import config", move |db| { + let mut stmt = db.prepare( + "SELECT * from config where + unixepoch(expires) - unixepoch() > 0 + order by expires, domain, source, name", + )?; + while let Ok(sqlite::State::Row) = stmt.next() { + let rule_hash: String = stmt.read("rule_hash")?; + let site_name: String = stmt.read("site_name")?; + let reason: String = stmt.read("reason")?; + let domain: String = stmt.read("domain")?; + let mx_rollup: i64 = stmt.read("mx_rollup")?; + let source: String = stmt.read("source")?; + let name: String = stmt.read("name")?; + let config_value: String = stmt.read("value")?; + let expires: String = stmt.read("expires")?; + + let config_value = serde_json::from_str(&config_value)?; + let config_value = json_to_toml_value(&config_value)?; + + let matching_scope = ActionHash::from_legacy_hash_and_site(&rule_hash, &site_name); + state.insert_config_override( + matching_scope, + ConfigurationOverride { + domain, + reason, + mx_rollup: mx_rollup != 0, + source, + option: EgressPathConfigValueUnchecked { + name, + value: config_value.into(), + }, + expires: expires.parse()?, + }, + ); + } + Ok(()) + }) + .await } async fn get_config_v1(_: TrustedIpRequired) -> Result { - let result = HISTORY.perform("get_config_v1", do_get_config).await?; + let result = TSA_STATE + .get() + .expect("tsa_state missing") + .export_config_override_toml(); Ok(result) } diff --git a/crates/tsa-daemon/src/state.rs b/crates/tsa-daemon/src/state.rs index 665a02ec..8b8c301f 100644 --- a/crates/tsa-daemon/src/state.rs +++ b/crates/tsa-daemon/src/state.rs @@ -1,13 +1,18 @@ -use crate::http_server::{Sha256Hasher, DB_PATH}; +use crate::http_server::{ + import_configs_from_sqlite, regex_list_to_string, toml_to_toml_edit_value, PreferRollup, + Sha256Hasher, DB_PATH, +}; use anyhow::Context; use chrono::{DateTime, Utc}; use dashmap::DashMap; -use kumo_api_types::shaping::Rule; +use kumo_api_types::shaping::{ + Action, EgressPathConfigValue, EgressPathConfigValueUnchecked, Rule, +}; use kumo_log_types::JsonLogRecord; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::hash::Hash; -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use std::time::Instant; pub static TSA_STATE: OnceLock = OnceLock::new(); @@ -33,6 +38,24 @@ impl SiteKey { } } +#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct ActionHash(#[serde(with = "serde_bytes")] [u8; 32], SiteKey); + +impl ActionHash { + pub fn from_rule_and_record(rule: &Rule, action: &Action, record: &JsonLogRecord) -> Self { + let mut hasher = Sha256Hasher::new(); + rule.hash(&mut hasher); + action.hash(&mut hasher); + Self(hasher.get_binary(), SiteKey::from_record(record)) + } + + pub fn from_legacy_hash_and_site(hash: &str, site: &str) -> Self { + let mut bytes = [0u8; 32]; + hex::decode_to_slice(hash, &mut bytes).ok(); + Self(bytes, SiteKey(site.to_string())) + } +} + #[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] pub struct MatchingScope(RuleHash, SiteKey); @@ -76,14 +99,30 @@ impl EventData { } } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ConfigurationOverride { + pub domain: String, + pub mx_rollup: bool, + pub source: String, + pub reason: String, + /// Explicitly store unchecked to accommodate version skew + /// where we might not know about a value yet + pub option: EgressPathConfigValueUnchecked, + pub expires: DateTime, +} + #[derive(Default)] pub struct TsaState { event_history: DashMap, + config_overrides: DashMap, } #[derive(Serialize, Deserialize)] struct SerializableState { + #[serde(default)] event_history: HashMap, + #[serde(default)] + config_overrides: HashMap, } impl TsaState { @@ -101,6 +140,115 @@ impl TsaState { series.insert_and_count(record) as u64 } + pub fn create_config_override( + &self, + scope: &ActionHash, + rule: &Rule, + record: &JsonLogRecord, + config: &EgressPathConfigValue, + domain: &str, + source: &str, + prefer_rollup: PreferRollup, + ) { + let reason = format!("automation rule: {}", regex_list_to_string(&rule.regex)); + self.insert_config_override( + scope.clone(), + ConfigurationOverride { + domain: domain.to_string(), + reason, + mx_rollup: match prefer_rollup { + PreferRollup::Yes => rule.was_rollup, + PreferRollup::No => false, + }, + source: source.to_string(), + option: config.clone().into(), + expires: record.timestamp + rule.duration, + }, + ); + } + + pub fn insert_config_override(&self, scope: ActionHash, over: ConfigurationOverride) { + if Utc::now() >= over.expires { + // Skip already expired entry + tracing::warn!("Skip insertion of {scope:?} {over:?} because it has already expired"); + return; + } + tracing::warn!("inserting {scope:?} {over:?}"); + self.config_overrides.insert(scope, over); + } + + pub fn export_config_override_toml(&self) -> String { + use toml_edit::{value, Item}; + let mut doc = toml_edit::DocumentMut::new(); + let now = Utc::now(); + + let mut entries = vec![]; + for entry in self.config_overrides.iter() { + let over = entry.value(); + if now >= over.expires { + continue; + } + entries.push(over.clone()); + } + + entries.sort_by_key(|over| { + ( + over.expires, + over.domain.clone(), + over.source.clone(), + over.option.name.clone(), + ) + }); + let num_entries = entries.len(); + + for over in entries { + let domain_entry = doc + .entry(&over.domain) + .or_insert_with(|| { + let mut tbl = toml_edit::Table::new(); + tbl["mx_rollup"] = value(over.mx_rollup); + Item::Table(tbl) + }) + .as_table_mut() + .unwrap(); + let sources = domain_entry + .entry("sources") + .or_insert_with(|| { + let tbl = toml_edit::Table::new(); + Item::Table(tbl) + }) + .as_table_mut() + .unwrap(); + let source_entry = sources + .entry(&over.source) + .or_insert_with(|| { + let tbl = toml_edit::Table::new(); + Item::Table(tbl) + }) + .as_table_mut() + .unwrap(); + + let item = toml_to_toml_edit_value(over.option.value.clone()); + source_entry.insert(&over.option.name, Item::Value(item)); + + if let Some(mut key) = source_entry.key_mut(&over.option.name) { + key.leaf_decor_mut().set_prefix(format!( + "# reason: {}\n# expires: {}\n", + over.reason, + over.expires.to_rfc3339() + )); + } + } + + format!( + "# Generated by tsa-daemon\n# Number of entries: {num_entries}\n\n\ + {}\n\n\ + # Generated by tsa-daemon\n# Number of entries: {num_entries}\n", + doc + ) + } + + /// Return a serializable version of the state fn serializable(&self) -> SerializableState { SerializableState { event_history: self @@ -108,6 +256,11 @@ impl TsaState { .iter() .map(|entry| (entry.key().clone(), entry.value().clone())) .collect(), + config_overrides: self + .config_overrides + .iter() + .map(|entry| (entry.key().clone(), entry.value().clone())) + .collect(), } } @@ -115,6 +268,44 @@ impl TsaState { let now = Utc::now(); let now_ts = to_unix_ts(&now); self.prune_events(now_ts).await; + self.prune_config_overrides(&now).await; + } + + async fn prune_config_overrides(&self, now: &DateTime) { + let mut visited = 0; + let start = Instant::now(); + + let is_prunable = |over: &ConfigurationOverride| *now >= over.expires; + + let keys_to_prune: Vec = self + .config_overrides + .iter() + .filter_map(|entry| { + visited += 1; + let over = entry.value(); + if is_prunable(over) { + Some(entry.key().clone()) + } else { + None + } + }) + .collect(); + + let mut num_pruned = 0; + for key in keys_to_prune { + let pruned = self + .config_overrides + .remove_if(&key, |_key, over| is_prunable(over)) + .is_some(); + if pruned { + num_pruned += 1; + } + } + tracing::debug!( + "visited {visited} and pruned {num_pruned} \ + config_overrides entries in {:?}", + start.elapsed() + ); } async fn prune_events(&self, now_ts: UnixTimeStamp) { @@ -202,7 +393,24 @@ pub async fn load_state() -> anyhow::Result<()> { } }; - TSA_STATE.set(state).ok(); + let import_holder = Arc::new(state); + + if import_holder.config_overrides.is_empty() { + // Import configs from the sqlite database + if let Err(err) = import_configs_from_sqlite(import_holder.clone()).await { + tracing::warn!( + "Failed to import legacy config entries from sqlite: {err:#}. Proceeding without them"); + } else { + tracing::info!( + "Imported {} config entries from sqlite", + import_holder.config_overrides.len() + ); + } + } + + TSA_STATE + .set(Arc::into_inner(import_holder).expect("only we hold a ref")) + .ok(); Ok(()) }