openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-10 21:00:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Owen Lin
2026-05-07 17:34:22 -07:00
parent 0c63bcda5e
commit 974ff5ba55
13 changed files with 593 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/otel/src/metrics/client.rs
-33
View File
@@ -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<String, String> = 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())
);
}
}
codex-rs/state/src/lib.rs
+10 -10
View File
@@ -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>,
codex-rs/state/src/runtime.rs
+10 -41
View File
@@ -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<sqlx::SqlitePool>,
logs_pool: Arc<sqlx::SqlitePool>,
state_db: InstrumentedDb,
logs_db: InstrumentedDb,
thread_updated_at_millis: Arc<AtomicI64>,
metrics: Option<DbMetricsRecorderHandle>,
}
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<DbMetricsRecorderHandle> {
self.metrics.clone()
}
pub(crate) async fn record_db_operation<T, F>(
&self,
db: DbKind,
operation: &'static str,
access: DbAccess,
future: F,
) -> anyhow::Result<T>
where
F: std::future::Future<Output = anyhow::Result<T>>,
{
crate::telemetry::record_operation(self.metrics(), db, operation, access, future).await
}
pub(crate) fn record_db_operation_result<T>(
&self,
db: DbKind,
operation: &'static str,
access: DbAccess,
started: Instant,
result: &anyhow::Result<T>,
) {
crate::telemetry::record_operation_result(
self.metrics(),
db,
operation,
access,
started.elapsed(),
result,
);
self.state_db.metrics_handle()
}
}
codex-rs/state/src/runtime/agent_jobs.rs
+17 -17
View File
@@ -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<AgentJobItemRow> = builder
.build_query_as::<AgentJobItemRow>()
.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")?;
codex-rs/state/src/runtime/backfill.rs
+40 -54
View File
@@ -2,19 +2,21 @@ use super::*;
impl StateRuntime {
pub async fn get_backfill_state(&self) -> anyhow::Result<crate::BackfillState> {
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<bool> {
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");
codex-rs/state/src/runtime/db.rs
+155
View File
@@ -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<SqlitePool>,
kind: DbKind,
metrics: Option<DbMetricsRecorderHandle>,
}
impl InstrumentedDb {
pub(super) fn new(
pool: Arc<SqlitePool>,
kind: DbKind,
metrics: Option<DbMetricsRecorderHandle>,
) -> 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<DbMetricsRecorderHandle> {
self.metrics.clone()
}
pub(super) async fn read<T, F>(&self, operation: DbOperation, future: F) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
self.record_operation(operation, DbAccess::Read, future)
.await
}
pub(super) async fn write<T, F>(&self, operation: DbOperation, future: F) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
self.record_operation(operation, DbAccess::Write, future)
.await
}
pub(super) async fn transaction<T, F>(
&self,
operation: DbOperation,
future: F,
) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
self.record_operation(operation, DbAccess::Transaction, future)
.await
}
pub(super) async fn maintenance<T, F>(
&self,
operation: DbOperation,
future: F,
) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
self.record_operation(operation, DbAccess::Maintenance, future)
.await
}
pub(super) fn record_result<T>(
&self,
operation: DbOperation,
access: DbAccess,
started: Instant,
result: &anyhow::Result<T>,
) {
crate::telemetry::record_operation_result(
self.metrics(),
self.kind,
operation.as_str(),
access,
started.elapsed(),
result,
);
}
async fn record_operation<T, F>(
&self,
operation: DbOperation,
access: DbAccess,
future: F,
) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
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",
}
}
}
codex-rs/state/src/runtime/goals.rs
+9 -9
View File
@@ -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(
codex-rs/state/src/runtime/logs.rs
+27 -28
View File
@@ -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::<Sqlite>::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<u64> {
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::<LogRow>()
.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::<Sqlite>::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<i64> = row.try_get("max_id")?;
Ok(max_id.unwrap_or(0))
}
codex-rs/state/src/runtime/memories.rs
+55 -55
View File
@@ -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<bool> {
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<bool> {
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::<i64, _>("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::<i64, _>("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::<i64, _>("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::<i64, _>("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::<i64, _>("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::<i64, _>("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");
codex-rs/state/src/runtime/metrics_tests.rs
-139
View File
@@ -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<String, String>,
}
#[derive(Default)]
struct TestMetrics {
events: Mutex<Vec<MetricEvent>>,
}
impl TestMetrics {
fn metric_points(&self, name: &str) -> Vec<BTreeMap<String, String>> {
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<TestMetrics> {
Arc::new(TestMetrics::default())
}
fn tags_to_map(tags: &[(&str, &str)]) -> BTreeMap<String, String> {
tags.iter()
.map(|(key, value)| ((*key).to_string(), (*value).to_string()))
.collect()
}
fn find_point(
points: &[BTreeMap<String, String>],
key: &str,
value: &str,
) -> BTreeMap<String, String> {
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;
}
codex-rs/state/src/runtime/remote_control.rs
+3 -3
View File
@@ -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())
}
codex-rs/state/src/runtime/threads.rs
+139 -151
View File
@@ -5,9 +5,11 @@ use std::sync::atomic::Ordering;
impl StateRuntime {
pub async fn get_thread(&self, id: ThreadId) -> anyhow::Result<Option<crate::ThreadMetadata>> {
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<Option<String>> {
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<Option<Vec<DynamicToolSpec>>> {
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::<Value>(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::<Value>(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::<String, _>("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::<String, _>("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<bool>,
) -> anyhow::Result<Option<PathBuf>> {
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::<Sqlite>::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::<String, _>("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<crate::ThreadsPage> {
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::<Sqlite>::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::<Sqlite>::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::<Result<Vec<_>, _>>()?;
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::<Result<Vec<_>, _>>()?;
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<Utc>,
) -> anyhow::Result<bool> {
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<u64> {
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");
codex-rs/state/src/telemetry.rs
+128 -97
View File
@@ -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<T, F>(
metrics: Option<&dyn DbMetricsRecorder>,
db: DbKind,
operation: &'static str,
access: DbAccess,
future: F,
) -> anyhow::Result<T>
where
F: Future<Output = anyhow::Result<T>>,
{
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<T>(
metrics: Option<&dyn DbMetricsRecorder>,
db: DbKind,
@@ -103,8 +83,7 @@ pub(crate) fn record_init_result<T>(
("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::<sqlx::migrate::MigrateError>()
.is_some()
{
return DbErrorTags {
error_class: "migration",
sqlite_code: "none".to_string(),
};
return DbErrorTags { error: "migration" };
}
if cause.downcast_ref::<serde_json::Error>().is_some() {
return DbErrorTags {
error_class: "serde",
sqlite_code: "none".to_string(),
};
return DbErrorTags { error: "serde" };
}
if cause.downcast_ref::<std::io::Error>().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::<i32>().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<T>(
("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::<serde_json::Error>() => {
DbErrorTags {
error_class: "serde",
sqlite_code: "none".to_string(),
}
DbErrorTags { error: "serde" }
}
sqlx::Error::Decode(source) if source.is::<serde_json::Error>() => 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::<serde_json::Error>() => {
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<Vec<MetricEvent>>,
}
#[derive(Debug, Eq, PartialEq)]
struct MetricEvent {
name: String,
tags: BTreeMap<String, String>,
}
impl TestMetrics {
fn events(&self) -> Vec<MetricEvent> {
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<String, String> {
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::<serde_json::Value>("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()),
]),
}]
);
}
}