From 5c554634e2312fdf7c6ee7e0a20e5a6aa0315a27 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 28 Aug 2026 02:08:01 -0700 Subject: [PATCH 1/2] feat(metrics): profile sqlite storage --- engine/packages/depot-client/src/database.rs | 16 +- engine/packages/depot-client/src/query.rs | 300 +++- engine/packages/depot-client/src/vfs.rs | 522 +++++- engine/packages/depot-client/src/worker.rs | 185 ++- .../packages/depot-client/tests/inline/vfs.rs | 115 +- .../rivetkit-core/src/actor/config.rs | 112 ++ .../rivetkit-core/src/actor/context.rs | 5 +- .../rivetkit-core/src/actor/metrics.rs | 1396 +++++++++++++++++ .../rivetkit-core/src/actor/sqlite.rs | 463 ++++++ .../src/actor/sqlite/profiling.rs | 238 +++ .../packages/rivetkit-core/src/lib.rs | 4 + .../packages/rivetkit-core/tests/config.rs | 45 +- .../packages/rivetkit-core/tests/metrics.rs | 525 ++++++- .../rivetkit-core/tests/metrics_helpers.rs | 9 +- rivetkit-rust/packages/rivetkit/src/lib.rs | 4 + .../packages/rivetkit/src/prelude.rs | 6 +- rivetkit-rust/packages/rivetkit/src/sqlite.rs | 54 + .../packages/rivetkit-napi/index.d.ts | 18 + .../rivetkit-napi/src/actor_factory.rs | 40 +- .../packages/rivetkit-wasm/src/lib.rs | 2 + .../rivetkit/src/common/database/config.ts | 23 + .../rivetkit/src/common/database/mod.ts | 15 +- .../packages/rivetkit/src/db/mod.ts | 1 + .../packages/rivetkit/src/registry/native.ts | 10 +- .../packages/rivetkit/src/registry/runtime.ts | 7 +- 25 files changed, 3999 insertions(+), 116 deletions(-) create mode 100644 rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/profiling.rs create mode 100644 rivetkit-rust/packages/rivetkit/src/sqlite.rs diff --git a/engine/packages/depot-client/src/database.rs b/engine/packages/depot-client/src/database.rs index 98657225c2..59dd481edf 100644 --- a/engine/packages/depot-client/src/database.rs +++ b/engine/packages/depot-client/src/database.rs @@ -13,7 +13,7 @@ use crate::{ SqliteVfsMetricsSnapshot, VfsConfig, VfsPreloadHintSnapshot, fetch_initial_pages_for_registration, }, - worker::{SqliteWorkerFatalError, SqliteWorkerHandle}, + worker::{SqliteWorkerFatalError, SqliteWorkerHandle, SqliteWorkerResult}, }; #[derive(Clone)] @@ -220,6 +220,11 @@ impl NativeDatabaseHandle { self.map_worker_result(self.worker.exec(sql).await) } + pub async fn exec_profiled(&self, sql: String) -> Result> { + self.check_fatal_error()?; + self.map_worker_result(self.worker.exec_profiled(sql).await) + } + pub async fn query(&self, sql: String, params: Option>) -> Result { self.execute(sql, params).await.map(|result| QueryResult { columns: result.columns, @@ -242,6 +247,15 @@ impl NativeDatabaseHandle { self.map_worker_result(self.worker.execute(sql, params).await) } + pub async fn execute_profiled( + &self, + sql: String, + params: Option>, + ) -> Result> { + self.check_fatal_error()?; + self.map_worker_result(self.worker.execute_profiled(sql, params).await) + } + pub async fn close(&self) -> Result<()> { match self.worker.close().await { Ok(()) => Ok(()), diff --git a/engine/packages/depot-client/src/query.rs b/engine/packages/depot-client/src/query.rs index ce342cadb6..9df1cb381e 100644 --- a/engine/packages/depot-client/src/query.rs +++ b/engine/packages/depot-client/src/query.rs @@ -13,6 +13,15 @@ use libsqlite3_sys::{ sqlite3_finalize, sqlite3_last_insert_rowid, sqlite3_prepare_v2, sqlite3_step, }; +#[derive(Clone, Copy, Debug, Default)] +pub struct SqliteQueryProfile { + pub bind_count: u64, + pub bind_logical_bytes: u64, + pub result_rows: u64, + pub result_columns: u64, + pub result_logical_bytes: u64, +} + pub fn execute_statement( db: *mut sqlite3, sql: &str, @@ -178,6 +187,86 @@ pub fn execute_single_statement( result } +pub fn execute_single_statement_profiled( + db: *mut sqlite3, + sql: &str, + params: Option<&[BindParam]>, +) -> Result<(ExecuteResult, SqliteQueryProfile)> { + let c_sql = CString::new(sql).map_err(|err| anyhow!(err.to_string()))?; + let mut stmt = ptr::null_mut(); + let mut tail = ptr::null(); + let rc = unsafe { sqlite3_prepare_v2(db, c_sql.as_ptr(), -1, &mut stmt, &mut tail) }; + if rc != SQLITE_OK { + return Err(sqlite_error( + db, + "failed to prepare sqlite execute statement", + )); + } + if has_non_whitespace_tail(tail) { + if !stmt.is_null() { + unsafe { + sqlite3_finalize(stmt); + } + } + return Err(anyhow!("sqlite execute only supports a single statement")); + } + if stmt.is_null() { + return Ok(( + ExecuteResult { + columns: Vec::new(), + rows: Vec::new(), + changes: 0, + last_insert_row_id: None, + }, + SqliteQueryProfile::default(), + )); + } + + let result = (|| { + let mut profile = SqliteQueryProfile::default(); + if let Some(params) = params { + bind_params_profiled(db, stmt, params, &mut profile)?; + } + + let columns = collect_columns(stmt); + profile.result_columns = columns.len() as u64; + let mut rows = Vec::new(); + loop { + let step_rc = unsafe { sqlite3_step(stmt) }; + if step_rc == SQLITE_DONE { + break; + } + if step_rc != SQLITE_ROW { + return Err(sqlite_error(db, "failed to step sqlite execute statement")); + } + + let mut row = Vec::with_capacity(columns.len()); + for index in 0..columns.len() { + row.push(column_value_profiled(stmt, index as i32, &mut profile)); + } + rows.push(row); + profile.result_rows = profile.result_rows.saturating_add(1); + } + + let changes = unsafe { sqlite3_changes(db) as i64 }; + Ok(( + ExecuteResult { + columns, + rows, + changes, + last_insert_row_id: (changes > 0).then(|| unsafe { sqlite3_last_insert_rowid(db) }), + }, + profile, + )) + })(); + + unsafe { + sqlite3_finalize(stmt); + } + + result +} + pub fn exec_statements(db: *mut sqlite3, sql: &str) -> Result { let c_sql = CString::new(sql).map_err(|err| anyhow!(err.to_string()))?; let mut remaining = c_sql.as_ptr(); @@ -242,40 +331,141 @@ pub fn exec_statements(db: *mut sqlite3, sql: &str) -> Result { Ok(final_result) } +pub fn exec_statements_profiled( + db: *mut sqlite3, + sql: &str, +) -> Result<(QueryResult, SqliteQueryProfile)> { + let c_sql = CString::new(sql).map_err(|err| anyhow!(err.to_string()))?; + let mut remaining = c_sql.as_ptr(); + let mut final_result = QueryResult { + columns: Vec::new(), + rows: Vec::new(), + }; + let mut final_profile = SqliteQueryProfile::default(); + + while unsafe { *remaining } != 0 { + let mut stmt = ptr::null_mut(); + let mut tail = ptr::null(); + let rc = unsafe { sqlite3_prepare_v2(db, remaining, -1, &mut stmt, &mut tail) }; + if rc != SQLITE_OK { + return Err(sqlite_error(db, "failed to prepare sqlite exec statement")); + } + + if stmt.is_null() { + if tail == remaining { + break; + } + remaining = tail; + continue; + } + + let result = (|| { + let columns = collect_columns(stmt); + let mut profile = SqliteQueryProfile { + result_columns: columns.len() as u64, + ..Default::default() + }; + let mut rows = Vec::new(); + loop { + let step_rc = unsafe { sqlite3_step(stmt) }; + if step_rc == SQLITE_DONE { + break; + } + if step_rc != SQLITE_ROW { + return Err(sqlite_error(db, "failed to step sqlite exec statement")); + } + + let mut row = Vec::with_capacity(columns.len()); + for index in 0..columns.len() { + row.push(column_value_profiled(stmt, index as i32, &mut profile)); + } + rows.push(row); + profile.result_rows = profile.result_rows.saturating_add(1); + } + + Ok((columns, rows, profile)) + })(); + + unsafe { + sqlite3_finalize(stmt); + } + + let (columns, rows, profile) = result?; + if !columns.is_empty() || !rows.is_empty() { + final_result = QueryResult { columns, rows }; + final_profile = profile; + } + + if tail == remaining { + break; + } + remaining = tail; + } + + Ok((final_result, final_profile)) +} + fn bind_params( db: *mut sqlite3, stmt: *mut libsqlite3_sys::sqlite3_stmt, params: &[BindParam], ) -> Result<()> { for (index, param) in params.iter().enumerate() { - let bind_index = (index + 1) as i32; - let rc = match param { - BindParam::Null => unsafe { sqlite3_bind_null(stmt, bind_index) }, - BindParam::Integer(value) => unsafe { sqlite3_bind_int64(stmt, bind_index, *value) }, - BindParam::Float(value) => unsafe { sqlite3_bind_double(stmt, bind_index, *value) }, - BindParam::Text(value) => unsafe { - sqlite3_bind_text( - stmt, - bind_index, - value.as_ptr() as *const c_char, - value.len() as i32, - SQLITE_TRANSIENT(), - ) - }, - BindParam::Blob(value) => unsafe { - sqlite3_bind_blob( - stmt, - bind_index, - value.as_ptr() as *const _, - value.len() as i32, - SQLITE_TRANSIENT(), - ) - }, - }; + bind_param(db, stmt, (index + 1) as i32, param)?; + } - if rc != SQLITE_OK { - return Err(sqlite_error(db, "failed to bind sqlite parameter")); - } + Ok(()) +} + +fn bind_params_profiled( + db: *mut sqlite3, + stmt: *mut libsqlite3_sys::sqlite3_stmt, + params: &[BindParam], + profile: &mut SqliteQueryProfile, +) -> Result<()> { + for (index, param) in params.iter().enumerate() { + profile.bind_count = profile.bind_count.saturating_add(1); + profile.bind_logical_bytes = profile + .bind_logical_bytes + .saturating_add(logical_bind_bytes(param)); + bind_param(db, stmt, (index + 1) as i32, param)?; + } + + Ok(()) +} + +fn bind_param( + db: *mut sqlite3, + stmt: *mut libsqlite3_sys::sqlite3_stmt, + bind_index: i32, + param: &BindParam, +) -> Result<()> { + let rc = match param { + BindParam::Null => unsafe { sqlite3_bind_null(stmt, bind_index) }, + BindParam::Integer(value) => unsafe { sqlite3_bind_int64(stmt, bind_index, *value) }, + BindParam::Float(value) => unsafe { sqlite3_bind_double(stmt, bind_index, *value) }, + BindParam::Text(value) => unsafe { + sqlite3_bind_text( + stmt, + bind_index, + value.as_ptr() as *const c_char, + value.len() as i32, + SQLITE_TRANSIENT(), + ) + }, + BindParam::Blob(value) => unsafe { + sqlite3_bind_blob( + stmt, + bind_index, + value.as_ptr() as *const _, + value.len() as i32, + SQLITE_TRANSIENT(), + ) + }, + }; + + if rc != SQLITE_OK { + return Err(sqlite_error(db, "failed to bind sqlite parameter")); } Ok(()) @@ -327,6 +517,62 @@ fn column_value(stmt: *mut libsqlite3_sys::sqlite3_stmt, index: i32) -> ColumnVa } } +fn column_value_profiled( + stmt: *mut libsqlite3_sys::sqlite3_stmt, + index: i32, + profile: &mut SqliteQueryProfile, +) -> ColumnValue { + let mut logical_bytes = 0_u64; + let value = match unsafe { sqlite3_column_type(stmt, index) } { + SQLITE_NULL => ColumnValue::Null, + SQLITE_INTEGER => { + logical_bytes = 8; + ColumnValue::Integer(unsafe { sqlite3_column_int64(stmt, index) }) + } + SQLITE_FLOAT => { + logical_bytes = 8; + ColumnValue::Float(unsafe { sqlite3_column_double(stmt, index) }) + } + SQLITE_TEXT => { + let text_ptr = unsafe { sqlite3_column_text(stmt, index) }; + if text_ptr.is_null() { + ColumnValue::Null + } else { + let text_len = unsafe { sqlite3_column_bytes(stmt, index) } as usize; + logical_bytes = text_len as u64; + let text = String::from_utf8_lossy(unsafe { + std::slice::from_raw_parts(text_ptr as *const u8, text_len) + }) + .into_owned(); + ColumnValue::Text(text) + } + } + SQLITE_BLOB => { + let blob_ptr = unsafe { sqlite3_column_blob(stmt, index) }; + if blob_ptr.is_null() { + ColumnValue::Null + } else { + let blob_len = unsafe { sqlite3_column_bytes(stmt, index) } as usize; + logical_bytes = blob_len as u64; + let blob = unsafe { std::slice::from_raw_parts(blob_ptr as *const u8, blob_len) }; + ColumnValue::Blob(blob.to_vec()) + } + } + _ => ColumnValue::Null, + }; + profile.result_logical_bytes = profile.result_logical_bytes.saturating_add(logical_bytes); + value +} + +fn logical_bind_bytes(param: &BindParam) -> u64 { + match param { + BindParam::Null => 0, + BindParam::Integer(_) | BindParam::Float(_) => 8, + BindParam::Text(value) => value.len() as u64, + BindParam::Blob(value) => value.len() as u64, + } +} + fn has_non_whitespace_tail(tail: *const c_char) -> bool { if tail.is_null() { return false; diff --git a/engine/packages/depot-client/src/vfs.rs b/engine/packages/depot-client/src/vfs.rs index ca3256f97d..93b14db220 100644 --- a/engine/packages/depot-client/src/vfs.rs +++ b/engine/packages/depot-client/src/vfs.rs @@ -7,7 +7,7 @@ use std::ffi::{CStr, CString, c_char, c_int, c_void}; use std::ptr; use std::slice; use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::time::{Duration, Instant}; use anyhow::Result; @@ -288,6 +288,147 @@ pub struct SqliteVfsMetricsSnapshot { pub db_size_pages: u64, } +pub const MAX_PROFILED_GET_PAGES_REQUESTS: usize = 16; +pub const MAX_PROFILED_TRANSACTION_STATEMENTS: usize = 32; + +/// Bounded details for one physical `get_pages` call made while a native SQLite +/// command is executing. +#[derive(Debug, Clone, Default)] +pub struct SqliteGetPagesProfile { + pub ordinal: u64, + pub duration_ns: u64, + pub demand_requested: u64, + pub prefetch_requested: u64, + pub response_present: u64, + pub response_missing: u64, + pub overflow_expansion_extra: u64, + pub response_bytes: u64, + pub success: bool, +} + +/// Fixed-size, connection-local counters collected by VFS callbacks for one +/// native SQLite command. The worker installs and removes this context around +/// each command, so VFS activity cannot be attributed to a neighboring query. +#[derive(Debug, Clone, Default)] +pub struct SqliteOperationProfile { + pub worker_wait_ns: u64, + pub sqlite_execution_ns: u64, + pub storage_ns: u64, + pub bind_count: u64, + pub bind_logical_bytes: u64, + pub result_rows: u64, + pub result_columns: u64, + pub result_logical_bytes: u64, + pub sqlite_requested_pages: u64, + pub cache_hit_pages: u64, + pub cache_miss_pages: u64, + pub depot_demand_requested_pages: u64, + pub vfs_prefetch_requested_pages: u64, + pub response_present_pages: u64, + pub response_missing_pages: u64, + pub overflow_expansion_extra_pages: u64, + pub btree_pages: u64, + pub non_btree_pages: u64, + pub prefetch_consumed_pages: u64, + pub prefetch_unused_pages: u64, + pub dirty_pages: u64, + pub storage_response_bytes: u64, + pub dirty_bytes: u64, + pub get_pages_round_trips: u64, + pub commit_round_trips: u64, + pub get_pages_requests: [Option; MAX_PROFILED_GET_PAGES_REQUESTS], + pub omitted_get_pages_requests: u64, +} + +#[derive(Debug, Clone)] +pub struct SqliteOperationMetric { + pub operation_type: &'static str, + pub fingerprint: String, + pub fingerprint_source: &'static str, + pub transaction_mode: &'static str, + pub storage_transport: &'static str, + pub outcome: &'static str, + pub sql_bytes: u64, + pub total_ns: u64, + pub transaction_wait_ns: u64, + pub profile: SqliteOperationProfile, +} + +#[derive(Debug, Clone)] +pub struct SqliteTransactionMetric { + pub fingerprint: String, + pub fingerprint_source: &'static str, + pub shape_fingerprint: String, + pub statement_fingerprint_hashes: [Option<[u8; 16]>; MAX_PROFILED_TRANSACTION_STATEMENTS], + pub omitted_statement_fingerprints: u64, + pub storage_transport: &'static str, + pub outcome: &'static str, + pub total_ns: u64, + pub transaction_wait_ns: u64, + pub worker_wait_ns: u64, + pub storage_ns: u64, + pub local_work_ns: u64, + pub application_time_ns: u64, + pub commit_ns: u64, + pub get_pages_round_trips: u64, + pub statement_count: u64, + pub dirty_pages: u64, + pub dirty_bytes: u64, +} + +impl SqliteOperationProfile { + fn push_get_pages(&mut self, request: SqliteGetPagesProfile, request_limit: usize) { + self.storage_ns = self.storage_ns.saturating_add(request.duration_ns); + self.get_pages_round_trips = self.get_pages_round_trips.saturating_add(1); + self.depot_demand_requested_pages = self + .depot_demand_requested_pages + .saturating_add(request.demand_requested); + self.vfs_prefetch_requested_pages = self + .vfs_prefetch_requested_pages + .saturating_add(request.prefetch_requested); + self.response_present_pages = self + .response_present_pages + .saturating_add(request.response_present); + self.response_missing_pages = self + .response_missing_pages + .saturating_add(request.response_missing); + self.overflow_expansion_extra_pages = self + .overflow_expansion_extra_pages + .saturating_add(request.overflow_expansion_extra); + self.storage_response_bytes = self + .storage_response_bytes + .saturating_add(request.response_bytes); + let index = self.get_pages_round_trips.saturating_sub(1) as usize; + if index < request_limit + && let Some(slot) = self.get_pages_requests.get_mut(index) + { + *slot = Some(request); + } else { + self.omitted_get_pages_requests = self.omitted_get_pages_requests.saturating_add(1); + } + } +} + +pub(crate) struct SqliteOperationProfileGuard<'a> { + ctx: &'a VfsContext, + active: bool, +} + +impl SqliteOperationProfileGuard<'_> { + pub(crate) fn finish(mut self) -> SqliteOperationProfile { + self.active = false; + self.ctx.finish_operation_profile() + } +} + +impl Drop for SqliteOperationProfileGuard<'_> { + fn drop(&mut self) { + if self.active { + let _ = self.ctx.finish_operation_profile(); + } + } +} + /// Cumulative count of network round trips the VFS has issued to the engine. /// /// `get_pages` counts `SqliteGetPagesRequest` fetches and `commit` counts @@ -332,69 +473,94 @@ impl SqliteOpenPhase { } pub trait SqliteVfsMetrics: Send + Sync { - fn record_resolve_pages(&self, _requested_pages: u64) {} + fn profiling_enabled(&self) -> bool; - fn record_resolve_cache_hits(&self, _pages: u64) {} + fn max_profiled_get_pages_requests(&self) -> usize; - fn record_resolve_cache_misses(&self, _pages: u64) {} + /// Records an operation and returns whether its original fingerprint was + /// admitted instead of being routed to the shared `other` series. + fn observe_operation_profile(&self, profile: &SqliteOperationMetric) -> bool; - fn record_get_pages_request(&self, _pages: u64, _prefetch_pages: u64, _page_size: u64) {} + /// Records a transaction and returns whether its original fingerprint was + /// admitted instead of being routed to the shared `other` series. + fn observe_transaction_profile(&self, profile: &SqliteTransactionMetric) -> bool; - fn observe_get_pages_duration(&self, _duration_ns: u64) {} + fn emit_operation_diagnostic_event( + &self, + actor_id: &str, + generation: Option, + profile: &SqliteOperationMetric, + ); - fn observe_open_phase( + fn emit_transaction_diagnostic_event( &self, - _phase: SqliteOpenPhase, - _outcome: &'static str, - _duration_ns: u64, - ) { - } + actor_id: &str, + generation: Option, + profile: &SqliteTransactionMetric, + ); + + fn record_fingerprint_catalog( + &self, + operation_type: &'static str, + fingerprint: &str, + identity: &str, + format_version: u8, + ); + + fn record_resolve_pages(&self, requested_pages: u64); - fn record_startup_preload_pages(&self, _kind: &'static str, _pages: u64) {} + fn record_resolve_cache_hits(&self, pages: u64); - fn record_commit(&self) {} + fn record_resolve_cache_misses(&self, pages: u64); + + fn record_get_pages_request(&self, pages: u64, prefetch_pages: u64, page_size: u64); + + fn observe_get_pages_duration(&self, duration_ns: u64); + + fn observe_open_phase(&self, phase: SqliteOpenPhase, outcome: &'static str, duration_ns: u64); + + fn record_startup_preload_pages(&self, kind: &'static str, pages: u64); + + fn record_commit(&self); fn observe_commit_phases( &self, - _request_build_ns: u64, - _serialize_ns: u64, - _transport_ns: u64, - _state_update_ns: u64, - _total_ns: u64, - ) { - } + request_build_ns: u64, + serialize_ns: u64, + transport_ns: u64, + state_update_ns: u64, + total_ns: u64, + ); + + fn set_worker_queue_depth(&self, depth: u64); - fn set_worker_queue_depth(&self, _depth: u64) {} + fn set_worker_active(&self, active: bool); - fn set_worker_active(&self, _active: bool) {} + fn set_worker_inflight(&self, active: bool); - fn record_worker_queue_overload(&self) {} + fn set_coordinator_queue_depth(&self, depth: u64); + + fn record_worker_queue_overload(&self); fn observe_worker_command_duration( &self, - _operation: &'static str, - _in_tx: bool, - _stmt_kind: &'static str, - _duration_ns: u64, - ) { - } + operation: &'static str, + in_tx: bool, + stmt_kind: &'static str, + duration_ns: u64, + ); - fn observe_transaction_round_trips( - &self, - _get_pages_round_trips: u64, - _commit_round_trips: u64, - ) { - } + fn observe_transaction_round_trips(&self, get_pages_round_trips: u64, commit_round_trips: u64); - fn record_worker_command_error(&self, _operation: &'static str, _code: &'static str) {} + fn record_worker_command_error(&self, operation: &'static str, code: &'static str); - fn observe_worker_close_duration(&self, _duration_ns: u64) {} + fn observe_worker_close_duration(&self, duration_ns: u64); - fn record_worker_close_timeout(&self) {} + fn record_worker_close_timeout(&self); - fn record_worker_crash(&self) {} + fn record_worker_crash(&self); - fn record_worker_unclean_close(&self) {} + fn record_worker_unclean_close(&self); } #[derive(Debug, Clone, Copy, Default)] @@ -437,6 +603,12 @@ pub struct VfsContext { pub commit_transport_ns: AtomicU64, pub commit_state_update_ns: AtomicU64, pub commit_duration_ns_total: AtomicU64, + // SQLite invokes VFS callbacks synchronously, so this profiling context cannot + // use an async mutex. The native worker releases each guard before returning + // to async code. + operation_profile_active: AtomicBool, + operation_profile: Mutex>, + operation_prefetch_unused_start: AtomicU64, metrics: Option>, } @@ -448,6 +620,9 @@ struct VfsState { page_cache: Cache>, committed_page_cache: Cache>, protected_page_cache: Arc>>, + profiling_enabled: bool, + prefetched_pages: Arc>, + prefetch_unused_total: Arc, write_buffer: WriteBuffer, predictor: ClassifiedPredictor, read_ahead: ClassifiedReadAhead, @@ -954,7 +1129,7 @@ fn push_coalesced_range(ranges: &mut VecDeque, range: VfsPr } impl VfsState { - fn new(config: &VfsConfig) -> Self { + fn new(config: &VfsConfig, profiling_enabled: bool) -> Self { let page_cache = build_page_cache(config); let committed_page_cache = build_page_cache(config); let mut state = Self { @@ -964,6 +1139,9 @@ impl VfsState { page_cache, committed_page_cache, protected_page_cache: Arc::new(SccHashMap::new()), + profiling_enabled, + prefetched_pages: Arc::new(scc::HashSet::new()), + prefetch_unused_total: Arc::new(AtomicU64::new(0)), write_buffer: WriteBuffer::default(), predictor: ClassifiedPredictor::default(), read_ahead: ClassifiedReadAhead::default(), @@ -991,23 +1169,32 @@ impl VfsState { config, &self.page_cache, &self.protected_page_cache, + &self.prefetched_pages, + &self.prefetch_unused_total, + self.profiling_enabled, kind, pgno, bytes, ); } - fn cached_page(&self, config: &VfsConfig, pgno: u32) -> Option> { + fn cached_page(&self, config: &VfsConfig, pgno: u32) -> Option<(Vec, bool)> { if !can_read_cached_page(config, pgno) { return None; } - self.committed_page_cache + let bytes = self + .committed_page_cache .get(&pgno) .or_else(|| { self.protected_page_cache .read_sync(&pgno, |_, bytes| bytes.clone()) }) - .or_else(|| self.page_cache.get(&pgno)) + .or_else(|| self.page_cache.get(&pgno)); + bytes.map(|bytes| { + let was_prefetched = + self.profiling_enabled && self.prefetched_pages.remove_sync(&pgno).is_some(); + (bytes, was_prefetched) + }) } /// Record target page accesses into the per-class predictor and read-ahead @@ -1124,6 +1311,11 @@ impl VfsState { } fn invalidate_page_cache(&mut self) { + if self.profiling_enabled { + self.prefetch_unused_total + .fetch_add(self.prefetched_pages.len() as u64, Ordering::Relaxed); + self.prefetched_pages.clear_sync(); + } self.page_cache.invalidate_all(); self.committed_page_cache.invalidate_all(); self.protected_page_cache.clear_sync(); @@ -1143,6 +1335,9 @@ fn cache_page( config: &VfsConfig, page_cache: &Cache>, _protected_page_cache: &SccHashMap>, + prefetched_pages: &scc::HashSet, + prefetch_unused_total: &AtomicU64, + profiling_enabled: bool, kind: PageCacheInsertKind, pgno: u32, bytes: Vec, @@ -1150,6 +1345,20 @@ fn cache_page( if !should_cache_page(config, kind, pgno) { return; } + if !profiling_enabled { + page_cache.insert(pgno, bytes); + return; + } + if kind == PageCacheInsertKind::Prefetch { + let capacity = config.cache_capacity_pages.max(1) as usize; + if prefetched_pages.len() >= capacity { + prefetch_unused_total.fetch_add(prefetched_pages.len() as u64, Ordering::Relaxed); + prefetched_pages.clear_sync(); + } + let _ = prefetched_pages.insert_sync(pgno); + } else if prefetched_pages.remove_sync(&pgno).is_some() { + prefetch_unused_total.fetch_add(1, Ordering::Relaxed); + } page_cache.insert(pgno, bytes); } @@ -1182,7 +1391,10 @@ impl VfsContext { initial_pages: impl Into, metrics: Option>, ) -> std::result::Result { - let mut state = VfsState::new(&config); + let profiling_enabled = metrics + .as_ref() + .is_some_and(|metrics| metrics.profiling_enabled()); + let mut state = VfsState::new(&config, profiling_enabled); let initial_pages = initial_pages.into(); state.head_txid = initial_pages.head_txid; for (pgno, page) in initial_pages.pages { @@ -1217,10 +1429,62 @@ impl VfsContext { commit_transport_ns: AtomicU64::new(0), commit_state_update_ns: AtomicU64::new(0), commit_duration_ns_total: AtomicU64::new(0), + operation_profile_active: AtomicBool::new(false), + operation_profile: Mutex::new(None), + operation_prefetch_unused_start: AtomicU64::new(0), metrics, }) } + fn begin_operation_profile(&self) { + if self + .metrics + .as_ref() + .is_some_and(|metrics| metrics.profiling_enabled()) + { + let unused = self + .state + .read() + .prefetch_unused_total + .load(Ordering::Relaxed); + self.operation_prefetch_unused_start + .store(unused, Ordering::Relaxed); + *self.operation_profile.lock() = Some(SqliteOperationProfile::default()); + self.operation_profile_active.store(true, Ordering::Release); + } + } + + fn finish_operation_profile(&self) -> SqliteOperationProfile { + if !self.operation_profile_active.swap(false, Ordering::AcqRel) { + return SqliteOperationProfile::default(); + } + let mut profile = self.operation_profile.lock().take().unwrap_or_default(); + let unused = self + .state + .read() + .prefetch_unused_total + .load(Ordering::Relaxed) + .saturating_sub(self.operation_prefetch_unused_start.load(Ordering::Relaxed)); + profile.prefetch_unused_pages = unused; + profile + } + + fn update_operation_profile(&self, update: impl FnOnce(&mut SqliteOperationProfile)) { + if !self.operation_profile_active.load(Ordering::Relaxed) { + return; + } + if let Some(profile) = self.operation_profile.lock().as_mut() { + update(profile); + } + } + + fn profile_request_limit(&self) -> usize { + self.metrics + .as_ref() + .map_or(0, |metrics| metrics.max_profiled_get_pages_requests()) + .min(MAX_PROFILED_GET_PAGES_REQUESTS) + } + fn clear_last_error(&self) { *self.last_error.lock() = None; } @@ -1256,6 +1520,12 @@ impl VfsContext { state_update_ns: u64, total_ns: u64, ) { + self.update_operation_profile(|profile| { + profile.storage_ns = profile + .storage_ns + .saturating_add(transport_metrics.transport_ns); + profile.commit_round_trips = profile.commit_round_trips.saturating_add(1); + }); if let Some(metrics) = &self.metrics { metrics.observe_commit_phases( request_build_ns, @@ -1417,6 +1687,11 @@ impl VfsContext { ) -> std::result::Result>>, GetPagesError> { use std::sync::atomic::Ordering::Relaxed; self.resolve_pages_total.fetch_add(1, Relaxed); + self.update_operation_profile(|profile| { + profile.sqlite_requested_pages = profile + .sqlite_requested_pages + .saturating_add(target_pgnos.len() as u64); + }); if let Some(metrics) = &self.metrics { metrics.record_resolve_pages(target_pgnos.len() as u64); } @@ -1424,6 +1699,7 @@ impl VfsContext { let mut resolved = HashMap::new(); let mut missing = Vec::new(); let mut seen = HashSet::new(); + let mut prefetch_consumed = 0_u64; { let state = self.state.read(); @@ -1441,13 +1717,21 @@ impl VfsContext { resolved.insert(pgno, Some(bytes.clone())); continue; } - if let Some(bytes) = state.cached_page(&self.config, pgno) { + if let Some((bytes, was_prefetched)) = state.cached_page(&self.config, pgno) { + prefetch_consumed = prefetch_consumed.saturating_add(u64::from(was_prefetched)); resolved.insert(pgno, Some(bytes)); continue; } missing.push(pgno); } } + if prefetch_consumed > 0 { + self.update_operation_profile(|profile| { + profile.prefetch_consumed_pages = profile + .prefetch_consumed_pages + .saturating_add(prefetch_consumed); + }); + } if missing.is_empty() { self.resolve_pages_cache_hits @@ -1467,6 +1751,11 @@ impl VfsContext { if let Some(metrics) = &self.metrics { metrics.record_resolve_cache_hits(target_pgnos.len() as u64); } + self.update_operation_profile(|profile| { + profile.cache_hit_pages = profile + .cache_hit_pages + .saturating_add(target_pgnos.len() as u64); + }); return Ok(resolved); } self.resolve_pages_cache_hits @@ -1475,6 +1764,14 @@ impl VfsContext { metrics.record_resolve_cache_hits((seen.len() - missing.len()) as u64); metrics.record_resolve_cache_misses(missing.len() as u64); } + self.update_operation_profile(|profile| { + profile.cache_hit_pages = profile + .cache_hit_pages + .saturating_add((seen.len() - missing.len()) as u64); + profile.cache_miss_pages = profile + .cache_miss_pages + .saturating_add(missing.len() as u64); + }); let ( to_fetch, @@ -1588,21 +1885,95 @@ impl VfsContext { expected_generation: None, expected_head_txid, })); + let get_pages_duration_ns = get_pages_start.elapsed().as_nanos() as u64; if let Some(metrics) = &self.metrics { - metrics.observe_get_pages_duration(get_pages_start.elapsed().as_nanos() as u64); - } - let response = response.map_err(|err| GetPagesError::Other(err.to_string()))?; + metrics.observe_get_pages_duration(get_pages_duration_ns); + } + let prefetch_count = to_fetch.len().saturating_sub(missing.len()) as u64; + let profile_active = self.operation_profile_active.load(Ordering::Relaxed); + let ordinal = if profile_active { + self.operation_profile + .lock() + .as_ref() + .map_or(1, |profile| profile.get_pages_round_trips.saturating_add(1)) + } else { + 1 + }; + let response = match response { + Ok(response) => response, + Err(err) => { + self.update_operation_profile(|profile| { + profile.push_get_pages( + SqliteGetPagesProfile { + ordinal, + duration_ns: get_pages_duration_ns, + demand_requested: missing.len() as u64, + prefetch_requested: prefetch_count, + ..Default::default() + }, + self.profile_request_limit(), + ); + }); + return Err(GetPagesError::Other(err.to_string())); + } + }; match response { protocol::SqliteGetPagesResponse::SqliteGetPagesOk(ok) => { + if profile_active { + let requested = to_fetch.iter().copied().collect::>(); + let mut request_profile = SqliteGetPagesProfile { + ordinal, + duration_ns: get_pages_duration_ns, + demand_requested: missing.len() as u64, + prefetch_requested: prefetch_count, + success: true, + ..Default::default() + }; + let mut btree_pages = 0_u64; + let mut non_btree_pages = 0_u64; + for fetched in &ok.pages { + if !requested.contains(&fetched.pgno) { + request_profile.overflow_expansion_extra = + request_profile.overflow_expansion_extra.saturating_add(1); + } + if let Some(bytes) = &fetched.bytes { + request_profile.response_present = + request_profile.response_present.saturating_add(1); + request_profile.response_bytes = request_profile + .response_bytes + .saturating_add(bytes.len() as u64); + match classify(fetched.pgno, bytes) { + PageClass::Btree => btree_pages = btree_pages.saturating_add(1), + PageClass::Overflow => { + non_btree_pages = non_btree_pages.saturating_add(1) + } + } + } else { + request_profile.response_missing = + request_profile.response_missing.saturating_add(1); + } + } + self.update_operation_profile(|profile| { + profile.btree_pages = profile.btree_pages.saturating_add(btree_pages); + profile.non_btree_pages = + profile.non_btree_pages.saturating_add(non_btree_pages); + profile.push_get_pages(request_profile, self.profile_request_limit()); + }); + } let response_head_txid = ok.head_txid; if let Some(head_txid) = response_head_txid { self.state.write().head_txid = Some(head_txid); } let missing_pages = missing.iter().copied().collect::>(); - let (page_cache, protected_page_cache) = { + let (page_cache, protected_page_cache, prefetched_pages, prefetch_unused_total) = { let state = self.state.read(); - (state.page_cache.clone(), state.protected_page_cache.clone()) + ( + state.page_cache.clone(), + state.protected_page_cache.clone(), + state.prefetched_pages.clone(), + state.prefetch_unused_total.clone(), + ) }; #[cfg(debug_assertions)] let mut returned_pgnos = HashSet::new(); @@ -1641,6 +2012,9 @@ impl VfsContext { &self.config, &page_cache, &protected_page_cache, + &prefetched_pages, + &prefetch_unused_total, + profile_active, kind, fetched.pgno, bytes.clone(), @@ -1683,6 +2057,18 @@ impl VfsContext { Ok(resolved) } protocol::SqliteGetPagesResponse::SqliteErrorResponse(error) => { + self.update_operation_profile(|profile| { + profile.push_get_pages( + SqliteGetPagesProfile { + ordinal, + duration_ns: get_pages_duration_ns, + demand_requested: missing.len() as u64, + prefetch_requested: prefetch_count, + ..Default::default() + }, + self.profile_request_limit(), + ); + }); if self.commit_total.load(Relaxed) == 0 && missing.contains(&1) && is_initial_main_page_missing(&error.message) @@ -1749,6 +2135,18 @@ impl VfsContext { } }; let request_build_ns = request_build_start.elapsed().as_nanos() as u64; + self.update_operation_profile(|profile| { + profile.dirty_pages = profile + .dirty_pages + .saturating_add(request.dirty_pages.len() as u64); + profile.dirty_bytes = profile.dirty_bytes.saturating_add( + request + .dirty_pages + .iter() + .map(|page| page.bytes.len() as u64) + .sum::(), + ); + }); let (outcome, transport_metrics) = // Transport rejection, including envoy shutdown while a VFS callback is @@ -1863,6 +2261,18 @@ impl VfsContext { } }; let request_build_ns = request_build_start.elapsed().as_nanos() as u64; + self.update_operation_profile(|profile| { + profile.dirty_pages = profile + .dirty_pages + .saturating_add(request.dirty_pages.len() as u64); + profile.dirty_bytes = profile.dirty_bytes.saturating_add( + request + .dirty_pages + .iter() + .map(|page| page.bytes.len() as u64) + .sum::(), + ); + }); let (outcome, transport_metrics) = match self.block_on_buffered_commit(request.clone(), timeout) { @@ -3309,6 +3719,14 @@ impl NativeDatabase { self._vfs.ctx.round_trip_counts() } + pub(crate) fn begin_operation_profile(&self) -> SqliteOperationProfileGuard<'_> { + self._vfs.ctx.begin_operation_profile(); + SqliteOperationProfileGuard { + ctx: &self._vfs.ctx, + active: true, + } + } + pub fn snapshot_preload_hints(&self) -> VfsPreloadHintSnapshot { self._vfs.snapshot_preload_hints() } diff --git a/engine/packages/depot-client/src/worker.rs b/engine/packages/depot-client/src/worker.rs index 492b484115..17d448beff 100644 --- a/engine/packages/depot-client/src/worker.rs +++ b/engine/packages/depot-client/src/worker.rs @@ -16,10 +16,14 @@ use parking_lot::Mutex; use tokio::sync::{Notify, oneshot}; use crate::{ - query::{BindParam, ExecuteResult, QueryResult, exec_statements, execute_single_statement}, + query::{ + BindParam, ExecuteResult, QueryResult, exec_statements, exec_statements_profiled, + execute_single_statement, execute_single_statement_profiled, + }, vfs::{ - NativeConnection, NativeVfsHandle, SqliteRoundTripCounts, SqliteVfsMetrics, - configure_connection_for_database, open_connection, verify_batch_atomic_writes, + NativeConnection, NativeVfsHandle, SqliteOperationProfile, SqliteRoundTripCounts, + SqliteVfsMetrics, configure_connection_for_database, open_connection, + verify_batch_atomic_writes, }, }; @@ -56,11 +60,13 @@ enum SqliteCommand { Execute { sql: String, params: Option>, - reply: oneshot::Sender>, + enqueued_at: Option, + reply: oneshot::Sender>>, }, Exec { sql: String, - reply: oneshot::Sender>, + enqueued_at: Option, + reply: oneshot::Sender>>, }, #[cfg(test)] Pause { @@ -71,8 +77,33 @@ enum SqliteCommand { Panic, } +#[derive(Debug)] +pub struct SqliteWorkerResult { + pub result: Result, + pub profile: SqliteOperationProfile, +} + struct CloseRequest; +struct WorkerInflightGuard<'a>(Option<&'a dyn SqliteVfsMetrics>); + +impl<'a> WorkerInflightGuard<'a> { + fn new(metrics: Option<&'a dyn SqliteVfsMetrics>) -> Self { + if let Some(metrics) = metrics { + metrics.set_worker_inflight(true); + } + Self(metrics) + } +} + +impl Drop for WorkerInflightGuard<'_> { + fn drop(&mut self) { + if let Some(metrics) = self.0 { + metrics.set_worker_inflight(false); + } + } +} + /// Tracks an in-progress SQLite transaction so its wall-clock duration and the /// network round trips it issued can be reported once it commits or rolls back. /// @@ -156,8 +187,24 @@ impl SqliteWorkerHandle { } pub async fn exec(&self, sql: String) -> Result { + self.exec_inner(sql, None).await?.result + } + + pub async fn exec_profiled(&self, sql: String) -> Result> { + self.exec_inner(sql, Some(Instant::now())).await + } + + async fn exec_inner( + &self, + sql: String, + enqueued_at: Option, + ) -> Result> { let (reply, result) = oneshot::channel(); - self.enqueue(SqliteCommand::Exec { sql, reply })?; + self.enqueue(SqliteCommand::Exec { + sql, + enqueued_at, + reply, + })?; result.await.map_err(|_| sqlite_worker_dead_error())? } @@ -166,8 +213,30 @@ impl SqliteWorkerHandle { sql: String, params: Option>, ) -> Result { + self.execute_inner(sql, params, None).await?.result + } + + pub async fn execute_profiled( + &self, + sql: String, + params: Option>, + ) -> Result> { + self.execute_inner(sql, params, Some(Instant::now())).await + } + + async fn execute_inner( + &self, + sql: String, + params: Option>, + enqueued_at: Option, + ) -> Result> { let (reply, result) = oneshot::channel(); - self.enqueue(SqliteCommand::Execute { sql, params, reply })?; + self.enqueue(SqliteCommand::Execute { + sql, + params, + enqueued_at, + reply, + })?; result.await.map_err(|_| sqlite_worker_dead_error())? } @@ -408,6 +477,7 @@ fn worker_main(mut ctx: WorkerContext) { fail_queued_sql(&ctx.sql_rx); break; } + let _inflight = WorkerInflightGuard::new(ctx.inner.metrics.as_deref()); run_command( &mut db, command, @@ -447,7 +517,12 @@ fn run_command( ) { let start = Instant::now(); match command { - SqliteCommand::Execute { sql, params, reply } => { + SqliteCommand::Execute { + sql, + params, + enqueued_at, + reply, + } => { if reply.is_closed() { return; } @@ -457,29 +532,95 @@ fn run_command( // behind (a BEGIN flips autocommit off, a COMMIT flips it back on). let in_tx = command_in_tx(db); let stmt_kind = classify_statement(&sql); - let result = execute_single_statement(db.as_ptr(), &sql, params.as_deref()); - record_command_metrics( - metrics, - "execute", - in_tx, - stmt_kind, - &result, - start.elapsed(), - ); + let worker_result = if let Some(enqueued_at) = enqueued_at { + let worker_wait_ns = enqueued_at.elapsed().as_nanos() as u64; + let operation_profile = db.begin_operation_profile(); + let execution_start = Instant::now(); + let result = + execute_single_statement_profiled(db.as_ptr(), &sql, params.as_deref()); + let execution_ns = execution_start.elapsed().as_nanos() as u64; + let mut profile = operation_profile.finish(); + profile.worker_wait_ns = worker_wait_ns; + profile.sqlite_execution_ns = execution_ns.saturating_sub(profile.storage_ns); + if let Ok((_, query_profile)) = &result { + profile.bind_count = query_profile.bind_count; + profile.bind_logical_bytes = query_profile.bind_logical_bytes; + profile.result_rows = query_profile.result_rows; + profile.result_columns = query_profile.result_columns; + profile.result_logical_bytes = query_profile.result_logical_bytes; + } + record_command_metrics( + metrics, + "execute", + in_tx, + stmt_kind, + &result, + start.elapsed(), + ); + SqliteWorkerResult { + result: result.map(|(value, _)| value), + profile, + } + } else { + let result = execute_single_statement(db.as_ptr(), &sql, params.as_deref()); + record_command_metrics( + metrics, + "execute", + in_tx, + stmt_kind, + &result, + start.elapsed(), + ); + SqliteWorkerResult { + result, + profile: SqliteOperationProfile::default(), + } + }; finalize_transaction_if_complete(db, metrics, file_name, transaction); - let _ = reply.send(result); + let _ = reply.send(Ok(worker_result)); } - SqliteCommand::Exec { sql, reply } => { + SqliteCommand::Exec { + sql, + enqueued_at, + reply, + } => { if reply.is_closed() { return; } begin_transaction_if_needed(db, transaction); let in_tx = command_in_tx(db); let stmt_kind = classify_statement(&sql); - let result = exec_statements(db.as_ptr(), &sql); - record_command_metrics(metrics, "exec", in_tx, stmt_kind, &result, start.elapsed()); + let worker_result = if let Some(enqueued_at) = enqueued_at { + let worker_wait_ns = enqueued_at.elapsed().as_nanos() as u64; + let operation_profile = db.begin_operation_profile(); + let execution_start = Instant::now(); + let result = exec_statements_profiled(db.as_ptr(), &sql); + let execution_ns = execution_start.elapsed().as_nanos() as u64; + let mut profile = operation_profile.finish(); + profile.worker_wait_ns = worker_wait_ns; + profile.sqlite_execution_ns = execution_ns.saturating_sub(profile.storage_ns); + if let Ok((_, query_profile)) = &result { + profile.bind_count = query_profile.bind_count; + profile.bind_logical_bytes = query_profile.bind_logical_bytes; + profile.result_rows = query_profile.result_rows; + profile.result_columns = query_profile.result_columns; + profile.result_logical_bytes = query_profile.result_logical_bytes; + } + record_command_metrics(metrics, "exec", in_tx, stmt_kind, &result, start.elapsed()); + SqliteWorkerResult { + result: result.map(|(value, _)| value), + profile, + } + } else { + let result = exec_statements(db.as_ptr(), &sql); + record_command_metrics(metrics, "exec", in_tx, stmt_kind, &result, start.elapsed()); + SqliteWorkerResult { + result, + profile: SqliteOperationProfile::default(), + } + }; finalize_transaction_if_complete(db, metrics, file_name, transaction); - let _ = reply.send(result); + let _ = reply.send(Ok(worker_result)); } #[cfg(test)] SqliteCommand::Pause { entered, resume } => { diff --git a/engine/packages/depot-client/tests/inline/vfs.rs b/engine/packages/depot-client/tests/inline/vfs.rs index 19f8bd2058..56e6202459 100644 --- a/engine/packages/depot-client/tests/inline/vfs.rs +++ b/engine/packages/depot-client/tests/inline/vfs.rs @@ -321,7 +321,7 @@ fn vfs_staging_cache_retains_only_speculative_pages() { staging_cache_ttl_ms: DEFAULT_VFS_STAGING_CACHE_TTL_MS, ..VfsConfig::default() }; - let mut state = VfsState::new(&config); + let mut state = VfsState::new(&config, true); state.cache_page( &config, @@ -358,6 +358,29 @@ fn vfs_staging_cache_retains_only_speculative_pages() { assert!(state.cached_page(&config, 4).is_none()); } +#[test] +fn disabled_profiling_skips_prefetch_tracking() { + let config = VfsConfig { + page_cache_mode: SqliteVfsPageCacheMode::All, + staging_cache_ttl_ms: DEFAULT_VFS_STAGING_CACHE_TTL_MS, + ..VfsConfig::default() + }; + let mut state = VfsState::new(&config, false); + + state.cache_page( + &config, + PageCacheInsertKind::Prefetch, + 2, + vec![2; DEFAULT_PAGE_SIZE], + ); + + assert!(state.prefetched_pages.is_empty()); + assert_eq!( + state.cached_page(&config, 2), + Some((vec![2; DEFAULT_PAGE_SIZE], false)) + ); +} + #[test] fn vfs_staging_cache_ttl_zero_disables_speculative_retention() { let config = VfsConfig { @@ -365,7 +388,7 @@ fn vfs_staging_cache_ttl_zero_disables_speculative_retention() { staging_cache_ttl_ms: 0, ..VfsConfig::default() }; - let mut state = VfsState::new(&config); + let mut state = VfsState::new(&config, true); state.cache_page( &config, @@ -578,10 +601,89 @@ struct WorkerTestMetrics { } impl SqliteVfsMetrics for WorkerTestMetrics { + fn profiling_enabled(&self) -> bool { + false + } + + fn max_profiled_get_pages_requests(&self) -> usize { + 0 + } + + fn observe_operation_profile(&self, _profile: &SqliteOperationMetric) -> bool { + false + } + + fn observe_transaction_profile(&self, _profile: &SqliteTransactionMetric) -> bool { + false + } + + fn emit_operation_diagnostic_event( + &self, + _actor_id: &str, + _generation: Option, + _profile: &SqliteOperationMetric, + ) { + } + + fn emit_transaction_diagnostic_event( + &self, + _actor_id: &str, + _generation: Option, + _profile: &SqliteTransactionMetric, + ) { + } + + fn record_fingerprint_catalog( + &self, + _operation_type: &'static str, + _fingerprint: &str, + _identity: &str, + _format_version: u8, + ) { + } + + fn record_resolve_pages(&self, _requested_pages: u64) {} + + fn record_resolve_cache_hits(&self, _pages: u64) {} + + fn record_resolve_cache_misses(&self, _pages: u64) {} + + fn record_get_pages_request(&self, _pages: u64, _prefetch_pages: u64, _page_size: u64) {} + + fn observe_get_pages_duration(&self, _duration_ns: u64) {} + + fn observe_open_phase( + &self, + _phase: SqliteOpenPhase, + _outcome: &'static str, + _duration_ns: u64, + ) { + } + + fn record_startup_preload_pages(&self, _kind: &'static str, _pages: u64) {} + + fn record_commit(&self) {} + + fn observe_commit_phases( + &self, + _request_build_ns: u64, + _serialize_ns: u64, + _transport_ns: u64, + _state_update_ns: u64, + _total_ns: u64, + ) { + } + fn set_worker_queue_depth(&self, depth: u64) { self.queue_depth.store(depth, Ordering::Release); } + fn set_worker_active(&self, _active: bool) {} + + fn set_worker_inflight(&self, _active: bool) {} + + fn set_coordinator_queue_depth(&self, _depth: u64) {} + fn record_worker_queue_overload(&self) { self.overloads.fetch_add(1, Ordering::AcqRel); } @@ -596,6 +698,13 @@ impl SqliteVfsMetrics for WorkerTestMetrics { self.command_durations.fetch_add(1, Ordering::AcqRel); } + fn observe_transaction_round_trips( + &self, + _get_pages_round_trips: u64, + _commit_round_trips: u64, + ) { + } + fn record_worker_command_error(&self, _operation: &'static str, _code: &'static str) { self.command_errors.fetch_add(1, Ordering::AcqRel); } @@ -604,6 +713,8 @@ impl SqliteVfsMetrics for WorkerTestMetrics { self.close_durations.fetch_add(1, Ordering::AcqRel); } + fn record_worker_close_timeout(&self) {} + fn record_worker_crash(&self) { self.crashes.fetch_add(1, Ordering::AcqRel); } diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs index 0abbe759c6..3460c708f3 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/config.rs @@ -53,6 +53,94 @@ pub struct ActionDefinition { pub name: String, } +/// Experimental SQLite profiling configuration. +/// +/// This entire configuration surface, including every field, is subject to +/// change without notice. +#[derive(Clone, Debug)] +pub struct SqliteProfilingConfig { + pub enabled: bool, + pub max_tracked_statement_fingerprints: usize, + pub max_tracked_transaction_fingerprints: usize, + pub max_prometheus_series: usize, + pub max_statements_per_transaction_trace: usize, + pub max_get_pages_requests_per_trace: usize, + pub max_transaction_name_bytes: usize, + pub slow_operation_threshold_ms: u64, + pub baseline_sample_rate: f64, + pub max_diagnostic_events_per_minute: usize, + pub diagnostic_event_queue_capacity: usize, +} + +impl Default for SqliteProfilingConfig { + fn default() -> Self { + Self { + enabled: true, + max_tracked_statement_fingerprints: 128, + max_tracked_transaction_fingerprints: 8, + max_prometheus_series: 25_000, + max_statements_per_transaction_trace: 32, + max_get_pages_requests_per_trace: 16, + max_transaction_name_bytes: 128, + slow_operation_threshold_ms: 10, + baseline_sample_rate: 0.001, + max_diagnostic_events_per_minute: 120, + diagnostic_event_queue_capacity: 256, + } + } +} + +/// Sparse experimental SQLite profiling configuration used at runtime +/// boundaries. +/// +/// This entire configuration surface, including every field, is subject to +/// change without notice. +#[derive(Clone, Debug, Default)] +pub struct SqliteProfilingConfigInput { + pub enabled: Option, + pub max_tracked_statement_fingerprints: Option, + pub max_tracked_transaction_fingerprints: Option, + pub max_prometheus_series: Option, + pub max_statements_per_transaction_trace: Option, + pub max_get_pages_requests_per_trace: Option, + pub max_transaction_name_bytes: Option, + pub slow_operation_threshold_ms: Option, + pub baseline_sample_rate: Option, + pub max_diagnostic_events_per_minute: Option, + pub diagnostic_event_queue_capacity: Option, +} + +impl SqliteProfilingConfig { + fn from_input(input: SqliteProfilingConfigInput) -> Self { + let mut config = Self::default(); + macro_rules! set_usize { + ($field:ident) => { + if let Some(value) = input.$field { + config.$field = value as usize; + } + }; + } + if let Some(value) = input.enabled { + config.enabled = value; + } + set_usize!(max_tracked_statement_fingerprints); + set_usize!(max_tracked_transaction_fingerprints); + set_usize!(max_prometheus_series); + set_usize!(max_statements_per_transaction_trace); + set_usize!(max_get_pages_requests_per_trace); + set_usize!(max_transaction_name_bytes); + if let Some(value) = input.slow_operation_threshold_ms { + config.slow_operation_threshold_ms = u64::from(value); + } + if let Some(value) = input.baseline_sample_rate { + config.baseline_sample_rate = value; + } + set_usize!(max_diagnostic_events_per_minute); + set_usize!(diagnostic_event_queue_capacity); + config + } +} + #[derive(Clone, Debug)] pub struct ActorConfig { pub name: Option, @@ -61,6 +149,7 @@ pub struct ActorConfig { /// on the TS side). Gates the inspector database tab. pub has_database: bool, pub remote_sqlite: bool, + pub sqlite_profiling: SqliteProfilingConfig, /// Whether the user declared actor state (`state: ...` or `createState`). /// Gates the inspector state tab and state-subscription messages. pub has_state: bool, @@ -98,6 +187,7 @@ pub struct ActorConfigInput { pub icon: Option, pub has_database: Option, pub remote_sqlite: Option, + pub sqlite_profiling: Option, pub has_state: Option, pub can_hibernate_websocket: Option, pub state_save_interval_ms: Option, @@ -129,6 +219,10 @@ impl ActorConfig { icon: config.icon, has_database: config.has_database.unwrap_or(false), remote_sqlite: config.remote_sqlite.unwrap_or(false), + sqlite_profiling: config + .sqlite_profiling + .map(SqliteProfilingConfig::from_input) + .unwrap_or_default(), has_state: config.has_state.unwrap_or(false), ..Self::default() }; @@ -215,6 +309,23 @@ impl ActorConfig { /// config so the actor never starts with garbage state. pub fn validate(&self) -> anyhow::Result<()> { crate::inspector::validate_inspector_tabs(&self.inspector_tabs)?; + anyhow::ensure!( + self.sqlite_profiling.baseline_sample_rate.is_finite() + && (0.0..=1.0).contains(&self.sqlite_profiling.baseline_sample_rate), + "SQLite profiling baselineSampleRate must be between 0 and 1" + ); + anyhow::ensure!( + self.sqlite_profiling.max_get_pages_requests_per_trace <= 16, + "SQLite profiling maxGetPagesRequestsPerTrace must be at most 16" + ); + anyhow::ensure!( + self.sqlite_profiling.max_statements_per_transaction_trace <= 32, + "SQLite profiling maxStatementsPerTransactionTrace must be at most 32" + ); + anyhow::ensure!( + self.sqlite_profiling.max_transaction_name_bytes > 0, + "SQLite profiling maxTransactionNameBytes must be greater than zero" + ); Ok(()) } } @@ -226,6 +337,7 @@ impl Default for ActorConfig { icon: None, has_database: false, remote_sqlite: false, + sqlite_profiling: SqliteProfilingConfig::default(), has_state: false, can_hibernate_websocket: CanHibernateWebSocket::default(), state_save_interval: DEFAULT_STATE_SAVE_INTERVAL, diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index 89e245de11..1646e9535d 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -230,10 +230,13 @@ impl ActorContext { kv: Kv, sql: SqliteDb, ) -> Self { - let metrics = ActorMetrics::new(name.clone()); #[cfg(feature = "sqlite-local")] let mut sql = sql; #[cfg(feature = "sqlite-local")] + sql.set_profiling_config(config.sqlite_profiling.clone()); + let metrics = + ActorMetrics::new_with_sqlite_profiling(name.clone(), config.sqlite_profiling.clone()); + #[cfg(feature = "sqlite-local")] sql.set_vfs_metrics(Arc::new(metrics.clone())); let diagnostics = ActorDiagnostics::new(actor_id.clone()); let state_save_interval = config.state_save_interval; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs index 167e4cbdf7..97892f8734 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/metrics.rs @@ -1,13 +1,21 @@ use std::collections::BTreeMap; use std::fmt; use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +#[cfg(feature = "sqlite-local")] +use std::sync::atomic::{AtomicU64, AtomicUsize}; use std::sync::{Arc, LazyLock}; +#[cfg(feature = "sqlite-local")] +use std::sync::{OnceLock, mpsc}; use std::time::Duration; +#[cfg(feature = "sqlite-local")] +use std::time::{SystemTime, UNIX_EPOCH}; use parking_lot::Mutex; use rivet_metrics::prometheus::{ CounterVec, HistogramOpts, HistogramVec, IntCounterVec, IntGaugeVec, Opts, Registry, }; +#[cfg(feature = "sqlite-local")] +use rivet_metrics::prometheus::{Histogram, IntCounter, IntGauge}; use crate::actor::task_types::{ShutdownKind, StateMutationReason, UserTaskKind}; use crate::time::Instant; @@ -124,6 +132,10 @@ pub(crate) struct StartupTimer { #[derive(Debug)] struct ActorMetricInner { labels: ActorMetricLabels, + #[cfg(feature = "sqlite-local")] + sqlite_profiling: crate::SqliteProfilingConfig, + #[cfg(feature = "sqlite-local")] + sqlite_profile_low_card_handles: OnceLock>>, state: Mutex, active: AtomicBool, startup_is_new: AtomicU8, @@ -231,6 +243,745 @@ struct ActorMetricCollectors { static METRICS: LazyLock = LazyLock::new(ActorMetricCollectors::new); +#[cfg(feature = "sqlite-local")] +struct SqliteProfileCollectors { + duration_seconds: HistogramVec, + phase_duration_seconds: HistogramVec, + get_pages_round_trips: HistogramVec, + transaction_statement_count: HistogramVec, + outcome_total: IntCounterVec, + local_pages_total: IntCounterVec, + local_bytes_total: IntCounterVec, + get_pages_duration_seconds: HistogramVec, + get_pages_pages: HistogramVec, + get_pages_response_bytes: HistogramVec, + get_pages_missing_pages_total: IntCounterVec, + fingerprint_overflow_total: IntCounterVec, + event_dropped_total: IntCounterVec, + worker_queue_depth: IntGaugeVec, + worker_inflight: IntGaugeVec, + coordinator_queue_depth: IntGaugeVec, +} + +#[cfg(feature = "sqlite-local")] +static SQLITE_PROFILE_METRICS: LazyLock = + LazyLock::new(SqliteProfileCollectors::new); + +#[cfg(feature = "sqlite-local")] +const SQLITE_PROFILE_PAGE_KINDS: [&str; 12] = [ + "sqlite_requested", + "cache_hit", + "cache_miss", + "depot_demand_requested", + "vfs_prefetch_requested", + "response_present", + "response_missing", + "overflow_expansion_extra", + "btree", + "non_btree", + "prefetch_consumed", + "prefetch_unused", +]; +#[cfg(feature = "sqlite-local")] +const SQLITE_PROFILE_BYTE_KINDS: [&str; 4] = [ + "bind_logical", + "result_logical", + "storage_response", + "dirty", +]; +#[cfg(feature = "sqlite-local")] +const SQLITE_PROFILE_REQUEST_ORDINALS: [&str; 6] = ["1", "2", "3", "4", "5-8", "9+"]; +#[cfg(feature = "sqlite-local")] +const SQLITE_PROFILE_REQUEST_PAGE_KINDS: [&str; 4] = [ + "demand_requested", + "prefetch_requested", + "response_present", + "overflow_expansion_extra", +]; +#[cfg(feature = "sqlite-local")] +const SQLITE_PROFILE_OUTCOMES: [&str; 7] = [ + "success", + "error", + "rollback", + "timeout", + "expired", + "connection_lost", + "cancelled", +]; + +#[cfg(feature = "sqlite-local")] +struct SqliteFingerprintMetricHandles { + duration: [Histogram; 2], + transaction_wait: Histogram, + worker_wait: Histogram, + storage: Histogram, + local_work: Histogram, + application_time: Option, + commit: Option, + get_pages_round_trips: Histogram, + transaction_statement_count: Option, + outcomes: [IntCounter; SQLITE_PROFILE_OUTCOMES.len()], +} + +#[cfg(feature = "sqlite-local")] +impl SqliteFingerprintMetricHandles { + fn new( + actor_name: &str, + operation_type: &'static str, + fingerprint: &str, + fingerprint_source: &'static str, + transaction_mode: &'static str, + storage_transport: &'static str, + ) -> Self { + let base = [ + actor_name, + operation_type, + fingerprint, + fingerprint_source, + transaction_mode, + storage_transport, + ]; + let phase = |name| { + SQLITE_PROFILE_METRICS + .phase_duration_seconds + .with_label_values(&[base[0], base[1], base[2], base[3], base[4], base[5], name]) + }; + let is_transaction = operation_type == "transaction"; + Self { + duration: std::array::from_fn(|index| { + SQLITE_PROFILE_METRICS.duration_seconds.with_label_values(&[ + base[0], + base[1], + base[2], + base[3], + base[4], + base[5], + if index == 0 { "success" } else { "non_success" }, + ]) + }), + transaction_wait: phase("transaction_wait"), + worker_wait: phase("worker_wait"), + storage: phase("storage"), + local_work: phase("local_work"), + application_time: is_transaction.then(|| phase("application_time")), + commit: is_transaction.then(|| phase("commit")), + get_pages_round_trips: SQLITE_PROFILE_METRICS + .get_pages_round_trips + .with_label_values(&base), + transaction_statement_count: is_transaction.then(|| { + SQLITE_PROFILE_METRICS + .transaction_statement_count + .with_label_values(&base) + }), + outcomes: std::array::from_fn(|index| { + SQLITE_PROFILE_METRICS.outcome_total.with_label_values(&[ + base[0], + base[1], + base[2], + base[3], + base[4], + base[5], + SQLITE_PROFILE_OUTCOMES[index], + ]) + }), + } + } + + fn observe_duration(&self, outcome: &str, duration: f64) { + self.duration[usize::from(outcome != "success")].observe(duration); + } + + fn observe_phase(&self, phase: &str, duration: f64) { + let handle = match phase { + "transaction_wait" => Some(&self.transaction_wait), + "worker_wait" => Some(&self.worker_wait), + "storage" => Some(&self.storage), + "local_work" => Some(&self.local_work), + "application_time" => self.application_time.as_ref(), + "commit" => self.commit.as_ref(), + _ => None, + }; + if let Some(handle) = handle { + handle.observe(duration); + } + } + + fn record_outcome(&self, outcome: &str) { + let index = match outcome { + "success" => 0, + "error" => 1, + "rollback" => 2, + "timeout" => 3, + "expired" => 4, + "connection_lost" => 5, + "cancelled" => 6, + _ => 1, + }; + self.outcomes[index].inc(); + } +} + +#[cfg(feature = "sqlite-local")] +struct SqliteRequestMetricHandles { + duration: [Histogram; 2], + pages: [Histogram; SQLITE_PROFILE_REQUEST_PAGE_KINDS.len()], + response_bytes: Histogram, + missing_pages: IntCounter, +} + +#[cfg(feature = "sqlite-local")] +struct SqliteLowCardMetricHandles { + local_pages: [[IntCounter; SQLITE_PROFILE_PAGE_KINDS.len()]; 2], + local_bytes: [[IntCounter; SQLITE_PROFILE_BYTE_KINDS.len()]; 2], + requests: [SqliteRequestMetricHandles; SQLITE_PROFILE_REQUEST_ORDINALS.len()], + event_dropped: [IntCounter; 2], + worker_queue_depth: IntGauge, + worker_inflight: IntGauge, + coordinator_queue_depth: IntGauge, +} + +#[cfg(feature = "sqlite-local")] +impl fmt::Debug for SqliteLowCardMetricHandles { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SqliteLowCardMetricHandles") + .finish_non_exhaustive() + } +} + +#[cfg(feature = "sqlite-local")] +impl SqliteLowCardMetricHandles { + fn new(actor_name: &str, storage_transport: &'static str) -> Self { + let operation_types = ["statement", "transaction"]; + Self { + local_pages: std::array::from_fn(|operation_index| { + std::array::from_fn(|kind_index| { + SQLITE_PROFILE_METRICS + .local_pages_total + .with_label_values(&[ + actor_name, + operation_types[operation_index], + SQLITE_PROFILE_PAGE_KINDS[kind_index], + storage_transport, + ]) + }) + }), + local_bytes: std::array::from_fn(|operation_index| { + std::array::from_fn(|kind_index| { + SQLITE_PROFILE_METRICS + .local_bytes_total + .with_label_values(&[ + actor_name, + operation_types[operation_index], + SQLITE_PROFILE_BYTE_KINDS[kind_index], + storage_transport, + ]) + }) + }), + requests: std::array::from_fn(|ordinal_index| { + let ordinal = SQLITE_PROFILE_REQUEST_ORDINALS[ordinal_index]; + SqliteRequestMetricHandles { + duration: std::array::from_fn(|outcome_index| { + SQLITE_PROFILE_METRICS + .get_pages_duration_seconds + .with_label_values(&[ + actor_name, + ordinal, + if outcome_index == 0 { + "success" + } else { + "non_success" + }, + storage_transport, + ]) + }), + pages: std::array::from_fn(|kind_index| { + SQLITE_PROFILE_METRICS.get_pages_pages.with_label_values(&[ + actor_name, + ordinal, + SQLITE_PROFILE_REQUEST_PAGE_KINDS[kind_index], + storage_transport, + ]) + }), + response_bytes: SQLITE_PROFILE_METRICS + .get_pages_response_bytes + .with_label_values(&[actor_name, ordinal, storage_transport]), + missing_pages: SQLITE_PROFILE_METRICS + .get_pages_missing_pages_total + .with_label_values(&[actor_name, ordinal, storage_transport]), + } + }), + event_dropped: std::array::from_fn(|index| { + SQLITE_PROFILE_METRICS + .event_dropped_total + .with_label_values(&[ + actor_name, + if index == 0 { + "rate_limit" + } else { + "backpressure" + }, + ]) + }), + worker_queue_depth: SQLITE_PROFILE_METRICS + .worker_queue_depth + .with_label_values(&[actor_name]), + worker_inflight: SQLITE_PROFILE_METRICS + .worker_inflight + .with_label_values(&[actor_name]), + coordinator_queue_depth: SQLITE_PROFILE_METRICS + .coordinator_queue_depth + .with_label_values(&[actor_name]), + } + } +} + +#[cfg(feature = "sqlite-local")] +struct AdmittedSqliteProfile<'a> { + fingerprint: String, + fingerprint_handles: Arc, + low_card_handles: &'a SqliteLowCardMetricHandles, +} + +#[cfg(feature = "sqlite-local")] +struct SqliteProfileAdmission { + statements: scc::HashSet, + transactions: scc::HashSet, + tuples: scc::HashMap>, + low_card_tuples: scc::HashMap>, + candidates: scc::HashMap, + statement_count: AtomicUsize, + transaction_count: AtomicUsize, + series: AtomicUsize, +} + +#[cfg(feature = "sqlite-local")] +static SQLITE_PROFILE_ADMISSION: LazyLock = + LazyLock::new(|| SqliteProfileAdmission { + statements: scc::HashSet::new(), + transactions: scc::HashSet::new(), + tuples: scc::HashMap::new(), + low_card_tuples: scc::HashMap::new(), + candidates: scc::HashMap::new(), + statement_count: AtomicUsize::new(0), + transaction_count: AtomicUsize::new(0), + series: AtomicUsize::new(0), + }); + +#[cfg(feature = "sqlite-local")] +#[derive(Debug)] +enum SqliteDiagnosticEvent { + Operation { + invocation_id: u64, + actor_id: String, + generation: Option, + actor_name: String, + profile: depot_client::vfs::SqliteOperationMetric, + }, + Transaction { + invocation_id: u64, + actor_id: String, + generation: Option, + actor_name: String, + profile: depot_client::vfs::SqliteTransactionMetric, + }, +} + +#[cfg(feature = "sqlite-local")] +static SQLITE_DIAGNOSTIC_SENDER: OnceLock>> = + OnceLock::new(); +#[cfg(feature = "sqlite-local")] +static SQLITE_DIAGNOSTIC_SAMPLE_SEQUENCE: AtomicU64 = AtomicU64::new(1); +#[cfg(feature = "sqlite-local")] +static SQLITE_DIAGNOSTIC_INVOCATION_ID: AtomicU64 = AtomicU64::new(1); +#[cfg(feature = "sqlite-local")] +static SQLITE_DIAGNOSTIC_RATE_STATE: AtomicU64 = AtomicU64::new(0); +#[cfg(feature = "sqlite-local")] +static SQLITE_PROFILE_CAPACITY_WARNINGS: AtomicU64 = AtomicU64::new(0); + +#[cfg(feature = "sqlite-local")] +fn sqlite_diagnostic_sender( + capacity: usize, +) -> Option<&'static mpsc::SyncSender> { + SQLITE_DIAGNOSTIC_SENDER + .get_or_init(|| { + let (sender, receiver) = mpsc::sync_channel(capacity.max(1)); + match std::thread::Builder::new() + .name("rivetkit-sqlite-diagnostics".to_owned()) + .spawn(move || { + while let Ok(event) = receiver.recv() { + match event { + SqliteDiagnosticEvent::Operation { + invocation_id, + actor_id, + generation, + actor_name, + profile, + } => tracing::info!( + target: "rivetkit_sqlite_profile", + event_type = "operation", + invocation_id, + actor_id, + generation, + actor_name, + operation_type = profile.operation_type, + fingerprint = profile.fingerprint, + fingerprint_source = profile.fingerprint_source, + transaction_mode = profile.transaction_mode, + storage_transport = profile.storage_transport, + outcome = profile.outcome, + sql_bytes = profile.sql_bytes, + total_ns = profile.total_ns, + transaction_wait_ns = profile.transaction_wait_ns, + worker_wait_ns = profile.profile.worker_wait_ns, + storage_ns = profile.profile.storage_ns, + sqlite_execution_ns = profile.profile.sqlite_execution_ns, + bind_count = profile.profile.bind_count, + bind_logical_bytes = profile.profile.bind_logical_bytes, + result_rows = profile.profile.result_rows, + result_columns = profile.profile.result_columns, + result_logical_bytes = profile.profile.result_logical_bytes, + sqlite_requested_pages = profile.profile.sqlite_requested_pages, + cache_hit_pages = profile.profile.cache_hit_pages, + cache_miss_pages = profile.profile.cache_miss_pages, + response_present_pages = profile.profile.response_present_pages, + response_missing_pages = profile.profile.response_missing_pages, + overflow_expansion_extra_pages = profile.profile.overflow_expansion_extra_pages, + prefetch_consumed_pages = profile.profile.prefetch_consumed_pages, + prefetch_unused_pages = profile.profile.prefetch_unused_pages, + dirty_pages = profile.profile.dirty_pages, + dirty_bytes = profile.profile.dirty_bytes, + get_pages_requests = ?profile.profile.get_pages_requests, + omitted_get_pages_requests = profile.profile.omitted_get_pages_requests, + "sampled SQLite operation profile" + ), + SqliteDiagnosticEvent::Transaction { + invocation_id, + actor_id, + generation, + actor_name, + profile, + } => tracing::info!( + target: "rivetkit_sqlite_profile", + event_type = "transaction", + invocation_id, + actor_id, + generation, + actor_name, + fingerprint = profile.fingerprint, + fingerprint_source = profile.fingerprint_source, + shape_fingerprint = profile.shape_fingerprint, + statement_fingerprint_hashes = ?profile.statement_fingerprint_hashes, + omitted_statement_fingerprints = profile.omitted_statement_fingerprints, + storage_transport = profile.storage_transport, + outcome = profile.outcome, + total_ns = profile.total_ns, + transaction_wait_ns = profile.transaction_wait_ns, + worker_wait_ns = profile.worker_wait_ns, + storage_ns = profile.storage_ns, + local_work_ns = profile.local_work_ns, + application_time_ns = profile.application_time_ns, + commit_ns = profile.commit_ns, + statement_count = profile.statement_count, + dirty_pages = profile.dirty_pages, + dirty_bytes = profile.dirty_bytes, + "sampled SQLite transaction profile" + ), + } + } + }) { + Ok(_) => Some(sender), + Err(error) => { + tracing::error!(%error, "failed to start SQLite diagnostic exporter"); + None + } + } + }) + .as_ref() +} + +#[cfg(feature = "sqlite-local")] +fn sqlite_baseline_sample_selected(rate: f64) -> bool { + if rate <= 0.0 { + return false; + } + if rate >= 1.0 { + return true; + } + let sequence = SQLITE_DIAGNOSTIC_SAMPLE_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let mixed = sequence.wrapping_mul(0x9e37_79b9_7f4a_7c15).rotate_left(17); + (mixed as f64 / u64::MAX as f64) < rate +} + +#[cfg(feature = "sqlite-local")] +fn try_acquire_sqlite_diagnostic_rate(max_events_per_minute: usize) -> bool { + const COUNT_BITS: u32 = 20; + const COUNT_MASK: u64 = (1 << COUNT_BITS) - 1; + let limit = (max_events_per_minute as u64).min(COUNT_MASK); + if limit == 0 { + return false; + } + let minute = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + / 60; + loop { + let current = SQLITE_DIAGNOSTIC_RATE_STATE.load(Ordering::Acquire); + let current_minute = current >> COUNT_BITS; + let current_count = current & COUNT_MASK; + let next_count = if current_minute == minute { + if current_count >= limit { + return false; + } + current_count + 1 + } else { + 1 + }; + let next = (minute << COUNT_BITS) | next_count; + if SQLITE_DIAGNOSTIC_RATE_STATE + .compare_exchange_weak(current, next, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + return true; + } + } +} + +#[cfg(feature = "sqlite-local")] +impl SqliteProfileCollectors { + fn new() -> Self { + let duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_duration_seconds", + "complete SQLite operation wall duration in seconds", + ) + .buckets(sqlite_worker_duration_buckets()), + &[ + "actor_name", + "type", + "fingerprint", + "fingerprint_source", + "transaction_mode", + "storage_transport", + "outcome_class", + ], + ) + .expect("create sqlite profiling duration histogram"); + let phase_duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_phase_duration_seconds", + "SQLite operation phase duration in seconds", + ) + .buckets(sqlite_worker_duration_buckets()), + &[ + "actor_name", + "type", + "fingerprint", + "fingerprint_source", + "transaction_mode", + "storage_transport", + "phase", + ], + ) + .expect("create sqlite profiling phase histogram"); + let get_pages_round_trips = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_get_pages_round_trips", + "physical get_pages calls per SQLite operation", + ) + .buckets(sqlite_round_trip_count_buckets()), + &[ + "actor_name", + "type", + "fingerprint", + "fingerprint_source", + "transaction_mode", + "storage_transport", + ], + ) + .expect("create sqlite profiling get_pages histogram"); + let transaction_statement_count = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_transaction_statement_count", + "user statements per explicit SQLite transaction", + ) + .buckets(sqlite_round_trip_count_buckets()), + &[ + "actor_name", + "type", + "fingerprint", + "fingerprint_source", + "transaction_mode", + "storage_transport", + ], + ) + .expect("create sqlite transaction statement count histogram"); + let outcome_total = IntCounterVec::new( + Opts::new("rivetkit_sqlite_outcome_total", "SQLite operation outcomes"), + &[ + "actor_name", + "type", + "fingerprint", + "fingerprint_source", + "transaction_mode", + "storage_transport", + "outcome", + ], + ) + .expect("create sqlite outcome counter"); + let local_pages_total = IntCounterVec::new( + Opts::new( + "rivetkit_sqlite_local_pages_total", + "local SQLite page totals", + ), + &["actor_name", "type", "page_kind", "storage_transport"], + ) + .expect("create sqlite local pages counter"); + let local_bytes_total = IntCounterVec::new( + Opts::new( + "rivetkit_sqlite_local_bytes_total", + "local SQLite logical byte totals", + ), + &["actor_name", "type", "byte_kind", "storage_transport"], + ) + .expect("create sqlite local bytes counter"); + let get_pages_duration_seconds = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_get_pages_duration_seconds", + "physical get_pages request duration in seconds", + ) + .buckets(sqlite_worker_duration_buckets()), + &[ + "actor_name", + "request_ordinal", + "outcome_class", + "storage_transport", + ], + ) + .expect("create sqlite request duration histogram"); + let get_pages_pages = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_get_pages_pages", + "pages per physical get_pages request", + ) + .buckets(sqlite_round_trip_count_buckets()), + &[ + "actor_name", + "request_ordinal", + "page_kind", + "storage_transport", + ], + ) + .expect("create sqlite request pages histogram"); + let get_pages_response_bytes = HistogramVec::new( + HistogramOpts::new( + "rivetkit_sqlite_get_pages_response_bytes", + "bytes per physical get_pages response", + ) + .buckets(vec![ + 512.0, + 4096.0, + 16_384.0, + 65_536.0, + 262_144.0, + 1_048_576.0, + ]), + &["actor_name", "request_ordinal", "storage_transport"], + ) + .expect("create sqlite request response bytes histogram"); + let get_pages_missing_pages_total = IntCounterVec::new( + Opts::new( + "rivetkit_sqlite_get_pages_missing_pages_total", + "missing pages in physical get_pages responses", + ), + &["actor_name", "request_ordinal", "storage_transport"], + ) + .expect("create sqlite missing pages counter"); + let fingerprint_overflow_total = IntCounterVec::new( + Opts::new( + "rivetkit_sqlite_fingerprint_overflow_total", + "SQLite observations routed to the shared other fingerprint", + ), + &["actor_name", "type", "reason"], + ) + .expect("create sqlite fingerprint overflow counter"); + let event_dropped_total = IntCounterVec::new( + Opts::new( + "rivetkit_sqlite_event_dropped_total", + "SQLite diagnostic events dropped before export", + ), + &["actor_name", "reason"], + ) + .expect("create sqlite diagnostic event drop counter"); + let worker_queue_depth = IntGaugeVec::new( + Opts::new( + "rivetkit_sqlite_worker_queue_depth", + "queued native SQLite commands", + ), + ACTOR_LABELS, + ) + .expect("create sqlite worker queue depth gauge"); + let worker_inflight = IntGaugeVec::new( + Opts::new( + "rivetkit_sqlite_worker_inflight", + "native SQLite commands executing", + ), + ACTOR_LABELS, + ) + .expect("create sqlite worker inflight gauge"); + let coordinator_queue_depth = IntGaugeVec::new( + Opts::new( + "rivetkit_sqlite_coordinator_queue_depth", + "operations waiting for SQLite transaction coordination", + ), + ACTOR_LABELS, + ) + .expect("create sqlite coordinator queue depth gauge"); + + register_metric(&rivet_metrics::REGISTRY, duration_seconds.clone()); + register_metric(&rivet_metrics::REGISTRY, phase_duration_seconds.clone()); + register_metric(&rivet_metrics::REGISTRY, get_pages_round_trips.clone()); + register_metric( + &rivet_metrics::REGISTRY, + transaction_statement_count.clone(), + ); + register_metric(&rivet_metrics::REGISTRY, outcome_total.clone()); + register_metric(&rivet_metrics::REGISTRY, local_pages_total.clone()); + register_metric(&rivet_metrics::REGISTRY, local_bytes_total.clone()); + register_metric(&rivet_metrics::REGISTRY, get_pages_duration_seconds.clone()); + register_metric(&rivet_metrics::REGISTRY, get_pages_pages.clone()); + register_metric(&rivet_metrics::REGISTRY, get_pages_response_bytes.clone()); + register_metric( + &rivet_metrics::REGISTRY, + get_pages_missing_pages_total.clone(), + ); + register_metric(&rivet_metrics::REGISTRY, fingerprint_overflow_total.clone()); + register_metric(&rivet_metrics::REGISTRY, event_dropped_total.clone()); + register_metric(&rivet_metrics::REGISTRY, worker_queue_depth.clone()); + register_metric(&rivet_metrics::REGISTRY, worker_inflight.clone()); + register_metric(&rivet_metrics::REGISTRY, coordinator_queue_depth.clone()); + + Self { + duration_seconds, + phase_duration_seconds, + get_pages_round_trips, + transaction_statement_count, + outcome_total, + local_pages_total, + local_bytes_total, + get_pages_duration_seconds, + get_pages_pages, + get_pages_response_bytes, + get_pages_missing_pages_total, + fingerprint_overflow_total, + event_dropped_total, + worker_queue_depth, + worker_inflight, + coordinator_queue_depth, + } + } +} + impl ActorMetricCollectors { fn new() -> Self { let actor_active_count = IntGaugeVec::new( @@ -833,9 +1584,20 @@ impl ActorMetricCollectors { impl ActorMetrics { pub(crate) fn new(actor_name: impl Into) -> Self { + Self::new_with_sqlite_profiling(actor_name, crate::SqliteProfilingConfig::default()) + } + + pub(crate) fn new_with_sqlite_profiling( + actor_name: impl Into, + _sqlite_profiling: crate::SqliteProfilingConfig, + ) -> Self { let labels = ActorMetricLabels { actor_name: actor_name.into(), }; + #[cfg(feature = "sqlite-local")] + if _sqlite_profiling.enabled { + let _ = sqlite_diagnostic_sender(_sqlite_profiling.diagnostic_event_queue_capacity); + } let metrics = &*METRICS; metrics .actor_active_count @@ -848,6 +1610,10 @@ impl ActorMetrics { Self { inner: Arc::new(ActorMetricInner { labels, + #[cfg(feature = "sqlite-local")] + sqlite_profiling: _sqlite_profiling, + #[cfg(feature = "sqlite-local")] + sqlite_profile_low_card_handles: OnceLock::new(), state: Mutex::new(ActorMetricState::default()), active: AtomicBool::new(true), startup_is_new: AtomicU8::new(STARTUP_KIND_UNKNOWN), @@ -1272,6 +2038,11 @@ impl ActorMetricInner { &METRICS.sqlite_workers_active, &labels, ); + if let Some(Some(handles)) = self.sqlite_profile_low_card_handles.get() { + handles.worker_queue_depth.set(0); + handles.worker_inflight.set(0); + handles.coordinator_queue_depth.set(0); + } } } } @@ -1288,8 +2059,600 @@ impl fmt::Debug for ActorMetrics { } } +#[cfg(feature = "sqlite-local")] +impl ActorMetrics { + const SQLITE_LOW_CARD_SERIES_COST: usize = 891; + + fn reserve_sqlite_series(&self, cost: usize) -> bool { + SQLITE_PROFILE_ADMISSION + .series + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current.saturating_add(cost) <= self.inner.sqlite_profiling.max_prometheus_series) + .then_some(current + cost) + }) + .is_ok() + } + + fn sqlite_low_card_handles( + &self, + storage_transport: &'static str, + ) -> Option<&SqliteLowCardMetricHandles> { + self.inner + .sqlite_profile_low_card_handles + .get_or_init(|| { + let tuple = format!("{}\0{storage_transport}", self.actor_labels()[0]); + if let Some(handles) = SQLITE_PROFILE_ADMISSION + .low_card_tuples + .read_sync(&tuple, |_, handles| Arc::clone(handles)) + { + return Some(handles); + } + if !self.reserve_sqlite_series(Self::SQLITE_LOW_CARD_SERIES_COST) { + self.warn_sqlite_profile_capacity("low_card_series_budget", "profile"); + return None; + } + let handles = Arc::new(SqliteLowCardMetricHandles::new( + self.actor_labels()[0], + storage_transport, + )); + match SQLITE_PROFILE_ADMISSION.low_card_tuples.entry_sync(tuple) { + scc::hash_map::Entry::Occupied(entry) => { + SQLITE_PROFILE_ADMISSION + .series + .fetch_sub(Self::SQLITE_LOW_CARD_SERIES_COST, Ordering::AcqRel); + Some(Arc::clone(entry.get())) + } + scc::hash_map::Entry::Vacant(entry) => { + entry.insert_entry(Arc::clone(&handles)); + Some(handles) + } + } + }) + .as_deref() + } + + fn warn_sqlite_profile_capacity(&self, reason: &'static str, operation_type: &'static str) { + let count = SQLITE_PROFILE_CAPACITY_WARNINGS.fetch_add(1, Ordering::Relaxed) + 1; + if count == 1 || count.is_power_of_two() { + tracing::warn!( + actor_name = self.actor_labels()[0], + operation_type, + reason, + overflow_observations = count, + "SQLite profiling capacity reached; observations are being aggregated or dropped" + ); + } + } + + fn record_sqlite_fingerprint_overflow( + &self, + operation_type: &'static str, + reason: &'static str, + ) { + SQLITE_PROFILE_METRICS + .fingerprint_overflow_total + .with_label_values(&[self.actor_labels()[0], operation_type, reason]) + .inc(); + self.warn_sqlite_profile_capacity(reason, operation_type); + } + + fn select_sqlite_fingerprint( + &self, + operation_type: &'static str, + fingerprint: &str, + total_ns: u64, + ) -> (String, bool) { + if fingerprint == "other" { + return ("other".to_owned(), false); + } + + let (set, count, cap) = if operation_type == "transaction" { + ( + &SQLITE_PROFILE_ADMISSION.transactions, + &SQLITE_PROFILE_ADMISSION.transaction_count, + self.inner + .sqlite_profiling + .max_tracked_transaction_fingerprints, + ) + } else { + ( + &SQLITE_PROFILE_ADMISSION.statements, + &SQLITE_PROFILE_ADMISSION.statement_count, + self.inner + .sqlite_profiling + .max_tracked_statement_fingerprints, + ) + }; + if set.contains_sync(fingerprint) { + return (fingerprint.to_owned(), false); + } + + let candidate_key = format!( + "{}\0{operation_type}\0{fingerprint}", + self.actor_labels()[0] + ); + let slow_threshold_ns = self + .inner + .sqlite_profiling + .slow_operation_threshold_ms + .saturating_mul(1_000_000); + if operation_type == "statement" && total_ns < slow_threshold_ns { + let candidate_cap = self + .inner + .sqlite_profiling + .max_tracked_statement_fingerprints + .saturating_add( + self.inner + .sqlite_profiling + .max_tracked_transaction_fingerprints, + ) + .saturating_mul(4) + .max(1); + if !SQLITE_PROFILE_ADMISSION + .candidates + .contains_sync(&candidate_key) + && SQLITE_PROFILE_ADMISSION.candidates.len() >= candidate_cap + { + self.record_sqlite_fingerprint_overflow(operation_type, "candidate_cap"); + return ("other".to_owned(), false); + } + let observation_count = *SQLITE_PROFILE_ADMISSION + .candidates + .entry_sync(candidate_key.clone()) + .and_modify(|count| *count = count.saturating_add(1)) + .or_insert(1) + .get(); + if observation_count < 2 { + return ("other".to_owned(), false); + } + } + if count + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + (current < cap).then_some(current + 1) + }) + .is_err() + { + self.record_sqlite_fingerprint_overflow(operation_type, "logical_cap"); + return ("other".to_owned(), false); + } + let newly_admitted = if set.insert_sync(fingerprint.to_owned()).is_err() { + count.fetch_sub(1, Ordering::AcqRel); + false + } else { + true + }; + let _ = SQLITE_PROFILE_ADMISSION + .candidates + .remove_sync(&candidate_key); + (fingerprint.to_owned(), newly_admitted) + } + + fn rollback_sqlite_logical_admission(&self, operation_type: &'static str, fingerprint: &str) { + let (set, count) = if operation_type == "transaction" { + ( + &SQLITE_PROFILE_ADMISSION.transactions, + &SQLITE_PROFILE_ADMISSION.transaction_count, + ) + } else { + ( + &SQLITE_PROFILE_ADMISSION.statements, + &SQLITE_PROFILE_ADMISSION.statement_count, + ) + }; + if set.remove_sync(fingerprint).is_some() { + count.fetch_sub(1, Ordering::AcqRel); + } + } + + fn admit_sqlite_fingerprint_tuple( + &self, + operation_type: &'static str, + fingerprint: &str, + fingerprint_source: &'static str, + transaction_mode: &'static str, + storage_transport: &'static str, + ) -> Option> { + let cost = if operation_type == "transaction" { + 207 + } else { + 147 + }; + let tuple = format!( + "{}\0{operation_type}\0{fingerprint}\0{fingerprint_source}\0{transaction_mode}\0{storage_transport}", + self.actor_labels()[0] + ); + if let Some(handles) = SQLITE_PROFILE_ADMISSION + .tuples + .read_sync(&tuple, |_, handles| Arc::clone(handles)) + { + return Some(handles); + } + if !self.reserve_sqlite_series(cost) { + return None; + } + let handles = Arc::new(SqliteFingerprintMetricHandles::new( + self.actor_labels()[0], + operation_type, + fingerprint, + fingerprint_source, + transaction_mode, + storage_transport, + )); + match SQLITE_PROFILE_ADMISSION.tuples.entry_sync(tuple) { + scc::hash_map::Entry::Occupied(entry) => { + SQLITE_PROFILE_ADMISSION + .series + .fetch_sub(cost, Ordering::AcqRel); + Some(Arc::clone(entry.get())) + } + scc::hash_map::Entry::Vacant(entry) => { + entry.insert_entry(Arc::clone(&handles)); + Some(handles) + } + } + } + + fn admitted_sqlite_fingerprint( + &self, + operation_type: &'static str, + fingerprint: &str, + fingerprint_source: &'static str, + transaction_mode: &'static str, + storage_transport: &'static str, + total_ns: u64, + ) -> Option> { + let low_card_handles = self.sqlite_low_card_handles(storage_transport)?; + let (selected, newly_admitted) = + self.select_sqlite_fingerprint(operation_type, fingerprint, total_ns); + if let Some(fingerprint_handles) = self.admit_sqlite_fingerprint_tuple( + operation_type, + &selected, + fingerprint_source, + transaction_mode, + storage_transport, + ) { + return Some(AdmittedSqliteProfile { + fingerprint: selected, + fingerprint_handles, + low_card_handles, + }); + } + if newly_admitted { + self.rollback_sqlite_logical_admission(operation_type, &selected); + } + + self.record_sqlite_fingerprint_overflow(operation_type, "series_budget"); + if selected != "other" { + let fingerprint_handles = self.admit_sqlite_fingerprint_tuple( + operation_type, + "other", + fingerprint_source, + transaction_mode, + storage_transport, + )?; + Some(AdmittedSqliteProfile { + fingerprint: "other".to_owned(), + fingerprint_handles, + low_card_handles, + }) + } else { + None + } + } + + fn observe_profile_common( + &self, + operation_type: &'static str, + fingerprint: &str, + fingerprint_source: &'static str, + transaction_mode: &'static str, + storage_transport: &'static str, + outcome: &'static str, + total_ns: u64, + phases: &[(&'static str, u64)], + get_pages_round_trips: u64, + ) -> Option> { + let admitted = self.admitted_sqlite_fingerprint( + operation_type, + fingerprint, + fingerprint_source, + transaction_mode, + storage_transport, + total_ns, + )?; + admitted + .fingerprint_handles + .observe_duration(outcome, ns_to_seconds(total_ns)); + for (phase, duration_ns) in phases { + admitted + .fingerprint_handles + .observe_phase(phase, ns_to_seconds(*duration_ns)); + } + admitted + .fingerprint_handles + .get_pages_round_trips + .observe(get_pages_round_trips as f64); + admitted.fingerprint_handles.record_outcome(outcome); + Some(admitted) + } + + fn record_local_profile_totals( + &self, + operation_type: &'static str, + handles: &SqliteLowCardMetricHandles, + profile: &depot_client::vfs::SqliteOperationProfile, + ) { + let operation_index = usize::from(operation_type == "transaction"); + for (index, pages) in [ + profile.sqlite_requested_pages, + profile.cache_hit_pages, + profile.cache_miss_pages, + profile.depot_demand_requested_pages, + profile.vfs_prefetch_requested_pages, + profile.response_present_pages, + profile.response_missing_pages, + profile.overflow_expansion_extra_pages, + profile.btree_pages, + profile.non_btree_pages, + profile.prefetch_consumed_pages, + profile.prefetch_unused_pages, + ] + .into_iter() + .enumerate() + { + if pages > 0 { + handles.local_pages[operation_index][index].inc_by(pages); + } + } + for (index, bytes) in [ + profile.bind_logical_bytes, + profile.result_logical_bytes, + profile.storage_response_bytes, + profile.dirty_bytes, + ] + .into_iter() + .enumerate() + { + if bytes > 0 { + handles.local_bytes[operation_index][index].inc_by(bytes); + } + } + } + + fn emit_sqlite_diagnostic_event(&self, event: SqliteDiagnosticEvent) { + let low_card_handles = self.sqlite_low_card_handles("proxy"); + if !try_acquire_sqlite_diagnostic_rate( + self.inner.sqlite_profiling.max_diagnostic_events_per_minute, + ) { + if let Some(handles) = low_card_handles { + handles.event_dropped[0].inc(); + } + return; + } + let Some(sender) = + sqlite_diagnostic_sender(self.inner.sqlite_profiling.diagnostic_event_queue_capacity) + else { + if let Some(handles) = low_card_handles { + handles.event_dropped[1].inc(); + } + return; + }; + if sender.try_send(event).is_err() + && let Some(handles) = low_card_handles + { + handles.event_dropped[1].inc(); + } + } + + fn operation_diagnostic_selected( + &self, + profile: &depot_client::vfs::SqliteOperationMetric, + ) -> bool { + let slow = profile.total_ns + >= self + .inner + .sqlite_profiling + .slow_operation_threshold_ms + .saturating_mul(1_000_000); + let page_amplification = profile.profile.response_present_pages + > profile + .profile + .sqlite_requested_pages + .saturating_mul(8) + .saturating_add(16); + let byte_amplification = profile.profile.storage_response_bytes > 64 * 1024 * 1024 + || profile.profile.dirty_bytes > 64 * 1024 * 1024; + slow || profile.outcome != "success" + || page_amplification + || byte_amplification + || sqlite_baseline_sample_selected(self.inner.sqlite_profiling.baseline_sample_rate) + } + + fn transaction_diagnostic_selected( + &self, + profile: &depot_client::vfs::SqliteTransactionMetric, + ) -> bool { + profile.total_ns + >= self + .inner + .sqlite_profiling + .slow_operation_threshold_ms + .saturating_mul(1_000_000) + || profile.outcome != "success" + || profile.dirty_bytes > 64 * 1024 * 1024 + || sqlite_baseline_sample_selected(self.inner.sqlite_profiling.baseline_sample_rate) + } +} + #[cfg(feature = "sqlite-local")] impl depot_client::vfs::SqliteVfsMetrics for ActorMetrics { + fn profiling_enabled(&self) -> bool { + self.inner.sqlite_profiling.enabled + } + + fn max_profiled_get_pages_requests(&self) -> usize { + self.inner.sqlite_profiling.max_get_pages_requests_per_trace + } + + fn observe_operation_profile( + &self, + profile: &depot_client::vfs::SqliteOperationMetric, + ) -> bool { + if !self.inner.sqlite_profiling.enabled { + return false; + } + let local_work_ns = profile + .total_ns + .saturating_sub(profile.transaction_wait_ns) + .saturating_sub(profile.profile.worker_wait_ns) + .saturating_sub(profile.profile.storage_ns); + let Some(admitted) = self.observe_profile_common( + profile.operation_type, + &profile.fingerprint, + profile.fingerprint_source, + profile.transaction_mode, + profile.storage_transport, + profile.outcome, + profile.total_ns, + &[ + ("transaction_wait", profile.transaction_wait_ns), + ("worker_wait", profile.profile.worker_wait_ns), + ("storage", profile.profile.storage_ns), + ("local_work", local_work_ns), + ], + profile.profile.get_pages_round_trips, + ) else { + return false; + }; + self.record_local_profile_totals( + profile.operation_type, + admitted.low_card_handles, + &profile.profile, + ); + for request in profile.profile.get_pages_requests.iter().flatten() { + let handles = + &admitted.low_card_handles.requests[sqlite_request_ordinal_index(request.ordinal)]; + handles.duration[usize::from(!request.success)] + .observe(ns_to_seconds(request.duration_ns)); + for (index, pages) in [ + request.demand_requested, + request.prefetch_requested, + request.response_present, + request.overflow_expansion_extra, + ] + .into_iter() + .enumerate() + { + handles.pages[index].observe(pages as f64); + } + handles + .response_bytes + .observe(request.response_bytes as f64); + if request.response_missing > 0 { + handles.missing_pages.inc_by(request.response_missing); + } + } + profile.fingerprint != "other" && admitted.fingerprint == profile.fingerprint + } + + fn observe_transaction_profile( + &self, + profile: &depot_client::vfs::SqliteTransactionMetric, + ) -> bool { + if !self.inner.sqlite_profiling.enabled { + return false; + } + let Some(admitted) = self.observe_profile_common( + "transaction", + &profile.fingerprint, + profile.fingerprint_source, + "explicit", + profile.storage_transport, + profile.outcome, + profile.total_ns, + &[ + ("application_time", profile.application_time_ns), + ("transaction_wait", profile.transaction_wait_ns), + ("worker_wait", profile.worker_wait_ns), + ("storage", profile.storage_ns), + ("local_work", profile.local_work_ns), + ("commit", profile.commit_ns), + ], + profile.get_pages_round_trips, + ) else { + return false; + }; + admitted + .fingerprint_handles + .transaction_statement_count + .as_ref() + .expect("transaction handles include statement count") + .observe(profile.statement_count as f64); + let mut aggregate = depot_client::vfs::SqliteOperationProfile::default(); + aggregate.dirty_pages = profile.dirty_pages; + aggregate.dirty_bytes = profile.dirty_bytes; + self.record_local_profile_totals("transaction", admitted.low_card_handles, &aggregate); + profile.fingerprint != "other" && admitted.fingerprint == profile.fingerprint + } + + fn emit_operation_diagnostic_event( + &self, + actor_id: &str, + generation: Option, + profile: &depot_client::vfs::SqliteOperationMetric, + ) { + if !self.inner.sqlite_profiling.enabled { + return; + } + if !self.operation_diagnostic_selected(profile) { + return; + } + self.emit_sqlite_diagnostic_event(SqliteDiagnosticEvent::Operation { + invocation_id: SQLITE_DIAGNOSTIC_INVOCATION_ID.fetch_add(1, Ordering::Relaxed), + actor_id: actor_id.to_owned(), + generation, + actor_name: self.actor_labels()[0].to_owned(), + profile: profile.clone(), + }); + } + + fn emit_transaction_diagnostic_event( + &self, + actor_id: &str, + generation: Option, + profile: &depot_client::vfs::SqliteTransactionMetric, + ) { + if !self.inner.sqlite_profiling.enabled { + return; + } + if !self.transaction_diagnostic_selected(profile) { + return; + } + self.emit_sqlite_diagnostic_event(SqliteDiagnosticEvent::Transaction { + invocation_id: SQLITE_DIAGNOSTIC_INVOCATION_ID.fetch_add(1, Ordering::Relaxed), + actor_id: actor_id.to_owned(), + generation, + actor_name: self.actor_labels()[0].to_owned(), + profile: profile.clone(), + }); + } + + fn record_fingerprint_catalog( + &self, + operation_type: &'static str, + fingerprint: &str, + identity: &str, + format_version: u8, + ) { + tracing::info!( + actor_name = self.actor_labels()[0], + operation_type, + fingerprint, + identity, + format_version, + "sqlite fingerprint catalog" + ); + } fn record_resolve_pages(&self, requested_pages: u64) { let labels = self.sqlite_vfs_labels(); METRICS @@ -1415,6 +2778,11 @@ impl depot_client::vfs::SqliteVfsMetrics for ActorMetrics { &METRICS.sqlite_worker_queue_depth, &labels, ); + if self.inner.sqlite_profiling.enabled + && let Some(handles) = self.sqlite_low_card_handles("proxy") + { + handles.worker_queue_depth.set(u64_to_i64(depth)); + } } fn set_worker_active(&self, active: bool) { @@ -1428,6 +2796,22 @@ impl depot_client::vfs::SqliteVfsMetrics for ActorMetrics { ); } + fn set_worker_inflight(&self, active: bool) { + if self.inner.sqlite_profiling.enabled + && let Some(handles) = self.sqlite_low_card_handles("proxy") + { + handles.worker_inflight.set(if active { 1 } else { 0 }); + } + } + + fn set_coordinator_queue_depth(&self, depth: u64) { + if self.inner.sqlite_profiling.enabled + && let Some(handles) = self.sqlite_low_card_handles("proxy") + { + handles.coordinator_queue_depth.set(u64_to_i64(depth)); + } + } + fn record_worker_queue_overload(&self) { METRICS .sqlite_worker_queue_overload_total @@ -1504,6 +2888,18 @@ fn ns_to_seconds(duration_ns: u64) -> f64 { Duration::from_nanos(duration_ns).as_secs_f64() } +#[cfg(feature = "sqlite-local")] +fn sqlite_request_ordinal_index(ordinal: u64) -> usize { + match ordinal { + 1 => 0, + 2 => 1, + 3 => 2, + 4 => 3, + 5..=8 => 4, + _ => 5, + } +} + #[cfg(feature = "sqlite-local")] fn sqlite_worker_duration_buckets() -> Vec { vec![ diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite.rs index 67a692506e..dfbf24deac 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite.rs @@ -20,8 +20,13 @@ use tokio::sync::Mutex as AsyncMutex; #[cfg(feature = "sqlite-local")] use tokio::task::JoinHandle; +#[cfg(feature = "sqlite-local")] +const DEFAULT_TRANSACTION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); + #[cfg(feature = "sqlite-local")] mod envoy_sqlite_transport; +#[cfg(feature = "sqlite-local")] +mod profiling; #[cfg(feature = "sqlite-local")] use crate::error::ActorLifecycle; @@ -36,6 +41,7 @@ use depot_client::{ worker::{ SQLITE_WORKER_QUEUE_CAPACITY, SqliteWorkerCloseTimeoutError, SqliteWorkerClosingError, SqliteWorkerDeadError, SqliteWorkerFatalError, SqliteWorkerOverloadedError, + SqliteWorkerResult, }, }; #[cfg(feature = "sqlite-local")] @@ -94,6 +100,167 @@ pub struct SqliteDb { worker_fatal_reported: Arc, #[cfg(feature = "sqlite-local")] vfs_metrics: Option>, + #[cfg(feature = "sqlite-local")] + profiling: Arc, + #[cfg(feature = "sqlite-local")] + transaction_lock: Arc>, +} + +#[cfg(feature = "sqlite-local")] +#[derive(Clone)] +pub struct SqliteTransaction { + db: SqliteDb, + state: Arc>, +} + +#[cfg(feature = "sqlite-local")] +struct SqliteTransactionState { + _guard: Option>, + profile: profiling::TransactionProfile, + finished: bool, + timeout_cancel: tokio_util::sync::CancellationToken, +} + +#[cfg(feature = "sqlite-local")] +impl SqliteTransaction { + pub async fn exec(&self, sql: impl Into) -> Result { + let sql = sql.into(); + let mut state = self.state.lock().await; + anyhow::ensure!(!state.finished, "sqlite transaction is already finished"); + let started_at = crate::time::Instant::now(); + let profiled = self.db.local_exec_profiled(sql.clone()).await?; + let outcome = if profiled.result.is_ok() { + "success" + } else { + "error" + }; + if let Some(observation) = self.db.observe_statement_profile( + &sql, + started_at, + Some(profiled.profile), + outcome, + "explicit", + ) { + state.profile.record_statement(&observation); + } + profiled.result + } + + pub async fn query( + &self, + sql: impl Into, + params: Option>, + ) -> Result { + self.execute(sql, params) + .await + .map(ExecuteResult::into_query_result) + } + + pub async fn execute( + &self, + sql: impl Into, + params: Option>, + ) -> Result { + let sql = sql.into(); + let mut state = self.state.lock().await; + anyhow::ensure!(!state.finished, "sqlite transaction is already finished"); + let started_at = crate::time::Instant::now(); + let profiled = self + .db + .local_execute_profiled(sql.clone(), params) + .await?; + let outcome = if profiled.result.is_ok() { + "success" + } else { + "error" + }; + if let Some(observation) = self.db.observe_statement_profile( + &sql, + started_at, + Some(profiled.profile), + outcome, + "explicit", + ) { + state.profile.record_statement(&observation); + } + profiled.result + } + + pub async fn commit(&self) -> Result<()> { + self.finish(true, None).await + } + + pub async fn rollback(&self) -> Result<()> { + self.finish(false, None).await + } + + async fn finish(&self, commit: bool, forced_outcome: Option<&'static str>) -> Result<()> { + let mut state = self.state.lock().await; + anyhow::ensure!(!state.finished, "sqlite transaction is already finished"); + state.timeout_cancel.cancel(); + let started_at = crate::time::Instant::now(); + let profiled = self + .db + .local_exec_profiled(if commit { + "COMMIT".to_owned() + } else { + "ROLLBACK".to_owned() + }) + .await?; + let outcome = forced_outcome.unwrap_or(if profiled.result.is_ok() { + if commit { "success" } else { "rollback" } + } else { + "error" + }); + state + .profile + .record_control(&profiled.profile, duration_ns(started_at.elapsed()), commit); + let total_ns = duration_ns(state.profile.started_at.elapsed()); + let application_time_ns = total_ns + .saturating_sub(state.profile.worker_wait_ns) + .saturating_sub(state.profile.storage_ns) + .saturating_sub(state.profile.local_work_ns); + if let Some(metrics) = &self.db.vfs_metrics { + let metric = depot_client::vfs::SqliteTransactionMetric { + fingerprint: state.profile.fingerprint.clone(), + fingerprint_source: "name", + shape_fingerprint: state.profile.fingerprint.clone(), + statement_fingerprint_hashes: state.profile.statement_fingerprint_hashes, + omitted_statement_fingerprints: state.profile.omitted_statement_fingerprints, + storage_transport: "proxy", + outcome, + total_ns, + transaction_wait_ns: 0, + worker_wait_ns: state.profile.worker_wait_ns, + storage_ns: state.profile.storage_ns, + local_work_ns: state.profile.local_work_ns, + application_time_ns, + commit_ns: state.profile.commit_ns, + get_pages_round_trips: state.profile.get_pages_round_trips, + statement_count: state.profile.statement_count, + dirty_pages: state.profile.dirty_pages, + dirty_bytes: state.profile.dirty_bytes, + }; + if metrics.observe_transaction_profile(&metric) + && self.db.profiling.mark_cataloged(&metric.fingerprint) + { + metrics.record_fingerprint_catalog( + "transaction", + &metric.fingerprint, + &state.profile.name, + profiling::FINGERPRINT_FORMAT_VERSION, + ); + } + metrics.emit_transaction_diagnostic_event( + self.db.actor_id.as_deref().unwrap_or("unknown"), + self.db.generation, + &metric, + ); + } + state.finished = true; + state._guard.take(); + profiled.result.map(|_| ()) + } } impl SqliteDb { @@ -125,6 +292,10 @@ impl SqliteDb { worker_fatal_reported: Default::default(), #[cfg(feature = "sqlite-local")] vfs_metrics: None, + #[cfg(feature = "sqlite-local")] + profiling: Arc::new(profiling::SqliteProfilingState::default()), + #[cfg(feature = "sqlite-local")] + transaction_lock: Default::default(), } } @@ -133,6 +304,93 @@ impl SqliteDb { self.vfs_metrics = Some(metrics); } + #[cfg(feature = "sqlite-local")] + pub(crate) fn set_profiling_config(&mut self, config: crate::SqliteProfilingConfig) { + self.profiling = Arc::new(profiling::SqliteProfilingState::new(config)); + } + + #[cfg(all(test, feature = "sqlite-local"))] + pub(crate) fn from_native_database_for_test( + actor_id: impl Into, + generation: u64, + native_db: NativeDatabaseHandle, + metrics: Arc, + profiling: crate::SqliteProfilingConfig, + ) -> Self { + Self { + actor_id: Some(actor_id.into()), + generation: Some(generation), + backend: SqliteBackend::LocalNative, + enabled: true, + db: Arc::new(Mutex::new(Some(native_db))), + vfs_metrics: Some(metrics), + profiling: Arc::new(profiling::SqliteProfilingState::new(profiling)), + ..Self::default() + } + } + + #[cfg(feature = "sqlite-local")] + pub async fn begin_transaction( + &self, + timeout: Option, + ) -> Result { + self.begin_named_transaction(None, timeout).await + } + + #[cfg(feature = "sqlite-local")] + pub async fn begin_named_transaction( + &self, + name: Option<&str>, + timeout: Option, + ) -> Result { + anyhow::ensure!( + self.backend == SqliteBackend::LocalNative, + "named transactions require local native SQLite" + ); + let name = name.unwrap_or("transaction"); + let timeout = timeout.unwrap_or(DEFAULT_TRANSACTION_TIMEOUT); + anyhow::ensure!(!timeout.is_zero(), "SQLite transaction timeout must be positive"); + anyhow::ensure!( + name.len() <= self.profiling.config.max_transaction_name_bytes, + "SQLite transaction name is too long" + ); + let started_at = crate::time::Instant::now(); + let guard = Arc::clone(&self.transaction_lock).lock_owned().await; + let control_started_at = crate::time::Instant::now(); + let profiled = self.local_exec_profiled("BEGIN".to_owned()).await?; + profiled.result?; + let mut profile = profiling::TransactionProfile::new( + name.to_owned(), + started_at, + self.profiling.config.max_statements_per_transaction_trace, + ); + profile.record_control( + &profiled.profile, + duration_ns(control_started_at.elapsed()), + false, + ); + let timeout_cancel = tokio_util::sync::CancellationToken::new(); + let transaction = SqliteTransaction { + db: self.clone(), + state: Arc::new(AsyncMutex::new(SqliteTransactionState { + _guard: Some(guard), + profile, + finished: false, + timeout_cancel: timeout_cancel.clone(), + })), + }; + let timeout_transaction = transaction.clone(); + RuntimeSpawner::spawn(async move { + tokio::select! { + _ = timeout_cancel.cancelled() => {} + _ = tokio::time::sleep(timeout) => { + let _ = timeout_transaction.finish(false, Some("timeout")).await; + } + } + }); + Ok(transaction) + } + pub fn is_enabled(&self) -> bool { self.enabled } @@ -210,6 +468,12 @@ impl SqliteDb { self.map_local_worker_result(self.native_db_handle()?.exec(sql).await) } + #[cfg(feature = "sqlite-local")] + async fn local_exec_profiled(&self, sql: String) -> Result> { + self.open().await?; + self.map_local_worker_result(self.native_db_handle()?.exec_profiled(sql).await) + } + #[cfg(not(feature = "sqlite-local"))] async fn local_exec(&self, _sql: String) -> Result { Err(SqliteRuntimeError::Unavailable.build()) @@ -255,6 +519,16 @@ impl SqliteDb { self.map_local_worker_result(self.native_db_handle()?.execute(sql, params).await) } + #[cfg(feature = "sqlite-local")] + async fn local_execute_profiled( + &self, + sql: String, + params: Option>, + ) -> Result> { + self.open().await?; + self.map_local_worker_result(self.native_db_handle()?.execute_profiled(sql, params).await) + } + #[cfg(not(feature = "sqlite-local"))] async fn local_execute( &self, @@ -264,14 +538,95 @@ impl SqliteDb { Err(SqliteRuntimeError::Unavailable.build()) } + #[cfg(feature = "sqlite-local")] + fn observe_statement_profile( + &self, + sql: &str, + started_at: crate::time::Instant, + profile: Option, + outcome: &'static str, + transaction_mode: &'static str, + ) -> Option { + if self.backend != SqliteBackend::LocalNative { + return None; + } + let Some(fingerprint) = self.profiling.statement_fingerprint(sql) else { + return None; + }; + let observation = profiling::StatementObservation { + fingerprint, + total_ns: duration_ns(started_at.elapsed()), + profile: profile.unwrap_or_default(), + }; + if let Some(metrics) = &self.vfs_metrics { + let metric = depot_client::vfs::SqliteOperationMetric { + operation_type: "statement", + fingerprint: observation.fingerprint.display.clone(), + fingerprint_source: "query", + transaction_mode, + storage_transport: "proxy", + outcome, + sql_bytes: sql.len().try_into().unwrap_or(u64::MAX), + total_ns: observation.total_ns, + transaction_wait_ns: 0, + profile: observation.profile.clone(), + }; + if metrics.observe_operation_profile(&metric) + && self.profiling.mark_cataloged(&metric.fingerprint) + { + metrics.record_fingerprint_catalog( + "statement", + &metric.fingerprint, + sql, + profiling::FINGERPRINT_FORMAT_VERSION, + ); + } + metrics.emit_operation_diagnostic_event( + self.actor_id.as_deref().unwrap_or("unknown"), + self.generation, + &metric, + ); + } + Some(observation) + } + pub async fn exec(&self, sql: impl Into) -> Result { + #[cfg(feature = "sqlite-local")] + let _transaction_guard = self.transaction_lock.lock().await; let sql = sql.into(); let sql_for_log = sql.clone(); + #[cfg(feature = "sqlite-local")] + let started_at = (self.backend == SqliteBackend::LocalNative + && self.profiling.config.enabled) + .then(crate::time::Instant::now); + #[cfg(feature = "sqlite-local")] + let (result, profile) = match self.backend { + SqliteBackend::LocalNative if started_at.is_some() => { + match self.local_exec_profiled(sql).await { + Ok(profiled) => (profiled.result, Some(profiled.profile)), + Err(error) => (Err(error), None), + } + } + SqliteBackend::LocalNative => (self.local_exec(sql).await, None), + SqliteBackend::RemoteEnvoy => (self.remote_exec(sql).await, None), + SqliteBackend::Unavailable => (Err(SqliteRuntimeError::Unavailable.build()), None), + }; + #[cfg(not(feature = "sqlite-local"))] let result = match self.backend { SqliteBackend::LocalNative => self.local_exec(sql).await, SqliteBackend::RemoteEnvoy => self.remote_exec(sql).await, SqliteBackend::Unavailable => Err(SqliteRuntimeError::Unavailable.build()), }; + #[cfg(feature = "sqlite-local")] + if let Some(started_at) = started_at { + let _ = self.observe_statement_profile( + &sql_for_log, + started_at, + profile, + if result.is_ok() { "success" } else { "error" }, + "autocommit", + ); + } match result { Ok(result) => Ok(result), Err(error) => { @@ -287,9 +642,36 @@ impl SqliteDb { sql: impl Into, params: Option>, ) -> Result { + #[cfg(feature = "sqlite-local")] + let _transaction_guard = self.transaction_lock.lock().await; let sql = sql.into(); let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); + #[cfg(feature = "sqlite-local")] + let started_at = (self.backend == SqliteBackend::LocalNative + && self.profiling.config.enabled) + .then(crate::time::Instant::now); + #[cfg(feature = "sqlite-local")] + let (result, profile) = match self.backend { + SqliteBackend::LocalNative if started_at.is_some() => { + match self.local_execute_profiled(sql, params).await { + Ok(profiled) => ( + profiled.result.map(ExecuteResult::into_query_result), + Some(profiled.profile), + ), + Err(error) => (Err(error), None), + } + } + SqliteBackend::LocalNative => (self.local_query(sql, params).await, None), + SqliteBackend::RemoteEnvoy => ( + self.remote_execute(sql, params) + .await + .map(ExecuteResult::into_query_result), + None, + ), + SqliteBackend::Unavailable => (Err(SqliteRuntimeError::Unavailable.build()), None), + }; + #[cfg(not(feature = "sqlite-local"))] let result = match self.backend { SqliteBackend::LocalNative => self.local_query(sql, params).await, SqliteBackend::RemoteEnvoy => self @@ -298,6 +680,16 @@ impl SqliteDb { .map(ExecuteResult::into_query_result), SqliteBackend::Unavailable => Err(SqliteRuntimeError::Unavailable.build()), }; + #[cfg(feature = "sqlite-local")] + if let Some(started_at) = started_at { + let _ = self.observe_statement_profile( + &sql_for_log, + started_at, + profile, + if result.is_ok() { "success" } else { "error" }, + "autocommit", + ); + } match result { Ok(result) => Ok(result), Err(error) => { @@ -313,9 +705,36 @@ impl SqliteDb { sql: impl Into, params: Option>, ) -> Result { + #[cfg(feature = "sqlite-local")] + let _transaction_guard = self.transaction_lock.lock().await; let sql = sql.into(); let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); + #[cfg(feature = "sqlite-local")] + let started_at = (self.backend == SqliteBackend::LocalNative + && self.profiling.config.enabled) + .then(crate::time::Instant::now); + #[cfg(feature = "sqlite-local")] + let (result, profile) = match self.backend { + SqliteBackend::LocalNative if started_at.is_some() => { + match self.local_execute_profiled(sql, params).await { + Ok(profiled) => ( + profiled.result.map(ExecuteResult::into_exec_result), + Some(profiled.profile), + ), + Err(error) => (Err(error), None), + } + } + SqliteBackend::LocalNative => (self.local_run(sql, params).await, None), + SqliteBackend::RemoteEnvoy => ( + self.remote_execute(sql, params) + .await + .map(ExecuteResult::into_exec_result), + None, + ), + SqliteBackend::Unavailable => (Err(SqliteRuntimeError::Unavailable.build()), None), + }; + #[cfg(not(feature = "sqlite-local"))] let result = match self.backend { SqliteBackend::LocalNative => self.local_run(sql, params).await, SqliteBackend::RemoteEnvoy => self @@ -324,6 +743,16 @@ impl SqliteDb { .map(ExecuteResult::into_exec_result), SqliteBackend::Unavailable => Err(SqliteRuntimeError::Unavailable.build()), }; + #[cfg(feature = "sqlite-local")] + if let Some(started_at) = started_at { + let _ = self.observe_statement_profile( + &sql_for_log, + started_at, + profile, + if result.is_ok() { "success" } else { "error" }, + "autocommit", + ); + } match result { Ok(result) => Ok(result), Err(error) => { @@ -339,14 +768,43 @@ impl SqliteDb { sql: impl Into, params: Option>, ) -> Result { + #[cfg(feature = "sqlite-local")] + let _transaction_guard = self.transaction_lock.lock().await; let sql = sql.into(); let sql_for_log = sql.clone(); let binding_count = bind_param_count(¶ms); + #[cfg(feature = "sqlite-local")] + let started_at = (self.backend == SqliteBackend::LocalNative + && self.profiling.config.enabled) + .then(crate::time::Instant::now); + #[cfg(feature = "sqlite-local")] + let (result, profile) = match self.backend { + SqliteBackend::LocalNative if started_at.is_some() => { + match self.local_execute_profiled(sql, params).await { + Ok(profiled) => (profiled.result, Some(profiled.profile)), + Err(error) => (Err(error), None), + } + } + SqliteBackend::LocalNative => (self.local_execute(sql, params).await, None), + SqliteBackend::RemoteEnvoy => (self.remote_execute(sql, params).await, None), + SqliteBackend::Unavailable => (Err(SqliteRuntimeError::Unavailable.build()), None), + }; + #[cfg(not(feature = "sqlite-local"))] let result = match self.backend { SqliteBackend::LocalNative => self.local_execute(sql, params).await, SqliteBackend::RemoteEnvoy => self.remote_execute(sql, params).await, SqliteBackend::Unavailable => Err(SqliteRuntimeError::Unavailable.build()), }; + #[cfg(feature = "sqlite-local")] + if let Some(started_at) = started_at { + let _ = self.observe_statement_profile( + &sql_for_log, + started_at, + profile, + if result.is_ok() { "success" } else { "error" }, + "autocommit", + ); + } match result { Ok(result) => Ok(result), Err(error) => { @@ -689,6 +1147,11 @@ fn bind_param_count(params: &Option>) -> usize { params.as_ref().map_or(0, Vec::len) } +#[cfg(feature = "sqlite-local")] +fn duration_ns(duration: std::time::Duration) -> u64 { + duration.as_nanos().try_into().unwrap_or(u64::MAX) +} + #[cfg(feature = "sqlite-local")] fn is_fatal_worker_error(error: &anyhow::Error) -> bool { error.downcast_ref::().is_some() diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/profiling.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/profiling.rs new file mode 100644 index 0000000000..f339d5e5a8 --- /dev/null +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/profiling.rs @@ -0,0 +1,238 @@ +use depot_client::vfs::SqliteOperationProfile; +use sha2::{Digest, Sha256}; + +use crate::SqliteProfilingConfig; + +pub(super) const FINGERPRINT_FORMAT_VERSION: u8 = 2; + +#[derive(Clone, Debug)] +pub(super) struct StatementFingerprint { + pub(super) display: String, + pub(super) hash: String, +} + +#[derive(Debug)] +pub(super) struct SqliteProfilingState { + pub(super) config: SqliteProfilingConfig, + cataloged: scc::HashSet, +} + +impl Default for SqliteProfilingState { + fn default() -> Self { + Self::new(SqliteProfilingConfig::default()) + } +} + +impl SqliteProfilingState { + pub(super) fn new(config: SqliteProfilingConfig) -> Self { + Self { + config, + cataloged: scc::HashSet::new(), + } + } + + pub(super) fn statement_fingerprint(&self, sql: &str) -> Option { + if !self.config.enabled { + return None; + } + let class = statement_class(sql)?; + let hash = fingerprint_hash(b"rivetkit-sqlite-statement", sql.as_bytes()); + Some(StatementFingerprint { + display: format!("{class}-{hash}"), + hash, + }) + } + + pub(super) fn mark_cataloged(&self, fingerprint: &str) -> bool { + self.cataloged.insert_sync(fingerprint.to_owned()).is_ok() + } +} + +#[derive(Clone, Debug)] +pub(super) struct StatementObservation { + pub(super) fingerprint: StatementFingerprint, + pub(super) total_ns: u64, + pub(super) profile: SqliteOperationProfile, +} + +#[derive(Debug)] +pub(super) struct TransactionProfile { + pub(super) started_at: crate::time::Instant, + pub(super) fingerprint: String, + pub(super) name: String, + pub(super) statement_fingerprint_hashes: + [Option<[u8; 16]>; depot_client::vfs::MAX_PROFILED_TRANSACTION_STATEMENTS], + pub(super) omitted_statement_fingerprints: u64, + statement_fingerprint_limit: usize, + pub(super) statement_count: u64, + pub(super) worker_wait_ns: u64, + pub(super) storage_ns: u64, + pub(super) local_work_ns: u64, + pub(super) get_pages_round_trips: u64, + pub(super) dirty_pages: u64, + pub(super) dirty_bytes: u64, + pub(super) commit_ns: u64, +} + +impl TransactionProfile { + pub(super) fn new( + name: String, + started_at: crate::time::Instant, + statement_fingerprint_limit: usize, + ) -> Self { + let fingerprint = format!( + "txn-{}", + fingerprint_hash(b"rivetkit-sqlite-transaction", name.as_bytes()) + ); + Self { + started_at, + fingerprint, + name, + statement_fingerprint_hashes: [None; + depot_client::vfs::MAX_PROFILED_TRANSACTION_STATEMENTS], + omitted_statement_fingerprints: 0, + statement_fingerprint_limit: statement_fingerprint_limit + .min(depot_client::vfs::MAX_PROFILED_TRANSACTION_STATEMENTS), + statement_count: 0, + worker_wait_ns: 0, + storage_ns: 0, + local_work_ns: 0, + get_pages_round_trips: 0, + dirty_pages: 0, + dirty_bytes: 0, + commit_ns: 0, + } + } + + pub(super) fn record_control( + &mut self, + profile: &SqliteOperationProfile, + duration_ns: u64, + is_commit: bool, + ) { + self.worker_wait_ns = self.worker_wait_ns.saturating_add(profile.worker_wait_ns); + self.storage_ns = self.storage_ns.saturating_add(profile.storage_ns); + self.local_work_ns = self.local_work_ns.saturating_add( + duration_ns + .saturating_sub(profile.worker_wait_ns) + .saturating_sub(profile.storage_ns), + ); + self.get_pages_round_trips = self + .get_pages_round_trips + .saturating_add(profile.get_pages_round_trips); + self.dirty_pages = self.dirty_pages.saturating_add(profile.dirty_pages); + self.dirty_bytes = self.dirty_bytes.saturating_add(profile.dirty_bytes); + if is_commit { + self.commit_ns = duration_ns; + } + } + + pub(super) fn record_statement(&mut self, observation: &StatementObservation) { + let statement_index = self.statement_count as usize; + if statement_index < self.statement_fingerprint_limit { + let mut hash = [0; 16]; + let source = observation.fingerprint.hash.as_bytes(); + let copy_len = source.len().min(hash.len()); + hash[..copy_len].copy_from_slice(&source[..copy_len]); + self.statement_fingerprint_hashes[statement_index] = Some(hash); + } else { + self.omitted_statement_fingerprints = + self.omitted_statement_fingerprints.saturating_add(1); + } + self.statement_count = self.statement_count.saturating_add(1); + self.worker_wait_ns = self + .worker_wait_ns + .saturating_add(observation.profile.worker_wait_ns); + self.storage_ns = self + .storage_ns + .saturating_add(observation.profile.storage_ns); + self.local_work_ns = self.local_work_ns.saturating_add( + observation + .total_ns + .saturating_sub(observation.profile.worker_wait_ns) + .saturating_sub(observation.profile.storage_ns), + ); + self.get_pages_round_trips = self + .get_pages_round_trips + .saturating_add(observation.profile.get_pages_round_trips); + self.dirty_pages = self + .dirty_pages + .saturating_add(observation.profile.dirty_pages); + self.dirty_bytes = self + .dirty_bytes + .saturating_add(observation.profile.dirty_bytes); + } +} + +fn fingerprint_hash(domain: &[u8], value: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(domain); + hasher.update([0, FINGERPRINT_FORMAT_VERSION, 0]); + hasher.update(value); + let digest = hasher.finalize(); + let mut output = String::with_capacity(16); + for byte in digest.iter().take(8) { + use std::fmt::Write; + let _ = write!(output, "{byte:02x}"); + } + output +} + +fn statement_class(sql: &str) -> Option<&'static str> { + let keyword = sql.split_ascii_whitespace().next()?; + if keyword.eq_ignore_ascii_case("select") || keyword.eq_ignore_ascii_case("values") { + Some("select") + } else if keyword.eq_ignore_ascii_case("insert") || keyword.eq_ignore_ascii_case("replace") { + Some("insert") + } else if keyword.eq_ignore_ascii_case("update") { + Some("update") + } else if keyword.eq_ignore_ascii_case("delete") { + Some("delete") + } else if keyword.eq_ignore_ascii_case("pragma") { + Some("pragma") + } else if ["begin", "commit", "end", "rollback", "savepoint", "release"] + .iter() + .any(|candidate| keyword.eq_ignore_ascii_case(candidate)) + { + None + } else if [ + "create", "alter", "drop", "vacuum", "reindex", "analyze", "attach", "detach", + ] + .iter() + .any(|candidate| keyword.eq_ignore_ascii_case(candidate)) + { + Some("ddl") + } else { + Some("other") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fingerprints_exact_query_text() { + let state = SqliteProfilingState::default(); + let first = state + .statement_fingerprint("SELECT * FROM orders WHERE id = ?") + .unwrap(); + let second = state + .statement_fingerprint("SELECT * FROM orders WHERE id = ?") + .unwrap(); + let differently_formatted = state + .statement_fingerprint("select * from orders where id=?") + .unwrap(); + + assert_eq!(first.display, second.display); + assert_ne!(first.display, differently_formatted.display); + assert!(first.display.starts_with("select-")); + } + + #[test] + fn transaction_control_is_not_profiled() { + let state = SqliteProfilingState::default(); + assert!(state.statement_fingerprint("BEGIN").is_none()); + assert!(state.statement_fingerprint("ROLLBACK").is_none()); + } +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/lib.rs b/rivetkit-rust/packages/rivetkit-core/src/lib.rs index 63e5be3cde..060e81e02f 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/lib.rs @@ -117,6 +117,7 @@ pub use actor::{kv, sqlite}; pub use actor::action::ActionDispatchError; pub use actor::config::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorConfigOverrides, CanHibernateWebSocket, + SqliteProfilingConfig, SqliteProfilingConfigInput, }; pub use actor::connection::ConnHandle; pub use actor::context::{ActorContext, ActorWorkRegion, KeepAwakeRegion, WebSocketCallbackRegion}; @@ -134,6 +135,8 @@ pub use actor::queue::{ pub use actor::sqlite::{ BindParam, ColumnValue, ExecResult, ExecuteResult, QueryResult, SqliteBackend, SqliteDb, }; +#[cfg(feature = "sqlite-local")] +pub use actor::sqlite::SqliteTransaction; pub use actor::state::RequestSaveOpts; pub use actor::task::{ ActionDispatchResult, ActorTask, DispatchCommand, HttpDispatchResult, LifecycleCommand, @@ -144,6 +147,7 @@ pub use actor::work_registry::{ActorWorkKind, ActorWorkPolicy}; pub use error::ActorLifecycle; pub use inspector::{Inspector, InspectorSnapshot}; pub use registry::{CoreRegistry, EngineSpawnMode, ServeConfig}; +pub use rivet_envoy_client::config::ResponseChunk; pub use runtime::{RuntimeBoxFuture, RuntimeSpawner, boxed_runtime_future}; pub use serverless::{CoreServerlessRuntime, ServerlessRequest, ServerlessResponse}; pub use types::{ diff --git a/rivetkit-rust/packages/rivetkit-core/tests/config.rs b/rivetkit-rust/packages/rivetkit-core/tests/config.rs index b550599730..9cdcf10473 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/config.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/config.rs @@ -3,7 +3,7 @@ use super::*; mod moved_tests { use std::time::Duration; - use super::{ActorConfig, ActorConfigInput}; + use super::{ActorConfig, ActorConfigInput, SqliteProfilingConfig, SqliteProfilingConfigInput}; #[test] fn actor_config_from_input_applies_overrides() { @@ -85,6 +85,49 @@ mod moved_tests { super::CanHibernateWebSocket::Bool(false), )); assert!(config.overrides.is_none()); + assert!(config.sqlite_profiling.enabled); + assert_eq!( + config.sqlite_profiling.max_tracked_statement_fingerprints, + 128 + ); + assert_eq!( + config.sqlite_profiling.max_tracked_transaction_fingerprints, + 8 + ); + assert_eq!(config.sqlite_profiling.max_prometheus_series, 25_000); + assert_eq!(config.sqlite_profiling.max_get_pages_requests_per_trace, 16); + assert_eq!(config.sqlite_profiling.slow_operation_threshold_ms, 10); + assert_eq!(config.sqlite_profiling.baseline_sample_rate, 0.001); + } + + #[test] + fn actor_config_applies_and_validates_sqlite_profiling_overrides() { + let config = ActorConfig::from_input(ActorConfigInput { + sqlite_profiling: Some(SqliteProfilingConfigInput { + enabled: Some(false), + max_tracked_statement_fingerprints: Some(7), + baseline_sample_rate: Some(0.25), + ..Default::default() + }), + ..Default::default() + }); + + assert!(!config.sqlite_profiling.enabled); + assert_eq!( + config.sqlite_profiling.max_tracked_statement_fingerprints, + 7 + ); + assert_eq!(config.sqlite_profiling.baseline_sample_rate, 0.25); + config.validate().expect("profiling config should be valid"); + + let invalid = ActorConfig { + sqlite_profiling: SqliteProfilingConfig { + baseline_sample_rate: 1.5, + ..Default::default() + }, + ..Default::default() + }; + assert!(invalid.validate().is_err()); } #[test] diff --git a/rivetkit-rust/packages/rivetkit-core/tests/metrics.rs b/rivetkit-rust/packages/rivetkit-core/tests/metrics.rs index 9728d1a742..2362fb3668 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/metrics.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/metrics.rs @@ -4,16 +4,80 @@ use super::*; mod metrics_helpers; mod moved_tests { + #[cfg(feature = "sqlite-local")] + use std::collections::BTreeMap; use std::panic::{AssertUnwindSafe, catch_unwind}; + #[cfg(feature = "sqlite-local")] + use std::sync::Arc; use std::time::Duration; use rivet_metrics::prometheus::{IntGauge, Opts, Registry}; use crate::actor::task_types::UserTaskKind; - use super::metrics_helpers::{metric_line_for_actor, render_global_metrics}; + use super::metrics_helpers::{ + metric_line_for_actor, metric_name_matches, render_global_metrics, + }; use super::*; + #[cfg(feature = "sqlite-local")] + #[derive(Default)] + struct InMemorySqliteTransport { + state: parking_lot::Mutex, + } + + #[cfg(feature = "sqlite-local")] + #[derive(Default)] + struct InMemorySqliteState { + pages: BTreeMap>, + head_txid: u64, + } + + #[cfg(feature = "sqlite-local")] + #[async_trait::async_trait] + impl depot_client::vfs::SqliteTransport for InMemorySqliteTransport { + async fn get_pages( + &self, + request: rivet_envoy_client::protocol::SqliteGetPagesRequest, + ) -> anyhow::Result { + let state = self.state.lock(); + Ok( + rivet_envoy_client::protocol::SqliteGetPagesResponse::SqliteGetPagesOk( + rivet_envoy_client::protocol::SqliteGetPagesOk { + pages: request + .pgnos + .into_iter() + .map(|pgno| rivet_envoy_client::protocol::SqliteFetchedPage { + pgno, + bytes: state.pages.get(&pgno).cloned(), + }) + .collect(), + head_txid: Some(state.head_txid), + }, + ), + ) + } + + async fn commit( + &self, + request: rivet_envoy_client::protocol::SqliteCommitRequest, + ) -> anyhow::Result { + let mut state = self.state.lock(); + for page in request.dirty_pages { + state.pages.insert(page.pgno, page.bytes); + } + state.pages.retain(|pgno, _| *pgno <= request.db_size_pages); + state.head_txid = state.head_txid.saturating_add(1); + Ok( + rivet_envoy_client::protocol::SqliteCommitResponse::SqliteCommitOk( + rivet_envoy_client::protocol::SqliteCommitOk { + head_txid: Some(state.head_txid), + }, + ), + ) + } + } + #[test] fn duplicate_metric_registration_uses_noop_fallback() { let registry = Registry::new(); @@ -174,6 +238,463 @@ mod moved_tests { ); } + #[cfg(feature = "sqlite-local")] + #[test] + fn sqlite_profile_reuses_pre_resolved_metric_handles() { + let actor_name = "counter-sqlite-pre-resolved-handles"; + let first = ActorMetrics::new(actor_name); + let second = ActorMetrics::new(actor_name); + + let first_low_card = first + .sqlite_low_card_handles("proxy") + .expect("low-cardinality handles should be admitted"); + let second_low_card = second + .sqlite_low_card_handles("proxy") + .expect("low-cardinality handles should be reused"); + assert!(std::ptr::eq(first_low_card, second_low_card)); + + let first_fingerprint = first + .admit_sqlite_fingerprint_tuple( + "statement", + "select-pre-resolved-handles", + "query", + "autocommit", + "proxy", + ) + .expect("fingerprint handles should be admitted"); + let second_fingerprint = second + .admit_sqlite_fingerprint_tuple( + "statement", + "select-pre-resolved-handles", + "query", + "autocommit", + "proxy", + ) + .expect("fingerprint handles should be reused"); + assert!(Arc::ptr_eq(&first_fingerprint, &second_fingerprint)); + } + + #[cfg(feature = "sqlite-local")] + #[test] + fn sqlite_profiling_metrics_render_statement_and_transaction_end_to_end() { + use depot_client::vfs::{ + SqliteGetPagesProfile, SqliteOperationMetric, SqliteOperationProfile, + SqliteTransactionMetric, SqliteVfsMetrics, + }; + + let actor_name = "counter-sqlite-profile-e2e"; + let metrics = ActorMetrics::new(actor_name); + let statement_profile = SqliteOperationProfile { + worker_wait_ns: 2_000_000, + storage_ns: 3_000_000, + sqlite_requested_pages: 3, + cache_hit_pages: 1, + cache_miss_pages: 2, + depot_demand_requested_pages: 2, + response_present_pages: 3, + overflow_expansion_extra_pages: 1, + storage_response_bytes: 12_288, + bind_count: 1, + bind_logical_bytes: 8, + result_rows: 3, + result_columns: 2, + result_logical_bytes: 48, + get_pages_round_trips: 1, + get_pages_requests: std::array::from_fn(|index| { + (index == 0).then_some(SqliteGetPagesProfile { + ordinal: 1, + duration_ns: 3_000_000, + demand_requested: 2, + response_present: 3, + overflow_expansion_extra: 1, + response_bytes: 12_288, + success: true, + ..Default::default() + }) + }), + ..Default::default() + }; + metrics.observe_operation_profile(&SqliteOperationMetric { + operation_type: "statement", + fingerprint: "select-e2e000000000001".to_owned(), + fingerprint_source: "query", + transaction_mode: "autocommit", + storage_transport: "proxy", + outcome: "success", + sql_bytes: 32, + total_ns: 10_000_000, + transaction_wait_ns: 1_000_000, + profile: statement_profile, + }); + metrics.observe_transaction_profile(&SqliteTransactionMetric { + fingerprint: "txn-e2e0000000000002".to_owned(), + fingerprint_source: "name", + shape_fingerprint: "shape-e2e000000000003".to_owned(), + statement_fingerprint_hashes: [None; + depot_client::vfs::MAX_PROFILED_TRANSACTION_STATEMENTS], + omitted_statement_fingerprints: 0, + storage_transport: "proxy", + outcome: "rollback", + total_ns: 25_000_000, + transaction_wait_ns: 2_000_000, + worker_wait_ns: 3_000_000, + storage_ns: 4_000_000, + local_work_ns: 5_000_000, + application_time_ns: 11_000_000, + commit_ns: 0, + get_pages_round_trips: 2, + statement_count: 3, + dirty_pages: 2, + dirty_bytes: 8_192, + }); + + let rendered = render_global_metrics(); + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_duration_seconds_count", + actor_name, + &[ + "type=\"statement\"", + "fingerprint=\"select-e2e000000000001\"", + "outcome_class=\"success\"", + ], + "1", + ); + for phase in ["transaction_wait", "worker_wait", "storage", "local_work"] { + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_phase_duration_seconds_count", + actor_name, + &["type=\"statement\"", &format!("phase=\"{phase}\"")], + "1", + ); + } + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_outcome_total", + actor_name, + &["type=\"transaction\"", "outcome=\"rollback\""], + "1", + ); + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_transaction_statement_count_count", + actor_name, + &["fingerprint=\"txn-e2e0000000000002\""], + "1", + ); + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_local_pages_total", + actor_name, + &["type=\"statement\"", "page_kind=\"cache_miss\""], + "2", + ); + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_local_bytes_total", + actor_name, + &["type=\"statement\"", "byte_kind=\"result_logical\""], + "48", + ); + assert_metric_value_with_labels( + &rendered, + "rivetkit_sqlite_get_pages_duration_seconds_count", + actor_name, + &["request_ordinal=\"1\"", "outcome_class=\"success\""], + "1", + ); + assert!( + std::mem::size_of::() < 2 * 1024, + "per-operation storage profile must stay below 2 KiB", + ); + } + + #[cfg(feature = "sqlite-local")] + #[test] + fn sqlite_fast_statement_requires_repetition_before_fingerprint_admission() { + use depot_client::vfs::{SqliteOperationMetric, SqliteVfsMetrics}; + + let actor_name = "counter-sqlite-profile-repetition"; + let metrics = ActorMetrics::new_with_sqlite_profiling( + actor_name, + crate::SqliteProfilingConfig { + slow_operation_threshold_ms: 100, + ..Default::default() + }, + ); + let profile = SqliteOperationMetric { + operation_type: "statement", + fingerprint: "select-repeat00000001".to_owned(), + fingerprint_source: "query", + transaction_mode: "autocommit", + storage_transport: "proxy", + outcome: "success", + sql_bytes: 8, + total_ns: 1_000_000, + transaction_wait_ns: 0, + profile: Default::default(), + }; + + assert!(!metrics.observe_operation_profile(&profile)); + assert!(metrics.observe_operation_profile(&profile)); + + let rendered = render_global_metrics(); + assert!(rendered.lines().any(|line| { + line.starts_with("rivetkit_sqlite_duration_seconds_count{") + && line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.contains("fingerprint=\"other\"") + })); + assert!(rendered.lines().any(|line| { + line.starts_with("rivetkit_sqlite_duration_seconds_count{") + && line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.contains("fingerprint=\"select-repeat00000001\"") + })); + } + + #[cfg(feature = "sqlite-local")] + #[test] + fn sqlite_disabled_profiling_records_no_profile_metrics_or_events() { + use depot_client::vfs::{SqliteOperationMetric, SqliteVfsMetrics}; + + let actor_name = "counter-sqlite-profile-disabled"; + let metrics = ActorMetrics::new_with_sqlite_profiling( + actor_name, + crate::SqliteProfilingConfig { + enabled: false, + ..Default::default() + }, + ); + let profile = SqliteOperationMetric { + operation_type: "statement", + fingerprint: "select-disabled000001".to_owned(), + fingerprint_source: "query", + transaction_mode: "autocommit", + storage_transport: "proxy", + outcome: "error", + sql_bytes: 8, + total_ns: 100_000_000, + transaction_wait_ns: 0, + profile: Default::default(), + }; + + assert!(!metrics.observe_operation_profile(&profile)); + metrics.set_worker_queue_depth(1); + metrics.set_worker_inflight(true); + metrics.set_coordinator_queue_depth(1); + metrics.emit_operation_diagnostic_event("private-actor-id", Some(1), &profile); + + let rendered = render_global_metrics(); + assert!(!rendered.lines().any(|line| { + line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.starts_with("rivetkit_sqlite_") + })); + } + + #[cfg(feature = "sqlite-local")] + #[test] + fn sqlite_diagnostic_rate_limit_reports_dropped_events() { + use depot_client::vfs::{SqliteOperationMetric, SqliteVfsMetrics}; + + let actor_name = "counter-sqlite-profile-event-drop"; + let metrics = ActorMetrics::new_with_sqlite_profiling( + actor_name, + crate::SqliteProfilingConfig { + slow_operation_threshold_ms: 0, + max_diagnostic_events_per_minute: 0, + ..Default::default() + }, + ); + metrics.emit_operation_diagnostic_event( + "private-actor-id", + Some(1), + &SqliteOperationMetric { + operation_type: "statement", + fingerprint: "select-drop0000000001".to_owned(), + fingerprint_source: "query", + transaction_mode: "autocommit", + storage_transport: "proxy", + outcome: "success", + sql_bytes: 8, + total_ns: 1, + transaction_wait_ns: 0, + profile: Default::default(), + }, + ); + + let rendered = render_global_metrics(); + assert_metric_value_with_label( + &rendered, + "rivetkit_sqlite_event_dropped_total", + actor_name, + "reason=\"rate_limit\"", + "1", + ); + } + + #[cfg(feature = "sqlite-local")] + #[tokio::test(flavor = "multi_thread")] + async fn sqlite_operations_report_profiles_through_the_native_stack() { + use depot_client::vfs::SqliteVfsMetrics; + + let actor_name = "counter-sqlite-native-stack"; + let actor_id = "counter-sqlite-native-stack-id"; + let profiling = crate::SqliteProfilingConfig { + slow_operation_threshold_ms: u64::MAX, + baseline_sample_rate: 0.0, + max_diagnostic_events_per_minute: 0, + ..Default::default() + }; + let metrics = Arc::new(ActorMetrics::new_with_sqlite_profiling( + actor_name, + profiling.clone(), + )); + let metric_sink: Arc = metrics.clone(); + let transport = Arc::new(InMemorySqliteTransport::default()); + let native_db = depot_client::database::open_database_from_transport( + transport.clone(), + actor_id.to_owned(), + 1, + tokio::runtime::Handle::current(), + Some(metric_sink.clone()), + ) + .await + .expect("native database should open"); + let db = crate::SqliteDb::from_native_database_for_test( + actor_id, + 1, + native_db, + metric_sink.clone(), + profiling, + ); + + db.exec("CREATE TABLE items (id INTEGER PRIMARY KEY, value TEXT NOT NULL)") + .await + .expect("table should be created"); + for (id, value) in [(1, "alpha"), (2, "beta")] { + db.execute( + "INSERT INTO items (id, value) VALUES (?, ?)", + Some(vec![ + crate::BindParam::Integer(id), + crate::BindParam::Text(value.to_owned()), + ]), + ) + .await + .expect("row should be inserted"); + } + let first = db + .query( + "SELECT value FROM items WHERE id = ?", + Some(vec![crate::BindParam::Integer(1)]), + ) + .await + .expect("first query should execute"); + let second = db + .query( + "SELECT value FROM items WHERE id = ?", + Some(vec![crate::BindParam::Integer(2)]), + ) + .await + .expect("second query should execute"); + assert_eq!(first.rows.len(), 1); + assert_eq!(second.rows.len(), 1); + db.exec( + "WITH RECURSIVE ids(id) AS (SELECT 3 UNION ALL SELECT id + 1 FROM ids WHERE id < 220) INSERT INTO items (id, value) SELECT id, zeroblob(4096) FROM ids", + ) + .await + .expect("database should grow beyond the preload window"); + + let transaction = db + .begin_named_transaction(Some("update-item"), None) + .await + .expect("named transaction should begin"); + transaction + .execute( + "UPDATE items SET value = ? WHERE id = ?", + Some(vec![ + crate::BindParam::Text("updated".to_owned()), + crate::BindParam::Integer(1), + ]), + ) + .await + .expect("transaction statement should execute"); + transaction + .commit() + .await + .expect("transaction should commit"); + db.close() + .await + .expect("first native database should close"); + + let reopened_native_db = depot_client::database::open_database_from_transport( + transport, + actor_id.to_owned(), + 2, + tokio::runtime::Handle::current(), + Some(metric_sink.clone()), + ) + .await + .expect("native database should reopen"); + let reopened = crate::SqliteDb::from_native_database_for_test( + actor_id, + 2, + reopened_native_db, + metric_sink, + crate::SqliteProfilingConfig { + slow_operation_threshold_ms: u64::MAX, + baseline_sample_rate: 0.0, + max_diagnostic_events_per_minute: 0, + ..Default::default() + }, + ); + let cold_result = reopened + .query("SELECT value FROM items WHERE id = 220", None) + .await + .expect("cold query should cross the VFS transport"); + assert_eq!(cold_result.rows.len(), 1); + reopened + .close() + .await + .expect("reopened native database should close"); + + let rendered = render_global_metrics(); + assert!( + rendered.lines().any(|line| { + line.starts_with("rivetkit_sqlite_duration_seconds_count{") + && line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.contains("type=\"statement\"") + && line.contains("fingerprint=\"select-") + }), + "repeated statements should render under one select fingerprint:\n{rendered}" + ); + assert!( + rendered.lines().any(|line| { + line.starts_with("rivetkit_sqlite_outcome_total{") + && line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.contains("type=\"transaction\"") + && line.contains("fingerprint=\"txn-") + && line.contains("outcome=\"success\"") + && line.ends_with(" 1") + }), + "named transaction outcome should render:\n{rendered}" + ); + for byte_kind in [ + "bind_logical", + "result_logical", + "storage_response", + "dirty", + ] { + assert!( + rendered.lines().any(|line| { + line.starts_with("rivetkit_sqlite_local_bytes_total{") + && line.contains(&format!("actor_name=\"{actor_name}\"")) + && line.contains(&format!("byte_kind=\"{byte_kind}\"")) + }), + "{byte_kind} should be reported through the native stack:\n{rendered}" + ); + } + } + #[test] fn actor_active_count_tracks_metric_lifetime() { let actor_name = "counter-active"; @@ -302,7 +823,7 @@ mod moved_tests { let line = metrics .lines() .find(|line| { - line.starts_with(name) + metric_name_matches(line, name) && line.contains(&format!("actor_name=\"{actor_name}\"")) && labels.iter().all(|label| line.contains(label)) }) diff --git a/rivetkit-rust/packages/rivetkit-core/tests/metrics_helpers.rs b/rivetkit-rust/packages/rivetkit-core/tests/metrics_helpers.rs index 1bff6771d9..ea80dc0e8b 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/metrics_helpers.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/metrics_helpers.rs @@ -11,5 +11,12 @@ pub(crate) fn render_global_metrics() -> String { } pub(crate) fn metric_line_for_actor(line: &str, name: &str, actor_name: &str) -> bool { - line.starts_with(name) && line.contains(&format!("actor_name=\"{actor_name}\"")) + metric_name_matches(line, name) && line.contains(&format!("actor_name=\"{actor_name}\"")) +} + +pub(crate) fn metric_name_matches(line: &str, name: &str) -> bool { + line.starts_with(name) + || line + .strip_prefix("rivet_") + .is_some_and(|unprefixed| unprefixed.starts_with(name)) } diff --git a/rivetkit-rust/packages/rivetkit/src/lib.rs b/rivetkit-rust/packages/rivetkit/src/lib.rs index 792bfe1295..7abf13fb2d 100644 --- a/rivetkit-rust/packages/rivetkit/src/lib.rs +++ b/rivetkit-rust/packages/rivetkit/src/lib.rs @@ -6,6 +6,8 @@ pub mod persist; pub mod prelude; pub mod queue; pub mod registry; +#[cfg(feature = "sqlite-local")] +pub mod sqlite; pub mod start; pub mod test; pub mod typed_client; @@ -23,6 +25,8 @@ pub use crate::{ start::{Events, Hibernated, Input, Snapshot, Start, run_actor}, typed_client::{IntoActorKey, TypedActorConnection, TypedActorHandle, TypedClientExt}, }; +#[cfg(feature = "sqlite-local")] +pub use crate::sqlite::{SqliteDbExt, SqliteTransactionOptions}; pub use rivetkit_client as client; pub use rivetkit_core::actor::state::OnStateChangeGuard; pub use rivetkit_core::metrics_endpoint::RenderedMetrics; diff --git a/rivetkit-rust/packages/rivetkit/src/prelude.rs b/rivetkit-rust/packages/rivetkit/src/prelude.rs index d09983da06..f3ff98984f 100644 --- a/rivetkit-rust/packages/rivetkit/src/prelude.rs +++ b/rivetkit-rust/packages/rivetkit/src/prelude.rs @@ -1,6 +1,8 @@ pub use anyhow::{Result, anyhow}; pub use crate::{ - Action, Actor, ConnCtx, Ctx, Event, Handles, Registry, RequestSaveOpts, RuntimeEvent, Start, - StateMut, StateRef, action, + Action, Actor, ConnCtx, Ctx, Event, Handles, Registry, RequestSaveOpts, RuntimeEvent, + Start, StateMut, StateRef, action, }; +#[cfg(feature = "sqlite-local")] +pub use crate::{SqliteDbExt, SqliteTransactionOptions}; diff --git a/rivetkit-rust/packages/rivetkit/src/sqlite.rs b/rivetkit-rust/packages/rivetkit/src/sqlite.rs new file mode 100644 index 0000000000..8d0526d9df --- /dev/null +++ b/rivetkit-rust/packages/rivetkit/src/sqlite.rs @@ -0,0 +1,54 @@ +use std::{future::Future, time::Duration}; + +use anyhow::Result; +use rivetkit_core::{SqliteDb, SqliteTransaction}; + +#[derive(Clone, Copy, Debug, Default)] +pub struct SqliteTransactionOptions<'a> { + pub name: Option<&'a str>, + pub timeout: Option, +} + +/// Ergonomic commit-on-success transaction helper for the high-level Rust API. +/// Coordination and profiling remain owned by `rivetkit-core`. +pub trait SqliteDbExt { + fn transaction<'a, T, F, Fut>( + &'a self, + callback: F, + options: SqliteTransactionOptions<'a>, + ) -> impl Future> + Send + 'a + where + T: Send + 'a, + F: FnOnce(SqliteTransaction) -> Fut + Send + 'a, + Fut: Future> + Send + 'a; +} + +impl SqliteDbExt for SqliteDb { + async fn transaction<'a, T, F, Fut>( + &'a self, + callback: F, + options: SqliteTransactionOptions<'a>, + ) -> Result + where + T: Send + 'a, + F: FnOnce(SqliteTransaction) -> Fut + Send + 'a, + Fut: Future> + Send + 'a, + { + let transaction = self + .begin_named_transaction(options.name, options.timeout) + .await?; + match callback(transaction.clone()).await { + Ok(value) => { + if let Err(error) = transaction.commit().await { + let _ = transaction.rollback().await; + return Err(error); + } + Ok(value) + } + Err(error) => { + let _ = transaction.rollback().await; + Err(error) + } + } + } +} diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 6d99414903..fac208ae90 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -67,11 +67,29 @@ export interface JsInspectorTabEntry { /** Set to true for HideBuiltin entries. */ hidden?: boolean } +/** + * Experimental SQLite profiling configuration. This entire configuration + * surface is subject to change without notice. + */ +export interface JsSqliteProfilingConfig { + enabled?: boolean + maxTrackedStatementFingerprints?: number + maxTrackedTransactionFingerprints?: number + maxPrometheusSeries?: number + maxStatementsPerTransactionTrace?: number + maxGetPagesRequestsPerTrace?: number + maxTransactionNameBytes?: number + slowOperationThresholdMs?: number + baselineSampleRate?: number + maxDiagnosticEventsPerMinute?: number + diagnosticEventQueueCapacity?: number +} export interface JsActorConfig { name?: string icon?: string hasDatabase?: boolean remoteSqlite?: boolean + sqliteProfiling?: JsSqliteProfilingConfig hasState?: boolean canHibernateWebsocket?: boolean stateSaveIntervalMs?: number diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs index c25e8c92ce..fa6ac8f1b9 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_factory.rs @@ -12,7 +12,7 @@ use rivetkit_core::inspector::InspectorTabEntry; use rivetkit_core::{ ActionDefinition, ActorConfig, ActorConfigInput, ActorContext as CoreActorContext, ActorFactory as CoreActorFactory, ConnHandle as CoreConnHandle, Request, Response, - WebSocket as CoreWebSocket, + SqliteProfilingConfigInput, WebSocket as CoreWebSocket, }; use crate::actor_context::{ActorContext, StateDeltaPayload}; @@ -78,6 +78,24 @@ pub struct JsInspectorTabEntry { pub hidden: Option, } +/// Experimental SQLite profiling configuration. This entire configuration +/// surface is subject to change without notice. +#[napi(object)] +#[derive(Clone, Default)] +pub struct JsSqliteProfilingConfig { + pub enabled: Option, + pub max_tracked_statement_fingerprints: Option, + pub max_tracked_transaction_fingerprints: Option, + pub max_prometheus_series: Option, + pub max_statements_per_transaction_trace: Option, + pub max_get_pages_requests_per_trace: Option, + pub max_transaction_name_bytes: Option, + pub slow_operation_threshold_ms: Option, + pub baseline_sample_rate: Option, + pub max_diagnostic_events_per_minute: Option, + pub diagnostic_event_queue_capacity: Option, +} + #[napi(object)] #[derive(Clone, Default)] pub struct JsActorConfig { @@ -85,6 +103,7 @@ pub struct JsActorConfig { pub icon: Option, pub has_database: Option, pub remote_sqlite: Option, + pub sqlite_profiling: Option, pub has_state: Option, pub can_hibernate_websocket: Option, pub state_save_interval_ms: Option, @@ -1021,6 +1040,7 @@ impl From for ActorConfigInput { icon: value.icon, has_database: value.has_database, remote_sqlite: value.remote_sqlite, + sqlite_profiling: value.sqlite_profiling.map(Into::into), has_state: value.has_state, can_hibernate_websocket: value.can_hibernate_websocket, state_save_interval_ms: value.state_save_interval_ms, @@ -1085,6 +1105,24 @@ impl From for ActorConfigInput { } } +impl From for SqliteProfilingConfigInput { + fn from(value: JsSqliteProfilingConfig) -> Self { + Self { + enabled: value.enabled, + max_tracked_statement_fingerprints: value.max_tracked_statement_fingerprints, + max_tracked_transaction_fingerprints: value.max_tracked_transaction_fingerprints, + max_prometheus_series: value.max_prometheus_series, + max_statements_per_transaction_trace: value.max_statements_per_transaction_trace, + max_get_pages_requests_per_trace: value.max_get_pages_requests_per_trace, + max_transaction_name_bytes: value.max_transaction_name_bytes, + slow_operation_threshold_ms: value.slow_operation_threshold_ms, + baseline_sample_rate: value.baseline_sample_rate, + max_diagnostic_events_per_minute: value.max_diagnostic_events_per_minute, + diagnostic_event_queue_capacity: value.diagnostic_event_queue_capacity, + } + } +} + // Test shim keeps moved tests in crate-root tests/ with private-module access. #[cfg(test)] #[path = "../tests/actor_factory.rs"] diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index b0ce74520b..e2bc793dd9 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -171,6 +171,8 @@ pub struct WasmActionDefinition { pub name: String, } +/// Experimental SQLite profiling configuration. This entire configuration +/// surface is subject to change without notice. #[derive(Clone, Default, serde::Deserialize)] #[serde(default, rename_all = "camelCase")] pub struct WasmActorConfig { diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts index 54a7d28ded..a204316a1a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/config.ts @@ -13,6 +13,25 @@ export type InferDatabaseClient = export type SqliteBindings = unknown[] | Record; +/** + * SQLite profiling controls. + * + * @experimental This entire configuration surface is subject to change. + */ +export interface SqliteProfilingOptions { + enabled?: boolean; + maxTrackedStatementFingerprints?: number; + maxTrackedTransactionFingerprints?: number; + maxPrometheusSeries?: number; + maxStatementsPerTransactionTrace?: number; + maxGetPagesRequestsPerTrace?: number; + maxTransactionNameBytes?: number; + slowOperationThresholdMs?: number; + baselineSampleRate?: number; + maxDiagnosticEventsPerMinute?: number; + diagnosticEventQueueCapacity?: number; +} + export interface SqliteQueryResult { columns: string[]; rows: unknown[][]; @@ -102,6 +121,10 @@ export interface DatabaseProviderContext { } export type DatabaseProvider = { + /** + * @experimental This entire configuration surface is subject to change. + */ + sqliteProfiling?: SqliteProfilingOptions; /** * Creates a new database client for the actor. * The result is passed to the actor context as `c.db`. diff --git a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts index 975511e6e1..bdd1bc262d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/common/database/mod.ts @@ -1,10 +1,19 @@ -import type { DatabaseProvider, RawAccess, SqliteDatabase } from "./config"; +import type { + DatabaseProvider, + RawAccess, + SqliteDatabase, + SqliteProfilingOptions, +} from "./config"; import { isSqliteBindingObject, toSqliteBindings } from "./shared"; export type { RawAccess } from "./config"; -interface DatabaseFactoryConfig { +export interface DatabaseFactoryConfig { onMigrate?: (db: RawAccess) => Promise | void; + /** + * @experimental This entire configuration surface is subject to change. + */ + profiling?: SqliteProfilingOptions; } function hasMultipleStatements(query: string): boolean { @@ -14,8 +23,10 @@ function hasMultipleStatements(query: string): boolean { export function db({ onMigrate, + profiling, }: DatabaseFactoryConfig = {}): DatabaseProvider { return { + sqliteProfiling: profiling, createClient: async (ctx) => { const nativeDatabaseProvider = ctx.nativeDatabaseProvider; if (!nativeDatabaseProvider) { diff --git a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts index c5daf6abd5..d0d157f252 100644 --- a/rivetkit-typescript/packages/rivetkit/src/db/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/db/mod.ts @@ -8,6 +8,7 @@ export type { SqliteBindings, SqliteDatabase, SqliteNativeMetrics, + SqliteProfilingOptions, SqliteQueryResult, } from "@/common/database/config"; export { db } from "@/common/database/mod"; diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 3c94653463..97357229fb 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -33,7 +33,10 @@ import { } from "@/client/client"; import { convertRegistryConfigToClientConfig } from "@/client/config"; import { HEADER_CONN_PARAMS } from "@/common/actor-router-consts"; -import type { AnyDatabaseProvider } from "@/common/database/config"; +import type { + AnyDatabaseProvider, + SqliteProfilingOptions, +} from "@/common/database/config"; import { wrapJsNativeDatabase } from "@/common/database/native-database"; import { assertJsonCompatValue, type JsonCompatValue } from "@/common/encoding"; import { decodeWorkflowHistoryTransport } from "@/common/inspector-transport"; @@ -3293,10 +3296,15 @@ function buildActorConfig( const options = (config.options ?? {}) as Record; const canHibernate = options.canHibernateWebSocket; + const sqliteProfiling = ( + config.db as { sqliteProfiling?: SqliteProfilingOptions } | undefined + )?.sqliteProfiling; + return { name: options.name as string | undefined, icon: options.icon as string | undefined, hasDatabase: config.db !== undefined, + sqliteProfiling, remoteSqlite: config.db !== undefined && sqliteBackendForConfig(registryConfig) === "remote", diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index b8dca80bb8..c82fb55e54 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -1,5 +1,8 @@ import { stringifyError } from "@/common/utils"; -import type { SqliteNativeMetrics } from "@/common/database/config"; +import type { + SqliteNativeMetrics, + SqliteProfilingOptions, +} from "@/common/database/config"; import type { RegistryConfig } from "./config"; import { logger } from "./log"; @@ -201,6 +204,8 @@ export interface RuntimeActorConfig { icon?: string; hasDatabase?: boolean; remoteSqlite?: boolean; + /** @experimental This entire configuration surface is subject to change. */ + sqliteProfiling?: SqliteProfilingOptions; hasState?: boolean; canHibernateWebsocket?: boolean; stateSaveIntervalMs?: number; From 4a9a623d79c62051147a58649cb19d05d6dfc3a4 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Fri, 28 Aug 2026 02:24:00 -0700 Subject: [PATCH 2/2] ci(release): enable npm oidc publishing --- .github/workflows/publish.yaml | 14 ++++++++++---- scripts/publish/src/ci/bin.ts | 5 +++++ scripts/publish/src/lib/npm.ts | 10 ++++++++++ scripts/publish/src/lib/version.ts | 31 ++++++++++++++++++++++++++++++ 4 files changed, 56 insertions(+), 4 deletions(-) diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index e28c13555e..03d727f2f8 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -55,7 +55,7 @@ jobs: - run: corepack enable - uses: actions/setup-node@v4 with: - node-version: "22" + node-version: "24" cache: pnpm - name: Install publish scripts run: pnpm install --frozen-lockfile --filter=publish @@ -363,7 +363,8 @@ jobs: !cancelled() && needs.build.result == 'success' && needs.docker-images.result == 'success' - runs-on: depot-ubuntu-24.04-8 + # npm trusted publishing currently requires a GitHub-hosted runner. + runs-on: ubuntu-24.04 permissions: contents: write # git tag + gh release (release only) id-token: write @@ -378,9 +379,11 @@ jobs: - run: corepack enable - uses: actions/setup-node@v4 with: - node-version: "22" + node-version: "24" registry-url: "https://registry.npmjs.org" cache: pnpm + - name: Install OIDC-capable npm + run: npm install --global npm@11.16.0 - run: pnpm install --frozen-lockfile - uses: ./.github/actions/docker-setup with: @@ -538,7 +541,10 @@ jobs: - name: Publish npm packages env: - NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} + # setup-node provides a dummy token value; clear it so token-only + # post-publish operations stay disabled for trusted publishing. + NODE_AUTH_TOKEN: "" + SKIP_WASM_BUILD: "1" run: | pnpm --filter=publish exec tsx src/ci/bin.ts publish-npm \ --tag ${{ needs.context.outputs.npm_tag }} \ diff --git a/scripts/publish/src/ci/bin.ts b/scripts/publish/src/ci/bin.ts index 17fc5ea3c4..e63aea36fe 100644 --- a/scripts/publish/src/ci/bin.ts +++ b/scripts/publish/src/ci/bin.ts @@ -186,6 +186,10 @@ program "--version-only", "Only rewrite package.json version fields without publish-time dependency injection", ) + .option( + "--repository ", + "GitHub repository recorded in publish-time package metadata (defaults to GITHUB_REPOSITORY)", + ) .option("--dry-run", "Do not write, only report") .action(async (opts) => { const repoRoot = findRepoRoot(); @@ -195,6 +199,7 @@ program dryRun: !!opts.dryRun, includeReleaseOnlyPackages: ctx.trigger === "release", versionOnly: !!opts.versionOnly, + repository: opts.repository ?? process.env.GITHUB_REPOSITORY, }); await bumpCargoVersions(repoRoot, version, { dryRun: !!opts.dryRun, diff --git a/scripts/publish/src/lib/npm.ts b/scripts/publish/src/lib/npm.ts index 9efcc1b0db..889057cc6e 100644 --- a/scripts/publish/src/lib/npm.ts +++ b/scripts/publish/src/lib/npm.ts @@ -220,6 +220,15 @@ export async function repairBranchPreviewLatestTags( opts: Required> & Pick, ): Promise { + // npm trusted publishing only authenticates `npm publish`; commands such as + // `npm dist-tag add` still require a traditional token. Preserve the repair + // for manual/token-authenticated runs, but do not make OIDC publishes fail + // after their packages have already been published successfully. + if (!process.env.NODE_AUTH_TOKEN) { + log.warn("skipping preview latest-tag repair during tokenless OIDC publish"); + return; + } + const previewPrefix = `0.0.0-${opts.tag}.`; const packages = discoverPackages(repoRoot, { includeReleaseOnly: opts.includeReleaseOnlyPackages, @@ -372,6 +381,7 @@ export async function publishAll( for (const r of results.filter((x) => x.status === "failed")) { log.error(` - ${r.pkg.name}: ${r.lastError}`); } + throw new Error(`${counts.failed} npm package(s) failed to publish`); } // In release mode, if *every* package was already published, treat it as diff --git a/scripts/publish/src/lib/version.ts b/scripts/publish/src/lib/version.ts index fd58d289b6..102fd2ad8d 100644 --- a/scripts/publish/src/lib/version.ts +++ b/scripts/publish/src/lib/version.ts @@ -32,6 +32,11 @@ const log = scoped("version"); interface PackageJson { name?: string; version?: string; + repository?: { + type: "git"; + url: string; + directory: string; + }; dependencies?: Record; devDependencies?: Record; peerDependencies?: Record; @@ -78,6 +83,26 @@ export interface BumpOptions { * the publish-time mode used by CI — never committed. */ versionOnly?: boolean; + /** GitHub repository slug recorded in publish-time package metadata. */ + repository?: string; +} + +export function githubRepositoryUrl(repository: string): string { + if (!/^[A-Za-z0-9_.-]+\/[A-Za-z0-9_.-]+$/.test(repository)) { + throw new Error( + `invalid GitHub repository ${JSON.stringify(repository)}; expected owner/repo`, + ); + } + return `https://github.com/${repository}.git`; +} + +function requirePublishRepository(repository: string | undefined): string { + if (!repository) { + throw new Error( + "publish-time package metadata requires a GitHub repository", + ); + } + return repository; } /** @@ -117,6 +142,12 @@ export async function bumpPackageJsons( pkgJson.version = version; if (!versionOnly) { + pkgJson.repository = { + type: "git", + url: githubRepositoryUrl(requirePublishRepository(opts.repository)), + directory: pkg.relDir, + }; + // Inject optionalDependencies on meta packages so end users get the // correct platform-specific binary via npm's os/cpu/libc resolution. const platformPkgs = metaPlatformMap.get(pkg.name);