From 974ff5ba5560feb6e9f5e83b19d372d8157efc2b Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 7 May 2026 17:34:22 -0700 Subject: [PATCH] update --- codex-rs/otel/src/metrics/client.rs | 33 --- codex-rs/state/src/lib.rs | 20 +- codex-rs/state/src/runtime.rs | 51 +--- codex-rs/state/src/runtime/agent_jobs.rs | 34 +-- codex-rs/state/src/runtime/backfill.rs | 94 +++--- codex-rs/state/src/runtime/db.rs | 155 ++++++++++ codex-rs/state/src/runtime/goals.rs | 18 +- codex-rs/state/src/runtime/logs.rs | 55 ++-- codex-rs/state/src/runtime/memories.rs | 110 +++---- codex-rs/state/src/runtime/metrics_tests.rs | 139 --------- codex-rs/state/src/runtime/remote_control.rs | 6 +- codex-rs/state/src/runtime/threads.rs | 290 +++++++++---------- codex-rs/state/src/telemetry.rs | 225 +++++++------- 13 files changed, 593 insertions(+), 637 deletions(-) create mode 100644 codex-rs/state/src/runtime/db.rs delete mode 100644 codex-rs/state/src/runtime/metrics_tests.rs diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index f304b154a7..32b050ae84 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -409,36 +409,3 @@ fn build_otlp_metric_exporter( } } } - -#[cfg(test)] -mod tests { - use super::*; - use opentelemetry_sdk::metrics::InMemoryMetricExporter; - use pretty_assertions::assert_eq; - use std::collections::BTreeMap; - - #[test] - fn metrics_resource_attributes_include_arch() { - let exporter = InMemoryMetricExporter::default(); - let metrics = MetricsClient::new( - MetricsConfig::in_memory("test", "codex-test", env!("CARGO_PKG_VERSION"), exporter) - .with_runtime_reader(), - ) - .expect("metrics client"); - - metrics - .counter("test.resource_attributes", /*inc*/ 1, &[]) - .expect("counter"); - let snapshot = metrics.snapshot().expect("snapshot"); - let attributes: BTreeMap = snapshot - .resource() - .iter() - .map(|(key, value)| (key.as_str().to_string(), value.as_str().to_string())) - .collect(); - - assert_eq!( - attributes.get(ARCH_ATTRIBUTE).map(String::as_str), - Some(sanitize_metric_tag_value(std::env::consts::ARCH).as_str()) - ); - } -} diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index bd9b2097dd..e891ad4f1e 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -74,18 +74,18 @@ pub const DB_ERROR_METRIC: &str = "codex.db.error"; pub const DB_METRIC_BACKFILL: &str = "codex.db.backfill"; /// Metrics on backfill duration. Tags: [status] pub const DB_METRIC_BACKFILL_DURATION_MS: &str = "codex.db.backfill.duration_ms"; -/// SQLite startup initialization attempts. Tags: [status, phase, db, error_class, sqlite_code] -pub const DB_INIT_METRIC: &str = "codex.db.init"; -/// SQLite startup initialization duration. Tags: [status, phase, db, error_class, sqlite_code] -pub const DB_INIT_DURATION_METRIC: &str = "codex.db.init.duration_ms"; -/// SQLite operation attempts. Tags: [status, db, operation, access, error_class, sqlite_code] -pub const DB_OPERATION_METRIC: &str = "codex.db.operation"; -/// SQLite operation duration. Tags: [status, db, operation, access, error_class, sqlite_code] -pub const DB_OPERATION_DURATION_METRIC: &str = "codex.db.operation.duration_ms"; +/// SQLite startup initialization attempts. Tags: [status, phase, db, error] +pub const DB_INIT_METRIC: &str = "codex.sqlite.init.count"; +/// SQLite startup initialization duration. Tags: [status, phase, db, error] +pub const DB_INIT_DURATION_METRIC: &str = "codex.sqlite.init.duration_ms"; +/// SQLite logical operation attempts. Tags: [status, db, operation, access, error] +pub const DB_OPERATION_METRIC: &str = "codex.sqlite.operation.count"; +/// SQLite logical operation duration. Tags: [status, db, operation, access, error] +pub const DB_OPERATION_DURATION_METRIC: &str = "codex.sqlite.operation.duration_ms"; /// Filesystem fallback after SQLite could not serve a request. Tags: [caller, reason] -pub const DB_FALLBACK_METRIC: &str = "codex.db.fallback"; +pub const DB_FALLBACK_METRIC: &str = "codex.sqlite.fallback.count"; /// SQLite log queue loss or flush failure. Tags: [event, reason] -pub const DB_LOG_QUEUE_METRIC: &str = "codex.db.log_queue"; +pub const DB_LOG_QUEUE_METRIC: &str = "codex.sqlite.log_queue.count"; pub fn record_db_fallback_metric( metrics: Option<&dyn DbMetricsRecorder>, diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 62ba16d723..cd452ed316 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -61,11 +61,10 @@ use tracing::warn; mod agent_jobs; mod backfill; +mod db; mod goals; mod logs; mod memories; -#[cfg(test)] -mod metrics_tests; mod remote_control; #[cfg(test)] mod test_support; @@ -77,6 +76,9 @@ pub use goals::ThreadGoalUpdate; pub use remote_control::RemoteControlEnrollmentRecord; pub use threads::ThreadFilterOptions; +use db::DbOperation; +use db::InstrumentedDb; + // "Partition" is the retained-log-content bucket we cap at 10 MiB: // - one bucket per non-null thread_id // - one bucket per threadless (thread_id IS NULL) non-null process_uuid @@ -90,10 +92,9 @@ const LOG_PARTITION_ROW_LIMIT: i64 = 1_000; pub struct StateRuntime { codex_home: PathBuf, default_provider: String, - pool: Arc, - logs_pool: Arc, + state_db: InstrumentedDb, + logs_db: InstrumentedDb, thread_updated_at_millis: Arc, - metrics: Option, } impl StateRuntime { @@ -186,12 +187,11 @@ impl StateRuntime { let thread_updated_at_millis = thread_updated_at_millis_result?; let thread_updated_at_millis = thread_updated_at_millis.unwrap_or(0); let runtime = Arc::new(Self { - pool, - logs_pool, + state_db: InstrumentedDb::new(pool, DbKind::State, metrics.clone()), + logs_db: InstrumentedDb::new(logs_pool, DbKind::Logs, metrics), codex_home, default_provider, thread_updated_at_millis: Arc::new(AtomicI64::new(thread_updated_at_millis)), - metrics, }); if let Err(err) = runtime.run_logs_startup_maintenance().await { warn!( @@ -208,42 +208,11 @@ impl StateRuntime { } pub(crate) fn metrics(&self) -> Option<&dyn DbMetricsRecorder> { - self.metrics.as_deref() + self.state_db.metrics() } pub(crate) fn metrics_handle(&self) -> Option { - self.metrics.clone() - } - - pub(crate) async fn record_db_operation( - &self, - db: DbKind, - operation: &'static str, - access: DbAccess, - future: F, - ) -> anyhow::Result - where - F: std::future::Future>, - { - crate::telemetry::record_operation(self.metrics(), db, operation, access, future).await - } - - pub(crate) fn record_db_operation_result( - &self, - db: DbKind, - operation: &'static str, - access: DbAccess, - started: Instant, - result: &anyhow::Result, - ) { - crate::telemetry::record_operation_result( - self.metrics(), - db, - operation, - access, - started.elapsed(), - result, - ); + self.state_db.metrics_handle() } } diff --git a/codex-rs/state/src/runtime/agent_jobs.rs b/codex-rs/state/src/runtime/agent_jobs.rs index 3f5526c58d..56bf818fdc 100644 --- a/codex-rs/state/src/runtime/agent_jobs.rs +++ b/codex-rs/state/src/runtime/agent_jobs.rs @@ -19,7 +19,7 @@ impl StateRuntime { .map(i64::try_from) .transpose() .map_err(|_| anyhow::anyhow!("invalid max_runtime_seconds value"))?; - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; sqlx::query( r#" INSERT INTO agent_jobs ( @@ -122,7 +122,7 @@ WHERE id = ? "#, ) .bind(job_id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; row.map(AgentJob::try_from).transpose() } @@ -166,7 +166,7 @@ WHERE job_id = } let rows: Vec = builder .build_query_as::() - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await?; rows.into_iter().map(AgentJobItem::try_from).collect() } @@ -199,7 +199,7 @@ WHERE job_id = ? AND item_id = ? ) .bind(job_id) .bind(item_id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; row.map(AgentJobItem::try_from).transpose() } @@ -222,7 +222,7 @@ WHERE id = ? .bind(now) .bind(now) .bind(job_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -240,7 +240,7 @@ WHERE id = ? .bind(now) .bind(now) .bind(job_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -263,7 +263,7 @@ WHERE id = ? .bind(now) .bind(error_message) .bind(job_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -288,7 +288,7 @@ WHERE id = ? AND status IN (?, ?) .bind(job_id) .bind(AgentJobStatus::Pending.as_str()) .bind(AgentJobStatus::Running.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -302,7 +302,7 @@ WHERE id = ? "#, ) .bind(job_id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; let Some(row) = row else { return Ok(false); @@ -334,7 +334,7 @@ WHERE job_id = ? AND item_id = ? AND status = ? .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Pending.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -364,7 +364,7 @@ WHERE job_id = ? AND item_id = ? AND status = ? .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Pending.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -393,7 +393,7 @@ WHERE job_id = ? AND item_id = ? AND status = ? .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -417,7 +417,7 @@ WHERE job_id = ? AND item_id = ? AND status = ? .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -458,7 +458,7 @@ WHERE .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) .bind(reporting_thread_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -490,7 +490,7 @@ WHERE .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -524,7 +524,7 @@ WHERE .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -547,7 +547,7 @@ WHERE job_id = ? .bind(AgentJobItemStatus::Completed.as_str()) .bind(AgentJobItemStatus::Failed.as_str()) .bind(job_id) - .fetch_one(self.pool.as_ref()) + .fetch_one(self.state_db.pool()) .await?; let total_items: i64 = row.try_get("total_items")?; diff --git a/codex-rs/state/src/runtime/backfill.rs b/codex-rs/state/src/runtime/backfill.rs index 8e91cbb88c..282ce508fd 100644 --- a/codex-rs/state/src/runtime/backfill.rs +++ b/codex-rs/state/src/runtime/backfill.rs @@ -2,19 +2,21 @@ use super::*; impl StateRuntime { pub async fn get_backfill_state(&self) -> anyhow::Result { - self.record_db_operation(DbKind::State, "get_backfill_state", DbAccess::Read, async { - let row = sqlx::query( - r#" + let pool = self.state_db.pool().clone(); + self.state_db + .read(DbOperation::GetBackfillState, async move { + let row = sqlx::query( + r#" SELECT status, last_watermark, last_success_at FROM backfill_state WHERE id = 1 "#, - ) - .fetch_one(self.pool.as_ref()) - .await?; - crate::BackfillState::try_from_row(&row) - }) - .await + ) + .fetch_one(&pool) + .await?; + crate::BackfillState::try_from_row(&row) + }) + .await } /// Attempt to claim ownership of rollout metadata backfill. @@ -23,12 +25,10 @@ WHERE id = 1 /// Returns `false` if backfill is already complete or currently owned by a /// non-expired worker. pub async fn try_claim_backfill(&self, lease_seconds: i64) -> anyhow::Result { - self.record_db_operation( - DbKind::State, - "try_claim_backfill", - DbAccess::Write, - async { - self.ensure_backfill_state_row().await?; + let pool = self.state_db.pool().clone(); + self.state_db + .write(DbOperation::TryClaimBackfill, async move { + ensure_backfill_state_row_in_pool(&pool).await?; let now = Utc::now().timestamp(); let lease_cutoff = now.saturating_sub(lease_seconds.max(0)); let result = sqlx::query( @@ -45,22 +45,19 @@ WHERE id = 1 .bind(crate::BackfillStatus::Complete.as_str()) .bind(crate::BackfillStatus::Running.as_str()) .bind(lease_cutoff) - .execute(self.pool.as_ref()) + .execute(&pool) .await?; Ok(result.rows_affected() == 1) - }, - ) - .await + }) + .await } /// Mark rollout metadata backfill as running. pub async fn mark_backfill_running(&self) -> anyhow::Result<()> { - self.record_db_operation( - DbKind::State, - "mark_backfill_running", - DbAccess::Write, - async { - self.ensure_backfill_state_row().await?; + let pool = self.state_db.pool().clone(); + self.state_db + .write(DbOperation::MarkBackfillRunning, async move { + ensure_backfill_state_row_in_pool(&pool).await?; sqlx::query( r#" UPDATE backfill_state @@ -70,22 +67,19 @@ WHERE id = 1 ) .bind(crate::BackfillStatus::Running.as_str()) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute(&pool) .await?; Ok(()) - }, - ) - .await + }) + .await } /// Persist rollout metadata backfill progress. pub async fn checkpoint_backfill(&self, watermark: &str) -> anyhow::Result<()> { - self.record_db_operation( - DbKind::State, - "checkpoint_backfill", - DbAccess::Write, - async { - self.ensure_backfill_state_row().await?; + let pool = self.state_db.pool().clone(); + self.state_db + .write(DbOperation::CheckpointBackfill, async move { + ensure_backfill_state_row_in_pool(&pool).await?; sqlx::query( r#" UPDATE backfill_state @@ -96,22 +90,19 @@ WHERE id = 1 .bind(crate::BackfillStatus::Running.as_str()) .bind(watermark) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute(&pool) .await?; Ok(()) - }, - ) - .await + }) + .await } /// Mark rollout metadata backfill as complete. pub async fn mark_backfill_complete(&self, last_watermark: Option<&str>) -> anyhow::Result<()> { - self.record_db_operation( - DbKind::State, - "mark_backfill_complete", - DbAccess::Write, - async { - self.ensure_backfill_state_row().await?; + let pool = self.state_db.pool().clone(); + self.state_db + .write(DbOperation::MarkBackfillComplete, async move { + ensure_backfill_state_row_in_pool(&pool).await?; let now = Utc::now().timestamp(); sqlx::query( r#" @@ -128,16 +119,11 @@ WHERE id = 1 .bind(last_watermark) .bind(now) .bind(now) - .execute(self.pool.as_ref()) + .execute(&pool) .await?; Ok(()) - }, - ) - .await - } - - async fn ensure_backfill_state_row(&self) -> anyhow::Result<()> { - ensure_backfill_state_row_in_pool(self.pool.as_ref()).await + }) + .await } } @@ -308,7 +294,7 @@ WHERE id = 1 ) .bind(crate::BackfillStatus::Running.as_str()) .bind(stale_updated_at) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("force stale backfill lease"); diff --git a/codex-rs/state/src/runtime/db.rs b/codex-rs/state/src/runtime/db.rs new file mode 100644 index 0000000000..84ca8bebfc --- /dev/null +++ b/codex-rs/state/src/runtime/db.rs @@ -0,0 +1,155 @@ +use std::future::Future; +use std::sync::Arc; +use std::time::Instant; + +use sqlx::SqlitePool; + +use crate::telemetry::DbAccess; +use crate::telemetry::DbKind; +use crate::telemetry::DbMetricsRecorder; +use crate::telemetry::DbMetricsRecorderHandle; + +/// SQLite pool plus Codex-level operation telemetry context. +#[derive(Clone)] +pub(super) struct InstrumentedDb { + pool: Arc, + kind: DbKind, + metrics: Option, +} + +impl InstrumentedDb { + pub(super) fn new( + pool: Arc, + kind: DbKind, + metrics: Option, + ) -> Self { + Self { + pool, + kind, + metrics, + } + } + + pub(super) fn pool(&self) -> &SqlitePool { + self.pool.as_ref() + } + + pub(super) fn metrics(&self) -> Option<&dyn DbMetricsRecorder> { + self.metrics.as_deref() + } + + pub(super) fn metrics_handle(&self) -> Option { + self.metrics.clone() + } + + pub(super) async fn read(&self, operation: DbOperation, future: F) -> anyhow::Result + where + F: Future>, + { + self.record_operation(operation, DbAccess::Read, future) + .await + } + + pub(super) async fn write(&self, operation: DbOperation, future: F) -> anyhow::Result + where + F: Future>, + { + self.record_operation(operation, DbAccess::Write, future) + .await + } + + pub(super) async fn transaction( + &self, + operation: DbOperation, + future: F, + ) -> anyhow::Result + where + F: Future>, + { + self.record_operation(operation, DbAccess::Transaction, future) + .await + } + + pub(super) async fn maintenance( + &self, + operation: DbOperation, + future: F, + ) -> anyhow::Result + where + F: Future>, + { + self.record_operation(operation, DbAccess::Maintenance, future) + .await + } + + pub(super) fn record_result( + &self, + operation: DbOperation, + access: DbAccess, + started: Instant, + result: &anyhow::Result, + ) { + crate::telemetry::record_operation_result( + self.metrics(), + self.kind, + operation.as_str(), + access, + started.elapsed(), + result, + ); + } + + async fn record_operation( + &self, + operation: DbOperation, + access: DbAccess, + future: F, + ) -> anyhow::Result + where + F: Future>, + { + let started = Instant::now(); + let result = future.await; + self.record_result(operation, access, started, &result); + result + } +} + +#[derive(Clone, Copy)] +pub(super) enum DbOperation { + CheckpointBackfill, + FindRolloutPathById, + GetBackfillState, + GetDynamicTools, + GetThread, + InsertLogs, + ListThreads, + LogsStartupMaintenance, + MarkBackfillComplete, + MarkBackfillRunning, + PersistDynamicTools, + TouchThreadUpdatedAt, + TryClaimBackfill, + UpsertThread, +} + +impl DbOperation { + fn as_str(self) -> &'static str { + match self { + Self::CheckpointBackfill => "checkpoint_backfill", + Self::FindRolloutPathById => "find_rollout_path_by_id", + Self::GetBackfillState => "get_backfill_state", + Self::GetDynamicTools => "get_dynamic_tools", + Self::GetThread => "get_thread", + Self::InsertLogs => "insert_logs", + Self::ListThreads => "list_threads", + Self::LogsStartupMaintenance => "logs_startup_maintenance", + Self::MarkBackfillComplete => "mark_backfill_complete", + Self::MarkBackfillRunning => "mark_backfill_running", + Self::PersistDynamicTools => "persist_dynamic_tools", + Self::TouchThreadUpdatedAt => "touch_thread_updated_at", + Self::TryClaimBackfill => "try_claim_backfill", + Self::UpsertThread => "upsert_thread", + } + } +} diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index 56613b7d92..60d8959498 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -42,7 +42,7 @@ WHERE thread_id = ? "#, ) .bind(thread_id.to_string()) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; row.map(|row| thread_goal_from_row(&row)).transpose() @@ -99,7 +99,7 @@ RETURNING .bind(token_budget) .bind(now_ms) .bind(now_ms) - .fetch_one(self.pool.as_ref()) + .fetch_one(self.state_db.pool()) .await?; thread_goal_from_row(&row) @@ -148,7 +148,7 @@ RETURNING .bind(token_budget) .bind(now_ms) .bind(now_ms) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; row.map(|row| thread_goal_from_row(&row)).transpose() @@ -196,7 +196,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? } (Some(status), None) => { @@ -224,7 +224,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? } (None, Some(token_budget)) => { @@ -250,7 +250,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? } (None, None) => { @@ -289,7 +289,7 @@ WHERE thread_id = ? .bind(crate::ThreadGoalStatus::Paused.as_str()) .bind(now_ms) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; if result.rows_affected() == 0 { @@ -307,7 +307,7 @@ WHERE thread_id = ? "#, ) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) @@ -392,7 +392,7 @@ RETURNING query = query.bind(expected_goal_id); } - let row = query.fetch_optional(self.pool.as_ref()).await?; + let row = query.fetch_optional(self.state_db.pool()).await?; let Some(row) = row else { return Ok(ThreadGoalAccountingOutcome::Unchanged( diff --git a/codex-rs/state/src/runtime/logs.rs b/codex-rs/state/src/runtime/logs.rs index cf79182d2f..822077e35d 100644 --- a/codex-rs/state/src/runtime/logs.rs +++ b/codex-rs/state/src/runtime/logs.rs @@ -15,7 +15,7 @@ impl StateRuntime { return Ok(()); } - let mut tx = self.logs_pool.begin().await?; + let mut tx = self.logs_db.pool().begin().await?; let mut builder = QueryBuilder::::new( "INSERT INTO logs (ts, ts_nanos, level, target, feedback_log_body, thread_id, process_uuid, module_path, file, line, estimated_bytes) ", ); @@ -48,9 +48,8 @@ impl StateRuntime { Ok(()) } .await; - self.record_db_operation_result( - DbKind::Logs, - "insert_logs", + self.logs_db.record_result( + DbOperation::InsertLogs, DbAccess::Transaction, started, &result, @@ -297,28 +296,28 @@ WHERE id IN ( Ok(()) } - pub(crate) async fn delete_logs_before(&self, cutoff_ts: i64) -> anyhow::Result { - let result = sqlx::query("DELETE FROM logs WHERE ts < ?") - .bind(cutoff_ts) - .execute(self.logs_pool.as_ref()) - .await?; - Ok(result.rows_affected()) - } - pub(crate) async fn run_logs_startup_maintenance(&self) -> anyhow::Result<()> { - let Some(cutoff) = - Utc::now().checked_sub_signed(chrono::Duration::days(LOG_RETENTION_DAYS)) - else { - return Ok(()); - }; - self.delete_logs_before(cutoff.timestamp()).await?; - // Startup cleanup should not wait behind or block foreground work. - // PASSIVE checkpoints copy whatever is immediately available and skip - // frames that would require waiting on active readers or writers. - sqlx::query("PRAGMA wal_checkpoint(PASSIVE)") - .execute(self.logs_pool.as_ref()) - .await?; - Ok(()) + let pool = self.logs_db.pool().clone(); + self.logs_db + .maintenance(DbOperation::LogsStartupMaintenance, async move { + let Some(cutoff) = + Utc::now().checked_sub_signed(chrono::Duration::days(LOG_RETENTION_DAYS)) + else { + return Ok(()); + }; + sqlx::query("DELETE FROM logs WHERE ts < ?") + .bind(cutoff.timestamp()) + .execute(&pool) + .await?; + // Startup cleanup should not wait behind or block foreground work. + // PASSIVE checkpoints copy whatever is immediately available and skip + // frames that would require waiting on active readers or writers. + sqlx::query("PRAGMA wal_checkpoint(PASSIVE)") + .execute(&pool) + .await?; + Ok(()) + }) + .await } /// Query logs with optional filters. @@ -338,7 +337,7 @@ WHERE id IN ( let rows = builder .build_query_as::() - .fetch_all(self.logs_pool.as_ref()) + .fetch_all(self.logs_db.pool()) .await?; Ok(rows) } @@ -410,7 +409,7 @@ ORDER BY ts DESC, ts_nanos DESC, id DESC } let rows = sql .bind(LOG_PARTITION_SIZE_LIMIT_BYTES) - .fetch_all(self.logs_pool.as_ref()) + .fetch_all(self.logs_db.pool()) .await?; let mut lines = Vec::new(); @@ -443,7 +442,7 @@ ORDER BY ts DESC, ts_nanos DESC, id DESC let mut builder = QueryBuilder::::new("SELECT MAX(id) AS max_id FROM logs WHERE 1 = 1"); push_log_filters(&mut builder, query); - let row = builder.build().fetch_one(self.logs_pool.as_ref()).await?; + let row = builder.build().fetch_one(self.logs_db.pool()).await?; let max_id: Option = row.try_get("max_id")?; Ok(max_id.unwrap_or(0)) } diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 186f2dd341..4206872e09 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -30,7 +30,7 @@ impl StateRuntime { /// stage-1 (`memory_stage1`) and phase-2 (`memory_consolidate_global`) /// memory pipelines. pub async fn clear_memory_data(&self) -> anyhow::Result<()> { - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; sqlx::query( r#" @@ -68,7 +68,7 @@ WHERE kind = ? OR kind = ? } let now = Utc::now().timestamp(); - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; let mut updated_rows = 0; for thread_id in thread_ids { @@ -209,7 +209,7 @@ LEFT JOIN jobs let items = builder .build() - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await? .into_iter() .map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) @@ -279,7 +279,7 @@ LIMIT ? "#, ) .bind(n as i64) - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await?; rows.into_iter() @@ -323,7 +323,7 @@ WHERE thread_id IN ( ) .bind(cutoff) .bind(limit as i64) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? .rows_affected(); @@ -400,7 +400,7 @@ ORDER BY selected.thread_id ASC .bind(cutoff) .bind(cutoff) .bind(n as i64) - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await?; let mut selected = Vec::with_capacity(current_rows.len()); @@ -421,7 +421,7 @@ ORDER BY selected.thread_id ASC ) -> anyhow::Result { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; let rows_affected = sqlx::query( r#" UPDATE threads @@ -489,7 +489,7 @@ WHERE thread_id = ? let thread_id = thread_id.to_string(); let worker_id = worker_id.to_string(); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.state_db.pool().begin_with("BEGIN IMMEDIATE").await?; let existing_output = sqlx::query( r#" @@ -673,7 +673,7 @@ WHERE kind = ? AND job_key = ? let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; let rows_affected = sqlx::query( r#" UPDATE jobs @@ -750,7 +750,7 @@ WHERE excluded.source_updated_at >= stage1_outputs.source_updated_at let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; let rows_affected = sqlx::query( r#" UPDATE jobs @@ -848,7 +848,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? .rows_affected(); @@ -863,7 +863,7 @@ WHERE kind = ? AND job_key = ? /// Phase 2 does not use this watermark as a dirty check; git workspace diffing /// decides whether consolidation work exists after the lock is claimed. pub async fn enqueue_global_consolidation(&self, input_watermark: i64) -> anyhow::Result<()> { - enqueue_global_consolidation_with_executor(self.pool.as_ref(), input_watermark).await + enqueue_global_consolidation_with_executor(self.state_db.pool(), input_watermark).await } /// Attempts to claim the global phase-2 consolidation lock. @@ -890,7 +890,7 @@ WHERE kind = ? AND job_key = ? let ownership_token = Uuid::new_v4().to_string(); let worker_id = worker_id.to_string(); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.state_db.pool().begin_with("BEGIN IMMEDIATE").await?; let existing_job = sqlx::query( r#" @@ -1035,7 +1035,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? .rows_affected(); @@ -1058,7 +1058,7 @@ WHERE kind = ? AND job_key = ? completed_watermark: i64, selected_outputs: &[Stage1Output], ) -> anyhow::Result { - let mut tx = self.pool.begin().await?; + let mut tx = self.state_db.pool().begin().await?; let rows_affected = mark_global_phase2_job_succeeded_row(&mut *tx, ownership_token, completed_watermark) .await?; @@ -1136,7 +1136,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? .rows_affected(); @@ -1178,7 +1178,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await? .rows_affected(); @@ -1300,7 +1300,7 @@ mod tests { .bind(Utc::now().timestamp() - PHASE2_SUCCESS_COOLDOWN_SECONDS - 1) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("age phase2 success beyond cooldown"); } @@ -1409,7 +1409,7 @@ mod tests { sqlx::query("UPDATE jobs SET lease_until = 0 WHERE kind = 'memory_stage1' AND job_key = ?") .bind(thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("force stale lease"); @@ -1789,7 +1789,7 @@ mod tests { .expect("upsert disabled thread"); sqlx::query("UPDATE threads SET memory_mode = 'disabled' WHERE id = ?") .bind(disabled_thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("disable thread memory mode"); @@ -1890,7 +1890,7 @@ mod tests { .expect("upsert disabled thread"); sqlx::query("UPDATE threads SET memory_mode = 'disabled' WHERE id = ?") .bind(disabled_thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("disable existing thread"); @@ -1900,7 +1900,7 @@ mod tests { .expect("clear memory data"); let stage1_outputs_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count stage1 outputs"); assert_eq!(stage1_outputs_count, 0); @@ -1909,7 +1909,7 @@ mod tests { sqlx::query_scalar("SELECT COUNT(*) FROM jobs WHERE kind = ? OR kind = ?") .bind(JOB_KIND_MEMORY_STAGE1) .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count memory jobs"); assert_eq!(memory_jobs_count, 0); @@ -1917,7 +1917,7 @@ mod tests { let enabled_memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(enabled_thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("read enabled thread memory mode"); assert_eq!(enabled_memory_mode, "enabled"); @@ -1925,7 +1925,7 @@ mod tests { let disabled_memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(disabled_thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("read disabled thread memory mode"); assert_eq!(disabled_memory_mode, "disabled"); @@ -2000,7 +2000,7 @@ INSERT INTO jobs ( .bind(lease_until) .bind(3) .bind(metadata.updated_at.timestamp()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("seed running stage1 job"); } @@ -2034,7 +2034,7 @@ WHERE kind = 'memory_stage1' "#, ) .bind(Utc::now().timestamp()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count running stage1 jobs") .try_get::("count") @@ -2191,7 +2191,7 @@ WHERE kind = 'memory_stage1' let count_before = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count before delete") .try_get::("count") @@ -2200,14 +2200,14 @@ WHERE kind = 'memory_stage1' sqlx::query("DELETE FROM threads WHERE id = ?") .bind(thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("delete thread"); let count_after = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count after delete") .try_get::("count") @@ -2258,7 +2258,7 @@ WHERE kind = 'memory_stage1' let output_row_count = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load stage1 output count") .try_get::("count") @@ -2279,7 +2279,7 @@ WHERE kind = 'memory_stage1' let global_job_row_count = sqlx::query("SELECT COUNT(*) AS count FROM jobs WHERE kind = ?") .bind("memory_consolidate_global") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load phase2 job row count") .try_get::("count") @@ -2383,7 +2383,7 @@ WHERE kind = 'memory_stage1' let output_row_count = sqlx::query("SELECT COUNT(*) AS count FROM stage1_outputs WHERE thread_id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load stage1 output count after delete") .try_get::("count") @@ -2494,7 +2494,7 @@ WHERE kind = 'memory_stage1' ) .bind("memory_stage1") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load stage1 job row after newer-source claim"); assert_eq!( @@ -2620,7 +2620,7 @@ WHERE kind = 'memory_stage1' sqlx::query("SELECT retry_remaining FROM jobs WHERE kind = ? AND job_key = ?") .bind("memory_consolidate_global") .bind("global") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load phase2 job row after retry exhaustion"); assert_eq!( @@ -2787,7 +2787,7 @@ VALUES (?, ?, ?, ?, ?) .bind("raw memory") .bind("summary") .bind(100_i64) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("insert non-empty stage1 output"); sqlx::query( @@ -2801,7 +2801,7 @@ VALUES (?, ?, ?, ?, ?) .bind("") .bind("") .bind(101_i64) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("insert empty stage1 output"); @@ -3292,7 +3292,7 @@ VALUES (?, ?, ?, ?, ?) "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load selected_for_phase2"); assert_eq!(selected_for_phase2, 1); @@ -3585,7 +3585,7 @@ VALUES (?, ?, ?, ?, ?) "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load selected snapshot after phase2"); assert_eq!(selected_for_phase2, 1); @@ -3698,7 +3698,7 @@ VALUES (?, ?, ?, ?, ?) "SELECT selected_for_phase2, selected_for_phase2_source_updated_at FROM stage1_outputs WHERE thread_id = ?", ) .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load selected_for_phase2"); assert_eq!(selected_for_phase2, 0); @@ -3802,13 +3802,13 @@ VALUES (?, ?, ?, ?, ?) let row_a = sqlx::query("SELECT usage_count, last_usage FROM stage1_outputs WHERE thread_id = ?") .bind(thread_a.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load stage1 usage row a"); let row_b = sqlx::query("SELECT usage_count, last_usage FROM stage1_outputs WHERE thread_id = ?") .bind(thread_b.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("load stage1 usage row b"); @@ -3908,7 +3908,7 @@ VALUES (?, ?, ?, ?, ?) .bind(usage_count) .bind(last_usage.timestamp()) .bind(thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("update usage metadata"); } @@ -4004,7 +4004,7 @@ VALUES (?, ?, ?, ?, ?) .bind(usage_count) .bind(last_usage.map(|value| value.timestamp())) .bind(thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("update usage metadata"); } @@ -4089,13 +4089,13 @@ VALUES (?, ?, ?, ?, ?) sqlx::query("UPDATE stage1_outputs SET generated_at = ? WHERE thread_id = ?") .bind(300_i64) .bind(older_thread.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("update older generated_at"); sqlx::query("UPDATE stage1_outputs SET generated_at = ? WHERE thread_id = ?") .bind(150_i64) .bind(newer_thread.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("update newer generated_at"); @@ -4201,14 +4201,14 @@ VALUES (?, ?, ?, ?, ?) .bind(3_i64) .bind(now - Duration::days(40).num_seconds()) .bind(stale_used.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("set stale used metadata"); sqlx::query( "UPDATE stage1_outputs SET selected_for_phase2 = 1, selected_for_phase2_source_updated_at = source_updated_at WHERE thread_id = ?", ) .bind(stale_selected.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("mark selected for phase2"); sqlx::query( @@ -4217,13 +4217,13 @@ VALUES (?, ?, ?, ?, ?) .bind(8_i64) .bind(now - Duration::days(2).num_seconds()) .bind(fresh_used.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("set fresh used metadata"); let before_jobs_count = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM jobs WHERE kind = 'memory_stage1'") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count stage1 jobs before prune"); @@ -4236,7 +4236,7 @@ VALUES (?, ?, ?, ?, ?) let remaining = sqlx::query_scalar::<_, String>( "SELECT thread_id FROM stage1_outputs ORDER BY thread_id", ) - .fetch_all(runtime.pool.as_ref()) + .fetch_all(runtime.state_db.pool()) .await .expect("load remaining stage1 outputs"); let mut expected_remaining = vec![fresh_used.to_string(), stale_selected.to_string()]; @@ -4245,7 +4245,7 @@ VALUES (?, ?, ?, ?, ?) let after_jobs_count = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM jobs WHERE kind = 'memory_stage1'") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count stage1 jobs after prune"); assert_eq!(after_jobs_count, before_jobs_count); @@ -4323,7 +4323,7 @@ VALUES (?, ?, ?, ?, ?) assert_eq!(pruned, 2); let remaining_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM stage1_outputs") - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("count remaining stage1 outputs"); assert_eq!(remaining_count, 1); @@ -4539,7 +4539,7 @@ VALUES (?, ?, ?, ?, ?) .bind(Utc::now().timestamp() - 1) .bind("memory_consolidate_global") .bind("global") - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("expire global consolidation lease"); @@ -4675,7 +4675,7 @@ VALUES (?, ?, ?, ?, ?) sqlx::query("UPDATE jobs SET ownership_token = NULL WHERE kind = ? AND job_key = ?") .bind("memory_consolidate_global") .bind("global") - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("clear ownership token"); diff --git a/codex-rs/state/src/runtime/metrics_tests.rs b/codex-rs/state/src/runtime/metrics_tests.rs deleted file mode 100644 index a9d65246b8..0000000000 --- a/codex-rs/state/src/runtime/metrics_tests.rs +++ /dev/null @@ -1,139 +0,0 @@ -use super::StateRuntime; -use super::test_support::unique_temp_dir; -use crate::DB_INIT_METRIC; -use crate::DB_OPERATION_METRIC; -use crate::DbMetricsRecorder; -use crate::DbMetricsRecorderHandle; -use pretty_assertions::assert_eq; -use std::collections::BTreeMap; -use std::sync::Arc; -use std::sync::Mutex; -use std::time::Duration; - -#[derive(Clone, Debug)] -struct MetricEvent { - name: String, - tags: BTreeMap, -} - -#[derive(Default)] -struct TestMetrics { - events: Mutex>, -} - -impl TestMetrics { - fn metric_points(&self, name: &str) -> Vec> { - self.events - .lock() - .expect("metrics lock") - .iter() - .filter(|event| event.name == name) - .map(|event| event.tags.clone()) - .collect() - } -} - -impl DbMetricsRecorder for TestMetrics { - fn counter(&self, name: &str, _inc: i64, tags: &[(&str, &str)]) { - self.events.lock().expect("metrics lock").push(MetricEvent { - name: name.to_string(), - tags: tags_to_map(tags), - }); - } - - fn record_duration(&self, name: &str, _duration: Duration, tags: &[(&str, &str)]) { - self.events.lock().expect("metrics lock").push(MetricEvent { - name: name.to_string(), - tags: tags_to_map(tags), - }); - } -} - -fn build_metrics() -> Arc { - Arc::new(TestMetrics::default()) -} - -fn tags_to_map(tags: &[(&str, &str)]) -> BTreeMap { - tags.iter() - .map(|(key, value)| ((*key).to_string(), (*value).to_string())) - .collect() -} - -fn find_point( - points: &[BTreeMap], - key: &str, - value: &str, -) -> BTreeMap { - points - .iter() - .find(|attrs| attrs.get(key).is_some_and(|actual| actual == value)) - .cloned() - .unwrap_or_else(|| panic!("missing point with {key}={value}: {points:?}")) -} - -#[tokio::test] -async fn init_records_success_metrics() { - let metrics = build_metrics(); - let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init_with_metrics( - codex_home.clone(), - "test-provider".to_string(), - Some(metrics.clone() as DbMetricsRecorderHandle), - ) - .await - .expect("initialize runtime"); - - let points = metrics.metric_points(DB_INIT_METRIC); - let open_state = find_point(&points, "phase", "open_state"); - assert_eq!( - open_state.get("status").map(String::as_str), - Some("success") - ); - assert_eq!(open_state.get("db").map(String::as_str), Some("state")); - assert_eq!( - open_state.get("error_class").map(String::as_str), - Some("none") - ); - assert_eq!( - open_state.get("sqlite_code").map(String::as_str), - Some("none") - ); - - runtime.pool.close().await; - runtime.logs_pool.close().await; - let _ = tokio::fs::remove_dir_all(codex_home).await; -} - -#[tokio::test] -async fn selected_operations_record_success_metrics() { - let metrics = build_metrics(); - let codex_home = unique_temp_dir(); - let runtime = StateRuntime::init_with_metrics( - codex_home.clone(), - "test-provider".to_string(), - Some(metrics.clone() as DbMetricsRecorderHandle), - ) - .await - .expect("initialize runtime"); - - runtime - .get_backfill_state() - .await - .expect("get backfill state"); - runtime - .checkpoint_backfill("sessions/test/rollout.jsonl") - .await - .expect("checkpoint backfill"); - - let points = metrics.metric_points(DB_OPERATION_METRIC); - let read = find_point(&points, "operation", "get_backfill_state"); - assert_eq!(read.get("status").map(String::as_str), Some("success")); - assert_eq!(read.get("access").map(String::as_str), Some("read")); - let write = find_point(&points, "operation", "checkpoint_backfill"); - assert_eq!(write.get("status").map(String::as_str), Some("success")); - assert_eq!(write.get("access").map(String::as_str), Some("write")); - - runtime.pool.close().await; - runtime.logs_pool.close().await; - let _ = tokio::fs::remove_dir_all(codex_home).await; -} diff --git a/codex-rs/state/src/runtime/remote_control.rs b/codex-rs/state/src/runtime/remote_control.rs index fa0b1823f8..ae0eb4ef60 100644 --- a/codex-rs/state/src/runtime/remote_control.rs +++ b/codex-rs/state/src/runtime/remote_control.rs @@ -44,7 +44,7 @@ WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? .bind(remote_control_app_server_client_name_key( app_server_client_name, )) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; row.map(|row| { @@ -92,7 +92,7 @@ ON CONFLICT(websocket_url, account_id, app_server_client_name) DO UPDATE SET .bind(&enrollment.environment_id) .bind(&enrollment.server_name) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -114,7 +114,7 @@ WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? .bind(remote_control_app_server_client_name_key( app_server_client_name, )) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected()) } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 0b60633de9..f7957c0cfc 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -5,9 +5,11 @@ use std::sync::atomic::Ordering; impl StateRuntime { pub async fn get_thread(&self, id: ThreadId) -> anyhow::Result> { - self.record_db_operation(DbKind::State, "get_thread", DbAccess::Read, async { - let row = sqlx::query( - r#" + let pool = self.state_db.pool().clone(); + self.state_db + .read(DbOperation::GetThread, async move { + let row = sqlx::query( + r#" SELECT threads.id, threads.rollout_path, @@ -35,20 +37,20 @@ SELECT FROM threads WHERE threads.id = ? "#, - ) - .bind(id.to_string()) - .fetch_optional(self.pool.as_ref()) - .await?; - row.map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) - .transpose() - }) - .await + ) + .bind(id.to_string()) + .fetch_optional(&pool) + .await?; + row.map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) + .transpose() + }) + .await } pub async fn get_thread_memory_mode(&self, id: ThreadId) -> anyhow::Result> { let row = sqlx::query("SELECT memory_mode FROM threads WHERE id = ?") .bind(id.to_string()) - .fetch_optional(self.pool.as_ref()) + .fetch_optional(self.state_db.pool()) .await?; Ok(row.and_then(|row| row.try_get("memory_mode").ok())) } @@ -58,36 +60,38 @@ WHERE threads.id = ? &self, thread_id: ThreadId, ) -> anyhow::Result>> { - self.record_db_operation(DbKind::State, "get_dynamic_tools", DbAccess::Read, async { - let rows = sqlx::query( - r#" + let pool = self.state_db.pool().clone(); + self.state_db + .read(DbOperation::GetDynamicTools, async move { + let rows = sqlx::query( + r#" SELECT namespace, name, description, input_schema, defer_loading FROM thread_dynamic_tools WHERE thread_id = ? ORDER BY position ASC "#, - ) - .bind(thread_id.to_string()) - .fetch_all(self.pool.as_ref()) - .await?; - if rows.is_empty() { - return Ok(None); - } - let mut tools = Vec::with_capacity(rows.len()); - for row in rows { - let input_schema: String = row.try_get("input_schema")?; - let input_schema = serde_json::from_str::(input_schema.as_str())?; - tools.push(DynamicToolSpec { - namespace: row.try_get("namespace")?, - name: row.try_get("name")?, - description: row.try_get("description")?, - input_schema, - defer_loading: row.try_get("defer_loading")?, - }); - } - Ok(Some(tools)) - }) - .await + ) + .bind(thread_id.to_string()) + .fetch_all(&pool) + .await?; + if rows.is_empty() { + return Ok(None); + } + let mut tools = Vec::with_capacity(rows.len()); + for row in rows { + let input_schema: String = row.try_get("input_schema")?; + let input_schema = serde_json::from_str::(input_schema.as_str())?; + tools.push(DynamicToolSpec { + namespace: row.try_get("namespace")?, + name: row.try_get("name")?, + description: row.try_get("description")?, + input_schema, + defer_loading: row.try_get("defer_loading")?, + }); + } + Ok(Some(tools)) + }) + .await } /// Persist or replace the directional parent-child edge for a spawned thread. @@ -112,7 +116,7 @@ ON CONFLICT(child_thread_id) DO UPDATE SET .bind(parent_thread_id.to_string()) .bind(child_thread_id.to_string()) .bind(status.as_ref()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -126,7 +130,7 @@ ON CONFLICT(child_thread_id) DO UPDATE SET sqlx::query("UPDATE thread_spawn_edges SET status = ? WHERE child_thread_id = ?") .bind(status.as_ref()) .bind(child_thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -192,7 +196,7 @@ LIMIT 2 ) .bind(parent_thread_id.to_string()) .bind(agent_path) - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await?; one_thread_id_from_rows(rows, agent_path) } @@ -224,7 +228,7 @@ LIMIT 2 ) .bind(root_thread_id.to_string()) .bind(agent_path) - .fetch_all(self.pool.as_ref()) + .fetch_all(self.state_db.pool()) .await?; one_thread_id_from_rows(rows, agent_path) } @@ -247,7 +251,7 @@ LIMIT 2 sql = sql.bind(status.to_string()); } - let rows = sql.fetch_all(self.pool.as_ref()).await?; + let rows = sql.fetch_all(self.state_db.pool()).await?; rows.into_iter() .map(|row| { ThreadId::try_from(row.try_get::("child_thread_id")?).map_err(Into::into) @@ -289,7 +293,7 @@ ORDER BY depth ASC, child_thread_id ASC sql = sql.bind(status.clone()).bind(status); } - let rows = sql.fetch_all(self.pool.as_ref()).await?; + let rows = sql.fetch_all(self.state_db.pool()).await?; rows.into_iter() .map(|row| { ThreadId::try_from(row.try_get::("child_thread_id")?).map_err(Into::into) @@ -315,7 +319,7 @@ ON CONFLICT(child_thread_id) DO NOTHING .bind(parent_thread_id.to_string()) .bind(child_thread_id.to_string()) .bind(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(()) } @@ -338,11 +342,9 @@ ON CONFLICT(child_thread_id) DO NOTHING id: ThreadId, archived_only: Option, ) -> anyhow::Result> { - self.record_db_operation( - DbKind::State, - "find_rollout_path_by_id", - DbAccess::Read, - async { + let pool = self.state_db.pool().clone(); + self.state_db + .read(DbOperation::FindRolloutPathById, async move { let mut builder = QueryBuilder::::new("SELECT rollout_path FROM threads WHERE id = "); builder.push_bind(id.to_string()); @@ -355,13 +357,12 @@ ON CONFLICT(child_thread_id) DO NOTHING } None => {} } - let row = builder.build().fetch_optional(self.pool.as_ref()).await?; + let row = builder.build().fetch_optional(&pool).await?; Ok(row .and_then(|r| r.try_get::("rollout_path").ok()) .map(PathBuf::from)) - }, - ) - .await + }) + .await } /// Find the newest thread whose user-facing title exactly matches `title`. @@ -403,7 +404,7 @@ ON CONFLICT(child_thread_id) DO NOTHING /*limit*/ 1, ); - let row = builder.build().fetch_optional(self.pool.as_ref()).await?; + let row = builder.build().fetch_optional(self.state_db.pool()).await?; row.map(|row| ThreadRow::try_from_row(&row).and_then(crate::ThreadMetadata::try_from)) .transpose() } @@ -414,38 +415,40 @@ ON CONFLICT(child_thread_id) DO NOTHING page_size: usize, filters: ThreadFilterOptions<'_>, ) -> anyhow::Result { - self.record_db_operation(DbKind::State, "list_threads", DbAccess::Read, async { - let limit = page_size.saturating_add(1); - let sort_key = filters.sort_key; - let sort_direction = filters.sort_direction; + let pool = self.state_db.pool().clone(); + self.state_db + .read(DbOperation::ListThreads, async move { + let limit = page_size.saturating_add(1); + let sort_key = filters.sort_key; + let sort_direction = filters.sort_direction; - let mut builder = QueryBuilder::::new(""); - push_thread_select_columns(&mut builder); - builder.push(" FROM threads"); - push_thread_filters(&mut builder, filters); - push_thread_order_and_limit(&mut builder, sort_key, sort_direction, limit); + let mut builder = QueryBuilder::::new(""); + push_thread_select_columns(&mut builder); + builder.push(" FROM threads"); + push_thread_filters(&mut builder, filters); + push_thread_order_and_limit(&mut builder, sort_key, sort_direction, limit); - let rows = builder.build().fetch_all(self.pool.as_ref()).await?; - let mut items = rows - .into_iter() - .map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) - .collect::, _>>()?; - let num_scanned_rows = items.len(); - let next_anchor = if items.len() > page_size { - items.pop(); - items - .last() - .and_then(|item| anchor_from_item(item, sort_key)) - } else { - None - }; - Ok(ThreadsPage { - items, - next_anchor, - num_scanned_rows, + let rows = builder.build().fetch_all(&pool).await?; + let mut items = rows + .into_iter() + .map(|row| ThreadRow::try_from_row(&row).and_then(ThreadMetadata::try_from)) + .collect::, _>>()?; + let num_scanned_rows = items.len(); + let next_anchor = if items.len() > page_size { + items.pop(); + items + .last() + .and_then(|item| anchor_from_item(item, sort_key)) + } else { + None + }; + Ok(ThreadsPage { + items, + next_anchor, + num_scanned_rows, + }) }) - }) - .await + .await } /// List thread ids using the underlying database (no rollout scanning). @@ -474,7 +477,7 @@ ON CONFLICT(child_thread_id) DO NOTHING ); push_thread_order_and_limit(&mut builder, sort_key, SortDirection::Desc, limit); - let rows = builder.build().fetch_all(self.pool.as_ref()).await?; + let rows = builder.build().fetch_all(self.state_db.pool()).await?; rows.into_iter() .map(|row| { let id: String = row.try_get("id")?; @@ -564,7 +567,7 @@ ON CONFLICT(id) DO NOTHING .bind(metadata.git_branch.as_deref()) .bind(metadata.git_origin_url.as_deref()) .bind("enabled") - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) .await?; @@ -579,7 +582,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET memory_mode = ? WHERE id = ?") .bind(memory_mode) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -592,7 +595,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET title = ? WHERE id = ?") .bind(title) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -602,24 +605,21 @@ ON CONFLICT(id) DO NOTHING thread_id: ThreadId, updated_at: DateTime, ) -> anyhow::Result { - self.record_db_operation( - DbKind::State, - "touch_thread_updated_at", - DbAccess::Write, - async { - let updated_at = self.allocate_thread_updated_at(updated_at)?; + let updated_at = self.allocate_thread_updated_at(updated_at)?; + let pool = self.state_db.pool().clone(); + self.state_db + .write(DbOperation::TouchThreadUpdatedAt, async move { let result = sqlx::query( "UPDATE threads SET updated_at = ?, updated_at_ms = ? WHERE id = ?", ) .bind(datetime_to_epoch_seconds(updated_at)) .bind(datetime_to_epoch_millis(updated_at)) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(&pool) .await?; Ok(result.rows_affected() > 0) - }, - ) - .await + }) + .await } /// Allocate a persisted `updated_at` value for thread-list cursor ordering. @@ -692,7 +692,7 @@ WHERE id = ? .bind(git_origin_url.is_some()) .bind(git_origin_url.flatten()) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected() > 0) } @@ -803,7 +803,7 @@ ON CONFLICT(id) DO UPDATE SET .bind(metadata.git_branch.as_deref()) .bind(metadata.git_origin_url.as_deref()) .bind(creation_memory_mode.unwrap_or("enabled")) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; self.insert_thread_spawn_edge_from_source_if_absent( metadata.id, @@ -813,13 +813,8 @@ ON CONFLICT(id) DO UPDATE SET Ok(()) } .await; - self.record_db_operation_result( - DbKind::State, - "upsert_thread", - DbAccess::Write, - started, - &result, - ); + self.state_db + .record_result(DbOperation::UpsertThread, DbAccess::Write, started, &result); result } @@ -832,21 +827,22 @@ ON CONFLICT(id) DO UPDATE SET thread_id: ThreadId, tools: Option<&[DynamicToolSpec]>, ) -> anyhow::Result<()> { - let started = Instant::now(); - let result: anyhow::Result<()> = async { - let Some(tools) = tools else { - return Ok(()); - }; - if tools.is_empty() { - return Ok(()); - } - let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; - for (idx, tool) in tools.iter().enumerate() { - let position = i64::try_from(idx).unwrap_or(i64::MAX); - let input_schema = serde_json::to_string(&tool.input_schema)?; - sqlx::query( - r#" + let pool = self.state_db.pool().clone(); + self.state_db + .transaction(DbOperation::PersistDynamicTools, async move { + let Some(tools) = tools else { + return Ok(()); + }; + if tools.is_empty() { + return Ok(()); + } + let thread_id = thread_id.to_string(); + let mut tx = pool.begin().await?; + for (idx, tool) in tools.iter().enumerate() { + let position = i64::try_from(idx).unwrap_or(i64::MAX); + let input_schema = serde_json::to_string(&tool.input_schema)?; + sqlx::query( + r#" INSERT INTO thread_dynamic_tools ( thread_id, position, @@ -858,29 +854,21 @@ INSERT INTO thread_dynamic_tools ( ) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(thread_id, position) DO NOTHING "#, - ) - .bind(thread_id.as_str()) - .bind(position) - .bind(tool.namespace.as_deref()) - .bind(tool.name.as_str()) - .bind(tool.description.as_str()) - .bind(input_schema) - .bind(tool.defer_loading) - .execute(&mut *tx) - .await?; - } - tx.commit().await?; - Ok(()) - } - .await; - self.record_db_operation_result( - DbKind::State, - "persist_dynamic_tools", - DbAccess::Transaction, - started, - &result, - ); - result + ) + .bind(thread_id.as_str()) + .bind(position) + .bind(tool.namespace.as_deref()) + .bind(tool.name.as_str()) + .bind(tool.description.as_str()) + .bind(input_schema) + .bind(tool.defer_loading) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(()) + }) + .await } /// Apply rollout items incrementally using the underlying database. @@ -990,7 +978,7 @@ ON CONFLICT(thread_id, position) DO NOTHING pub async fn delete_thread(&self, thread_id: ThreadId) -> anyhow::Result { let result = sqlx::query("DELETE FROM threads WHERE id = ?") .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute(self.state_db.pool()) .await?; Ok(result.rows_affected()) } @@ -1224,7 +1212,7 @@ mod tests { let memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("memory mode should be readable"); assert_eq!(memory_mode, "disabled"); @@ -1238,7 +1226,7 @@ mod tests { let memory_mode: String = sqlx::query_scalar("SELECT memory_mode FROM threads WHERE id = ?") .bind(thread_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("memory mode should remain readable"); assert_eq!(memory_mode, "disabled"); @@ -1592,7 +1580,7 @@ mod tests { .bind(123_i64) .bind("newer preview") .bind(thread_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("concurrent metadata write should succeed"); @@ -1792,7 +1780,7 @@ mod tests { "SELECT created_at, updated_at, created_at_ms, updated_at_ms FROM threads WHERE id = ?", ) .bind(second_id.to_string()) - .fetch_one(runtime.pool.as_ref()) + .fetch_one(runtime.state_db.pool()) .await .expect("thread timestamp row should load"); assert_eq!( @@ -1826,7 +1814,7 @@ mod tests { sqlx::query("UPDATE threads SET updated_at = ? WHERE id = ?") .bind(1_700_001_112_i64) .bind(first_id.to_string()) - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("legacy timestamp write should succeed"); let legacy = runtime @@ -2038,7 +2026,7 @@ INSERT INTO thread_spawn_edges ( .bind(parent_thread_id.to_string()) .bind(future_child_thread_id.to_string()) .bind("future") - .execute(runtime.pool.as_ref()) + .execute(runtime.state_db.pool()) .await .expect("future-status child edge insert should succeed"); diff --git a/codex-rs/state/src/telemetry.rs b/codex-rs/state/src/telemetry.rs index ed598aa738..8bb0ce699f 100644 --- a/codex-rs/state/src/telemetry.rs +++ b/codex-rs/state/src/telemetry.rs @@ -1,8 +1,6 @@ use std::borrow::Cow; -use std::future::Future; use std::sync::Arc; use std::time::Duration; -use std::time::Instant; use crate::DB_FALLBACK_METRIC; use crate::DB_INIT_DURATION_METRIC; @@ -48,6 +46,7 @@ pub(crate) enum DbAccess { Read, Write, Transaction, + Maintenance, } impl DbAccess { @@ -56,41 +55,22 @@ impl DbAccess { Self::Read => "read", Self::Write => "write", Self::Transaction => "transaction", + Self::Maintenance => "maintenance", } } } #[derive(Debug, Clone, Eq, PartialEq)] pub(crate) struct DbErrorTags { - pub error_class: &'static str, - pub sqlite_code: String, + pub error: &'static str, } impl DbErrorTags { fn none() -> Self { - Self { - error_class: "none", - sqlite_code: "none".to_string(), - } + Self { error: "none" } } } -pub(crate) async fn record_operation( - metrics: Option<&dyn DbMetricsRecorder>, - db: DbKind, - operation: &'static str, - access: DbAccess, - future: F, -) -> anyhow::Result -where - F: Future>, -{ - let started = Instant::now(); - let result = future.await; - record_operation_result(metrics, db, operation, access, started.elapsed(), &result); - result -} - pub(crate) fn record_init_result( metrics: Option<&dyn DbMetricsRecorder>, db: DbKind, @@ -103,8 +83,7 @@ pub(crate) fn record_init_result( ("status", outcome.status), ("phase", phase), ("db", db.as_str()), - ("error_class", outcome.error.error_class), - ("sqlite_code", outcome.error.sqlite_code.as_str()), + ("error", outcome.error.error), ]; record_counter(metrics, DB_INIT_METRIC, &tags); record_duration(metrics, DB_INIT_DURATION_METRIC, duration, &tags); @@ -145,43 +124,31 @@ pub(crate) fn classify_error(err: &anyhow::Error) -> DbErrorTags { .downcast_ref::() .is_some() { - return DbErrorTags { - error_class: "migration", - sqlite_code: "none".to_string(), - }; + return DbErrorTags { error: "migration" }; } if cause.downcast_ref::().is_some() { - return DbErrorTags { - error_class: "serde", - sqlite_code: "none".to_string(), - }; + return DbErrorTags { error: "serde" }; } if cause.downcast_ref::().is_some() { - return DbErrorTags { - error_class: "io", - sqlite_code: "none".to_string(), - }; + return DbErrorTags { error: "io" }; } } - DbErrorTags { - error_class: "unknown", - sqlite_code: "none".to_string(), - } + DbErrorTags { error: "unknown" } } pub(crate) fn classify_sqlite_code(code: &str) -> &'static str { let primary_code = code.parse::().ok().map(|code| code & 0xff); match primary_code { - Some(5) => "sqlite_busy", - Some(6) => "sqlite_locked", - Some(8) => "sqlite_readonly", - Some(10) => "sqlite_ioerr", - Some(11) => "sqlite_corrupt", - Some(13) => "sqlite_full", - Some(14) => "sqlite_cantopen", - Some(19) => "sqlite_constraint", - Some(17) => "sqlite_schema", + Some(5) => "busy", + Some(6) => "locked", + Some(8) => "readonly", + Some(10) => "io", + Some(11) => "corrupt", + Some(13) => "full", + Some(14) => "cantopen", + Some(19) => "constraint", + Some(17) => "schema", _ => "unknown", } } @@ -200,8 +167,7 @@ pub(crate) fn record_operation_result( ("db", db.as_str()), ("operation", operation), ("access", access.as_str()), - ("error_class", outcome.error.error_class), - ("sqlite_code", outcome.error.sqlite_code.as_str()), + ("error", outcome.error.error), ]; record_counter(metrics, DB_OPERATION_METRIC, &tags); record_duration(metrics, DB_OPERATION_DURATION_METRIC, duration, &tags); @@ -235,32 +201,20 @@ fn classify_sqlx_error(err: &sqlx::Error) -> DbErrorTags { .unwrap_or(Cow::Borrowed("none")) .to_string(); DbErrorTags { - error_class: classify_sqlite_code(code.as_str()), - sqlite_code: code, + error: classify_sqlite_code(code.as_str()), } } sqlx::Error::PoolTimedOut => DbErrorTags { - error_class: "pool_timeout", - sqlite_code: "none".to_string(), - }, - sqlx::Error::Io(_) => DbErrorTags { - error_class: "io", - sqlite_code: "none".to_string(), + error: "pool_timeout", }, + sqlx::Error::Io(_) => DbErrorTags { error: "io" }, sqlx::Error::ColumnDecode { source, .. } if source.is::() => { - DbErrorTags { - error_class: "serde", - sqlite_code: "none".to_string(), - } + DbErrorTags { error: "serde" } } - sqlx::Error::Decode(source) if source.is::() => DbErrorTags { - error_class: "serde", - sqlite_code: "none".to_string(), - }, - _ => DbErrorTags { - error_class: "unknown", - sqlite_code: "none".to_string(), - }, + sqlx::Error::Decode(source) if source.is::() => { + DbErrorTags { error: "serde" } + } + _ => DbErrorTags { error: "unknown" }, } } @@ -284,45 +238,76 @@ fn record_duration( #[cfg(test)] mod tests { use super::*; + use crate::DB_FALLBACK_METRIC; + use crate::DB_OPERATION_METRIC; use pretty_assertions::assert_eq; + use std::collections::BTreeMap; + use std::sync::Mutex; + + #[derive(Default)] + struct TestMetrics { + events: Mutex>, + } + + #[derive(Debug, Eq, PartialEq)] + struct MetricEvent { + name: String, + tags: BTreeMap, + } + + impl TestMetrics { + fn events(&self) -> Vec { + self.events + .lock() + .expect("metrics lock") + .iter() + .map(|event| MetricEvent { + name: event.name.clone(), + tags: event.tags.clone(), + }) + .collect() + } + } + + impl DbMetricsRecorder for TestMetrics { + fn counter(&self, name: &str, _inc: i64, tags: &[(&str, &str)]) { + self.events.lock().expect("metrics lock").push(MetricEvent { + name: name.to_string(), + tags: tags_to_map(tags), + }); + } + + fn record_duration(&self, _name: &str, _duration: Duration, _tags: &[(&str, &str)]) {} + } + + fn tags_to_map(tags: &[(&str, &str)]) -> BTreeMap { + tags.iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect() + } #[test] fn classifies_sqlite_primary_codes() { - assert_eq!(classify_sqlite_code("5"), "sqlite_busy"); - assert_eq!(classify_sqlite_code("6"), "sqlite_locked"); - assert_eq!(classify_sqlite_code("14"), "sqlite_cantopen"); - assert_eq!(classify_sqlite_code("2067"), "sqlite_constraint"); + assert_eq!(classify_sqlite_code("5"), "busy"); + assert_eq!(classify_sqlite_code("6"), "locked"); + assert_eq!(classify_sqlite_code("14"), "cantopen"); + assert_eq!(classify_sqlite_code("2067"), "constraint"); } #[test] fn classifies_non_sqlite_errors() { let io_error = anyhow::Error::new(std::io::Error::new(std::io::ErrorKind::NotFound, "missing")); - assert_eq!( - classify_error(&io_error), - DbErrorTags { - error_class: "io", - sqlite_code: "none".to_string() - } - ); + assert_eq!(classify_error(&io_error), DbErrorTags { error: "io" }); let serde_error = anyhow::Error::new(serde_json::from_str::("not-json").unwrap_err()); - assert_eq!( - classify_error(&serde_error), - DbErrorTags { - error_class: "serde", - sqlite_code: "none".to_string() - } - ); + assert_eq!(classify_error(&serde_error), DbErrorTags { error: "serde" }); let unknown_error = anyhow::anyhow!("plain failure"); assert_eq!( classify_error(&unknown_error), - DbErrorTags { - error_class: "unknown", - sqlite_code: "none".to_string() - } + DbErrorTags { error: "unknown" } ); } @@ -332,9 +317,55 @@ mod tests { assert_eq!( classify_error(&err), DbErrorTags { - error_class: "pool_timeout", - sqlite_code: "none".to_string() + error: "pool_timeout" } ); } + + #[test] + fn records_operation_metric_with_stable_tags() { + let metrics = TestMetrics::default(); + let result: anyhow::Result<()> = Ok(()); + + record_operation_result( + Some(&metrics), + DbKind::State, + "list_threads", + DbAccess::Read, + Duration::from_millis(3), + &result, + ); + + assert_eq!( + metrics.events(), + vec![MetricEvent { + name: DB_OPERATION_METRIC.to_string(), + tags: BTreeMap::from([ + ("access".to_string(), "read".to_string()), + ("db".to_string(), "state".to_string()), + ("error".to_string(), "none".to_string()), + ("operation".to_string(), "list_threads".to_string()), + ("status".to_string(), "success".to_string()), + ]), + }] + ); + } + + #[test] + fn records_fallback_metric_with_reason() { + let metrics = TestMetrics::default(); + + record_fallback(Some(&metrics), "list_threads", "db_error"); + + assert_eq!( + metrics.events(), + vec![MetricEvent { + name: DB_FALLBACK_METRIC.to_string(), + tags: BTreeMap::from([ + ("caller".to_string(), "list_threads".to_string()), + ("reason".to_string(), "db_error".to_string()), + ]), + }] + ); + } }