1use std::collections::{BTreeMap, HashMap, HashSet};
6use std::ffi::{CStr, CString, c_char, c_int, c_void};
7use std::ptr;
8use std::slice;
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
11use std::time::{Duration, Instant};
12
13use crate::page_buffer::{PageBudget, PageBuffer, PageTrace};
14use anyhow::Result;
15use async_trait::async_trait;
16use futures_util::stream::{self, StreamExt};
17use libsqlite3_sys::*;
18use parking_lot::{Mutex, RwLock};
19use rivet_envoy_protocol as protocol;
20use sqlite_types::{PageBytes, is_head_fence_mismatch};
21use tokio::runtime::Handle;
22use tracing::Instrument;
23
24use crate::optimization_flags::{
25 SqliteOptimizationFlags, SqliteVfsPageCacheMode, sqlite_optimization_flags,
26};
27use crate::sqlite_page::{PageClass, classify};
28
29const DEFAULT_MAX_PAGES_PER_STAGE: usize = 4_000;
30const DEFAULT_PAGE_SIZE: usize = 4096;
31const STAGE_PIPELINE_BYTE_BUDGET: usize = 8 * 1024 * 1024;
32const STAGE_MAX_SEGMENT_BYTES: usize = sqlite_types::SHARD_SIZE as usize
33 * sqlite_types::COMMIT_SEGMENT_MAX_SHARDS as usize
34 * DEFAULT_PAGE_SIZE;
35const STAGE_PIPELINE_DEPTH: usize = STAGE_PIPELINE_BYTE_BUDGET / STAGE_MAX_SEGMENT_BYTES;
36const NATIVE_DATABASE_DROP_FLUSH_TIMEOUT: Duration = Duration::from_millis(250);
37const MAX_PATHNAME: c_int = 64;
38const TEMP_AUX_PATH_PREFIX: &str = "__sqlite_temp__";
39const SQLITE_HEADER_MAGIC: &[u8; 16] = b"SQLite format 3\0";
40const EMPTY_DB_PAGE_HEADER_PREFIX: [u8; 108] = [
41 83, 81, 76, 105, 116, 101, 32, 102, 111, 114, 109, 97, 116, 32, 51, 0, 16, 0, 1, 1, 0, 64, 32,
42 32, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
43 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
44 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 46, 138, 17, 13, 0, 0, 0, 0, 16, 0, 0,
45];
46
47static NEXT_TEMP_AUX_ID: AtomicU64 = AtomicU64::new(1);
48
49unsafe extern "C" {
50 fn sqlite3_close_v2(db: *mut sqlite3) -> c_int;
51}
52
53fn empty_db_page() -> Vec<u8> {
54 let mut page = vec![0u8; DEFAULT_PAGE_SIZE];
55 page[..EMPTY_DB_PAGE_HEADER_PREFIX.len()].copy_from_slice(&EMPTY_DB_PAGE_HEADER_PREFIX);
56 page
57}
58
59fn sqlite_header_page_size(page: &[u8]) -> Option<usize> {
60 if page.len() < 100 || &page[..SQLITE_HEADER_MAGIC.len()] != SQLITE_HEADER_MAGIC {
61 return None;
62 }
63
64 let raw = u16::from_be_bytes([page[16], page[17]]);
65 let page_size = if raw == 1 { 65_536 } else { usize::from(raw) };
66
67 if (512..=65_536).contains(&page_size) && page_size.is_power_of_two() {
68 Some(page_size)
69 } else {
70 None
71 }
72}
73
74fn sqlite_header_db_size_pages(page: &[u8]) -> Option<u32> {
75 if page.len() < 100 || &page[..SQLITE_HEADER_MAGIC.len()] != SQLITE_HEADER_MAGIC {
76 return None;
77 }
78
79 let db_size_pages = u32::from_be_bytes([page[28], page[29], page[30], page[31]]);
80 Some(db_size_pages.max(1))
81}
82
83fn panic_message(payload: &Box<dyn std::any::Any + Send>) -> String {
84 if let Some(message) = payload.downcast_ref::<&str>() {
85 message.to_string()
86 } else if let Some(message) = payload.downcast_ref::<String>() {
87 message.clone()
88 } else {
89 "unknown panic".to_string()
90 }
91}
92
93macro_rules! vfs_catch_unwind {
94 ($err_val:expr, $body:expr) => {
95 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| $body)) {
96 Ok(result) => result,
97 Err(panic) => {
98 tracing::error!(msg = panic_message(&panic), "sqlite callback panicked");
99 $err_val
100 }
101 }
102 };
103}
104
105#[async_trait]
106pub trait SqliteTransport: Send + Sync {
107 async fn get_pages(
108 &self,
109 request: protocol::SqliteGetPagesRequest,
110 ) -> Result<protocol::SqliteGetPagesResponse>;
111
112 async fn commit(
113 &self,
114 request: protocol::SqliteCommitRequest,
115 ) -> Result<protocol::SqliteCommitResponse>;
116
117 async fn commit_stage_begin(
124 &self,
125 request: protocol::SqliteCommitStageBeginRequest,
126 ) -> Result<protocol::SqliteCommitStageBeginResponse>;
127
128 async fn commit_stage_segment(
130 &self,
131 request: protocol::SqliteCommitStageSegmentRequest,
132 ) -> Result<protocol::SqliteCommitStageSegmentResponse>;
133
134 async fn commit_finalize(
136 &self,
137 request: protocol::SqliteCommitFinalizeRequest,
138 ) -> Result<protocol::SqliteCommitFinalizeResponse>;
139}
140
141pub type SqliteTransportHandle = Arc<dyn SqliteTransport>;
142fn sqlite_now_ms() -> Result<i64> {
143 use std::time::{SystemTime, UNIX_EPOCH};
144
145 Ok(SystemTime::now()
146 .duration_since(UNIX_EPOCH)?
147 .as_millis()
148 .try_into()?)
149}
150
151#[derive(Debug, Clone)]
152pub struct VfsConfig {
153 pub page_budget: Arc<PageBudget>,
156 pub page_cache_mode: SqliteVfsPageCacheMode,
157 pub staging_cache_ttl_ms: u64,
158 pub max_pages_per_stage: usize,
159 pub idle_release_ms: u64,
161 pub max_commit_dirty_pages: usize,
164}
165
166impl Default for VfsConfig {
167 fn default() -> Self {
168 Self::from_optimization_flags(*sqlite_optimization_flags())
169 }
170}
171
172impl VfsConfig {
173 pub fn commit_page_cap(&self) -> usize {
175 self.max_commit_dirty_pages
176 .min(sqlite_types::MAX_COMMIT_DIRTY_PAGES)
177 }
178
179 pub fn from_optimization_flags(flags: SqliteOptimizationFlags) -> Self {
180 let caches_pages = flags.vfs_page_cache_mode.caches_any_pages();
181 Self {
182 page_budget: PageBudget::new(
183 if caches_pages {
184 flags.vfs_page_buffer_process_pages
185 } else {
186 0
187 },
188 if caches_pages {
189 flags.vfs_staging_cache_ttl_ms
190 } else {
191 0
192 },
193 None,
194 ),
195 page_cache_mode: flags.vfs_page_cache_mode,
196 staging_cache_ttl_ms: if caches_pages {
197 flags.vfs_staging_cache_ttl_ms
198 } else {
199 0
200 },
201 max_pages_per_stage: DEFAULT_MAX_PAGES_PER_STAGE,
202 idle_release_ms: flags.idle_release_ms,
203 max_commit_dirty_pages: sqlite_types::MAX_COMMIT_DIRTY_PAGES,
204 }
205 }
206}
207
208#[derive(Debug, Clone, Default)]
209pub(crate) struct InitialPages {
210 pub pages: Vec<(u32, PageBytes)>,
211 pub head_txid: Option<u64>,
212 pub requested_page_count: u32,
213}
214
215impl From<Vec<(u32, Vec<u8>)>> for InitialPages {
216 fn from(pages: Vec<(u32, Vec<u8>)>) -> Self {
217 let requested_page_count = pages.len().try_into().unwrap_or(u32::MAX);
218 Self {
219 pages: pages
220 .into_iter()
221 .map(|(pgno, bytes)| (pgno, bytes.into()))
222 .collect(),
223 head_txid: None,
224 requested_page_count,
225 }
226 }
227}
228
229#[derive(Debug, Clone, PartialEq, Eq)]
230pub enum CommitPath {
231 Fast,
232 Slow,
233 Staged,
236}
237
238#[derive(Debug, Clone)]
239pub struct BufferedCommitRequest {
240 pub actor_id: String,
241 pub new_db_size_pages: u32,
242 pub dirty_pages: Vec<BufferedDirtyPage>,
243 pub expected_head_txid: Option<u64>,
244}
245
246#[derive(Debug, Clone)]
247pub struct BufferedDirtyPage {
248 pub pgno: u32,
249 pub bytes: PageBytes,
250}
251
252#[derive(Debug, Clone)]
253pub struct BufferedCommitOutcome {
254 pub path: CommitPath,
255 pub db_size_pages: u32,
256 pub head_txid: Option<u64>,
257}
258
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub enum CommitBufferError {
261 FenceMismatch(String),
262 Indeterminate(String),
263 Storage(SqliteStorageError),
266 Other(String),
267}
268
269#[derive(Debug, Clone, PartialEq, Eq)]
273pub enum SqliteStorageError {
274 Overloaded {
277 retry_after_ms: u64,
278 message: String,
279 },
280 Timeout { operation: String },
283}
284
285impl std::fmt::Display for SqliteStorageError {
286 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
287 match self {
288 Self::Overloaded {
289 retry_after_ms,
290 message,
291 } => write!(
292 f,
293 "sqlite storage is overloaded, retry after {retry_after_ms} ms: {message}"
294 ),
295 Self::Timeout { operation } => {
296 write!(
297 f,
298 "sqlite storage {operation} timed out before storage answered"
299 )
300 }
301 }
302 }
303}
304
305impl std::error::Error for SqliteStorageError {}
306
307#[derive(Debug)]
311pub struct SqliteStorageTimeoutError {
312 pub operation: String,
313}
314
315impl std::fmt::Display for SqliteStorageTimeoutError {
316 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
317 write!(f, "sqlite {} request timed out", self.operation)
318 }
319}
320
321impl std::error::Error for SqliteStorageTimeoutError {}
322
323fn storage_refusal(error: &protocol::SqliteErrorResponse) -> Option<SqliteStorageError> {
325 let retry_after_ms =
326 sqlite_types::storage_overload_retry_after_ms(&error.group, &error.code, &error.message)?;
327 Some(SqliteStorageError::Overloaded {
328 retry_after_ms,
329 message: error.message.to_string(),
330 })
331}
332
333fn storage_timeout(error: &anyhow::Error) -> Option<SqliteStorageError> {
335 error
336 .downcast_ref::<SqliteStorageTimeoutError>()
337 .map(|timeout| SqliteStorageError::Timeout {
338 operation: timeout.operation.clone(),
339 })
340}
341
342#[derive(Debug)]
344pub struct SqliteCommitIndeterminateError(pub String);
345
346impl std::fmt::Display for SqliteCommitIndeterminateError {
347 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
348 write!(f, "sqlite {} result is indeterminate", self.0)
349 }
350}
351
352impl std::error::Error for SqliteCommitIndeterminateError {}
353
354#[derive(Debug, Clone, PartialEq, Eq)]
359pub struct SqliteTransactionTooLargeError {
360 pub max_pages: usize,
361 pub page_size: usize,
362}
363
364impl std::fmt::Display for SqliteTransactionTooLargeError {
365 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
366 let max_kib = self.max_pages * self.page_size / 1024;
367 let max_size = if max_kib >= 1024 {
368 format!("{} MiB", max_kib / 1024)
369 } else {
370 format!("{max_kib} KiB")
371 };
372 write!(
373 f,
374 "sqlite transaction changes more than {} pages ({}), the most one commit can hold, \
375 so the statement failed and its changes were rolled back; split the work into smaller \
376 transactions. VACUUM rewrites the whole \
377 database in one transaction, so it only runs while the database's in-use pages fit under \
378 this limit",
379 self.max_pages, max_size,
380 )
381 }
382}
383
384impl std::error::Error for SqliteTransactionTooLargeError {}
385
386pub(crate) const FCNTL_COMMIT_PAGE_CAP: c_int = 0x5256_0001;
389
390#[derive(Debug, Clone, Copy, Default)]
391pub struct SqliteVfsMetricsSnapshot {
392 pub request_build_ns: u64,
393 pub serialize_ns: u64,
394 pub transport_ns: u64,
395 pub state_update_ns: u64,
396 pub total_ns: u64,
397 pub commit_count: u64,
398 pub page_cache_entries: u64,
399 pub page_cache_weighted_size: u64,
400 pub page_cache_capacity_pages: u64,
401 pub write_buffer_dirty_pages: u64,
402 pub db_size_pages: u64,
403}
404
405pub const MAX_PROFILED_GET_PAGES_REQUESTS: usize = 16;
406pub const MAX_PROFILED_TRANSACTION_STATEMENTS: usize = 32;
407
408#[derive(Debug, Clone, Default)]
411pub struct SqliteGetPagesProfile {
412 pub ordinal: u64,
413 pub duration_ns: u64,
414 pub demand_requested: u64,
415 pub response_present: u64,
416 pub response_missing: u64,
417 pub overflow_expansion_extra: u64,
418 pub response_bytes: u64,
419 pub success: bool,
420}
421
422#[derive(Debug, Clone, Default)]
426pub struct SqliteOperationProfile {
427 pub worker_wait_ns: u64,
428 pub sqlite_execution_ns: u64,
429 pub storage_ns: u64,
430 pub bind_count: u64,
431 pub bind_logical_bytes: u64,
432 pub result_rows: u64,
433 pub result_columns: u64,
434 pub result_logical_bytes: u64,
435 pub sqlite_requested_pages: u64,
436 pub cache_hit_pages: u64,
437 pub cache_miss_pages: u64,
438 pub depot_demand_requested_pages: u64,
439 pub response_present_pages: u64,
440 pub response_missing_pages: u64,
441 pub overflow_expansion_extra_pages: u64,
442 pub btree_pages: u64,
443 pub non_btree_pages: u64,
444 pub prefetch_consumed_pages: u64,
445 pub prefetch_unused_pages: u64,
446 pub dirty_pages: u64,
447 pub storage_response_bytes: u64,
448 pub dirty_bytes: u64,
449 pub get_pages_round_trips: u64,
450 pub commit_round_trips: u64,
451 pub get_pages_requests: [Option<SqliteGetPagesProfile>; MAX_PROFILED_GET_PAGES_REQUESTS],
452 pub omitted_get_pages_requests: u64,
453}
454
455#[derive(Debug, Clone)]
456pub struct SqliteOperationMetric {
457 pub operation_type: &'static str,
458 pub fingerprint: String,
459 pub fingerprint_source: &'static str,
460 pub transaction_mode: &'static str,
461 pub storage_transport: &'static str,
462 pub outcome: &'static str,
463 pub sql_bytes: u64,
464 pub total_ns: u64,
465 pub transaction_wait_ns: u64,
466 pub profile: SqliteOperationProfile,
467}
468
469#[derive(Debug, Clone)]
470pub struct SqliteTransactionMetric {
471 pub fingerprint: String,
472 pub fingerprint_source: &'static str,
473 pub shape_fingerprint: String,
474 pub statement_fingerprint_hashes: [Option<[u8; 16]>; MAX_PROFILED_TRANSACTION_STATEMENTS],
475 pub omitted_statement_fingerprints: u64,
476 pub storage_transport: &'static str,
477 pub outcome: &'static str,
478 pub total_ns: u64,
479 pub transaction_wait_ns: u64,
480 pub worker_wait_ns: u64,
481 pub storage_ns: u64,
482 pub local_work_ns: u64,
483 pub application_time_ns: u64,
484 pub commit_ns: u64,
485 pub get_pages_round_trips: u64,
486 pub statement_count: u64,
487 pub dirty_pages: u64,
488 pub dirty_bytes: u64,
489}
490
491impl SqliteOperationProfile {
492 fn push_get_pages(&mut self, request: SqliteGetPagesProfile, request_limit: usize) {
493 self.storage_ns = self.storage_ns.saturating_add(request.duration_ns);
494 self.get_pages_round_trips = self.get_pages_round_trips.saturating_add(1);
495 self.depot_demand_requested_pages = self
496 .depot_demand_requested_pages
497 .saturating_add(request.demand_requested);
498 self.response_present_pages = self
499 .response_present_pages
500 .saturating_add(request.response_present);
501 self.response_missing_pages = self
502 .response_missing_pages
503 .saturating_add(request.response_missing);
504 self.overflow_expansion_extra_pages = self
505 .overflow_expansion_extra_pages
506 .saturating_add(request.overflow_expansion_extra);
507 self.storage_response_bytes = self
508 .storage_response_bytes
509 .saturating_add(request.response_bytes);
510 let index = self.get_pages_round_trips.saturating_sub(1) as usize;
511 if index < request_limit
512 && let Some(slot) = self.get_pages_requests.get_mut(index)
513 {
514 *slot = Some(request);
515 } else {
516 self.omitted_get_pages_requests = self.omitted_get_pages_requests.saturating_add(1);
517 }
518 }
519}
520
521pub(crate) struct SqliteOperationProfileGuard<'a> {
522 ctx: &'a VfsContext,
523 active: bool,
524}
525
526impl SqliteOperationProfileGuard<'_> {
527 pub(crate) fn finish(mut self) -> SqliteOperationProfile {
528 self.active = false;
529 self.ctx.finish_operation_profile()
530 }
531}
532
533impl Drop for SqliteOperationProfileGuard<'_> {
534 fn drop(&mut self) {
535 if self.active {
536 let _ = self.ctx.finish_operation_profile();
537 }
538 }
539}
540
541#[derive(Debug, Clone, Copy, Default)]
547pub struct SqliteRoundTripCounts {
548 pub get_pages: u64,
549 pub commit: u64,
550}
551
552impl SqliteRoundTripCounts {
553 pub fn total(&self) -> u64 {
554 self.get_pages.saturating_add(self.commit)
555 }
556
557 pub fn since(&self, earlier: SqliteRoundTripCounts) -> SqliteRoundTripCounts {
558 SqliteRoundTripCounts {
559 get_pages: self.get_pages.saturating_sub(earlier.get_pages),
560 commit: self.commit.saturating_sub(earlier.commit),
561 }
562 }
563}
564
565#[derive(Debug, Clone, Copy)]
566pub enum SqliteOpenPhase {
567 InitialPreload,
568 VfsRegister,
569 WorkerReady,
570 Total,
571}
572
573impl SqliteOpenPhase {
574 pub fn as_label(self) -> &'static str {
575 match self {
576 SqliteOpenPhase::InitialPreload => "initial_preload",
577 SqliteOpenPhase::VfsRegister => "vfs_register",
578 SqliteOpenPhase::WorkerReady => "worker_ready",
579 SqliteOpenPhase::Total => "total",
580 }
581 }
582}
583
584pub trait SqliteVfsMetrics: Send + Sync {
585 fn profiling_enabled(&self) -> bool;
586
587 fn max_profiled_get_pages_requests(&self) -> usize;
588
589 fn observe_operation_profile(&self, profile: &SqliteOperationMetric) -> bool;
591
592 fn observe_transaction_profile(&self, profile: &SqliteTransactionMetric) -> bool;
594
595 fn operation_diagnostic_selected(&self, _profile: &SqliteOperationMetric) -> bool {
597 true
598 }
599
600 fn transaction_diagnostic_selected(&self, _profile: &SqliteTransactionMetric) -> bool {
601 true
602 }
603
604 fn emit_operation_diagnostic_event(
606 &self,
607 actor_id: &str,
608 generation: Option<u64>,
609 profile: &SqliteOperationMetric,
610 );
611
612 fn emit_transaction_diagnostic_event(
613 &self,
614 actor_id: &str,
615 generation: Option<u64>,
616 profile: &SqliteTransactionMetric,
617 );
618
619 fn record_fingerprint_catalog(
620 &self,
621 operation_type: &'static str,
622 fingerprint: &str,
623 identity: &str,
624 format_version: u8,
625 );
626
627 fn record_resolve_pages(&self, requested_pages: u64);
628
629 fn record_resolve_cache_hits(&self, pages: u64);
630
631 fn record_resolve_cache_misses(&self, pages: u64);
632
633 fn record_get_pages_request(&self, pages: u64, prefetch_pages: u64, page_size: u64);
634
635 fn observe_get_pages_duration(&self, duration_ns: u64);
636
637 fn observe_open_phase(&self, phase: SqliteOpenPhase, outcome: &'static str, duration_ns: u64);
638
639 fn record_startup_preload_pages(&self, kind: &'static str, pages: u64);
640
641 fn record_commit(&self, dirty_pages: u64);
643
644 fn observe_commit_phases(
645 &self,
646 request_build_ns: u64,
647 serialize_ns: u64,
648 transport_ns: u64,
649 state_update_ns: u64,
650 total_ns: u64,
651 );
652
653 fn set_worker_queue_depth(&self, depth: u64);
654
655 fn set_worker_active(&self, active: bool);
656
657 fn set_worker_inflight(&self, active: bool);
658
659 fn set_vfs_cached_bytes(&self, bytes: u64);
662 fn record_prefetch_unused_pages(&self, _pages: u64) {}
664
665 fn set_coordinator_queue_depth(&self, depth: u64);
666
667 fn record_worker_queue_overload(&self);
668
669 fn observe_worker_command_duration(
670 &self,
671 operation: &'static str,
672 in_tx: bool,
673 stmt_kind: &'static str,
674 duration_ns: u64,
675 );
676
677 fn observe_transaction_round_trips(&self, get_pages_round_trips: u64, commit_round_trips: u64);
678
679 fn record_worker_command_error(&self, operation: &'static str, code: &'static str);
680
681 fn observe_worker_close_duration(&self, duration_ns: u64);
682
683 fn record_worker_close_timeout(&self);
684
685 fn record_worker_crash(&self);
686
687 fn record_worker_unclean_close(&self);
688}
689
690#[derive(Debug, Clone, Copy, Default)]
691struct CommitTransportMetrics {
692 serialize_ns: u64,
693 transport_ns: u64,
694}
695
696enum CommitWait<T> {
697 Completed(T),
698 TimedOut,
699}
700
701pub struct VfsContext {
702 actor_id: String,
703 generation: Option<u64>,
704 runtime: Handle,
705 transport: SqliteTransportHandle,
706 config: VfsConfig,
707 state: RwLock<VfsState>,
708 aux_files: RwLock<BTreeMap<String, Arc<AuxFileState>>>,
709 last_error: Mutex<Option<String>>,
710 indeterminate_result: Mutex<Option<String>>,
711 transient_commit_error: Mutex<Option<String>>,
712 transaction_too_large: Mutex<Option<SqliteTransactionTooLargeError>>,
714 storage_error: Mutex<Option<SqliteStorageError>>,
717 fatal_error: RwLock<Option<String>>,
718 commit_atomic_count: AtomicU64,
719 trace_commit_attempts: AtomicU64,
721 io_methods: Box<sqlite3_io_methods>,
722 pub resolve_pages_total: AtomicU64,
724 pub resolve_pages_cache_hits: AtomicU64,
725 pub resolve_pages_fetches: AtomicU64,
726 pub pages_fetched_total: AtomicU64,
727 pub prefetch_pages_total: AtomicU64,
728 pub prefetch_consumed_total: AtomicU64,
730 pub commit_total: AtomicU64,
731 pub commit_request_build_ns: AtomicU64,
732 pub commit_serialize_ns: AtomicU64,
733 pub commit_transport_ns: AtomicU64,
734 pub commit_state_update_ns: AtomicU64,
735 pub commit_duration_ns_total: AtomicU64,
736 operation_profile_active: AtomicBool,
740 operation_profile: Mutex<Option<SqliteOperationProfile>>,
741 operation_prefetch_unused_start: AtomicU64,
742 reported_prefetch_unused: AtomicU64,
743 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
744}
745
746#[derive(Debug, Clone)]
747struct VfsState {
748 db_size_pages: u32,
749 committed_db_size_pages: u32,
753 head_txid: Option<u64>,
754 page_size: usize,
755 prefetch_buffer: PageBuffer,
757 committed_buffer: PageBuffer,
760 prefetch_unused_total: Arc<AtomicU64>,
761 write_buffer: WriteBuffer,
762 corked: Option<BTreeMap<u32, PageBytes>>,
766 dead: bool,
767 corked_db_size_pages: u32,
770 refused_write_pending: bool,
775 closing: bool,
781}
782
783#[derive(Debug, Clone, Copy, PartialEq, Eq)]
784enum CorkOutcome {
785 NotCorked,
786 Corked,
787 Overflow,
788}
789
790#[derive(Debug, Clone, Default)]
791struct WriteBuffer {
792 in_atomic_write: bool,
793 saved_db_size: u32,
794 dirty: BTreeMap<u32, PageBytes>,
795}
796
797impl WriteBuffer {
798 fn snapshot(&self) -> Vec<BufferedDirtyPage> {
799 self.dirty
800 .iter()
801 .map(|(pgno, bytes)| BufferedDirtyPage {
802 pgno: *pgno,
803 bytes: bytes.clone(),
804 })
805 .collect()
806 }
807}
808
809#[derive(Debug)]
810enum GetPagesError {
811 FenceMismatch(String),
812 Storage(SqliteStorageError),
813 Other(String),
814}
815
816#[repr(C)]
817struct VfsFile {
818 base: sqlite3_file,
819 ctx: *const VfsContext,
820 aux: *mut AuxFileHandle,
821}
822
823#[derive(Default)]
824struct AuxFileState {
825 bytes: Mutex<Vec<u8>>,
826 journal_records: AtomicUsize,
829}
830
831struct AuxFileHandle {
832 path: String,
833 state: Arc<AuxFileState>,
834 delete_on_close: bool,
835 is_main_journal: bool,
837}
838
839#[derive(Debug, Clone, Copy, PartialEq, Eq)]
840enum PageCacheInsertKind {
841 Target,
842 PartialTarget,
843 Prefetch,
844 Startup,
845}
846
847unsafe impl Send for VfsContext {}
848unsafe impl Sync for VfsContext {}
849
850pub struct SqliteVfs {
851 _registration: SqliteVfsRegistration,
852 _name: CString,
853 ctx: Box<VfsContext>,
854}
855
856unsafe impl Send for SqliteVfs {}
857unsafe impl Sync for SqliteVfs {}
858
859struct SqliteVfsRegistration {
860 vfs_ptr: *mut sqlite3_vfs,
861}
862
863pub struct NativeDatabase {
864 db: *mut sqlite3,
865 _vfs: NativeVfsHandle,
866}
867
868unsafe impl Send for NativeDatabase {}
869
870pub type NativeVfsHandle = Arc<SqliteVfs>;
871pub type NativeConnection = NativeDatabase;
872
873impl VfsState {
874 fn committed_size_pages_before_transaction(&self) -> u32 {
876 if self.corked.is_some() {
877 self.corked_db_size_pages
878 } else {
879 self.committed_db_size_pages
880 }
881 }
882
883 fn discard_uncommitted_pages(&mut self) {
888 self.write_buffer.dirty.clear();
889 self.write_buffer.in_atomic_write = false;
890 self.db_size_pages = self.committed_size_pages_before_transaction();
891 self.refused_write_pending = false;
892 }
893
894 fn new(config: &VfsConfig, _profiling_enabled: bool) -> Self {
895 let prefetch_unused_total = Arc::new(AtomicU64::new(0));
896 let prefetch_buffer =
897 PageBuffer::with_unused_counter(&config.page_budget, prefetch_unused_total.clone());
898 let committed_buffer = build_page_buffer(config);
899 let mut state = Self {
900 db_size_pages: 1,
901 committed_db_size_pages: 1,
902 head_txid: None,
903 page_size: DEFAULT_PAGE_SIZE,
904 prefetch_buffer,
905 committed_buffer,
906 prefetch_unused_total,
907 write_buffer: WriteBuffer::default(),
908 corked: None,
909 dead: false,
910 corked_db_size_pages: 0,
911 closing: false,
912 refused_write_pending: false,
913 };
914 state.cache_page(
915 config,
916 PageCacheInsertKind::Startup,
917 1,
918 empty_db_page().into(),
919 );
920 state
921 }
922
923 fn cache_page(
924 &mut self,
925 config: &VfsConfig,
926 kind: PageCacheInsertKind,
927 pgno: u32,
928 bytes: PageBytes,
929 ) {
930 if !should_cache_page(config, kind, pgno) {
931 return;
932 }
933 cache_page(config, &self.prefetch_buffer, kind, pgno, bytes);
934 }
935
936 fn take_buffered_page(
942 &self,
943 config: &VfsConfig,
944 pgno: u32,
945 partial: bool,
946 ) -> Option<(PageBytes, bool)> {
947 if !can_read_cached_page(config, pgno) {
948 return None;
949 }
950 if pgno == 1 || partial {
951 let bytes = self
952 .committed_buffer
953 .get(&pgno)
954 .or_else(|| self.prefetch_buffer.get(&pgno))?;
955 return Some((bytes, false));
956 }
957
958 if let Some(bytes) = self.committed_buffer.take(&pgno) {
984 self.prefetch_buffer.invalidate(&pgno);
988 Some((bytes, false))
989 } else {
990 self.prefetch_buffer.take_with_status(&pgno)
991 }
992 }
993
994 fn corked_page(&self, pgno: u32) -> Option<&PageBytes> {
996 self.corked.as_ref().and_then(|corked| corked.get(&pgno))
997 }
998
999 fn cork_dirty_pages(&mut self, max_pages: usize) -> CorkOutcome {
1006 let Some(corked) = self.corked.as_mut() else {
1007 return CorkOutcome::NotCorked;
1008 };
1009 let added_pages = self
1010 .write_buffer
1011 .dirty
1012 .keys()
1013 .filter(|pgno| !corked.contains_key(pgno))
1014 .count();
1015 if corked.len() + added_pages > max_pages {
1016 return CorkOutcome::Overflow;
1017 }
1018 corked.append(&mut self.write_buffer.dirty);
1019 self.corked_db_size_pages = self.db_size_pages;
1020 CorkOutcome::Corked
1021 }
1022
1023 fn note_committed_page(&mut self, config: &VfsConfig, pgno: u32, bytes: &PageBytes) {
1027 self.prefetch_buffer.invalidate(&pgno);
1028 if pgno != 1 {}
1029 if config.staging_cache_ttl_ms == 0 || !config.page_cache_mode.caches_any_pages() {
1030 return;
1031 }
1032 self.committed_buffer.insert(pgno, bytes.clone());
1033 }
1034
1035 fn seed_page(
1036 &mut self,
1037 config: &VfsConfig,
1038 kind: PageCacheInsertKind,
1039 pgno: u32,
1040 page: PageBytes,
1041 ) {
1042 if pgno == 1 {
1043 self.seed_main_page(config, kind, page);
1044 } else {
1045 self.cache_page(config, kind, pgno, page);
1046 }
1047 }
1048
1049 fn seed_main_page(&mut self, config: &VfsConfig, kind: PageCacheInsertKind, page: PageBytes) {
1050 if let Some(page_size) = sqlite_header_page_size(&page) {
1051 self.page_size = page_size;
1052 }
1053 if let Some(db_size_pages) = sqlite_header_db_size_pages(&page) {
1054 self.db_size_pages = db_size_pages;
1055 self.committed_db_size_pages = db_size_pages;
1056 }
1057 self.cache_page(config, kind, 1, page);
1058 }
1059
1060 fn invalidate_page_buffers(&mut self) {
1061 self.prefetch_buffer.invalidate_all();
1062 self.committed_buffer.invalidate_all();
1063 }
1064}
1065
1066fn build_page_buffer(config: &VfsConfig) -> PageBuffer {
1067 PageBuffer::new(&config.page_budget)
1068}
1069
1070fn cache_page(
1071 config: &VfsConfig,
1072 prefetch_buffer: &PageBuffer,
1073 kind: PageCacheInsertKind,
1074 pgno: u32,
1075 bytes: PageBytes,
1076) {
1077 if !should_cache_page(config, kind, pgno) {
1078 return;
1079 }
1080 if kind == PageCacheInsertKind::Prefetch || (kind == PageCacheInsertKind::Startup && pgno != 1)
1081 {
1082 prefetch_buffer.insert_prefetched(pgno, bytes);
1083 } else {
1084 prefetch_buffer.insert(pgno, bytes);
1085 }
1086}
1087
1088fn should_cache_page(config: &VfsConfig, kind: PageCacheInsertKind, pgno: u32) -> bool {
1089 match kind {
1090 PageCacheInsertKind::Target => false,
1091 PageCacheInsertKind::PartialTarget => true,
1092 PageCacheInsertKind::Prefetch => {
1093 config.staging_cache_ttl_ms > 0 && config.page_cache_mode.caches_prefetched_pages()
1094 }
1095 PageCacheInsertKind::Startup => {
1096 pgno == 1
1097 || (config.staging_cache_ttl_ms > 0
1098 && config.page_cache_mode.caches_startup_preloaded_pages())
1099 }
1100 }
1101}
1102
1103fn can_read_cached_page(config: &VfsConfig, pgno: u32) -> bool {
1104 pgno == 1 || config.page_cache_mode.caches_any_pages()
1105}
1106
1107impl VfsContext {
1108 fn new(
1109 actor_id: String,
1110 generation: Option<u64>,
1111 runtime: Handle,
1112 transport: SqliteTransportHandle,
1113 config: VfsConfig,
1114 io_methods: sqlite3_io_methods,
1115 initial_pages: impl Into<InitialPages>,
1116 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
1117 ) -> std::result::Result<Self, String> {
1118 let profiling_enabled = metrics
1119 .as_ref()
1120 .is_some_and(|metrics| metrics.profiling_enabled());
1121 let mut state = VfsState::new(&config, profiling_enabled);
1122 let initial_pages = initial_pages.into();
1123 state.head_txid = initial_pages.head_txid;
1124 for (pgno, page) in initial_pages.pages {
1125 state.seed_page(&config, PageCacheInsertKind::Startup, pgno, page);
1126 }
1127
1128 Ok(Self {
1129 actor_id,
1130 generation,
1131 runtime,
1132 transport,
1133 config: config.clone(),
1134 state: RwLock::new(state),
1135 aux_files: RwLock::new(BTreeMap::new()),
1136 last_error: Mutex::new(None),
1137 indeterminate_result: Mutex::new(None),
1138 transient_commit_error: Mutex::new(None),
1139 transaction_too_large: Mutex::new(None),
1140 storage_error: Mutex::new(None),
1141 fatal_error: RwLock::new(None),
1142 commit_atomic_count: AtomicU64::new(0),
1143 trace_commit_attempts: AtomicU64::new(0),
1144 io_methods: Box::new(io_methods),
1145 resolve_pages_total: AtomicU64::new(0),
1146 resolve_pages_cache_hits: AtomicU64::new(0),
1147 resolve_pages_fetches: AtomicU64::new(0),
1148 pages_fetched_total: AtomicU64::new(0),
1149 prefetch_pages_total: AtomicU64::new(0),
1150 prefetch_consumed_total: AtomicU64::new(0),
1151 commit_total: AtomicU64::new(0),
1152 commit_request_build_ns: AtomicU64::new(0),
1153 commit_serialize_ns: AtomicU64::new(0),
1154 commit_transport_ns: AtomicU64::new(0),
1155 commit_state_update_ns: AtomicU64::new(0),
1156 commit_duration_ns_total: AtomicU64::new(0),
1157 operation_profile_active: AtomicBool::new(false),
1158 operation_profile: Mutex::new(None),
1159 operation_prefetch_unused_start: AtomicU64::new(0),
1160 reported_prefetch_unused: AtomicU64::new(0),
1161 metrics,
1162 })
1163 }
1164
1165 fn begin_operation_profile(&self) {
1166 if self
1167 .metrics
1168 .as_ref()
1169 .is_some_and(|metrics| metrics.profiling_enabled())
1170 {
1171 let unused = self
1172 .state
1173 .read()
1174 .prefetch_unused_total
1175 .load(Ordering::Relaxed);
1176 self.operation_prefetch_unused_start
1177 .store(unused, Ordering::Relaxed);
1178 *self.operation_profile.lock() = Some(SqliteOperationProfile::default());
1179 self.operation_profile_active.store(true, Ordering::Release);
1180 }
1181 }
1182
1183 fn report_prefetch_unused(&self) {
1184 let total = self
1185 .state
1186 .read()
1187 .prefetch_unused_total
1188 .load(Ordering::Relaxed);
1189 let previous = self
1190 .reported_prefetch_unused
1191 .fetch_max(total, Ordering::Relaxed);
1192 if total > previous
1193 && let Some(metrics) = &self.metrics
1194 {
1195 metrics.record_prefetch_unused_pages(total - previous);
1196 }
1197 }
1198
1199 fn finish_operation_profile(&self) -> SqliteOperationProfile {
1200 self.report_prefetch_unused();
1201 if !self.operation_profile_active.swap(false, Ordering::AcqRel) {
1202 return SqliteOperationProfile::default();
1203 }
1204 let mut profile = self.operation_profile.lock().take().unwrap_or_default();
1205 let unused = self
1206 .state
1207 .read()
1208 .prefetch_unused_total
1209 .load(Ordering::Relaxed)
1210 .saturating_sub(self.operation_prefetch_unused_start.load(Ordering::Relaxed));
1211 profile.prefetch_unused_pages = unused;
1212 profile
1213 }
1214
1215 fn update_operation_profile(&self, update: impl FnOnce(&mut SqliteOperationProfile)) {
1216 if !self.operation_profile_active.load(Ordering::Relaxed) {
1217 return;
1218 }
1219 if let Some(profile) = self.operation_profile.lock().as_mut() {
1220 update(profile);
1221 }
1222 }
1223
1224 fn profile_request_limit(&self) -> usize {
1225 self.metrics
1226 .as_ref()
1227 .map_or(0, |metrics| metrics.max_profiled_get_pages_requests())
1228 .min(MAX_PROFILED_GET_PAGES_REQUESTS)
1229 }
1230
1231 fn clear_last_error(&self) {
1232 *self.last_error.lock() = None;
1233 }
1234
1235 fn set_last_error(&self, message: String) {
1236 *self.last_error.lock() = Some(message);
1237 }
1238
1239 fn mark_indeterminate(&self, operation: String) {
1240 *self.indeterminate_result.lock() = Some(operation.clone());
1241 self.mark_fatal(SqliteCommitIndeterminateError(operation).to_string());
1242 }
1243
1244 fn take_indeterminate_result(&self) -> Option<String> {
1245 self.indeterminate_result.lock().take()
1246 }
1247
1248 fn clone_last_error(&self) -> Option<String> {
1249 self.last_error.lock().clone()
1250 }
1251
1252 fn clone_fatal_error(&self) -> Option<String> {
1253 self.fatal_error.read().clone()
1254 }
1255
1256 fn defer_transient_commit_error(&self, message: String) {
1257 *self.transient_commit_error.lock() = Some(message);
1258 }
1259
1260 fn take_transient_commit_error(&self) -> Option<String> {
1261 self.transient_commit_error.lock().take()
1262 }
1263
1264 fn transaction_too_large_error(&self) -> SqliteTransactionTooLargeError {
1265 SqliteTransactionTooLargeError {
1266 max_pages: self.config.commit_page_cap(),
1267 page_size: self.page_size(),
1268 }
1269 }
1270
1271 fn refuse_oversize_transaction(&self) -> c_int {
1275 let error = self.transaction_too_large_error();
1276 tracing::warn!(
1277 actor_id = %self.actor_id,
1278 max_pages = error.max_pages,
1279 "sqlite transaction exceeds the commit page cap; the application must split it into smaller transactions"
1280 );
1281 self.set_last_error(error.to_string());
1282 *self.transaction_too_large.lock() = Some(error);
1283 SQLITE_FULL
1284 }
1285
1286 fn transaction_too_large_commit_error(&self) -> CommitBufferError {
1289 CommitBufferError::Other(self.transaction_too_large_error().to_string())
1290 }
1291
1292 fn commit_cork_alone(&self) -> std::result::Result<(), CommitBufferError> {
1297 let request = {
1298 let state = self.state.read();
1299 let Some(corked) = state.corked.as_ref() else {
1300 return Ok(());
1301 };
1302 BufferedCommitRequest {
1303 actor_id: self.actor_id.clone(),
1304 new_db_size_pages: state.corked_db_size_pages,
1305 expected_head_txid: state.head_txid,
1306 dirty_pages: corked
1307 .iter()
1308 .map(|(pgno, bytes)| BufferedDirtyPage {
1309 pgno: *pgno,
1310 bytes: bytes.clone(),
1311 })
1312 .collect(),
1313 }
1314 };
1315
1316 let outcome = match self.block_on_buffered_commit(request.clone(), None) {
1317 Ok(CommitWait::Completed((outcome, _))) => outcome,
1318 Ok(CommitWait::TimedOut) => {
1319 return Err(CommitBufferError::Other(
1320 "sqlite commit timed out".to_string(),
1321 ));
1322 }
1323 Err(err) => {
1324 tracing::error!(
1325 actor_id = %self.actor_id,
1326 dirty_pages = request.dirty_pages.len(),
1327 ?err,
1328 "sqlite commit of a full cork failed"
1329 );
1330 handle_non_finalize_commit_error(self, &err);
1331 return Err(err);
1332 }
1333 };
1334 self.commit_total.fetch_add(1, Ordering::Relaxed);
1335 if let Some(metrics) = &self.metrics {
1336 metrics.record_commit(request.dirty_pages.len() as u64);
1337 }
1338
1339 let mut state = self.state.write();
1340 state.committed_db_size_pages = outcome.db_size_pages;
1341 state.head_txid = outcome
1342 .head_txid
1343 .or_else(|| state.head_txid.map(|head_txid| head_txid.saturating_add(1)));
1344 for dirty_page in &request.dirty_pages {
1345 state.note_committed_page(&self.config, dirty_page.pgno, &dirty_page.bytes);
1346 }
1347 if let Some(corked) = state.corked.as_mut() {
1348 corked.clear();
1349 }
1350 Ok(())
1351 }
1352
1353 pub(crate) fn take_transaction_too_large(&self) -> Option<SqliteTransactionTooLargeError> {
1354 self.transaction_too_large.lock().take()
1355 }
1356
1357 fn set_storage_error(&self, error: SqliteStorageError) {
1359 self.set_last_error(error.to_string());
1360 *self.storage_error.lock() = Some(error);
1361 }
1362
1363 pub(crate) fn take_storage_error(&self) -> Option<SqliteStorageError> {
1364 self.storage_error.lock().take()
1365 }
1366
1367 fn discard_refused_transaction(&self) {
1371 let mut state = self.state.write();
1372 if !state.refused_write_pending {
1373 return;
1374 }
1375 tracing::debug!(
1376 actor_id = %self.actor_id,
1377 dirty_pages = state.write_buffer.dirty.len(),
1378 "discarding the pages of an abandoned sqlite transaction at close"
1379 );
1380 state.discard_uncommitted_pages();
1381 }
1382
1383 fn abandon_write_buffer(&self) {
1387 self.state.write().refused_write_pending = true;
1388 }
1389
1390 fn discard_open_transaction(&self) {
1393 let mut state = self.state.write();
1394 if state.write_buffer.dirty.is_empty()
1395 && !state.write_buffer.in_atomic_write
1396 && state.db_size_pages == state.committed_size_pages_before_transaction()
1397 {
1398 return;
1399 }
1400 tracing::debug!(
1401 actor_id = %self.actor_id,
1402 dirty_pages = state.write_buffer.dirty.len(),
1403 "discarding the pages of an uncommitted sqlite transaction at close"
1404 );
1405 state.discard_uncommitted_pages();
1406 }
1407
1408 fn discard_unsent_commits(&self) -> usize {
1413 let mut state = self.state.write();
1414 let dirty_pages = state.write_buffer.dirty.len();
1415 state.discard_uncommitted_pages();
1416 dirty_pages
1417 }
1418
1419 pub(crate) fn take_last_error(&self) -> Option<String> {
1420 self.last_error.lock().take()
1421 }
1422
1423 fn add_commit_phase_metrics(
1424 &self,
1425 request_build_ns: u64,
1426 transport_metrics: CommitTransportMetrics,
1427 state_update_ns: u64,
1428 total_ns: u64,
1429 ) {
1430 self.update_operation_profile(|profile| {
1431 profile.storage_ns = profile
1432 .storage_ns
1433 .saturating_add(transport_metrics.transport_ns);
1434 profile.commit_round_trips = profile.commit_round_trips.saturating_add(1);
1435 });
1436 if let Some(metrics) = &self.metrics {
1437 metrics.observe_commit_phases(
1438 request_build_ns,
1439 transport_metrics.serialize_ns,
1440 transport_metrics.transport_ns,
1441 state_update_ns,
1442 total_ns,
1443 );
1444 }
1445 self.commit_request_build_ns
1446 .fetch_add(request_build_ns, Ordering::Relaxed);
1447 self.commit_serialize_ns
1448 .fetch_add(transport_metrics.serialize_ns, Ordering::Relaxed);
1449 self.commit_transport_ns
1450 .fetch_add(transport_metrics.transport_ns, Ordering::Relaxed);
1451 self.commit_state_update_ns
1452 .fetch_add(state_update_ns, Ordering::Relaxed);
1453 self.commit_duration_ns_total
1454 .fetch_add(total_ns, Ordering::Relaxed);
1455 }
1456
1457 fn sqlite_vfs_metrics(&self) -> SqliteVfsMetricsSnapshot {
1458 let state = self.state.read();
1459
1460 SqliteVfsMetricsSnapshot {
1461 request_build_ns: self.commit_request_build_ns.load(Ordering::Relaxed),
1462 serialize_ns: self.commit_serialize_ns.load(Ordering::Relaxed),
1463 transport_ns: self.commit_transport_ns.load(Ordering::Relaxed),
1464 state_update_ns: self.commit_state_update_ns.load(Ordering::Relaxed),
1465 total_ns: self.commit_duration_ns_total.load(Ordering::Relaxed),
1466 commit_count: self.commit_total.load(Ordering::Relaxed),
1467 page_cache_entries: state
1468 .prefetch_buffer
1469 .entry_count()
1470 .saturating_add(state.committed_buffer.entry_count()),
1471 page_cache_weighted_size: state.prefetch_buffer.weighted_size(),
1472 page_cache_capacity_pages: self.config.page_budget.capacity_pages(),
1473 write_buffer_dirty_pages: state.write_buffer.dirty.len() as u64,
1474 db_size_pages: state.db_size_pages as u64,
1475 }
1476 }
1477
1478 fn block_on_buffered_commit(
1479 &self,
1480 request: BufferedCommitRequest,
1481 timeout: Option<Duration>,
1482 ) -> std::result::Result<
1483 CommitWait<(BufferedCommitOutcome, CommitTransportMetrics)>,
1484 CommitBufferError,
1485 > {
1486 let trace_pages =
1487 tracing::enabled!(target: "sqlite_prefetch::trace", tracing::Level::TRACE).then(|| {
1488 (
1489 request
1490 .dirty_pages
1491 .iter()
1492 .map(|page| page.pgno)
1493 .collect::<Vec<_>>(),
1494 request.new_db_size_pages,
1495 )
1496 });
1497 let attempt = trace_pages
1498 .as_ref()
1499 .map(|_| self.trace_commit_attempts.fetch_add(1, Ordering::Relaxed));
1500 let span = tracing::trace_span!(target: "sqlite_prefetch::trace", "native_sqlite_commit",
1501 stream_id = %format_args!("sqlite-{:p}", self), actor_id = %self.actor_id,
1502 generation = self.generation, commit_attempt = attempt,
1503 local_request_id = %format_args!("c{}", self.commit_total.load(Ordering::Relaxed)));
1504 let commit =
1505 commit_buffered_pages(&*self.transport, request, self.config.commit_page_cap())
1506 .instrument(span.clone());
1507 let result = if let Some(timeout) = timeout {
1508 match self
1509 .runtime
1510 .block_on(async { tokio::time::timeout(timeout, commit).await })
1511 {
1512 Ok(result) => result,
1513 Err(_) => {
1514 span.in_scope(|| tracing::trace!(target: "sqlite_prefetch::trace", event = "commit_result", outcome = "timeout"));
1515 return Ok(CommitWait::TimedOut);
1516 }
1517 }
1518 } else {
1519 self.runtime.block_on(commit)
1520 };
1521
1522 let _span = span.enter();
1523 tracing::trace!(target: "sqlite_prefetch::trace", event = "commit_result",
1524 outcome = if result.is_ok() { "acknowledged" } else { "error" });
1525
1526 if result.is_ok()
1527 && let Some((pages, size_pages)) = trace_pages
1528 {
1529 let request_id = self.commit_total.load(Ordering::Relaxed);
1530 for pgno in pages {
1531 tracing::trace!(target: "sqlite_prefetch::trace", event = "access", op = "write",
1532 stream_id = %format_args!("sqlite-{:p}", self),
1533 pgno, hit = false, request_id = %format_args!("c{request_id}"));
1534 }
1535 tracing::trace!(target: "sqlite_prefetch::trace", event = "commit",
1536 stream_id = %format_args!("sqlite-{:p}", self),
1537 request_id = %format_args!("c{request_id}"), size_pages);
1538 }
1539 result.map(CommitWait::Completed)
1540 }
1541
1542 fn round_trip_counts(&self) -> SqliteRoundTripCounts {
1543 SqliteRoundTripCounts {
1544 get_pages: self.resolve_pages_fetches.load(Ordering::Relaxed),
1545 commit: self.commit_total.load(Ordering::Relaxed),
1546 }
1547 }
1548
1549 fn page_size(&self) -> usize {
1550 self.state.read().page_size.max(DEFAULT_PAGE_SIZE)
1551 }
1552
1553 fn open_aux_file(&self, path: &str) -> Arc<AuxFileState> {
1554 let mut aux_files = self.aux_files.write();
1555 aux_files
1556 .entry(path.to_string())
1557 .or_insert_with(|| Arc::new(AuxFileState::default()))
1558 .clone()
1559 }
1560
1561 fn aux_file_exists(&self, path: &str) -> bool {
1562 self.aux_files.read().contains_key(path)
1563 }
1564
1565 fn read_aux_file(&self, path: &str) -> Option<Vec<u8>> {
1566 let state = self.aux_files.read().get(path).cloned()?;
1567 let bytes = state.bytes.lock().clone();
1568 Some(bytes)
1569 }
1570
1571 fn delete_aux_file(&self, path: &str) {
1572 self.aux_files.write().remove(path);
1573 }
1574
1575 fn is_dead(&self) -> bool {
1576 self.state.read().dead
1577 }
1578
1579 fn mark_fatal(&self, message: String) {
1580 self.set_last_error(message.clone());
1581 self.state.write().dead = true;
1582 let mut fatal_error = self.fatal_error.write();
1583 if fatal_error.is_none() {
1584 *fatal_error = Some(message);
1585 }
1586 }
1587
1588 fn resolve_pages(
1591 &self,
1592 target_pgnos: &[u32],
1593 partial: bool,
1594 ) -> std::result::Result<HashMap<u32, Option<PageBytes>>, GetPagesError> {
1595 use std::sync::atomic::Ordering::Relaxed;
1596 let request_id = self.resolve_pages_total.fetch_add(1, Relaxed);
1597 self.update_operation_profile(|profile| {
1598 profile.sqlite_requested_pages = profile
1599 .sqlite_requested_pages
1600 .saturating_add(target_pgnos.len() as u64);
1601 });
1602 if let Some(metrics) = &self.metrics {
1603 metrics.record_resolve_pages(target_pgnos.len() as u64);
1604 }
1605
1606 let mut resolved = HashMap::new();
1607 let mut missing = Vec::new();
1608 let mut seen = HashSet::new();
1609 let mut prefetch_consumed = 0_u64;
1610
1611 {
1612 let state = self.state.read();
1613 if state.dead {
1614 return Err(GetPagesError::Other(
1615 "sqlite actor lost its fence".to_string(),
1616 ));
1617 }
1618
1619 for pgno in target_pgnos.iter().copied() {
1620 if !seen.insert(pgno) {
1621 continue;
1622 }
1623 if let Some(bytes) = state.write_buffer.dirty.get(&pgno) {
1624 resolved.insert(pgno, Some(bytes.clone()));
1625 continue;
1626 }
1627 if let Some(bytes) = state.corked_page(pgno) {
1628 resolved.insert(pgno, Some(bytes.clone()));
1629 continue;
1630 }
1631 if let Some((bytes, was_prefetched)) =
1632 state.take_buffered_page(&self.config, pgno, partial)
1633 {
1634 prefetch_consumed = prefetch_consumed.saturating_add(u64::from(was_prefetched));
1635 resolved.insert(pgno, Some(bytes));
1636 continue;
1637 }
1638 missing.push(pgno);
1639 }
1640 }
1641 if tracing::enabled!(target: "sqlite_prefetch::trace", tracing::Level::TRACE) {
1642 for pgno in target_pgnos {
1643 tracing::trace!(target: "sqlite_prefetch::trace", event = "access", op = "read",
1644 stream_id = %format_args!("sqlite-{:p}", self),
1645 pgno, hit = resolved.contains_key(pgno), request_id);
1646 }
1647 }
1648 if prefetch_consumed > 0 {
1649 self.prefetch_consumed_total
1650 .fetch_add(prefetch_consumed, Relaxed);
1651 self.update_operation_profile(|profile| {
1652 profile.prefetch_consumed_pages = profile
1653 .prefetch_consumed_pages
1654 .saturating_add(prefetch_consumed);
1655 });
1656 }
1657
1658 if missing.is_empty() {
1659 self.resolve_pages_cache_hits
1660 .fetch_add(target_pgnos.len() as u64, Relaxed);
1661 if let Some(metrics) = &self.metrics {
1662 metrics.record_resolve_cache_hits(target_pgnos.len() as u64);
1663 }
1664 self.update_operation_profile(|profile| {
1665 profile.cache_hit_pages = profile
1666 .cache_hit_pages
1667 .saturating_add(target_pgnos.len() as u64);
1668 });
1669 return Ok(resolved);
1670 }
1671 self.resolve_pages_cache_hits
1672 .fetch_add((seen.len() - missing.len()) as u64, Relaxed);
1673 if let Some(metrics) = &self.metrics {
1674 metrics.record_resolve_cache_hits((seen.len() - missing.len()) as u64);
1675 metrics.record_resolve_cache_misses(missing.len() as u64);
1676 }
1677 self.update_operation_profile(|profile| {
1678 profile.cache_hit_pages = profile
1679 .cache_hit_pages
1680 .saturating_add((seen.len() - missing.len()) as u64);
1681 profile.cache_miss_pages = profile
1682 .cache_miss_pages
1683 .saturating_add(missing.len() as u64);
1684 });
1685
1686 let (to_fetch, page_size, db_size_pages, expected_head_txid) = {
1687 let state = self.state.read();
1688 let to_fetch = missing.clone();
1690 (
1691 to_fetch,
1692 state.page_size.max(1),
1693 state.db_size_pages,
1694 state.head_txid,
1695 )
1696 };
1697
1698 {
1699 self.resolve_pages_fetches.fetch_add(1, Relaxed);
1700 self.pages_fetched_total
1701 .fetch_add(to_fetch.len() as u64, Relaxed);
1702 tracing::debug!(
1703 actor_id = %self.actor_id,
1704 generation = ?self.generation,
1705 requested_pages = ?target_pgnos,
1706 missing_pages = ?missing,
1707 total_fetch_pages = to_fetch.len(),
1708 total_fetch_bytes = to_fetch.len().saturating_mul(page_size),
1709 db_size_pages,
1710 expected_head_txid,
1711 "vfs get_pages fetch"
1712 );
1713 }
1714
1715 let trace_span = tracing::trace_span!(target: "sqlite_prefetch::trace", "native_sqlite_read",
1716 stream_id = %format_args!("sqlite-{:p}", self), actor_id = %self.actor_id,
1717 generation = self.generation, local_request_id = request_id);
1718 let get_pages_start = Instant::now();
1719 let response = self.runtime.block_on(
1723 self.transport
1724 .get_pages(protocol::SqliteGetPagesRequest {
1725 actor_id: self.actor_id.clone().into(),
1726 pgnos: to_fetch.clone(),
1727 expected_generation: None,
1728 expected_head_txid,
1729 })
1730 .instrument(trace_span.clone()),
1731 );
1732 let _trace_guard = trace_span.enter();
1735 let get_pages_duration_ns = get_pages_start.elapsed().as_nanos() as u64;
1736 if let Some(metrics) = &self.metrics {
1737 metrics.observe_get_pages_duration(get_pages_duration_ns);
1738 }
1739 let profile_active = self.operation_profile_active.load(Ordering::Relaxed);
1740 let ordinal = if profile_active {
1741 self.operation_profile
1742 .lock()
1743 .as_ref()
1744 .map_or(1, |profile| profile.get_pages_round_trips.saturating_add(1))
1745 } else {
1746 1
1747 };
1748 let response = match response {
1749 Ok(response) => response,
1750 Err(err) => {
1751 self.update_operation_profile(|profile| {
1752 profile.push_get_pages(
1753 SqliteGetPagesProfile {
1754 ordinal,
1755 duration_ns: get_pages_duration_ns,
1756 demand_requested: missing.len() as u64,
1757 ..Default::default()
1758 },
1759 self.profile_request_limit(),
1760 );
1761 });
1762 if let Some(error) = storage_timeout(&err) {
1763 return Err(GetPagesError::Storage(error));
1764 }
1765 return Err(GetPagesError::Other(err.to_string()));
1766 }
1767 };
1768
1769 match response {
1770 protocol::SqliteGetPagesResponse::SqliteGetPagesOk(ok) => {
1771 if profile_active {
1772 let requested = to_fetch.iter().copied().collect::<HashSet<_>>();
1773 let mut request_profile = SqliteGetPagesProfile {
1774 ordinal,
1775 duration_ns: get_pages_duration_ns,
1776 demand_requested: missing.len() as u64,
1777 success: true,
1778 ..Default::default()
1779 };
1780 let mut btree_pages = 0_u64;
1781 let mut non_btree_pages = 0_u64;
1782 for fetched in &ok.pages {
1783 if !requested.contains(&fetched.pgno) {
1784 request_profile.overflow_expansion_extra =
1785 request_profile.overflow_expansion_extra.saturating_add(1);
1786 }
1787 if let Some(bytes) = &fetched.bytes {
1788 request_profile.response_present =
1789 request_profile.response_present.saturating_add(1);
1790 request_profile.response_bytes = request_profile
1791 .response_bytes
1792 .saturating_add(bytes.len() as u64);
1793 match classify(fetched.pgno, bytes) {
1794 PageClass::Btree => btree_pages = btree_pages.saturating_add(1),
1795 PageClass::Overflow => {
1796 non_btree_pages = non_btree_pages.saturating_add(1)
1797 }
1798 }
1799 } else {
1800 request_profile.response_missing =
1801 request_profile.response_missing.saturating_add(1);
1802 }
1803 }
1804 self.update_operation_profile(|profile| {
1805 profile.btree_pages = profile.btree_pages.saturating_add(btree_pages);
1806 profile.non_btree_pages =
1807 profile.non_btree_pages.saturating_add(non_btree_pages);
1808 profile.push_get_pages(request_profile, self.profile_request_limit());
1809 });
1810 }
1811 let response_head_txid = ok.head_txid;
1812 if let Some(head_txid) = response_head_txid {
1813 self.state.write().head_txid = Some(head_txid);
1814 }
1815 let missing_pages = missing.iter().copied().collect::<HashSet<_>>();
1816 let prefetch_buffer = self.state.read().prefetch_buffer.clone();
1817 #[cfg(debug_assertions)]
1818 let mut returned_pgnos = HashSet::new();
1819 #[cfg(debug_assertions)]
1820 let mut returned_missing_pages = Vec::new();
1821 #[cfg(debug_assertions)]
1822 let mut returned_missing_in_range_pages = Vec::new();
1823 let mut server_prefetched = 0_u64;
1824 for (ordinal, fetched) in ok.pages.into_iter().enumerate() {
1825 tracing::trace!(target: "sqlite_prefetch::trace", event = "page_received",
1826 ordinal, pgno = fetched.pgno, present = fetched.bytes.is_some(),
1827 bytes = fetched.bytes.as_ref().map_or(0, |bytes| bytes.len()),
1828 demand = missing_pages.contains(&fetched.pgno));
1829 let delivered = fetched
1831 .bytes
1832 .as_ref()
1833 .and_then(|bytes| PageTrace::capture(ordinal, fetched.pgno, bytes.len()));
1834 #[cfg(debug_assertions)]
1835 {
1836 returned_pgnos.insert(fetched.pgno);
1837 if fetched.bytes.is_none() {
1838 returned_missing_pages.push(fetched.pgno);
1839 if fetched.pgno <= db_size_pages {
1840 returned_missing_in_range_pages.push(fetched.pgno);
1841 }
1842 }
1843 }
1844 let bytes = if fetched.bytes.is_none()
1845 && self.commit_total.load(Relaxed) == 0
1846 && missing_pages.contains(&fetched.pgno)
1847 && fetched.pgno == 1
1848 {
1849 self.state.write().head_txid = Some(response_head_txid.unwrap_or(0));
1852 Some(PageBytes::from(empty_db_page()))
1853 } else {
1854 fetched.bytes.map(PageBytes::from)
1855 };
1856 if let Some(bytes) = &bytes {
1857 let kind = if missing_pages.contains(&fetched.pgno) {
1860 if partial {
1861 PageCacheInsertKind::PartialTarget
1862 } else {
1863 PageCacheInsertKind::Target
1864 }
1865 } else {
1866 PageCacheInsertKind::Prefetch
1867 };
1868 if kind == PageCacheInsertKind::Prefetch {
1869 server_prefetched += 1;
1870 if should_cache_page(&self.config, kind, fetched.pgno) {
1871 prefetch_buffer.insert_prefetched_traced(
1872 fetched.pgno,
1873 bytes.clone(),
1874 delivered,
1875 );
1876 } else if let Some(trace) = delivered {
1877 trace.disposition("unused");
1878 }
1879 } else {
1880 if let Some(trace) = delivered {
1883 trace.disposition("consumed");
1884 }
1885 cache_page(
1886 &self.config,
1887 &prefetch_buffer,
1888 kind,
1889 fetched.pgno,
1890 bytes.clone(),
1891 );
1892 }
1893 }
1894 resolved.insert(fetched.pgno, bytes);
1895 }
1896 self.prefetch_pages_total
1897 .fetch_add(server_prefetched, Relaxed);
1898 if let Some(metrics) = &self.metrics {
1899 metrics.record_get_pages_request(
1900 to_fetch.len() as u64,
1901 server_prefetched,
1902 page_size as u64,
1903 );
1904 }
1905 #[cfg(debug_assertions)]
1906 {
1907 let absent_response_pages = to_fetch
1908 .iter()
1909 .copied()
1910 .filter(|pgno| !returned_pgnos.contains(pgno))
1911 .collect::<Vec<_>>();
1912 let absent_in_range_pages = absent_response_pages
1913 .iter()
1914 .copied()
1915 .filter(|pgno| *pgno <= db_size_pages)
1916 .collect::<Vec<_>>();
1917 if !returned_missing_in_range_pages.is_empty()
1918 || !absent_in_range_pages.is_empty()
1919 {
1920 tracing::warn!(
1921 actor_id = %self.actor_id,
1922 requested_pages = ?target_pgnos,
1923 missing_pages = ?missing,
1924 fetch_pages = ?to_fetch,
1925 db_size_pages,
1926 returned_missing_pages = ?returned_missing_pages,
1927 returned_missing_in_range_pages = ?returned_missing_in_range_pages,
1928 absent_response_pages = ?absent_response_pages,
1929 absent_in_range_pages = ?absent_in_range_pages,
1930 "sqlite get_pages returned missing pages within declared db size"
1931 );
1932 }
1933 }
1934 for pgno in missing {
1935 resolved.entry(pgno).or_insert(None);
1936 }
1937 Ok(resolved)
1938 }
1939 protocol::SqliteGetPagesResponse::SqliteErrorResponse(error) => {
1940 self.update_operation_profile(|profile| {
1941 profile.push_get_pages(
1942 SqliteGetPagesProfile {
1943 ordinal,
1944 duration_ns: get_pages_duration_ns,
1945 demand_requested: missing.len() as u64,
1946 ..Default::default()
1947 },
1948 self.profile_request_limit(),
1949 );
1950 });
1951 if self.commit_total.load(Relaxed) == 0
1952 && missing.contains(&1)
1953 && is_initial_main_page_missing(&error.message)
1954 {
1955 for pgno in missing {
1956 let bytes = if pgno == 1 {
1957 Some(PageBytes::from(empty_db_page()))
1958 } else {
1959 None
1960 };
1961 resolved.entry(pgno).or_insert(bytes);
1962 }
1963 return Ok(resolved);
1964 }
1965 if is_head_fence_mismatch_response(&error) {
1966 return Err(GetPagesError::FenceMismatch(error.message.to_string()));
1967 }
1968 if let Some(refusal) = storage_refusal(&error) {
1969 return Err(GetPagesError::Storage(refusal));
1970 }
1971 Err(GetPagesError::Other(error.message.to_string()))
1972 }
1973 }
1974 }
1975
1976 fn flush_dirty_pages(
1977 &self,
1978 ) -> std::result::Result<Option<BufferedCommitOutcome>, CommitBufferError> {
1979 match self.flush_dirty_pages_with_timeout(None)? {
1980 CommitWait::Completed(outcome) => Ok(outcome),
1981 CommitWait::TimedOut => Err(CommitBufferError::Other(
1982 "sqlite commit timed out".to_string(),
1983 )),
1984 }
1985 }
1986
1987 fn flush_dirty_pages_with_timeout(
1988 &self,
1989 timeout: Option<Duration>,
1990 ) -> std::result::Result<CommitWait<Option<BufferedCommitOutcome>>, CommitBufferError> {
1991 let total_start = Instant::now();
1992 let request_build_start = Instant::now();
1993 let request = {
1994 let mut state = self.state.write();
1995 if state.dead {
1996 return Err(CommitBufferError::Other(
1997 "sqlite actor lost its fence".to_string(),
1998 ));
1999 }
2000 if state.closing {
2001 state.discard_uncommitted_pages();
2002 return Ok(CommitWait::Completed(None));
2003 }
2004 if state.write_buffer.in_atomic_write
2005 || (state.write_buffer.dirty.is_empty()
2006 && state.db_size_pages == state.committed_db_size_pages)
2007 {
2008 return Ok(CommitWait::Completed(None));
2009 }
2010 match state.cork_dirty_pages(self.config.commit_page_cap()) {
2012 CorkOutcome::NotCorked => {}
2013 CorkOutcome::Corked => {
2014 state.refused_write_pending = false;
2015 return Ok(CommitWait::Completed(None));
2016 }
2017 CorkOutcome::Overflow => {
2018 drop(state);
2019 self.commit_cork_alone()?;
2020 let mut state = self.state.write();
2021 return match state.cork_dirty_pages(self.config.commit_page_cap()) {
2022 CorkOutcome::Corked | CorkOutcome::NotCorked => {
2023 state.refused_write_pending = false;
2024 Ok(CommitWait::Completed(None))
2025 }
2026 CorkOutcome::Overflow => Err(self.transaction_too_large_commit_error()),
2027 };
2028 }
2029 }
2030
2031 BufferedCommitRequest {
2032 actor_id: self.actor_id.clone(),
2033 new_db_size_pages: state.db_size_pages,
2034 expected_head_txid: state.head_txid,
2035 dirty_pages: state.write_buffer.snapshot(),
2036 }
2037 };
2038 let request_build_ns = request_build_start.elapsed().as_nanos() as u64;
2039 self.update_operation_profile(|profile| {
2040 profile.dirty_pages = profile
2041 .dirty_pages
2042 .saturating_add(request.dirty_pages.len() as u64);
2043 profile.dirty_bytes = profile.dirty_bytes.saturating_add(
2044 request
2045 .dirty_pages
2046 .iter()
2047 .map(|page| page.bytes.len() as u64)
2048 .sum::<u64>(),
2049 );
2050 });
2051
2052 let (outcome, transport_metrics) =
2053 match self.block_on_buffered_commit(request.clone(), timeout) {
2057 Ok(CommitWait::Completed(outcome)) => outcome,
2058 Ok(CommitWait::TimedOut) => return Ok(CommitWait::TimedOut),
2059 Err(err) => {
2060 tracing::error!(
2061 actor_id = %self.actor_id,
2062 new_db_size_pages = request.new_db_size_pages,
2063 dirty_pages = request.dirty_pages.len(),
2064 ?err,
2065 "sqlite flush commit failed"
2066 );
2067 handle_non_finalize_commit_error(self, &err);
2068 return Err(err);
2069 }
2070 };
2071 self.commit_total
2072 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2073 if let Some(metrics) = &self.metrics {
2074 metrics.record_commit(request.dirty_pages.len() as u64);
2075 }
2076 tracing::debug!(
2077 dirty_pages = request.dirty_pages.len(),
2078 path = ?outcome.path,
2079 requested_db_size_pages = request.new_db_size_pages,
2080 db_size_pages = outcome.db_size_pages,
2081 request_build_ns,
2082 serialize_ns = transport_metrics.serialize_ns,
2083 transport_ns = transport_metrics.transport_ns,
2084 "vfs commit complete (flush)"
2085 );
2086 #[cfg(debug_assertions)]
2087 {
2088 if outcome.db_size_pages != request.new_db_size_pages {
2089 tracing::warn!(
2090 actor_id = %self.actor_id,
2091 dirty_pages = request.dirty_pages.len(),
2092 path = ?outcome.path,
2093 requested_db_size_pages = request.new_db_size_pages,
2094 outcome_db_size_pages = outcome.db_size_pages,
2095 "sqlite flush commit returned db size different from request"
2096 );
2097 }
2098 }
2099 let state_update_start = Instant::now();
2100 let mut state = self.state.write();
2101 state.db_size_pages = outcome.db_size_pages;
2102 state.committed_db_size_pages = outcome.db_size_pages;
2103 state.head_txid = outcome
2104 .head_txid
2105 .or_else(|| state.head_txid.map(|head_txid| head_txid.saturating_add(1)));
2106 for dirty_page in &request.dirty_pages {
2107 state.note_committed_page(&self.config, dirty_page.pgno, &dirty_page.bytes);
2108 }
2109 state.write_buffer.dirty.clear();
2110 state.refused_write_pending = false;
2111 let state_update_ns = state_update_start.elapsed().as_nanos() as u64;
2112 self.add_commit_phase_metrics(
2113 request_build_ns,
2114 transport_metrics,
2115 state_update_ns,
2116 total_start.elapsed().as_nanos() as u64,
2117 );
2118 Ok(CommitWait::Completed(Some(outcome)))
2119 }
2120
2121 fn commit_atomic_write(&self) -> std::result::Result<(), CommitBufferError> {
2122 match self.commit_atomic_write_with_timeout(None)? {
2123 CommitWait::Completed(()) => Ok(()),
2124 CommitWait::TimedOut => Err(CommitBufferError::Other(
2125 "sqlite commit timed out".to_string(),
2126 )),
2127 }
2128 }
2129
2130 fn commit_atomic_write_with_timeout(
2131 &self,
2132 timeout: Option<Duration>,
2133 ) -> std::result::Result<CommitWait<()>, CommitBufferError> {
2134 let total_start = Instant::now();
2135 let request_build_start = Instant::now();
2136 let request = {
2137 let mut state = self.state.write();
2138 if state.dead {
2139 return Err(CommitBufferError::Other(
2140 "sqlite actor lost its fence".to_string(),
2141 ));
2142 }
2143 if !state.write_buffer.in_atomic_write {
2144 return Ok(CommitWait::Completed(()));
2145 }
2146 if state.closing {
2147 state.discard_uncommitted_pages();
2148 return Ok(CommitWait::Completed(()));
2149 }
2150 if state.write_buffer.dirty.is_empty()
2151 && state.db_size_pages == state.committed_db_size_pages
2152 {
2153 state.write_buffer.in_atomic_write = false;
2154 return Ok(CommitWait::Completed(()));
2155 }
2156 match state.cork_dirty_pages(self.config.commit_page_cap()) {
2157 CorkOutcome::NotCorked => {}
2158 CorkOutcome::Corked => {
2159 state.write_buffer.in_atomic_write = false;
2160 state.refused_write_pending = false;
2161 return Ok(CommitWait::Completed(()));
2162 }
2163 CorkOutcome::Overflow => {
2164 drop(state);
2165 self.commit_cork_alone()?;
2166 let mut state = self.state.write();
2167 return match state.cork_dirty_pages(self.config.commit_page_cap()) {
2168 CorkOutcome::Corked | CorkOutcome::NotCorked => {
2169 state.write_buffer.in_atomic_write = false;
2170 state.refused_write_pending = false;
2171 Ok(CommitWait::Completed(()))
2172 }
2173 CorkOutcome::Overflow => Err(self.transaction_too_large_commit_error()),
2174 };
2175 }
2176 }
2177
2178 BufferedCommitRequest {
2179 actor_id: self.actor_id.clone(),
2180 new_db_size_pages: state.db_size_pages,
2181 expected_head_txid: state.head_txid,
2182 dirty_pages: state.write_buffer.snapshot(),
2183 }
2184 };
2185 let request_build_ns = request_build_start.elapsed().as_nanos() as u64;
2186 self.update_operation_profile(|profile| {
2187 profile.dirty_pages = profile
2188 .dirty_pages
2189 .saturating_add(request.dirty_pages.len() as u64);
2190 profile.dirty_bytes = profile.dirty_bytes.saturating_add(
2191 request
2192 .dirty_pages
2193 .iter()
2194 .map(|page| page.bytes.len() as u64)
2195 .sum::<u64>(),
2196 );
2197 });
2198
2199 let (outcome, transport_metrics) =
2200 match self.block_on_buffered_commit(request.clone(), timeout) {
2201 Ok(CommitWait::Completed(outcome)) => outcome,
2202 Ok(CommitWait::TimedOut) => return Ok(CommitWait::TimedOut),
2203 Err(err) => {
2204 tracing::error!(
2205 actor_id = %self.actor_id,
2206 new_db_size_pages = request.new_db_size_pages,
2207 dirty_pages = request.dirty_pages.len(),
2208 ?err,
2209 "sqlite atomic commit failed"
2210 );
2211 handle_non_finalize_commit_error(self, &err);
2212 return Err(err);
2213 }
2214 };
2215 self.commit_total
2216 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2217 if let Some(metrics) = &self.metrics {
2218 metrics.record_commit(request.dirty_pages.len() as u64);
2219 }
2220 tracing::debug!(
2221 dirty_pages = request.dirty_pages.len(),
2222 path = ?outcome.path,
2223 requested_db_size_pages = request.new_db_size_pages,
2224 db_size_pages = outcome.db_size_pages,
2225 request_build_ns,
2226 serialize_ns = transport_metrics.serialize_ns,
2227 transport_ns = transport_metrics.transport_ns,
2228 "vfs commit complete (atomic)"
2229 );
2230 #[cfg(debug_assertions)]
2231 {
2232 if outcome.db_size_pages != request.new_db_size_pages {
2233 tracing::warn!(
2234 actor_id = %self.actor_id,
2235 dirty_pages = request.dirty_pages.len(),
2236 path = ?outcome.path,
2237 requested_db_size_pages = request.new_db_size_pages,
2238 outcome_db_size_pages = outcome.db_size_pages,
2239 "sqlite atomic commit returned db size different from request"
2240 );
2241 }
2242 }
2243 self.clear_last_error();
2244 let state_update_start = Instant::now();
2245 let mut state = self.state.write();
2246 state.db_size_pages = outcome.db_size_pages;
2247 state.committed_db_size_pages = outcome.db_size_pages;
2248 state.head_txid = outcome
2249 .head_txid
2250 .or_else(|| state.head_txid.map(|head_txid| head_txid.saturating_add(1)));
2251 for dirty_page in &request.dirty_pages {
2252 state.note_committed_page(&self.config, dirty_page.pgno, &dirty_page.bytes);
2253 }
2254 state.write_buffer.dirty.clear();
2255 state.write_buffer.in_atomic_write = false;
2256 state.refused_write_pending = false;
2257 let state_update_ns = state_update_start.elapsed().as_nanos() as u64;
2258 self.add_commit_phase_metrics(
2259 request_build_ns,
2260 transport_metrics,
2261 state_update_ns,
2262 total_start.elapsed().as_nanos() as u64,
2263 );
2264 Ok(CommitWait::Completed(()))
2265 }
2266
2267 fn truncate_main_file(&self, size: sqlite3_int64) -> bool {
2276 let page_size = self.page_size() as i64;
2277 let truncated_pages = ((size + page_size - 1) / page_size) as u32;
2278 let mut state = self.state.write();
2279 state.db_size_pages = truncated_pages;
2280 state
2281 .write_buffer
2282 .dirty
2283 .retain(|pgno, _| *pgno <= truncated_pages);
2284 let between_transactions =
2285 !state.write_buffer.in_atomic_write && state.write_buffer.dirty.is_empty();
2286 if let Some(corked) = state.corked.as_mut() {
2287 corked.retain(|pgno, _| *pgno <= truncated_pages);
2288 }
2289 if state.corked.is_some() && between_transactions {
2291 state.corked_db_size_pages = truncated_pages;
2292 }
2293 state.invalidate_page_buffers();
2294 state.corked.is_none()
2296 && !state.write_buffer.in_atomic_write
2297 && state.write_buffer.dirty.is_empty()
2298 && state.db_size_pages != state.committed_db_size_pages
2299 }
2300
2301 fn cork(&self) {
2305 let mut state = self.state.write();
2306 if state.corked.is_none() {
2307 state.corked = Some(BTreeMap::new());
2308 state.corked_db_size_pages = state.committed_db_size_pages;
2312 }
2313 }
2314
2315 fn uncork(&self) -> std::result::Result<(), CommitBufferError> {
2320 {
2321 let mut state = self.state.write();
2322 if state.corked.is_none() {
2323 return Ok(());
2324 }
2325 if state.write_buffer.in_atomic_write || !state.write_buffer.dirty.is_empty() {
2326 return Err(CommitBufferError::Other(
2327 "cannot uncork sqlite commits inside a transaction".to_string(),
2328 ));
2329 }
2330 state.write_buffer.dirty = state.corked.take().unwrap_or_default();
2331 }
2332 self.flush_dirty_pages().map(|_| ())
2333 }
2334}
2335
2336impl Drop for VfsContext {
2337 fn drop(&mut self) {
2338 self.state.write().invalidate_page_buffers();
2339 self.report_prefetch_unused();
2340 let state = self.state.read();
2341 let page_cache_entries = state
2342 .prefetch_buffer
2343 .entry_count()
2344 .saturating_add(state.committed_buffer.entry_count());
2345 let page_cache_weighted_size = state.prefetch_buffer.weighted_size();
2346 tracing::debug!(
2347 actor_id = %self.actor_id,
2348 generation = ?self.generation,
2349 resolve_pages_calls = self.resolve_pages_total.load(Ordering::Relaxed),
2350 resolve_pages_cache_hits = self.resolve_pages_cache_hits.load(Ordering::Relaxed),
2351 get_pages_round_trips = self.resolve_pages_fetches.load(Ordering::Relaxed),
2352 pages_fetched_total = self.pages_fetched_total.load(Ordering::Relaxed),
2353 prefetch_pages_total = self.prefetch_pages_total.load(Ordering::Relaxed),
2354 prefetch_consumed_total = self.prefetch_consumed_total.load(Ordering::Relaxed),
2355 commit_count = self.commit_total.load(Ordering::Relaxed),
2356 db_size_pages = state.db_size_pages,
2357 head_txid = state.head_txid,
2358 page_cache_entries,
2359 page_cache_weighted_size,
2360 page_cache_capacity_pages = self.config.page_budget.capacity_pages(),
2361 "sqlite vfs close summary"
2362 );
2363 }
2364}
2365
2366fn cleanup_batch_atomic_probe(db: *mut sqlite3) {
2367 if let Err(err) = sqlite_exec(db, "DROP TABLE IF EXISTS __rivet_batch_probe;") {
2368 tracing::warn!(%err, "failed to clean up sqlite batch atomic probe table");
2369 }
2370}
2371
2372fn assert_batch_atomic_probe(db: *mut sqlite3, vfs: &SqliteVfs) -> std::result::Result<(), String> {
2373 let commit_atomic_before = vfs.commit_atomic_count();
2374 let probe_sql = "\
2375 BEGIN IMMEDIATE;\
2376 CREATE TABLE IF NOT EXISTS __rivet_batch_probe(x INTEGER);\
2377 INSERT INTO __rivet_batch_probe VALUES(1);\
2378 DELETE FROM __rivet_batch_probe;\
2379 DROP TABLE IF EXISTS __rivet_batch_probe;\
2380 COMMIT;\
2381 ";
2382
2383 if let Err(err) = sqlite_exec(db, probe_sql) {
2384 let last_error = vfs.clone_last_error();
2385 tracing::error!(
2386 %err,
2387 last_error = ?last_error,
2388 commit_atomic_before,
2389 "sqlite batch atomic probe failed"
2390 );
2391 cleanup_batch_atomic_probe(db);
2392 if let Some(last_error) = last_error {
2393 return Err(format!(
2394 "batch atomic probe failed: {err}; vfs last_error: {last_error}"
2395 ));
2396 }
2397 return Err(format!("batch atomic probe failed: {err}"));
2398 }
2399
2400 let commit_atomic_after = vfs.commit_atomic_count();
2401 if commit_atomic_after == commit_atomic_before {
2402 tracing::error!(
2403 commit_atomic_before,
2404 commit_atomic_after,
2405 last_error = ?vfs.clone_last_error(),
2406 "batch atomic writes not active for sqlite, SQLITE_ENABLE_BATCH_ATOMIC_WRITE may be missing"
2407 );
2408 cleanup_batch_atomic_probe(db);
2409 return Err(
2410 "batch atomic writes not active for sqlite, SQLITE_ENABLE_BATCH_ATOMIC_WRITE may be missing"
2411 .to_string(),
2412 );
2413 }
2414
2415 Ok(())
2416}
2417
2418struct BatchAtomicProbeTransport;
2424
2425#[async_trait]
2426impl SqliteTransport for BatchAtomicProbeTransport {
2427 async fn get_pages(
2428 &self,
2429 request: protocol::SqliteGetPagesRequest,
2430 ) -> Result<protocol::SqliteGetPagesResponse> {
2431 Ok(protocol::SqliteGetPagesResponse::SqliteGetPagesOk(
2432 protocol::SqliteGetPagesOk {
2433 pages: request
2434 .pgnos
2435 .into_iter()
2436 .map(|pgno| protocol::SqliteFetchedPage { pgno, bytes: None })
2437 .collect(),
2438 head_txid: Some(0),
2439 },
2440 ))
2441 }
2442
2443 async fn commit(
2444 &self,
2445 _request: protocol::SqliteCommitRequest,
2446 ) -> Result<protocol::SqliteCommitResponse> {
2447 Ok(protocol::SqliteCommitResponse::SqliteCommitOk(
2448 protocol::SqliteCommitOk { head_txid: None },
2449 ))
2450 }
2451
2452 async fn commit_stage_begin(
2453 &self,
2454 _request: protocol::SqliteCommitStageBeginRequest,
2455 ) -> Result<protocol::SqliteCommitStageBeginResponse> {
2456 anyhow::bail!("batch-atomic startup probe unexpectedly attempted a staged commit")
2457 }
2458
2459 async fn commit_stage_segment(
2460 &self,
2461 _request: protocol::SqliteCommitStageSegmentRequest,
2462 ) -> Result<protocol::SqliteCommitStageSegmentResponse> {
2463 anyhow::bail!("batch-atomic startup probe unexpectedly attempted a staged commit")
2464 }
2465
2466 async fn commit_finalize(
2467 &self,
2468 _request: protocol::SqliteCommitFinalizeRequest,
2469 ) -> Result<protocol::SqliteCommitFinalizeResponse> {
2470 anyhow::bail!("batch-atomic startup probe unexpectedly attempted a staged commit")
2471 }
2472}
2473
2474fn verify_batch_atomic_writes_at_process_startup_blocking(
2475 runtime: Handle,
2476 config: VfsConfig,
2477) -> std::result::Result<(), String> {
2478 let probe_id = NEXT_TEMP_AUX_ID.fetch_add(1, Ordering::Relaxed);
2479 let file_name = format!("batch-atomic-startup-probe-{probe_id}.db");
2480 let vfs = SqliteVfs::register_with_transport_and_initial_pages(
2481 &format!("batch-atomic-startup-probe-{probe_id}"),
2482 Arc::new(BatchAtomicProbeTransport),
2483 file_name.clone(),
2484 runtime,
2485 config,
2486 InitialPages::default(),
2487 None,
2488 )?;
2489 let connection = open_connection(
2490 Arc::new(vfs),
2491 &file_name,
2492 SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE,
2493 )?;
2494 configure_connection_for_database(connection.as_ptr(), &connection._vfs, &file_name)
2495 .map_err(|error| error.to_string())?;
2496 assert_batch_atomic_probe(connection.as_ptr(), &connection._vfs)
2497}
2498
2499pub async fn verify_batch_atomic_writes_at_process_startup() -> Result<()> {
2502 verify_batch_atomic_writes_at_process_startup_with_config(VfsConfig::default()).await
2503}
2504
2505async fn verify_batch_atomic_writes_at_process_startup_with_config(
2506 config: VfsConfig,
2507) -> Result<()> {
2508 let runtime = Handle::try_current().map_err(|err| {
2509 anyhow::anyhow!("batch-atomic startup probe requires a Tokio runtime: {err}")
2510 })?;
2511 let result = tokio::task::spawn_blocking(move || {
2512 verify_batch_atomic_writes_at_process_startup_blocking(runtime, config)
2513 })
2514 .await
2515 .map_err(|err| anyhow::anyhow!("batch-atomic startup probe task failed: {err}"))?;
2516 result.map_err(anyhow::Error::msg)
2517}
2518
2519fn handle_non_finalize_commit_error(ctx: &VfsContext, err: &CommitBufferError) {
2520 match err {
2521 CommitBufferError::FenceMismatch(message) => ctx.mark_fatal(message.clone()),
2522 CommitBufferError::Indeterminate(operation) => {
2523 ctx.mark_indeterminate(operation.clone());
2524 }
2525 CommitBufferError::Storage(error) => ctx.set_storage_error(error.clone()),
2526 CommitBufferError::Other(message) => ctx.set_last_error(message.clone()),
2527 }
2528}
2529
2530fn handle_finalize_fence_error(ctx: &VfsContext, err: &CommitBufferError) {
2531 if let CommitBufferError::FenceMismatch(reason) = err {
2532 ctx.mark_fatal(reason.clone());
2533 }
2534}
2535
2536pub(crate) async fn fetch_initial_pages_for_registration(
2537 transport: SqliteTransportHandle,
2538 actor_id: &str,
2539 generation: u64,
2540 _config: &VfsConfig,
2541) -> std::result::Result<InitialPages, String> {
2542 fetch_initial_pages(transport, actor_id.to_string(), generation, 1).await
2543}
2544
2545async fn fetch_initial_pages(
2546 transport: SqliteTransportHandle,
2547 actor_id: String,
2548 generation: u64,
2549 page_count: u32,
2550) -> std::result::Result<InitialPages, String> {
2551 tracing::debug!(
2552 actor_id = %actor_id,
2553 generation,
2554 page_count,
2555 "sqlite initial page preload request"
2556 );
2557
2558 let request_actor_id = actor_id.clone();
2559 let response = transport
2560 .get_pages(protocol::SqliteGetPagesRequest {
2561 actor_id: request_actor_id.into(),
2562 pgnos: (1..=page_count).collect(),
2563 expected_generation: None,
2564 expected_head_txid: None,
2565 })
2566 .await;
2567
2568 match response {
2569 Ok(protocol::SqliteGetPagesResponse::SqliteGetPagesOk(ok)) => {
2570 let head_txid = ok.head_txid;
2571 let pages: Vec<_> = ok
2572 .pages
2573 .into_iter()
2574 .filter_map(|page| page.bytes.map(|bytes| (page.pgno, bytes.into())))
2575 .collect();
2576 tracing::debug!(
2577 actor_id = %actor_id,
2578 generation,
2579 page_count,
2580 loaded_pages = pages.len(),
2581 head_txid,
2582 "sqlite initial page preload result"
2583 );
2584 Ok(InitialPages {
2585 pages,
2586 head_txid,
2587 requested_page_count: page_count,
2588 })
2589 }
2590 Ok(protocol::SqliteGetPagesResponse::SqliteErrorResponse(error)) => {
2591 if !is_initial_main_page_missing(&error.message) {
2592 return Err(format!(
2593 "sqlite initial page fetch failed: {}",
2594 error.message
2595 ));
2596 }
2597 tracing::info!(
2598 actor_id = %actor_id,
2599 generation,
2600 page_count,
2601 error = %error.message,
2602 "sqlite initial page fetch did not find persisted data"
2603 );
2604 Ok(InitialPages {
2605 pages: Vec::new(),
2606 head_txid: Some(0),
2607 requested_page_count: page_count,
2608 })
2609 }
2610 Err(err) => Err(format!("sqlite initial page fetch failed: {err}")),
2611 }
2612}
2613
2614fn is_initial_main_page_missing(message: &str) -> bool {
2615 message.contains("sqlite database was not found in this bucket branch")
2616 || message.contains("sqlite meta missing for get_pages")
2617 || message == "actor does not exist"
2618}
2619
2620fn next_temp_aux_path() -> String {
2621 format!(
2622 "{TEMP_AUX_PATH_PREFIX}-{}",
2623 NEXT_TEMP_AUX_ID.fetch_add(1, Ordering::Relaxed)
2624 )
2625}
2626
2627unsafe fn get_aux_state(file: &VfsFile) -> Option<&AuxFileHandle> {
2628 (!file.aux.is_null()).then(|| &*file.aux)
2629}
2630
2631const COMMIT_STAGE_THRESHOLD_PAGES: usize = 320;
2642
2643async fn commit_buffered_pages(
2644 transport: &dyn SqliteTransport,
2645 request: BufferedCommitRequest,
2646 max_dirty_pages: usize,
2647) -> std::result::Result<(BufferedCommitOutcome, CommitTransportMetrics), CommitBufferError> {
2648 if request.dirty_pages.len() > max_dirty_pages {
2652 return Err(CommitBufferError::Other(format!(
2653 "commit of {} dirty pages exceeds the {} page maximum",
2654 request.dirty_pages.len(),
2655 max_dirty_pages,
2656 )));
2657 }
2658
2659 if request.dirty_pages.len() > COMMIT_STAGE_THRESHOLD_PAGES {
2660 return commit_staged_pages(transport, request).await;
2661 }
2662
2663 let mut metrics = CommitTransportMetrics::default();
2664 let serialize_start = Instant::now();
2665 let BufferedCommitRequest {
2666 actor_id,
2667 new_db_size_pages,
2668 dirty_pages,
2669 expected_head_txid,
2670 } = request;
2671 let commit_request = protocol::SqliteCommitRequest {
2672 actor_id: actor_id.into(),
2673 dirty_pages: dirty_pages
2674 .into_iter()
2675 .map(|page| protocol::SqliteDirtyPage {
2676 pgno: page.pgno,
2677 bytes: page.bytes.into_bytes(),
2678 })
2679 .collect(),
2680 db_size_pages: new_db_size_pages,
2681 now_ms: sqlite_now_ms().map_err(|err| CommitBufferError::Other(err.to_string()))?,
2682 expected_generation: None,
2683 expected_head_txid,
2684 };
2685 metrics.serialize_ns += serialize_start.elapsed().as_nanos() as u64;
2686 let transport_start = Instant::now();
2687 match transport
2688 .commit(commit_request)
2689 .await
2690 .map_err(commit_transport_error)?
2691 {
2692 protocol::SqliteCommitResponse::SqliteCommitOk(ok) => {
2693 metrics.transport_ns += transport_start.elapsed().as_nanos() as u64;
2694 Ok((
2695 BufferedCommitOutcome {
2696 path: CommitPath::Fast,
2697 db_size_pages: new_db_size_pages,
2698 head_txid: ok.head_txid,
2699 },
2700 metrics,
2701 ))
2702 }
2703 protocol::SqliteCommitResponse::SqliteErrorResponse(error) => {
2704 if is_head_fence_mismatch_response(&error) {
2705 Err(CommitBufferError::FenceMismatch(error.message.to_string()))
2706 } else if let Some(refusal) = storage_refusal(&error) {
2707 Err(CommitBufferError::Storage(refusal))
2708 } else {
2709 Err(CommitBufferError::Other(error.message.to_string()))
2710 }
2711 }
2712 }
2713}
2714
2715async fn commit_staged_pages(
2722 transport: &dyn SqliteTransport,
2723 request: BufferedCommitRequest,
2724) -> std::result::Result<(BufferedCommitOutcome, CommitTransportMetrics), CommitBufferError> {
2725 let mut metrics = CommitTransportMetrics::default();
2726 let serialize_start = Instant::now();
2727 let mut dirty_pages = request.dirty_pages;
2728 dirty_pages.sort_by_key(|page| page.pgno);
2731 metrics.serialize_ns += serialize_start.elapsed().as_nanos() as u64;
2732
2733 let transport_start = Instant::now();
2734 let txid = match transport
2735 .commit_stage_begin(protocol::SqliteCommitStageBeginRequest {
2736 actor_id: request.actor_id.clone().into(),
2737 expected_generation: None,
2738 expected_head_txid: request.expected_head_txid,
2739 })
2740 .await
2741 .map_err(commit_transport_error)?
2742 {
2743 protocol::SqliteCommitStageBeginResponse::SqliteCommitStageBeginOk(ok) => ok.txid,
2744 protocol::SqliteCommitStageBeginResponse::SqliteErrorResponse(error) => {
2745 return Err(staged_commit_error(error));
2746 }
2747 };
2748
2749 let segments = sqlite_types::cut_page_segments(&dirty_pages, |page| page.pgno);
2750 let segment_first_pgnos = segments.iter().map(|(first_pgno, _)| *first_pgno).collect();
2751 let responses = stream::iter(segments)
2754 .map(|(first_pgno, segment)| {
2755 let dirty_pages = segment
2756 .iter()
2757 .map(|page| protocol::SqliteDirtyPage {
2758 pgno: page.pgno,
2759 bytes: page.bytes.clone().into_bytes(),
2760 })
2761 .collect();
2762 let actor_id = request.actor_id.clone();
2763 async move {
2764 match transport
2765 .commit_stage_segment(protocol::SqliteCommitStageSegmentRequest {
2766 actor_id: actor_id.into(),
2767 expected_generation: None,
2768 txid,
2769 first_pgno,
2770 dirty_pages,
2771 })
2772 .await
2773 .map_err(commit_transport_error)?
2774 {
2775 protocol::SqliteCommitStageSegmentResponse::SqliteCommitStageSegmentOk(_) => {
2776 Ok(())
2777 }
2778 protocol::SqliteCommitStageSegmentResponse::SqliteErrorResponse(error) => {
2779 Err(staged_commit_error(error))
2780 }
2781 }
2782 }
2783 })
2784 .buffered(STAGE_PIPELINE_DEPTH)
2785 .collect::<Vec<_>>()
2786 .await;
2787 for response in responses {
2788 response?;
2789 }
2790
2791 let head_txid = match transport
2792 .commit_finalize(protocol::SqliteCommitFinalizeRequest {
2793 actor_id: request.actor_id.clone().into(),
2794 expected_generation: None,
2795 txid,
2796 new_db_size_pages: request.new_db_size_pages,
2797 now_ms: sqlite_now_ms().map_err(|err| CommitBufferError::Other(err.to_string()))?,
2798 segment_first_pgnos,
2799 })
2800 .await
2801 .map_err(commit_transport_error)?
2802 {
2803 protocol::SqliteCommitFinalizeResponse::SqliteCommitFinalizeOk(ok) => ok.head_txid,
2804 protocol::SqliteCommitFinalizeResponse::SqliteErrorResponse(error) => {
2805 return Err(staged_commit_error(error));
2806 }
2807 };
2808 metrics.transport_ns += transport_start.elapsed().as_nanos() as u64;
2809
2810 Ok((
2811 BufferedCommitOutcome {
2812 path: CommitPath::Staged,
2813 db_size_pages: request.new_db_size_pages,
2814 head_txid,
2815 },
2816 metrics,
2817 ))
2818}
2819
2820fn commit_transport_error(error: anyhow::Error) -> CommitBufferError {
2821 if let Some(timeout) = storage_timeout(&error) {
2822 return CommitBufferError::Storage(timeout);
2823 }
2824 match error.downcast::<SqliteCommitIndeterminateError>() {
2825 Ok(error) => CommitBufferError::Indeterminate(error.0),
2826 Err(error) => CommitBufferError::Other(error.to_string()),
2827 }
2828}
2829
2830fn staged_commit_error(error: protocol::SqliteErrorResponse) -> CommitBufferError {
2831 if is_head_fence_mismatch_response(&error) {
2832 CommitBufferError::FenceMismatch(error.message.to_string())
2833 } else if let Some(refusal) = storage_refusal(&error) {
2834 CommitBufferError::Storage(refusal)
2835 } else {
2836 CommitBufferError::Other(error.message.to_string())
2837 }
2838}
2839
2840fn is_head_fence_mismatch_response(error: &protocol::SqliteErrorResponse) -> bool {
2841 is_head_fence_mismatch(&error.group, &error.code)
2842}
2843
2844unsafe fn get_file(p: *mut sqlite3_file) -> &'static mut VfsFile {
2845 &mut *(p as *mut VfsFile)
2846}
2847
2848unsafe fn get_vfs_ctx(p: *mut sqlite3_vfs) -> &'static VfsContext {
2849 &*((*p).pAppData as *const VfsContext)
2850}
2851
2852fn sqlite_error_message(db: *mut sqlite3) -> String {
2853 unsafe {
2854 if db.is_null() {
2855 "unknown sqlite error".to_string()
2856 } else {
2857 CStr::from_ptr(sqlite3_errmsg(db))
2858 .to_string_lossy()
2859 .into_owned()
2860 }
2861 }
2862}
2863
2864pub(crate) fn sqlite_exec(db: *mut sqlite3, sql: &str) -> std::result::Result<(), String> {
2865 let c_sql = CString::new(sql).map_err(|err| err.to_string())?;
2866 let rc = unsafe { sqlite3_exec(db, c_sql.as_ptr(), None, ptr::null_mut(), ptr::null_mut()) };
2867 if rc != SQLITE_OK {
2868 return Err(format!(
2869 "`{sql}` failed with code {rc}: {}",
2870 sqlite_error_message(db)
2871 ));
2872 }
2873 Ok(())
2874}
2875
2876fn page_span(offset: i64, length: usize, page_size: usize) -> std::result::Result<Vec<u32>, ()> {
2877 if offset < 0 {
2878 return Err(());
2879 }
2880 if length == 0 {
2881 return Ok(Vec::new());
2882 }
2883
2884 let start = offset as usize / page_size + 1;
2885 let end = (offset as usize + length - 1) / page_size + 1;
2886 Ok((start as u32..=end as u32).collect())
2887}
2888
2889unsafe extern "C" fn io_close(p_file: *mut sqlite3_file) -> c_int {
2890 vfs_catch_unwind!(SQLITE_IOERR, {
2891 if p_file.is_null() {
2892 return SQLITE_OK;
2893 }
2894 let file = get_file(p_file);
2895 let result = if !file.aux.is_null() {
2896 let aux = Box::from_raw(file.aux);
2897 if aux.delete_on_close {
2898 let ctx = &*file.ctx;
2899 ctx.delete_aux_file(&aux.path);
2900 }
2901 file.aux = ptr::null_mut();
2902 Ok(())
2903 } else {
2904 let ctx = &*file.ctx;
2905 ctx.discard_refused_transaction();
2906 let should_flush = {
2907 let state = ctx.state.read();
2908 state.write_buffer.in_atomic_write
2909 || !state.write_buffer.dirty.is_empty()
2910 || state.db_size_pages != state.committed_db_size_pages
2911 || state.corked.is_some()
2912 };
2913 if should_flush {
2914 if ctx.state.read().write_buffer.in_atomic_write {
2915 ctx.commit_atomic_write().map(|_| ())
2916 } else {
2917 ctx.uncork()
2918 .and_then(|()| ctx.flush_dirty_pages().map(|_| ()))
2919 }
2920 } else {
2921 Ok(())
2922 }
2923 };
2924 file.base.pMethods = ptr::null();
2925 match result {
2926 Ok(()) => SQLITE_OK,
2927 Err(err) => {
2928 let ctx = &*file.ctx;
2929 handle_finalize_fence_error(ctx, &err);
2930 SQLITE_IOERR
2931 }
2932 }
2933 })
2934}
2935
2936unsafe extern "C" fn io_read(
2937 p_file: *mut sqlite3_file,
2938 p_buf: *mut c_void,
2939 i_amt: c_int,
2940 i_offset: sqlite3_int64,
2941) -> c_int {
2942 vfs_catch_unwind!(SQLITE_IOERR_READ, {
2943 if i_amt <= 0 {
2944 return SQLITE_OK;
2945 }
2946
2947 let file = get_file(p_file);
2948 if let Some(aux) = get_aux_state(file) {
2949 if i_offset < 0 {
2950 return SQLITE_IOERR_READ;
2951 }
2952
2953 let offset = i_offset as usize;
2954 let requested = i_amt as usize;
2955 let buf = slice::from_raw_parts_mut(p_buf.cast::<u8>(), requested);
2956 buf.fill(0);
2957
2958 let bytes = aux.state.bytes.lock();
2959 if offset >= bytes.len() {
2960 return SQLITE_IOERR_SHORT_READ;
2961 }
2962
2963 let copy_len = requested.min(bytes.len() - offset);
2964 buf[..copy_len].copy_from_slice(&bytes[offset..offset + copy_len]);
2965 return if copy_len < requested {
2966 SQLITE_IOERR_SHORT_READ
2967 } else {
2968 SQLITE_OK
2969 };
2970 }
2971
2972 let ctx = &*file.ctx;
2973 if ctx.is_dead() {
2974 return SQLITE_IOERR_READ;
2975 }
2976
2977 let buf = slice::from_raw_parts_mut(p_buf.cast::<u8>(), i_amt as usize);
2978 let requested_pages = match page_span(i_offset, i_amt as usize, ctx.page_size()) {
2979 Ok(pages) => pages,
2980 Err(_) => return SQLITE_IOERR_READ,
2981 };
2982 let page_size = ctx.page_size();
2983 let (file_size, db_size_pages) = {
2984 let state = ctx.state.read();
2985 (
2986 state.db_size_pages as usize * state.page_size,
2987 state.db_size_pages,
2988 )
2989 };
2990
2991 let partial = i_offset as usize % page_size != 0 || i_amt as usize % page_size != 0;
2992 let resolved = match ctx.resolve_pages(&requested_pages, partial) {
2993 Ok(pages) => pages,
2994 Err(GetPagesError::FenceMismatch(message)) => {
2995 tracing::error!(
2996 actor_id = %ctx.actor_id,
2997 requested_pages = ?requested_pages,
2998 error = %message,
2999 "sqlite xRead hit fatal sqlite error"
3000 );
3001 ctx.mark_fatal(message);
3002 return SQLITE_IOERR_READ;
3003 }
3004 Err(GetPagesError::Storage(error)) => {
3005 tracing::warn!(
3006 actor_id = %ctx.actor_id,
3007 requested_pages = ?requested_pages,
3008 %error,
3009 "sqlite xRead failed without changing storage; the statement fails with a typed storage error"
3010 );
3011 ctx.set_storage_error(error);
3012 return SQLITE_IOERR_READ;
3013 }
3014 Err(GetPagesError::Other(message)) => {
3015 tracing::error!(
3016 actor_id = %ctx.actor_id,
3017 requested_pages = ?requested_pages,
3018 error = %message,
3019 "sqlite xRead failed to resolve pages"
3020 );
3021 ctx.set_last_error(message);
3022 return SQLITE_IOERR_READ;
3023 }
3024 };
3025 ctx.clear_last_error();
3026
3027 #[cfg(debug_assertions)]
3028 {
3029 let missing_in_range_pages = requested_pages
3030 .iter()
3031 .copied()
3032 .filter(|pgno| *pgno <= db_size_pages)
3033 .filter(|pgno| !matches!(resolved.get(pgno), Some(Some(_))))
3034 .collect::<Vec<_>>();
3035 if !missing_in_range_pages.is_empty() {
3036 tracing::warn!(
3037 actor_id = %ctx.actor_id,
3038 offset = i_offset,
3039 amount = i_amt,
3040 page_size,
3041 db_size_pages,
3042 file_size,
3043 requested_pages = ?requested_pages,
3044 missing_in_range_pages = ?missing_in_range_pages,
3045 "sqlite xRead would zero-fill pages within declared db size"
3046 );
3047 }
3048 }
3049
3050 buf.fill(0);
3051 for pgno in requested_pages.iter().copied() {
3052 let Some(Some(bytes)) = resolved.get(&pgno) else {
3053 continue;
3054 };
3055 let page_start = (pgno as usize - 1) * page_size;
3056 let copy_start = page_start.max(i_offset as usize);
3057 let copy_end = (page_start + page_size).min(i_offset as usize + i_amt as usize);
3058 if copy_start >= copy_end {
3059 continue;
3060 }
3061 let page_offset = copy_start - page_start;
3062 let dest_offset = copy_start - i_offset as usize;
3063 let copy_len = copy_end - copy_start;
3064 buf[dest_offset..dest_offset + copy_len]
3065 .copy_from_slice(&bytes[page_offset..page_offset + copy_len]);
3066 }
3067
3068 if i_offset as usize + i_amt as usize > file_size {
3069 return SQLITE_IOERR_SHORT_READ;
3070 }
3071
3072 SQLITE_OK
3073 })
3074}
3075
3076fn journal_records_added(offset: usize, len: usize, file_len: usize, page_size: usize) -> usize {
3085 if len == page_size {
3086 return usize::from(offset % DEFAULT_PAGE_SIZE != 0);
3087 }
3088 if len == 4 {
3089 return 0;
3090 }
3091 let appended = (offset + len).saturating_sub(file_len.max(offset));
3092 appended.div_ceil(page_size + 8)
3093}
3094
3095unsafe extern "C" fn io_write(
3096 p_file: *mut sqlite3_file,
3097 p_buf: *const c_void,
3098 i_amt: c_int,
3099 i_offset: sqlite3_int64,
3100) -> c_int {
3101 vfs_catch_unwind!(SQLITE_IOERR_WRITE, {
3102 if i_amt <= 0 {
3103 return SQLITE_OK;
3104 }
3105
3106 let file = get_file(p_file);
3107 if let Some(aux) = get_aux_state(file) {
3108 if i_offset < 0 {
3109 return SQLITE_IOERR_WRITE;
3110 }
3111
3112 let offset = i_offset as usize;
3113 let source = slice::from_raw_parts(p_buf.cast::<u8>(), i_amt as usize);
3114 let end = offset + source.len();
3115 let mut bytes = aux.state.bytes.lock();
3116 if aux.is_main_journal {
3117 let ctx = &*file.ctx;
3118 let added =
3119 journal_records_added(offset, source.len(), bytes.len(), ctx.page_size());
3120 let records = aux.state.journal_records.load(Ordering::Relaxed);
3121 if added > 0 && records + added > ctx.config.commit_page_cap() {
3122 drop(bytes);
3123 return ctx.refuse_oversize_transaction();
3124 }
3125 aux.state
3126 .journal_records
3127 .store(records + added, Ordering::Relaxed);
3128 }
3129 if bytes.len() < end {
3130 bytes.resize(end, 0);
3131 }
3132 bytes[offset..end].copy_from_slice(source);
3133 return SQLITE_OK;
3134 }
3135
3136 let ctx = &*file.ctx;
3137 if ctx.is_dead() {
3138 return SQLITE_IOERR_WRITE;
3139 }
3140
3141 let page_size = ctx.page_size();
3142 let source = slice::from_raw_parts(p_buf.cast::<u8>(), i_amt as usize);
3143 let target_pages = match page_span(i_offset, i_amt as usize, page_size) {
3144 Ok(pages) => pages,
3145 Err(_) => return SQLITE_IOERR_WRITE,
3146 };
3147
3148 {
3152 let state = ctx.state.read();
3153 let new_pages = target_pages
3154 .iter()
3155 .filter(|pgno| !state.write_buffer.dirty.contains_key(pgno))
3156 .count();
3157 if state.write_buffer.dirty.len() + new_pages > ctx.config.commit_page_cap() {
3158 drop(state);
3159 ctx.state.write().refused_write_pending = true;
3160 return ctx.refuse_oversize_transaction();
3161 }
3162 }
3163
3164 let offset = i_offset as usize;
3169 let amt = i_amt as usize;
3170 let is_aligned_full_page = offset % page_size == 0 && amt % page_size == 0;
3171
3172 let mut resolved = HashMap::new();
3173 if !is_aligned_full_page {
3174 let (db_size_pages, pages_to_resolve): (u32, Vec<u32>) = {
3175 let state = ctx.state.read();
3176 let known_max = state.db_size_pages;
3177 (
3178 known_max,
3179 target_pages
3180 .iter()
3181 .copied()
3182 .filter(|pgno| {
3183 *pgno <= known_max
3184 && !state.write_buffer.dirty.contains_key(pgno)
3185 && state.corked_page(*pgno).is_none()
3186 })
3187 .collect(),
3188 )
3189 };
3190
3191 if !pages_to_resolve.is_empty() {
3192 resolved = match ctx.resolve_pages(&pages_to_resolve, false) {
3193 Ok(pages) => pages,
3194 Err(GetPagesError::FenceMismatch(message)) => {
3195 ctx.mark_fatal(message);
3196 return SQLITE_IOERR_WRITE;
3197 }
3198 Err(GetPagesError::Storage(error)) => {
3199 ctx.set_storage_error(error);
3200 return SQLITE_IOERR_WRITE;
3201 }
3202 Err(GetPagesError::Other(message)) => {
3203 ctx.set_last_error(message);
3204 return SQLITE_IOERR_WRITE;
3205 }
3206 };
3207 }
3208 for pgno in &target_pages {
3209 if *pgno > db_size_pages {
3210 resolved.entry(*pgno).or_insert(None);
3211 }
3212 }
3213 #[cfg(debug_assertions)]
3214 {
3215 let missing_existing_pages = target_pages
3216 .iter()
3217 .copied()
3218 .filter(|pgno| *pgno <= db_size_pages)
3219 .filter(|pgno| pages_to_resolve.contains(pgno))
3220 .filter(|pgno| !matches!(resolved.get(pgno), Some(Some(_))))
3221 .collect::<Vec<_>>();
3222 if !missing_existing_pages.is_empty() {
3223 tracing::warn!(
3224 actor_id = %ctx.actor_id,
3225 offset = i_offset,
3226 amount = i_amt,
3227 page_size,
3228 db_size_pages,
3229 target_pages = ?target_pages,
3230 missing_existing_pages = ?missing_existing_pages,
3231 "sqlite xWrite partial update would synthesize existing pages from zeros"
3232 );
3233 }
3234 }
3235 }
3236
3237 let mut state = ctx.state.write();
3238 for pgno in target_pages {
3239 let page_start = (pgno as usize - 1) * page_size;
3240 let patch_start = page_start.max(offset);
3241 let patch_end = (page_start + page_size).min(offset + amt);
3242 let Some(copy_len) = patch_end.checked_sub(patch_start) else {
3243 continue;
3244 };
3245 if copy_len == 0 {
3246 continue;
3247 }
3248
3249 let page_offset = patch_start - page_start;
3250 let source_offset = patch_start - offset;
3251 if is_aligned_full_page {
3252 state.write_buffer.dirty.insert(
3253 pgno,
3254 source[source_offset..source_offset + copy_len]
3255 .to_vec()
3256 .into(),
3257 );
3258 continue;
3259 }
3260
3261 let mut page = state
3262 .write_buffer
3263 .dirty
3264 .get_mut(&pgno)
3265 .map(std::mem::take)
3266 .map(PageBytes::into_vec)
3267 .unwrap_or_else(|| {
3268 resolved
3269 .remove(&pgno)
3270 .flatten()
3271 .map(PageBytes::into_vec)
3272 .unwrap_or_else(|| vec![0; page_size])
3273 });
3274 if page.len() < page_size {
3275 page.resize(page_size, 0);
3276 }
3277 page[page_offset..page_offset + copy_len]
3278 .copy_from_slice(&source[source_offset..source_offset + copy_len]);
3279 state.write_buffer.dirty.insert(pgno, page.into());
3280 }
3281
3282 let end_page = ((offset + amt) + page_size - 1) / page_size;
3283 state.db_size_pages = state.db_size_pages.max(end_page as u32);
3284 ctx.clear_last_error();
3285 SQLITE_OK
3286 })
3287}
3288
3289unsafe extern "C" fn io_truncate(p_file: *mut sqlite3_file, size: sqlite3_int64) -> c_int {
3290 vfs_catch_unwind!(SQLITE_IOERR_TRUNCATE, {
3291 if size < 0 {
3292 return SQLITE_IOERR_TRUNCATE;
3293 }
3294 let file = get_file(p_file);
3295 if let Some(aux) = get_aux_state(file) {
3296 let mut bytes = aux.state.bytes.lock();
3297 bytes.truncate(size as usize);
3298 bytes.shrink_to_fit();
3299 if size == 0 {
3300 aux.state.journal_records.store(0, Ordering::Relaxed);
3301 }
3302 return SQLITE_OK;
3303 }
3304 let ctx = &*file.ctx;
3305 if !ctx.truncate_main_file(size) {
3306 return SQLITE_OK;
3307 }
3308 match ctx.flush_dirty_pages() {
3309 Ok(_) => SQLITE_OK,
3310 Err(err) => {
3311 tracing::error!(
3312 actor_id = %ctx.actor_id,
3313 last_error = ?ctx.clone_last_error(),
3314 ?err,
3315 "sqlite truncate commit failed"
3316 );
3317 handle_finalize_fence_error(ctx, &err);
3318 ctx.abandon_write_buffer();
3319 SQLITE_IOERR_TRUNCATE
3320 }
3321 }
3322 })
3323}
3324
3325unsafe extern "C" fn io_sync(p_file: *mut sqlite3_file, _flags: c_int) -> c_int {
3329 vfs_catch_unwind!(SQLITE_IOERR_FSYNC, {
3330 let file = get_file(p_file);
3331 if get_aux_state(file).is_some() {
3332 return SQLITE_OK;
3333 }
3334 let ctx = &*file.ctx;
3335 if let Some(message) = ctx.take_transient_commit_error() {
3336 ctx.set_last_error(message);
3337 ctx.abandon_write_buffer();
3338 return SQLITE_IOERR_FSYNC;
3339 }
3340 match ctx.flush_dirty_pages() {
3341 Ok(_) => SQLITE_OK,
3342 Err(err) => {
3343 tracing::error!(
3344 actor_id = %ctx.actor_id,
3345 last_error = ?ctx.clone_last_error(),
3346 ?err,
3347 "sqlite sync failed"
3348 );
3349 handle_finalize_fence_error(ctx, &err);
3350 ctx.abandon_write_buffer();
3351 SQLITE_IOERR_FSYNC
3352 }
3353 }
3354 })
3355}
3356
3357unsafe extern "C" fn io_file_size(p_file: *mut sqlite3_file, p_size: *mut sqlite3_int64) -> c_int {
3358 vfs_catch_unwind!(SQLITE_IOERR_FSTAT, {
3359 let file = get_file(p_file);
3360 if let Some(aux) = get_aux_state(file) {
3361 *p_size = aux.state.bytes.lock().len() as sqlite3_int64;
3362 return SQLITE_OK;
3363 }
3364 let ctx = &*file.ctx;
3365 let state = ctx.state.read();
3366 *p_size = (state.db_size_pages as usize * state.page_size) as sqlite3_int64;
3367 SQLITE_OK
3368 })
3369}
3370
3371unsafe extern "C" fn io_lock(_p_file: *mut sqlite3_file, _level: c_int) -> c_int {
3378 vfs_catch_unwind!(SQLITE_IOERR_LOCK, SQLITE_OK)
3379}
3380
3381unsafe extern "C" fn io_unlock(_p_file: *mut sqlite3_file, _level: c_int) -> c_int {
3382 vfs_catch_unwind!(SQLITE_IOERR_UNLOCK, SQLITE_OK)
3383}
3384
3385unsafe extern "C" fn io_check_reserved_lock(
3386 _p_file: *mut sqlite3_file,
3387 p_res_out: *mut c_int,
3388) -> c_int {
3389 vfs_catch_unwind!(SQLITE_IOERR, {
3390 *p_res_out = 0;
3391 SQLITE_OK
3392 })
3393}
3394
3395unsafe extern "C" fn io_file_control(
3396 p_file: *mut sqlite3_file,
3397 op: c_int,
3398 p_arg: *mut c_void,
3399) -> c_int {
3400 vfs_catch_unwind!(SQLITE_IOERR, {
3401 let file = get_file(p_file);
3402 if get_aux_state(file).is_some() {
3403 return SQLITE_NOTFOUND;
3404 }
3405 let ctx = &*file.ctx;
3406
3407 match op {
3408 SQLITE_FCNTL_BEGIN_ATOMIC_WRITE => {
3409 let mut state = ctx.state.write();
3410 state.write_buffer.in_atomic_write = true;
3411 state.write_buffer.saved_db_size = state.db_size_pages;
3412 state.write_buffer.dirty.clear();
3413 state.refused_write_pending = false;
3414 SQLITE_OK
3415 }
3416 SQLITE_FCNTL_COMMIT_ATOMIC_WRITE => match ctx.commit_atomic_write() {
3417 Ok(()) => {
3418 ctx.commit_atomic_count.fetch_add(1, Ordering::Relaxed);
3419 SQLITE_OK
3420 }
3421 Err(err) => {
3422 tracing::error!(
3423 actor_id = %ctx.actor_id,
3424 last_error = ?ctx.clone_last_error(),
3425 ?err,
3426 "sqlite atomic write file control failed"
3427 );
3428 match &err {
3429 CommitBufferError::Other(message) => {
3430 ctx.defer_transient_commit_error(message.clone());
3431 }
3432 CommitBufferError::Storage(error) => {
3433 ctx.defer_transient_commit_error(error.to_string());
3434 }
3435 CommitBufferError::FenceMismatch(_)
3436 | CommitBufferError::Indeterminate(_) => {}
3437 }
3438 handle_finalize_fence_error(ctx, &err);
3439 SQLITE_IOERR
3440 }
3441 },
3442 SQLITE_FCNTL_ROLLBACK_ATOMIC_WRITE => {
3443 let mut state = ctx.state.write();
3444 state.write_buffer.dirty.clear();
3445 state.write_buffer.in_atomic_write = false;
3446 state.db_size_pages = state.write_buffer.saved_db_size;
3447 state.refused_write_pending = false;
3448 SQLITE_OK
3449 }
3450 SQLITE_FCNTL_OVERWRITE => {
3453 if p_arg.is_null() {
3454 return SQLITE_NOTFOUND;
3455 }
3456 let bytes = *(p_arg as *const i64);
3457 let pages = (bytes.max(0) as usize).div_ceil(ctx.page_size());
3458 if pages > ctx.config.commit_page_cap() {
3459 ctx.refuse_oversize_transaction()
3460 } else {
3461 SQLITE_NOTFOUND
3462 }
3463 }
3464 FCNTL_COMMIT_PAGE_CAP => {
3465 if p_arg.is_null() {
3466 return SQLITE_MISUSE;
3467 }
3468 *(p_arg as *mut i64) = ctx.config.commit_page_cap() as i64;
3469 SQLITE_OK
3470 }
3471 _ => SQLITE_NOTFOUND,
3472 }
3473 })
3474}
3475
3476unsafe extern "C" fn io_sector_size(_p_file: *mut sqlite3_file) -> c_int {
3477 vfs_catch_unwind!(DEFAULT_PAGE_SIZE as c_int, DEFAULT_PAGE_SIZE as c_int)
3478}
3479
3480unsafe extern "C" fn io_device_characteristics(p_file: *mut sqlite3_file) -> c_int {
3481 vfs_catch_unwind!(0, {
3482 let file = get_file(p_file);
3483 if get_aux_state(file).is_some() {
3484 0
3485 } else {
3486 SQLITE_IOCAP_BATCH_ATOMIC
3487 }
3488 })
3489}
3490
3491unsafe extern "C" fn vfs_open(
3492 p_vfs: *mut sqlite3_vfs,
3493 z_name: *const c_char,
3494 p_file: *mut sqlite3_file,
3495 flags: c_int,
3496 p_out_flags: *mut c_int,
3497) -> c_int {
3498 vfs_catch_unwind!(SQLITE_CANTOPEN, {
3499 let ctx = get_vfs_ctx(p_vfs);
3500 let delete_on_close = (flags & SQLITE_OPEN_DELETEONCLOSE) != 0;
3501 let path = if z_name.is_null() {
3502 if delete_on_close {
3503 next_temp_aux_path()
3504 } else {
3505 return SQLITE_CANTOPEN;
3506 }
3507 } else {
3508 match CStr::from_ptr(z_name).to_str() {
3509 Ok(path) => path.to_string(),
3510 Err(_) => return SQLITE_CANTOPEN,
3511 }
3512 };
3513 let is_main =
3514 path == ctx.actor_id && !delete_on_close && (flags & SQLITE_OPEN_MAIN_DB) != 0;
3515
3516 let base = sqlite3_file {
3517 pMethods: ctx.io_methods.as_ref(),
3518 };
3519 let aux = if is_main {
3520 ptr::null_mut()
3521 } else {
3522 Box::into_raw(Box::new(AuxFileHandle {
3523 is_main_journal: (flags & SQLITE_OPEN_MAIN_JOURNAL) != 0
3524 && path.strip_suffix("-journal") == Some(ctx.actor_id.as_str()),
3525 path: path.clone(),
3526 state: ctx.open_aux_file(&path),
3527 delete_on_close,
3528 }))
3529 };
3530 ptr::write(
3531 p_file.cast::<VfsFile>(),
3532 VfsFile {
3533 base,
3534 ctx: ctx as *const VfsContext,
3535 aux,
3536 },
3537 );
3538
3539 if !p_out_flags.is_null() {
3540 *p_out_flags = flags;
3541 }
3542
3543 SQLITE_OK
3544 })
3545}
3546
3547unsafe extern "C" fn vfs_delete(
3548 p_vfs: *mut sqlite3_vfs,
3549 z_name: *const c_char,
3550 _sync_dir: c_int,
3551) -> c_int {
3552 vfs_catch_unwind!(SQLITE_IOERR_DELETE, {
3553 if z_name.is_null() {
3554 return SQLITE_OK;
3555 }
3556
3557 let ctx = get_vfs_ctx(p_vfs);
3558 let path = match CStr::from_ptr(z_name).to_str() {
3559 Ok(path) => path,
3560 Err(_) => return SQLITE_OK,
3561 };
3562 if path == ctx.actor_id {
3563 ctx.set_last_error("main database deletion is unsupported".to_string());
3565 return SQLITE_IOERR_DELETE;
3566 } else {
3567 ctx.delete_aux_file(path);
3568 }
3569 SQLITE_OK
3570 })
3571}
3572
3573unsafe extern "C" fn vfs_access(
3574 p_vfs: *mut sqlite3_vfs,
3575 z_name: *const c_char,
3576 _flags: c_int,
3577 p_res_out: *mut c_int,
3578) -> c_int {
3579 vfs_catch_unwind!(SQLITE_IOERR_ACCESS, {
3580 if z_name.is_null() {
3581 *p_res_out = 0;
3582 return SQLITE_OK;
3583 }
3584
3585 let ctx = get_vfs_ctx(p_vfs);
3586 let path = match CStr::from_ptr(z_name).to_str() {
3587 Ok(path) => path,
3588 Err(_) => {
3589 *p_res_out = 0;
3590 return SQLITE_OK;
3591 }
3592 };
3593
3594 *p_res_out = if path == ctx.actor_id || ctx.aux_file_exists(path) {
3595 1
3596 } else {
3597 0
3598 };
3599 SQLITE_OK
3600 })
3601}
3602
3603unsafe extern "C" fn vfs_full_pathname(
3604 _p_vfs: *mut sqlite3_vfs,
3605 z_name: *const c_char,
3606 n_out: c_int,
3607 z_out: *mut c_char,
3608) -> c_int {
3609 vfs_catch_unwind!(SQLITE_IOERR, {
3610 if z_name.is_null() || z_out.is_null() || n_out <= 0 {
3611 return SQLITE_IOERR;
3612 }
3613
3614 let name = CStr::from_ptr(z_name);
3615 let bytes = name.to_bytes_with_nul();
3616 if bytes.len() >= n_out as usize {
3617 return SQLITE_IOERR;
3618 }
3619
3620 ptr::copy_nonoverlapping(bytes.as_ptr().cast::<c_char>(), z_out, bytes.len());
3621 SQLITE_OK
3622 })
3623}
3624
3625unsafe extern "C" fn vfs_randomness(
3626 _p_vfs: *mut sqlite3_vfs,
3627 n_byte: c_int,
3628 z_out: *mut c_char,
3629) -> c_int {
3630 vfs_catch_unwind!(0, {
3631 let buf = slice::from_raw_parts_mut(z_out.cast::<u8>(), n_byte as usize);
3632 match getrandom::getrandom(buf) {
3633 Ok(()) => n_byte,
3634 Err(_) => 0,
3635 }
3636 })
3637}
3638
3639unsafe extern "C" fn vfs_sleep(_p_vfs: *mut sqlite3_vfs, microseconds: c_int) -> c_int {
3640 vfs_catch_unwind!(0, {
3641 std::thread::sleep(std::time::Duration::from_micros(microseconds as u64));
3642 microseconds
3643 })
3644}
3645
3646unsafe extern "C" fn vfs_current_time(_p_vfs: *mut sqlite3_vfs, p_time_out: *mut f64) -> c_int {
3647 vfs_catch_unwind!(SQLITE_IOERR, {
3648 let now = std::time::SystemTime::now()
3649 .duration_since(std::time::UNIX_EPOCH)
3650 .unwrap_or_default();
3651 *p_time_out = 2440587.5 + (now.as_secs_f64() / 86400.0);
3652 SQLITE_OK
3653 })
3654}
3655
3656unsafe extern "C" fn vfs_get_last_error(
3657 p_vfs: *mut sqlite3_vfs,
3658 n_byte: c_int,
3659 z_err_msg: *mut c_char,
3660) -> c_int {
3661 vfs_catch_unwind!(SQLITE_IOERR, {
3662 if n_byte <= 0 || z_err_msg.is_null() {
3663 return 0;
3664 }
3665
3666 let ctx = get_vfs_ctx(p_vfs);
3667 let Some(message) = ctx.clone_last_error() else {
3668 *z_err_msg = 0;
3669 return 0;
3670 };
3671
3672 let bytes = message.as_bytes();
3673 let max_len = (n_byte as usize).saturating_sub(1);
3674 let copy_len = bytes.len().min(max_len);
3675 let dst = z_err_msg.cast::<u8>();
3676 ptr::copy_nonoverlapping(bytes.as_ptr(), dst, copy_len);
3677 *dst.add(copy_len) = 0;
3678 0
3679 })
3680}
3681
3682impl SqliteVfs {
3683 pub(crate) fn take_last_error(&self) -> Option<String> {
3684 self.ctx.take_last_error()
3685 }
3686
3687 pub(crate) fn take_indeterminate_result(&self) -> Option<String> {
3688 self.ctx.take_indeterminate_result()
3689 }
3690
3691 pub(crate) fn take_transaction_too_large(&self) -> Option<SqliteTransactionTooLargeError> {
3692 self.ctx.take_transaction_too_large()
3693 }
3694
3695 pub(crate) fn take_storage_error(&self) -> Option<SqliteStorageError> {
3696 self.ctx.take_storage_error()
3697 }
3698
3699 fn clone_last_error(&self) -> Option<String> {
3700 self.ctx.clone_last_error()
3701 }
3702
3703 pub fn clone_fatal_error(&self) -> Option<String> {
3704 self.ctx.clone_fatal_error()
3705 }
3706
3707 pub(crate) fn read_aux_file(&self, path: &str) -> Option<Vec<u8>> {
3711 self.ctx.read_aux_file(path)
3712 }
3713
3714 pub(crate) fn delete_aux_file(&self, path: &str) {
3715 self.ctx.delete_aux_file(path);
3716 }
3717
3718 pub(crate) fn sqlite_vfs_metrics(&self) -> SqliteVfsMetricsSnapshot {
3719 self.ctx.sqlite_vfs_metrics()
3720 }
3721
3722 pub(crate) fn register_with_transport_and_initial_pages(
3723 name: &str,
3724 transport: SqliteTransportHandle,
3725 actor_id: String,
3726 runtime: Handle,
3727 config: VfsConfig,
3728 initial_pages: InitialPages,
3729 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
3730 ) -> std::result::Result<Self, String> {
3731 let mut io_methods: sqlite3_io_methods = unsafe { std::mem::zeroed() };
3732 io_methods.iVersion = 1;
3733 io_methods.xClose = Some(io_close);
3734 io_methods.xRead = Some(io_read);
3735 io_methods.xWrite = Some(io_write);
3736 io_methods.xTruncate = Some(io_truncate);
3737 io_methods.xSync = Some(io_sync);
3738 io_methods.xFileSize = Some(io_file_size);
3739 io_methods.xLock = Some(io_lock);
3740 io_methods.xUnlock = Some(io_unlock);
3741 io_methods.xCheckReservedLock = Some(io_check_reserved_lock);
3742 io_methods.xFileControl = Some(io_file_control);
3743 io_methods.xSectorSize = Some(io_sector_size);
3744 io_methods.xDeviceCharacteristics = Some(io_device_characteristics);
3745
3746 let generation = name
3747 .rsplit_once("-g")
3748 .and_then(|(_, generation)| generation.parse::<u64>().ok());
3749 let mut ctx = Box::new(VfsContext::new(
3750 actor_id,
3751 generation,
3752 runtime,
3753 transport,
3754 config,
3755 io_methods,
3756 initial_pages,
3757 metrics,
3758 )?);
3759 tracing::trace!(target: "sqlite_prefetch::trace", event = "open",
3760 stream_id = %format_args!("sqlite-{:p}", &*ctx), actor_id = %ctx.actor_id, generation = ctx.generation);
3761 let ctx_ptr = (&mut *ctx) as *mut VfsContext;
3762 let name_cstring = CString::new(name).map_err(|err| err.to_string())?;
3763
3764 let mut vfs: sqlite3_vfs = unsafe { std::mem::zeroed() };
3765 vfs.iVersion = 1;
3766 vfs.szOsFile = std::mem::size_of::<VfsFile>() as c_int;
3767 vfs.mxPathname = MAX_PATHNAME;
3768 vfs.zName = name_cstring.as_ptr();
3769 vfs.pAppData = ctx_ptr.cast::<c_void>();
3770 vfs.xOpen = Some(vfs_open);
3771 vfs.xDelete = Some(vfs_delete);
3772 vfs.xAccess = Some(vfs_access);
3773 vfs.xFullPathname = Some(vfs_full_pathname);
3774 vfs.xRandomness = Some(vfs_randomness);
3775 vfs.xSleep = Some(vfs_sleep);
3776 vfs.xCurrentTime = Some(vfs_current_time);
3777 vfs.xGetLastError = Some(vfs_get_last_error);
3778
3779 let registration = SqliteVfsRegistration::register(vfs)?;
3780
3781 Ok(Self {
3782 _registration: registration,
3783 _name: name_cstring,
3784 ctx,
3785 })
3786 }
3787
3788 pub fn name_ptr(&self) -> *const c_char {
3789 self._name.as_ptr()
3790 }
3791
3792 pub fn cached_page_bytes(&self) -> u64 {
3795 let state = self.ctx.state.read();
3796 state
3797 .prefetch_buffer
3798 .entry_count()
3799 .saturating_add(state.committed_buffer.entry_count())
3800 .saturating_mul(state.page_size as u64)
3801 }
3802
3803 pub fn idle_release(&self) -> Option<std::time::Duration> {
3805 match self.ctx.config.idle_release_ms {
3806 0 => None,
3807 ms => Some(std::time::Duration::from_millis(ms)),
3808 }
3809 }
3810
3811 fn ctx(&self) -> &VfsContext {
3812 &self.ctx
3813 }
3814
3815 fn commit_atomic_count(&self) -> u64 {
3816 self.ctx.commit_atomic_count.load(Ordering::Relaxed)
3817 }
3818}
3819
3820impl SqliteVfsRegistration {
3821 fn register(vfs: sqlite3_vfs) -> std::result::Result<Self, String> {
3822 let vfs_ptr = Box::into_raw(Box::new(vfs));
3823 let rc = unsafe { sqlite3_vfs_register(vfs_ptr, 0) };
3824 if rc != SQLITE_OK {
3825 unsafe {
3826 drop(Box::from_raw(vfs_ptr));
3827 }
3828 return Err(format!("sqlite3_vfs_register failed with code {rc}"));
3829 }
3830
3831 Ok(Self { vfs_ptr })
3832 }
3833}
3834
3835impl Drop for SqliteVfsRegistration {
3836 fn drop(&mut self) {
3837 unsafe {
3838 sqlite3_vfs_unregister(self.vfs_ptr);
3839 drop(Box::from_raw(self.vfs_ptr));
3840 }
3841 }
3842}
3843
3844impl NativeDatabase {
3845 pub fn as_ptr(&self) -> *mut sqlite3 {
3846 self.db
3847 }
3848
3849 pub fn take_last_kv_error(&self) -> Option<String> {
3850 self._vfs.take_last_error()
3851 }
3852
3853 pub fn sqlite_vfs_metrics(&self) -> SqliteVfsMetricsSnapshot {
3854 self._vfs.ctx.sqlite_vfs_metrics()
3855 }
3856
3857 pub fn round_trip_counts(&self) -> SqliteRoundTripCounts {
3858 self._vfs.ctx.round_trip_counts()
3859 }
3860
3861 pub(crate) fn begin_operation_profile(&self) -> SqliteOperationProfileGuard<'_> {
3862 self._vfs.ctx.begin_operation_profile();
3863 SqliteOperationProfileGuard {
3864 ctx: &self._vfs.ctx,
3865 active: true,
3866 }
3867 }
3868
3869 pub fn cached_page_bytes(&self) -> u64 {
3871 self._vfs.cached_page_bytes()
3872 }
3873
3874 pub fn cork_commits(&self) {
3882 self._vfs.ctx().cork();
3883 }
3884
3885 pub fn uncork_commits(&self) -> Result<()> {
3892 let autocommit = unsafe { sqlite3_get_autocommit(self.as_ptr()) };
3894 if autocommit == 0 {
3895 anyhow::bail!("cannot uncork sqlite commits inside a transaction");
3896 }
3897 self._vfs
3898 .ctx()
3899 .uncork()
3900 .map_err(|error| anyhow::anyhow!("uncork sqlite commits: {error:?}"))
3901 }
3902
3903 pub fn release_idle_memory(&self) -> bool {
3904 {
3905 let state = self._vfs.ctx.state.read();
3906 if state.write_buffer.in_atomic_write
3907 || !state.write_buffer.dirty.is_empty()
3908 || state.corked.is_some()
3909 {
3910 return false;
3911 }
3912 }
3913 unsafe {
3914 sqlite3_db_release_memory(self.db);
3915 }
3916 self._vfs.ctx.state.write().invalidate_page_buffers();
3917 self._vfs.ctx.report_prefetch_unused();
3918 true
3919 }
3920}
3921
3922impl Drop for NativeDatabase {
3923 fn drop(&mut self) {
3924 if !self.db.is_null() {
3925 let ctx = self._vfs.ctx();
3926 ctx.discard_refused_transaction();
3927 if write_transaction_open(self.db) {
3930 ctx.discard_open_transaction();
3931 }
3932 {
3934 let mut state = ctx.state.write();
3935 if let Some(mut corked) = state.corked.take() {
3936 corked.append(&mut state.write_buffer.dirty);
3937 state.write_buffer.dirty = corked;
3938 }
3939 }
3940 let should_flush = {
3941 let state = ctx.state.read();
3942 state.write_buffer.in_atomic_write
3943 || !state.write_buffer.dirty.is_empty()
3944 || state.db_size_pages != state.committed_db_size_pages
3945 };
3946 if should_flush {
3947 let result = if ctx.state.read().write_buffer.in_atomic_write {
3948 ctx.commit_atomic_write_with_timeout(Some(NATIVE_DATABASE_DROP_FLUSH_TIMEOUT))
3949 } else {
3950 ctx.flush_dirty_pages_with_timeout(Some(NATIVE_DATABASE_DROP_FLUSH_TIMEOUT))
3951 .map(|wait| match wait {
3952 CommitWait::Completed(_) => CommitWait::Completed(()),
3953 CommitWait::TimedOut => CommitWait::TimedOut,
3954 })
3955 };
3956 match result {
3961 Ok(CommitWait::Completed(())) => {}
3962 Ok(CommitWait::TimedOut) => {
3963 let discarded_pages = ctx.discard_unsent_commits();
3964 tracing::error!(
3965 actor_id = %ctx.actor_id,
3966 timeout_ms = NATIVE_DATABASE_DROP_FLUSH_TIMEOUT.as_millis(),
3967 discarded_pages,
3968 "timed out flushing sqlite database before close, discarding its unsent commits"
3969 );
3970 }
3971 Err(err) => {
3972 handle_non_finalize_commit_error(ctx, &err);
3973 let discarded_pages = ctx.discard_unsent_commits();
3974 tracing::error!(
3975 actor_id = %ctx.actor_id,
3976 ?err,
3977 discarded_pages,
3978 "failed to flush sqlite database before close, discarding its unsent commits"
3979 );
3980 }
3981 }
3982 }
3983
3984 ctx.state.write().closing = true;
3989 let rc = unsafe { sqlite3_close_v2(self.db) };
3990 if rc != SQLITE_OK {
3991 tracing::warn!(
3992 rc,
3993 error = sqlite_error_message(self.db),
3994 "failed to close sqlite database"
3995 );
3996 }
3997 self.db = ptr::null_mut();
3998 }
3999 }
4000}
4001
4002fn write_transaction_open(db: *mut sqlite3) -> bool {
4005 let autocommit = unsafe { sqlite3_get_autocommit(db) };
4007 let txn_state = unsafe { sqlite3_txn_state(db, ptr::null()) };
4008 autocommit == 0 || txn_state == SQLITE_TXN_WRITE
4009}
4010
4011pub fn open_database(vfs: SqliteVfs, file_name: &str) -> Result<NativeDatabase> {
4012 open_connection(
4013 Arc::new(vfs),
4014 file_name,
4015 SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE,
4016 )
4017 .map_err(anyhow::Error::msg)
4018 .and_then(|connection| {
4019 configure_connection_for_database(connection.as_ptr(), &connection._vfs, file_name)?;
4020 Ok(connection)
4021 })
4022}
4023
4024pub fn open_connection(
4025 vfs: NativeVfsHandle,
4026 file_name: &str,
4027 flags: c_int,
4028) -> std::result::Result<NativeConnection, String> {
4029 let c_name = CString::new(file_name).map_err(|err| err.to_string())?;
4030 let mut db: *mut sqlite3 = ptr::null_mut();
4031
4032 let rc = unsafe { sqlite3_open_v2(c_name.as_ptr(), &mut db, flags, vfs.name_ptr()) };
4033 if rc != SQLITE_OK {
4034 let message = sqlite_error_message(db);
4035 tracing::error!(
4036 file_name,
4037 rc,
4038 %message,
4039 last_error = ?vfs.clone_last_error(),
4040 "failed to open sqlite database with custom VFS"
4041 );
4042 if !db.is_null() {
4043 unsafe {
4044 sqlite3_close(db);
4045 }
4046 }
4047 return Err(format!("sqlite3_open_v2 failed with code {rc}: {message}"));
4048 }
4049
4050 Ok(NativeDatabase { db, _vfs: vfs })
4051}
4052
4053#[derive(rivet_error::RivetError)]
4054#[error(
4055 "sqlite",
4056 "invalid_compile_options",
4057 "SQLite requires SQLITE_ENABLE_BATCH_ATOMIC_WRITE and SQLITE_DIRECT_OVERFLOW_READ=0. Rebuild with LIBSQLITE3_FLAGS=\"SQLITE_ENABLE_BATCH_ATOMIC_WRITE SQLITE_DIRECT_OVERFLOW_READ=0\"."
4058)]
4059struct SqliteInvalidCompileOptions;
4060
4061pub fn configure_connection_for_database(
4062 db: *mut sqlite3,
4063 vfs: &SqliteVfs,
4064 file_name: &str,
4065) -> Result<()> {
4066 let batch_atomic = unsafe { sqlite3_compileoption_used(c"ENABLE_BATCH_ATOMIC_WRITE".as_ptr()) };
4068 let direct_overflow = unsafe { sqlite3_compileoption_used(c"DIRECT_OVERFLOW_READ".as_ptr()) };
4069 if batch_atomic == 0 || direct_overflow != 0 {
4070 return Err(SqliteInvalidCompileOptions.build());
4071 }
4072
4073 let cache_size_kib = sqlite_optimization_flags().pager_cache_size_kib;
4079 let cache_size_pragma = format!("PRAGMA cache_size = -{cache_size_kib};");
4080
4081 let pragmas = [
4082 "PRAGMA page_size = 4096;",
4083 "PRAGMA journal_mode = DELETE;",
4084 "PRAGMA synchronous = NORMAL;",
4085 "PRAGMA temp_store = MEMORY;",
4086 "PRAGMA auto_vacuum = NONE;",
4087 "PRAGMA locking_mode = EXCLUSIVE;",
4088 "PRAGMA journal_size_limit = 0;",
4091 cache_size_pragma.as_str(),
4092 ];
4093 for pragma in &pragmas {
4094 if let Err(err) = sqlite_exec(db, pragma) {
4095 tracing::error!(
4096 file_name,
4097 pragma,
4098 %err,
4099 last_error = ?vfs.clone_last_error(),
4100 "failed to configure sqlite database"
4101 );
4102 return Err(anyhow::Error::msg(err));
4103 }
4104 }
4105
4106 let rc = unsafe { sqlite3_set_authorizer(db, Some(authorize_journal_mode), ptr::null_mut()) };
4107 if rc != SQLITE_OK {
4108 anyhow::bail!("failed to install the sqlite authorizer: {rc}");
4109 }
4110
4111 Ok(())
4112}
4113
4114extern "C" fn authorize_journal_mode(
4123 _context: *mut c_void,
4124 action: c_int,
4125 first: *const c_char,
4126 second: *const c_char,
4127 _database: *const c_char,
4128 _trigger: *const c_char,
4129) -> c_int {
4130 if action != SQLITE_PRAGMA || first.is_null() || second.is_null() {
4131 return SQLITE_OK;
4132 }
4133 let name = unsafe { CStr::from_ptr(first) }.to_string_lossy();
4134 if name.eq_ignore_ascii_case("journal_size_limit") {
4137 return SQLITE_DENY;
4138 }
4139 let value = unsafe { CStr::from_ptr(second) }.to_string_lossy();
4140 let value = value.trim_matches(|c| matches!(c, '"' | '\''));
4141 if name.eq_ignore_ascii_case("synchronous") {
4142 let syncs_every_commit = [
4143 "normal", "full", "extra", "1", "2", "3", "on", "yes", "true",
4144 ]
4145 .iter()
4146 .any(|level| value.eq_ignore_ascii_case(level));
4147 return if syncs_every_commit {
4148 SQLITE_OK
4149 } else {
4150 SQLITE_DENY
4151 };
4152 }
4153 if !name.eq_ignore_ascii_case("journal_mode") {
4154 return SQLITE_OK;
4155 }
4156 if value.eq_ignore_ascii_case("delete") {
4157 SQLITE_OK
4158 } else {
4159 SQLITE_DENY
4160 }
4161}