1use std::{
10 collections::{BTreeMap, BTreeSet, VecDeque},
11 error::Error,
12 fmt,
13 sync::{
14 Arc, Weak,
15 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
16 },
17 thread::JoinHandle,
18 time::{Duration, Instant},
19};
20
21use anyhow::{Result, anyhow};
22use libsqlite3_sys::{
23 SQLITE_OPEN_CREATE, SQLITE_OPEN_READWRITE, sqlite3_compileoption_used, sqlite3_get_autocommit,
24};
25use parking_lot::{Condvar, Mutex};
26use tokio::sync::{Notify, oneshot, watch};
27
28use crate::{
29 guard::{ConnectionGuard, InstalledGuard},
30 optimization_flags::sqlite_optimization_flags,
31 page_buffer::PageBudget,
32 query::{
33 BindParam, ExecuteResult, QueryResult, exec_statements, exec_statements_profiled,
34 execute_single_statement, execute_single_statement_profiled,
35 },
36 vfs::{
37 NativeConnection, NativeVfsHandle, SqliteCommitIndeterminateError, SqliteOperationProfile,
38 SqliteRoundTripCounts, SqliteVfsMetrics, configure_connection_for_database,
39 open_connection, sqlite_exec,
40 },
41};
42
43pub const SQLITE_WORKER_QUEUE_CAPACITY: usize = 128;
46const SQLITE_WORKER_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
47const SQLITE_SLOW_TRANSACTION_THRESHOLD: Duration = Duration::from_millis(2500);
51
52pub const SQLITE_WORKER_THREADS_ENV: &str = "RIVETKIT_SQLITE_WORKER_THREADS";
54pub const DEFAULT_SQLITE_WORKER_THREADS: usize = 256;
55const MAX_SQLITE_WORKER_THREADS: usize = 4096;
56const SQLITE_WORKER_STACK_BYTES: usize = 512 * 1024;
58const SQLITE_POOL_SATURATED_WARN: Duration = Duration::from_secs(1);
60const ALLOCATOR_TRIM_INTERVAL: Duration = Duration::from_secs(5);
64const BACKGROUND_TURN_EVERY: usize = 4;
66const IDLE_RELEASES_PER_WAKE: usize = 64;
68
69const STATE_RUNNING: u8 = 0;
70const STATE_CLOSING: u8 = 1;
71const STATE_CLOSED: u8 = 2;
72const STATE_DEAD: u8 = 3;
73
74pub trait SqlitePoolMetrics: Send + Sync {
76 fn set_pool_threads(&self, threads: u64);
78
79 fn set_pool_busy(&self, busy: u64);
81
82 fn set_pool_runnable(&self, runnable: u64);
84
85 fn observe_pool_queue_wait(&self, wait_ns: u64);
87
88 fn observe_pool_lock_wait(&self, lock: &'static str, wait_ns: u64);
90
91 fn set_page_buffer_capacity_pages(&self, _pages: u64) {}
93
94 fn set_page_buffer_pages(&self, _pages: u64) {}
96
97 fn set_pager_cache_capacity_bytes(&self, _bytes: u64) {}
99
100 fn set_pager_cache_bytes(&self, _bytes: u64) {}
102
103 fn set_pager_cache_pages(&self, _pages: u64) {}
105
106 fn set_pager_cache_unbudgeted_bytes(&self, _bytes: u64) {}
109}
110
111#[derive(Clone)]
115pub struct SqlitePool {
116 owner: Arc<PoolOwner>,
117 page_budget: Arc<PageBudget>,
118}
119
120struct PoolOwner {
123 shared: Arc<PoolShared>,
124}
125
126struct PoolShared {
127 max_threads: usize,
128 state: Mutex<PoolState>,
130 idle: Mutex<IdleQueue<DatabaseSlot>>,
133 epoch: Instant,
135 wake: Condvar,
136 metrics: Option<Arc<dyn SqlitePoolMetrics>>,
137
138}
139
140struct IdleQueue<T> {
145 queue: BTreeSet<(Instant, u64)>,
146 entries: BTreeMap<u64, (Instant, Weak<T>)>,
147}
148
149impl<T> Default for IdleQueue<T> {
150 fn default() -> Self {
151 Self {
152 queue: BTreeSet::new(),
153 entries: BTreeMap::new(),
154 }
155 }
156}
157
158impl<T> IdleQueue<T> {
159 fn insert(&mut self, id: u64, at: Instant, item: Weak<T>) -> bool {
162 if let Some((old, _)) = self.entries.insert(id, (at, item)) {
163 self.queue.remove(&(old, id));
164 }
165 self.queue.insert((at, id));
166 self.queue.first() == Some(&(at, id))
167 }
168
169 fn cancel(&mut self, id: u64) {
170 if let Some((at, _)) = self.entries.remove(&id) {
171 self.queue.remove(&(at, id));
172 }
173 }
174
175 fn head(&self) -> Option<Instant> {
176 self.queue.first().map(|(at, _)| *at)
177 }
178
179
180
181 fn pop_due(
184 &mut self,
185 now: Instant,
186 limit: usize,
187 due_at: impl Fn(&Arc<T>) -> Option<Instant>,
188 ) -> Vec<Arc<T>> {
189 let mut due = Vec::new();
190 while due.len() < limit {
191 let Some(&(at, id)) = self.queue.first() else {
192 break;
193 };
194 if at > now {
195 break;
196 }
197 self.queue.pop_first();
198 let Some((_, item)) = self.entries.remove(&id) else {
199 continue;
200 };
201 let Some(item) = item.upgrade() else {
202 continue;
203 };
204 match due_at(&item) {
205 Some(later) if later > now => {
206 self.entries.insert(id, (later, Arc::downgrade(&item)));
207 self.queue.insert((later, id));
208 }
209 Some(_) | None => due.push(item),
210 }
211 }
212 due
213 }
214}
215
216struct RunQueue<T> {
221 foreground: VecDeque<(T, Instant)>,
222 background: VecDeque<(T, Instant)>,
223 foreground_streak: usize,
225}
226
227impl<T> Default for RunQueue<T> {
228 fn default() -> Self {
229 Self {
230 foreground: VecDeque::new(),
231 background: VecDeque::new(),
232 foreground_streak: 0,
233 }
234 }
235}
236
237impl<T> RunQueue<T> {
238 fn push(&mut self, item: T, background: bool, now: Instant) {
239 if background {
240 self.background.push_back((item, now));
241 } else {
242 self.foreground.push_back((item, now));
243 }
244 }
245
246 fn pop(&mut self) -> Option<(T, Instant)> {
247 let background_turn = !self.background.is_empty()
248 && (self.foreground.is_empty() || self.foreground_streak >= BACKGROUND_TURN_EVERY);
249 if background_turn {
250 self.foreground_streak = 0;
251 return self.background.pop_front();
252 }
253 let next = self.foreground.pop_front()?;
254 if self.background.is_empty() {
255 self.foreground_streak = 0;
256 } else {
257 self.foreground_streak += 1;
258 }
259 Some(next)
260 }
261
262 fn len(&self) -> usize {
263 self.foreground.len() + self.background.len()
264 }
265
266 fn is_empty(&self) -> bool {
267 self.foreground.is_empty() && self.background.is_empty()
268 }
269}
270
271#[derive(Default)]
272struct PoolState {
273 runnable: RunQueue<Arc<DatabaseSlot>>,
274 orphans: Vec<NativeConnection>,
277 threads: Vec<JoinHandle<()>>,
278 waiting: usize,
280 busy: usize,
281 saturation: Saturation,
282 next_database: u64,
283 shutdown: bool,
284 trim_due_at: Option<Instant>,
286 last_trim_at: Option<Instant>,
288}
289
290#[derive(Default)]
294struct Saturation {
295 since: Option<Instant>,
296 warned: bool,
297}
298
299impl Saturation {
300 fn freed(&mut self) {
302 self.since = None;
303 self.warned = false;
304 }
305
306 fn work_waiting(&mut self, saturated: bool, now: Instant) -> bool {
310 if !saturated {
311 return false;
312 }
313 let since = *self.since.get_or_insert(now);
314 if self.warned || now.duration_since(since) < SQLITE_POOL_SATURATED_WARN {
315 return false;
316 }
317 self.warned = true;
318 true
319 }
320}
321
322struct DatabaseSlot {
325 opening_trace: Option<WorkerTrace>,
326 id: u64,
327 inner: Arc<SqliteWorkerInner>,
328 shared: Arc<PoolShared>,
329 vfs: NativeVfsHandle,
330 file_name: String,
331 idle_release: Option<Duration>,
332 last_used_ns: AtomicU64,
334 idle_queued: AtomicBool,
336 background: AtomicBool,
339 guard: Option<Box<InstalledGuard>>,
341 work: Mutex<SlotWork>,
342}
343
344#[derive(Clone)]
347struct WorkerTrace {
348 span: tracing::Span,
349 dispatch: tracing::Dispatch,
350}
351impl Drop for WorkerTrace {
352 fn drop(&mut self) {
353 tracing::dispatcher::with_default(&self.dispatch, || {
356 drop(std::mem::replace(&mut self.span, tracing::Span::none()));
357 });
358 }
359}
360
361impl WorkerTrace {
362 fn capture() -> Option<Self> {
363 tracing::enabled!(target: "sqlite_prefetch::trace", tracing::Level::TRACE).then(|| Self {
364 span: tracing::Span::current(),
365 dispatch: tracing::dispatcher::get_default(Clone::clone),
366 })
367 }
368}
369struct QueuedCommand {
370 command: SqliteCommand,
371 trace: Option<WorkerTrace>,
372}
373
374struct SlotWork {
375 commands: VecDeque<QueuedCommand>,
376 scheduled: bool,
378 idle_due: bool,
380 connection: ConnectionSlot,
381 close_requested_at: Option<Instant>,
382}
383
384enum ConnectionSlot {
385 Unopened(Option<oneshot::Sender<Result<()>>>),
387 Parked(OpenConnection),
389 Running,
391 Gone,
393}
394
395struct OpenConnection {
396 db: NativeConnection,
397 transaction: Option<TransactionTracker>,
398 pending_uncork: Option<oneshot::Sender<Result<()>>>,
401}
402
403enum Action {
405 Open(Option<oneshot::Sender<Result<()>>>),
406 Run(OpenConnection, QueuedCommand),
407 Release(OpenConnection),
408 Close(Option<OpenConnection>),
409}
410
411enum PoolTask {
413 Database(Arc<DatabaseSlot>, Instant),
414 IdleReleases(Vec<Arc<DatabaseSlot>>),
415 Orphan(NativeConnection),
416 TrimAllocator,
417}
418
419fn process_page_budget(metrics: Option<Arc<dyn SqlitePoolMetrics>>) -> Arc<PageBudget> {
422 let flags = sqlite_optimization_flags();
423 let caches_pages = flags.vfs_page_cache_mode.caches_any_pages();
424 let (capacity, ttl_ms) = if caches_pages {
425 (
426 flags.vfs_page_buffer_process_pages,
427 flags.vfs_staging_cache_ttl_ms,
428 )
429 } else {
430 (0, 0)
431 };
432 PageBudget::new(capacity, ttl_ms, metrics)
433}
434
435impl SqlitePool {
436 pub fn from_env(metrics: Option<Arc<dyn SqlitePoolMetrics>>) -> Self {
438 let threads = std::env::var(SQLITE_WORKER_THREADS_ENV)
439 .ok()
440 .and_then(|value| value.trim().parse::<usize>().ok())
441 .filter(|threads| *threads > 0)
442 .unwrap_or(DEFAULT_SQLITE_WORKER_THREADS)
443 .min(MAX_SQLITE_WORKER_THREADS);
444 Self::new(threads, metrics)
445 }
446
447 pub fn new(max_threads: usize, metrics: Option<Arc<dyn SqlitePoolMetrics>>) -> Self {
449 let max_threads = max_threads.max(1);
450 if let Some(metrics) = &metrics {
451 metrics.set_pool_threads(max_threads as u64);
452 }
453 let page_budget = process_page_budget(metrics.clone());
454 crate::pager_cache::install(
457 sqlite_optimization_flags()
458 .pager_process_cache_kib
459 .saturating_mul(1024),
460 metrics.clone(),
461 );
462 warn_if_direct_overflow_reads();
463 Self {
464 page_budget,
465 owner: Arc::new(PoolOwner {
466 shared: Arc::new(PoolShared {
467 max_threads,
468 state: Mutex::new(PoolState::default()),
469 idle: Mutex::new(IdleQueue::default()),
470 epoch: Instant::now(),
471 wake: Condvar::new(),
472 metrics,
473
474 }),
475 }),
476 }
477 }
478
479 pub fn page_budget(&self) -> &Arc<PageBudget> {
481 &self.page_budget
482 }
483
484 pub fn request_allocator_trim(&self) {
488 self.owner.shared.request_allocator_trim();
489 }
490
491 pub fn live_threads(&self) -> usize {
492 self.owner.shared.state.lock().threads.len()
493 }
494
495 pub fn max_threads(&self) -> usize {
496 self.owner.shared.max_threads
497 }
498
499
500
501
502}
503
504impl Drop for PoolOwner {
505 fn drop(&mut self) {
506 let threads = {
507 let mut state = self.shared.state.lock();
508 state.shutdown = true;
509 std::mem::take(&mut state.threads)
510 };
511 self.shared.wake.notify_all();
512 if threads.is_empty() {
513 return;
514 }
515 let joiner = std::thread::Builder::new()
517 .name("sqlite-pool-join".to_owned())
518 .stack_size(64 * 1024)
519 .spawn(move || {
520 for thread in threads {
521 if thread.join().is_err() {
522 tracing::error!("sqlite pool thread panicked during shutdown");
523 }
524 }
525 });
526 if let Err(error) = joiner {
527 tracing::error!(
528 ?error,
529 "failed to spawn sqlite pool joiner; threads exit unjoined"
530 );
531 }
532 }
533}
534
535impl PoolShared {
536 fn schedule(self: &Arc<Self>, slot: Arc<DatabaseSlot>) {
538 let mut state = self.lock_state();
539 let now = Instant::now();
540 let background = slot.background.load(Ordering::Acquire);
541 state.runnable.push(slot, background, now);
542 if let Some(metrics) = &self.metrics {
543 metrics.set_pool_runnable(state.runnable.len() as u64);
544 }
545 let mut failed_slots = Vec::new();
546 if state.waiting == 0 && state.threads.len() < self.max_threads && !state.shutdown {
547 if !self.spawn_thread(&mut state) && state.threads.is_empty() {
548 while let Some((slot, _)) = state.runnable.pop() {
550 failed_slots.push(slot);
551 }
552 if let Some(metrics) = &self.metrics {
553 metrics.set_pool_runnable(0);
554 }
555 }
556 }
557 self.note_waiting_work(&mut state, now);
558 drop(state);
559 for slot in failed_slots {
560 slot.fail_without_worker();
561 }
562 self.wake.notify_one();
563 }
564
565 fn note_waiting_work(&self, state: &mut PoolState, now: Instant) {
568 let saturated = state.waiting == 0 && state.threads.len() >= self.max_threads;
569 if state.saturation.work_waiting(saturated, now) {
570 tracing::warn!(
571 threads = self.max_threads,
572 runnable = state.runnable.len(),
573 "every sqlite pool thread has been busy for over a second while work waited; raise RIVETKIT_SQLITE_WORKER_THREADS or look for slow VFS round trips"
574 );
575 }
576 }
577
578 fn spawn_thread(self: &Arc<Self>, state: &mut PoolState) -> bool {
579 let shared = Arc::clone(self);
580 let index = state.threads.len();
581
582 let spawned = std::thread::Builder::new()
583 .name(format!("sqlite-pool-{index}"))
584 .stack_size(SQLITE_WORKER_STACK_BYTES)
585 .spawn(move || pool_thread(shared));
586 match spawned {
587 Ok(thread) => {
588 state.threads.push(thread);
589 true
590 }
591 Err(error) => {
592 tracing::error!(
594 ?error,
595 threads = state.threads.len(),
596 "failed to spawn sqlite pool thread"
597 );
598 false
599 }
600 }
601 }
602
603 fn lock_state(&self) -> parking_lot::MutexGuard<'_, PoolState> {
605 self.timed_lock(&self.state, "run_queue")
606 }
607
608 fn lock_idle(&self) -> parking_lot::MutexGuard<'_, IdleQueue<DatabaseSlot>> {
609 self.timed_lock(&self.idle, "idle")
610 }
611
612 fn timed_lock<'a, T>(
613 &self,
614 mutex: &'a Mutex<T>,
615 lock: &'static str,
616 ) -> parking_lot::MutexGuard<'a, T> {
617 if let Some(guard) = mutex.try_lock() {
618 return guard;
619 }
620 let started = Instant::now();
621 let guard = mutex.lock();
622 if let Some(metrics) = &self.metrics {
623 metrics.observe_pool_lock_wait(lock, started.elapsed().as_nanos() as u64);
624 }
625 guard
626 }
627
628 fn now_ns(&self) -> u64 {
629 self.epoch.elapsed().as_nanos().min(u64::MAX as u128) as u64
630 }
631
632 fn note_used(&self, slot: &Arc<DatabaseSlot>, idle: Duration) {
635 slot.last_used_ns.store(self.now_ns(), Ordering::SeqCst);
636 if slot.idle_queued.swap(true, Ordering::SeqCst) {
637 return;
638 }
639 let at = Instant::now() + idle;
640 let head = self.lock_idle().insert(slot.id, at, Arc::downgrade(slot));
641 if head {
642 drop(self.lock_state());
645 self.wake.notify_one();
646 }
647 }
648
649 fn idle_due_at(&self, slot: &DatabaseSlot) -> Option<Instant> {
651 let idle = slot.idle_release?;
652 let last_used = Duration::from_nanos(slot.last_used_ns.load(Ordering::SeqCst));
653 Some(self.epoch + last_used + idle)
654 }
655
656 fn cancel_idle(&self, id: u64) {
657 self.lock_idle().cancel(id);
658 }
659
660 fn due_idle_releases(&self, now: Instant) -> Vec<Arc<DatabaseSlot>> {
663 let mut idle = self.lock_idle();
664 let due = idle.pop_due(now, IDLE_RELEASES_PER_WAKE, |slot| self.idle_due_at(slot));
665 let mut released = Vec::with_capacity(due.len());
666 for slot in due {
667 slot.idle_queued.store(false, Ordering::SeqCst);
668 match self.idle_due_at(&slot) {
671 Some(at) if at > now => {
672 if !slot.idle_queued.swap(true, Ordering::SeqCst) {
673 idle.insert(slot.id, at, Arc::downgrade(&slot));
674 }
675 }
676 Some(_) | None => released.push(slot),
677 }
678 }
679 released
680 }
681
682 fn orphan(&self, connection: NativeConnection) {
683 let mut state = self.lock_state();
684 state.orphans.push(connection);
685 drop(state);
686 self.wake.notify_one();
687 }
688
689 fn next_task(&self) -> Option<PoolTask> {
691 let mut state = self.lock_state();
692 loop {
693 if let Some(connection) = state.orphans.pop() {
694 return Some(PoolTask::Orphan(connection));
695 }
696 if let Some((slot, queued_at)) = state.runnable.pop() {
697 if let Some(metrics) = &self.metrics {
698 metrics.set_pool_runnable(state.runnable.len() as u64);
699 }
700 if !state.runnable.is_empty() {
701 self.note_waiting_work(&mut state, Instant::now());
702 }
703 return Some(PoolTask::Database(slot, queued_at));
704 }
705 let due = self.due_idle_releases(Instant::now());
706 if !due.is_empty() {
707 return Some(PoolTask::IdleReleases(due));
708 }
709 if state.trim_due_at.is_some_and(|at| at <= Instant::now()) {
710 state.trim_due_at = None;
711 return Some(PoolTask::TrimAllocator);
712 }
713 if state.shutdown {
714 return None;
715 }
716 state.waiting += 1;
717 state.saturation.freed();
718 let idle_head = self.lock_idle().head();
719 let head = match (idle_head, state.trim_due_at) {
720 (Some(idle), Some(trim)) => Some(idle.min(trim)),
721 (idle, trim) => idle.or(trim),
722 };
723 match head {
724 Some(at) => {
725 self.wake.wait_until(&mut state, at);
726 }
727 None => self.wake.wait(&mut state),
728 }
729 state.waiting -= 1;
730 }
731 }
732
733 fn request_allocator_trim(self: &Arc<Self>) {
736 let mut state = self.lock_state();
737 if state.trim_due_at.is_some() || state.shutdown {
738 return;
739 }
740 let now = Instant::now();
741 let due = state
742 .last_trim_at
743 .map_or(now, |last| (last + ALLOCATOR_TRIM_INTERVAL).max(now));
744 state.trim_due_at = Some(due);
745 if state.waiting == 0 && state.threads.is_empty() {
746 self.spawn_thread(&mut state);
747 }
748 drop(state);
749 self.wake.notify_one();
751 }
752
753 fn set_busy(&self, delta: isize) {
755 let mut state = self.lock_state();
756 state.busy = state.busy.saturating_add_signed(delta);
757 if let Some(metrics) = &self.metrics {
758 metrics.set_pool_busy(state.busy as u64);
759 }
760 }
761}
762
763#[cfg(all(unix, target_env = "gnu"))]
766fn trim_allocator() {
767 unsafe extern "C" {
768 fn malloc_trim(pad: usize) -> i32;
769 }
770
771 let started = Instant::now();
772 let rc = unsafe { malloc_trim(0) };
773 tracing::debug!(
774 rc,
775 duration_us = started.elapsed().as_micros() as u64,
776 "trimmed native allocator"
777 );
778}
779
780#[cfg(not(all(unix, target_env = "gnu")))]
781fn trim_allocator() {}
782
783fn pool_thread(shared: Arc<PoolShared>) {
784 while let Some(task) = shared.next_task() {
785 shared.set_busy(1);
786 match task {
787 PoolTask::Database(slot, queued_at) => {
788 if let Some(metrics) = &shared.metrics {
789 metrics.observe_pool_queue_wait(queued_at.elapsed().as_nanos() as u64);
790 }
791 slot.run_one();
792 }
793 PoolTask::IdleReleases(slots) => {
794 for slot in slots {
795 slot.mark_idle_due();
796 }
797 }
798 PoolTask::Orphan(connection) => drop(connection),
799 PoolTask::TrimAllocator => {
800 trim_allocator();
801 shared.lock_state().last_trim_at = Some(Instant::now());
802 }
803 }
804 shared.set_busy(-1);
805 }
806}
807
808impl DatabaseSlot {
809 fn fail_without_worker(&self) {
811 self.inner.state.store(STATE_DEAD, Ordering::Release);
812 self.inner.explicit_transaction.send_replace(false);
813 let mut work = self.work.lock();
814 debug_assert!(matches!(work.connection, ConnectionSlot::Unopened(_)));
815 work.connection = ConnectionSlot::Gone;
816 work.scheduled = false;
817 fail_queued_commands(&mut work.commands);
818 drop(work);
819 self.inner.closed.notify_waiters();
820 }
821
822 fn enqueue(self: &Arc<Self>, command: SqliteCommand) -> Result<()> {
824 match self.inner.state.load(Ordering::Acquire) {
825 STATE_RUNNING => {}
826 STATE_CLOSING | STATE_CLOSED => return Err(sqlite_closing_error()),
827 STATE_DEAD => return Err(sqlite_worker_dead_error()),
828 other => return Err(anyhow!("unknown sqlite worker state {other}")),
829 }
830 let mut work = self.work.lock();
831 if work.commands.len() >= SQLITE_WORKER_QUEUE_CAPACITY {
832 drop(work);
833 if let Some(metrics) = &self.inner.metrics {
834 metrics.record_worker_queue_overload();
835 }
836 return Err(SqliteWorkerOverloadedError.into());
839 }
840 work.commands.push_back(QueuedCommand {
841 command,
842 trace: WorkerTrace::capture(),
843 });
844 self.inner.record_queue_depth(work.commands.len() as u64);
845 self.schedule_locked(work);
846 Ok(())
847 }
848
849 fn request_close(self: &Arc<Self>) {
851 let mut work = self.work.lock();
852 work.close_requested_at.get_or_insert_with(Instant::now);
853 self.schedule_locked(work);
854 }
855
856 fn mark_idle_due(self: &Arc<Self>) {
857 let mut work = self.work.lock();
858 if !matches!(work.connection, ConnectionSlot::Parked(_)) {
859 return;
860 }
861 work.idle_due = true;
862 self.schedule_locked(work);
863 }
864
865 fn schedule_locked(self: &Arc<Self>, mut work: parking_lot::MutexGuard<'_, SlotWork>) {
868 if work.scheduled {
869 return;
870 }
871 work.scheduled = true;
872 drop(work);
873 self.shared.schedule(Arc::clone(self));
874 }
875
876 fn run_one(self: &Arc<Self>) {
879 let action = {
880 let mut work = self.work.lock();
881 let closing = self.inner.state.load(Ordering::Acquire) == STATE_CLOSING;
882 let connection = std::mem::replace(&mut work.connection, ConnectionSlot::Running);
883 match connection {
884 ConnectionSlot::Unopened(ready) if closing => {
885 drop(ready);
886 Some(Action::Close(None))
887 }
888 ConnectionSlot::Unopened(ready) => Some(Action::Open(ready)),
889 ConnectionSlot::Parked(open) if closing => Some(Action::Close(Some(open))),
890 ConnectionSlot::Parked(open) => match work.commands.pop_front() {
891 Some(command) => {
892 self.inner.record_queue_depth(work.commands.len() as u64);
893 Some(Action::Run(open, command))
894 }
895 None if work.idle_due => Some(Action::Release(open)),
896 None => {
897 work.connection = ConnectionSlot::Parked(open);
898 None
899 }
900 },
901 ConnectionSlot::Running => {
902 tracing::error!("sqlite database scheduled while a thread holds it");
904 None
905 }
906 ConnectionSlot::Gone => {
907 work.connection = ConnectionSlot::Gone;
908 fail_queued_commands(&mut work.commands);
909 None
910 }
911 }
912 };
913 let Some(action) = action else {
914 self.finish_idle();
915 return;
916 };
917
918 let ran_command = matches!(action, Action::Run(..));
919 let released = matches!(action, Action::Release(_));
920 let outcome =
921 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.perform(action)));
922 let next = match outcome {
923 Ok(next) => next,
924 Err(panic) => {
925 self.inner.state.store(STATE_DEAD, Ordering::Release);
926 self.inner.explicit_transaction.send_replace(false);
927 if let Some(metrics) = &self.inner.metrics {
928 metrics.record_worker_crash();
929 metrics.set_vfs_cached_bytes(0);
930 }
931 self.inner.closed.notify_waiters();
932 tracing::error!(msg = panic_message(&panic), "sqlite worker panicked");
933 ConnectionSlot::Gone
934 }
935 };
936 let gone = matches!(next, ConnectionSlot::Gone);
937
938 let requeue = {
939 let mut work = self.work.lock();
940 work.connection = next;
941 if released {
942 work.idle_due = false;
943 }
944 if gone {
945 fail_queued_commands(&mut work.commands);
946 }
947 let more = !gone
948 && (!work.commands.is_empty()
949 || work.idle_due
950 || self.inner.state.load(Ordering::Acquire) == STATE_CLOSING);
951 if !more {
952 work.scheduled = false;
953 }
954 more
955 };
956 if gone {
957 self.shared.cancel_idle(self.id);
958 } else if ran_command && let Some(idle) = self.idle_release {
959 self.shared.note_used(self, idle);
960 }
961 if requeue {
962 self.shared.schedule(Arc::clone(self));
963 }
964 }
965
966 fn finish_idle(self: &Arc<Self>) {
968 let requeue = {
969 let mut work = self.work.lock();
970 let more = !work.commands.is_empty()
971 || (matches!(work.connection, ConnectionSlot::Parked(_))
972 && self.inner.state.load(Ordering::Acquire) == STATE_CLOSING);
973 if !more {
974 work.scheduled = false;
975 }
976 more
977 };
978 if requeue {
979 self.shared.schedule(Arc::clone(self));
980 }
981 }
982
983 fn perform(&self, action: Action) -> ConnectionSlot {
985 let trace = match &action {
986 Action::Run(_, command) => command.trace.clone(),
987 _ => self.opening_trace.clone(),
988 };
989 let _dispatch = trace
992 .as_ref()
993 .map(|trace| tracing::dispatcher::set_default(&trace.dispatch));
994 let _span = trace.as_ref().map(|trace| trace.span.enter());
995 let metrics = self.inner.metrics.as_deref();
996 match action {
997 Action::Open(ready) => match open_database_connection(&self.vfs, &self.file_name)
998 .and_then(|db| {
999 if let Some(guard) = &self.guard {
1000 guard.install(db.as_ptr())?;
1001 }
1002 Ok(db)
1003 }) {
1004 Ok(db) => {
1005 if let Some(ready) = ready {
1006 let _ = ready.send(Ok(()));
1007 }
1008 ConnectionSlot::Parked(OpenConnection {
1009 db,
1010 transaction: None,
1011 pending_uncork: None,
1012 })
1013 }
1014 Err(error) => {
1015 if let Some(ready) = ready {
1016 let _ = ready.send(Err(error));
1017 }
1018 self.inner.state.store(STATE_DEAD, Ordering::Release);
1019 self.inner.closed.notify_waiters();
1020 ConnectionSlot::Gone
1021 }
1022 },
1023 Action::Run(mut open, command) => {
1024 if let Some(guard) = &self.guard {
1025 guard.reset();
1026 }
1027 let _inflight = WorkerInflightGuard::new(metrics);
1028 let max_result_bytes = self.guard.as_ref().map(|guard| guard.max_result_bytes());
1029 run_command(
1030 &mut open.db,
1031 &self.vfs,
1032 command.command,
1033 max_result_bytes,
1034 metrics,
1035 &self.file_name,
1036 &mut open.transaction,
1037 &mut open.pending_uncork,
1038 &self.inner.explicit_transaction,
1039 );
1040 report_vfs_memory(&open.db, metrics);
1041 ConnectionSlot::Parked(open)
1042 }
1043 Action::Release(mut open) => {
1044 release_idle_memory(&mut open.db, metrics);
1045 ConnectionSlot::Parked(open)
1046 }
1047 Action::Close(open) => {
1048 let requested_at = {
1049 let mut work = self.work.lock();
1050 fail_queued_commands(&mut work.commands);
1051 work.close_requested_at
1052 };
1053 drop(open);
1054 if let Some(metrics) = metrics {
1055 metrics.set_vfs_cached_bytes(0);
1056 if let Some(requested_at) = requested_at {
1057 metrics
1058 .observe_worker_close_duration(requested_at.elapsed().as_nanos() as u64);
1059 }
1060 }
1061 self.inner.explicit_transaction.send_replace(false);
1062 self.inner.state.store(STATE_CLOSED, Ordering::Release);
1063 self.inner.closed.notify_waiters();
1064 ConnectionSlot::Gone
1065 }
1066 }
1067 }
1068}
1069
1070impl Drop for DatabaseSlot {
1071 fn drop(&mut self) {
1072 let work = self.work.get_mut();
1073 fail_queued_commands(&mut work.commands);
1074 if let ConnectionSlot::Parked(open) =
1075 std::mem::replace(&mut work.connection, ConnectionSlot::Gone)
1076 {
1077 self.shared.orphan(open.db);
1078 }
1079 self.shared.cancel_idle(self.id);
1080 if self.inner.state.load(Ordering::Acquire) == STATE_RUNNING {
1081 if let Some(metrics) = &self.inner.metrics {
1082 metrics.record_worker_unclean_close();
1083 metrics.set_vfs_cached_bytes(0);
1084 }
1085 tracing::error!("sqlite worker handle dropped without clean close");
1087 self.inner.explicit_transaction.send_replace(false);
1088 self.inner.state.store(STATE_CLOSED, Ordering::Release);
1089 self.inner.closed.notify_waiters();
1090 }
1091 }
1092}
1093
1094fn fail_queued_commands(commands: &mut VecDeque<QueuedCommand>) {
1095 for command in commands.drain(..) {
1096 fail_command(command.command);
1097 }
1098}
1099
1100#[derive(Clone)]
1101pub struct SqliteWorkerHandle {
1102 slot: Arc<DatabaseSlot>,
1103 _pool: SqlitePool,
1104}
1105
1106struct SqliteWorkerInner {
1107 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
1108 state: AtomicU8,
1109 closed: Notify,
1110 ready: Mutex<Option<oneshot::Receiver<Result<()>>>>,
1111 explicit_transaction: watch::Sender<bool>,
1115}
1116
1117enum SqliteCommand {
1118 Execute {
1119 sql: String,
1120 params: Option<Vec<BindParam>>,
1121 enqueued_at: Option<Instant>,
1122 reply: oneshot::Sender<Result<SqliteWorkerResult<ExecuteResult>>>,
1123 },
1124 Exec {
1125 sql: String,
1126 enqueued_at: Option<Instant>,
1127 reply: oneshot::Sender<Result<SqliteWorkerResult<QueryResult>>>,
1128 },
1129 ExecuteBatch {
1131 statements: Vec<(String, Option<Vec<BindParam>>)>,
1132 reply: oneshot::Sender<Result<SqliteWorkerResult<Vec<ExecuteResult>>>>,
1133 },
1134 Cork { reply: oneshot::Sender<Result<()>> },
1136 Uncork { reply: oneshot::Sender<Result<()>> },
1138
1139
1140}
1141
1142#[derive(Debug)]
1143pub struct SqliteWorkerResult<T> {
1144 pub result: Result<T>,
1145 pub profile: SqliteOperationProfile,
1146}
1147
1148struct WorkerInflightGuard<'a>(Option<&'a dyn SqliteVfsMetrics>);
1149
1150impl<'a> WorkerInflightGuard<'a> {
1151 fn new(metrics: Option<&'a dyn SqliteVfsMetrics>) -> Self {
1152 if let Some(metrics) = metrics {
1153 metrics.set_worker_inflight(true);
1154 }
1155 Self(metrics)
1156 }
1157}
1158
1159impl Drop for WorkerInflightGuard<'_> {
1160 fn drop(&mut self) {
1161 if let Some(metrics) = self.0 {
1162 metrics.set_worker_inflight(false);
1163 }
1164 }
1165}
1166
1167struct TransactionTracker {
1174 started_at: Instant,
1175 start_round_trips: SqliteRoundTripCounts,
1176 statement_count: u64,
1177}
1178
1179fn release_idle_memory(db: &mut NativeConnection, metrics: Option<&dyn SqliteVfsMetrics>) {
1182 if command_in_tx(db) {
1183 return;
1184 }
1185 if db.release_idle_memory() {
1186 report_vfs_memory(db, metrics);
1187 }
1188}
1189
1190fn report_vfs_memory(db: &NativeConnection, metrics: Option<&dyn SqliteVfsMetrics>) {
1193 if let Some(metrics) = metrics {
1194 metrics.set_vfs_cached_bytes(db.cached_page_bytes());
1195 }
1196}
1197
1198impl SqliteWorkerHandle {
1199 pub fn start(
1201 vfs: NativeVfsHandle,
1202 file_name: String,
1203 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
1204 ) -> Result<Self> {
1205 Self::start_on(&SqlitePool::new(1, None), vfs, file_name, metrics)
1206 }
1207
1208 pub fn start_on(
1211 pool: &SqlitePool,
1212 vfs: NativeVfsHandle,
1213 file_name: String,
1214 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
1215 ) -> Result<Self> {
1216 Self::start_on_with_guard(pool, vfs, file_name, metrics, None)
1217 }
1218
1219 pub fn start_on_with_guard(
1222 pool: &SqlitePool,
1223 vfs: NativeVfsHandle,
1224 file_name: String,
1225 metrics: Option<Arc<dyn SqliteVfsMetrics>>,
1226 guard: Option<ConnectionGuard>,
1227 ) -> Result<Self> {
1228 let (ready_tx, ready_rx) = oneshot::channel();
1229 let inner = Arc::new(SqliteWorkerInner {
1230 metrics,
1231 state: AtomicU8::new(STATE_RUNNING),
1232 closed: Notify::new(),
1233 ready: Mutex::new(Some(ready_rx)),
1234 explicit_transaction: watch::Sender::new(false),
1235 });
1236 let shared = &pool.owner.shared;
1237 let id = {
1238 let mut state = shared.state.lock();
1239 if state.shutdown {
1240 return Err(anyhow!("sqlite pool is shutting down"));
1241 }
1242 state.next_database += 1;
1243 state.next_database
1244 };
1245 let idle_release = vfs.idle_release();
1246 let slot = Arc::new(DatabaseSlot {
1247 opening_trace: WorkerTrace::capture(),
1248 id,
1249 inner,
1250 shared: Arc::clone(shared),
1251 vfs,
1252 file_name,
1253 idle_release,
1254 last_used_ns: AtomicU64::new(0),
1255 idle_queued: AtomicBool::new(false),
1256 background: AtomicBool::new(false),
1257 guard: guard.map(InstalledGuard::new),
1258 work: Mutex::new(SlotWork {
1259 commands: VecDeque::new(),
1260 scheduled: false,
1261 idle_due: false,
1262 connection: ConnectionSlot::Unopened(Some(ready_tx)),
1263 close_requested_at: None,
1264 }),
1265 });
1266 let work = slot.work.lock();
1267 slot.schedule_locked(work);
1268 Ok(Self {
1269 slot,
1270 _pool: pool.clone(),
1271 })
1272 }
1273
1274 fn inner(&self) -> &SqliteWorkerInner {
1275 &self.slot.inner
1276 }
1277
1278 pub async fn wait_ready(&self) -> Result<()> {
1279 let ready = self.inner().ready.lock().take();
1280 let Some(ready) = ready else {
1281 return Ok(());
1282 };
1283 ready
1284 .await
1285 .map_err(|_| sqlite_worker_dead_error())?
1286 .map_err(|err| anyhow!("failed to initialize sqlite worker: {err}"))
1287 }
1288
1289 pub fn set_background(&self, background: bool) {
1294 self.slot.background.store(background, Ordering::Release);
1295 }
1296
1297 pub fn explicit_transaction_open(&self) -> bool {
1298 *self.inner().explicit_transaction.borrow()
1299 }
1300
1301 pub fn subscribe_explicit_transaction(&self) -> watch::Receiver<bool> {
1304 self.inner().explicit_transaction.subscribe()
1305 }
1306
1307 pub async fn exec(&self, sql: String) -> Result<QueryResult> {
1308 self.exec_inner(sql, None).await?.result
1309 }
1310
1311 pub async fn exec_profiled(&self, sql: String) -> Result<SqliteWorkerResult<QueryResult>> {
1312 self.exec_inner(sql, Some(Instant::now())).await
1313 }
1314
1315 async fn exec_inner(
1316 &self,
1317 sql: String,
1318 enqueued_at: Option<Instant>,
1319 ) -> Result<SqliteWorkerResult<QueryResult>> {
1320 let (reply, result) = oneshot::channel();
1321 self.slot.enqueue(SqliteCommand::Exec {
1322 sql,
1323 enqueued_at,
1324 reply,
1325 })?;
1326 result.await.map_err(|_| sqlite_worker_dead_error())?
1327 }
1328
1329 pub async fn execute(
1330 &self,
1331 sql: String,
1332 params: Option<Vec<BindParam>>,
1333 ) -> Result<ExecuteResult> {
1334 self.execute_inner(sql, params, None).await?.result
1335 }
1336
1337 pub async fn execute_profiled(
1338 &self,
1339 sql: String,
1340 params: Option<Vec<BindParam>>,
1341 ) -> Result<SqliteWorkerResult<ExecuteResult>> {
1342 self.execute_inner(sql, params, Some(Instant::now())).await
1343 }
1344
1345 async fn execute_inner(
1346 &self,
1347 sql: String,
1348 params: Option<Vec<BindParam>>,
1349 enqueued_at: Option<Instant>,
1350 ) -> Result<SqliteWorkerResult<ExecuteResult>> {
1351 let (reply, result) = oneshot::channel();
1352 self.slot.enqueue(SqliteCommand::Execute {
1353 sql,
1354 params,
1355 enqueued_at,
1356 reply,
1357 })?;
1358 result.await.map_err(|_| sqlite_worker_dead_error())?
1359 }
1360
1361 pub async fn execute_batch(
1365 &self,
1366 statements: Vec<(String, Option<Vec<BindParam>>)>,
1367 ) -> Result<Vec<ExecuteResult>> {
1368 let (reply, result) = oneshot::channel();
1369 self.slot
1370 .enqueue(SqliteCommand::ExecuteBatch { statements, reply })?;
1371 result
1372 .await
1373 .map_err(|_| sqlite_worker_dead_error())??
1374 .result
1375 }
1376
1377 pub async fn cork_commits(&self) -> Result<()> {
1381 let (reply, result) = oneshot::channel();
1382 self.slot.enqueue(SqliteCommand::Cork { reply })?;
1383 result.await.map_err(|_| sqlite_worker_dead_error())?
1384 }
1385
1386 pub async fn uncork_commits(&self) -> Result<()> {
1389 let (reply, result) = oneshot::channel();
1390 self.slot.enqueue(SqliteCommand::Uncork { reply })?;
1391 result.await.map_err(|_| sqlite_worker_dead_error())?
1392 }
1393
1394 pub async fn close(&self) -> Result<()> {
1395 self.close_with_timeout(Some(SQLITE_WORKER_CLOSE_TIMEOUT))
1396 .await
1397 }
1398
1399 pub async fn close_and_wait(&self) -> Result<()> {
1401 self.close_with_timeout(None).await
1402 }
1403
1404 async fn close_with_timeout(&self, close_timeout: Option<Duration>) -> Result<()> {
1405 let start = Instant::now();
1406 if self.inner().mark_closing() {
1409 self.slot.request_close();
1410 }
1411
1412 let wait_closed = async {
1413 loop {
1414 let closed = self.inner().closed.notified();
1415 match self.inner().state.load(Ordering::Acquire) {
1416 STATE_CLOSED => return Ok(()),
1417 STATE_DEAD => return Err(sqlite_worker_dead_error()),
1418 STATE_RUNNING | STATE_CLOSING => closed.await,
1419 other => return Err(anyhow!("unknown sqlite worker state {other}")),
1420 }
1421 }
1422 };
1423
1424 match close_timeout {
1425 Some(close_timeout) => match tokio::time::timeout(close_timeout, wait_closed).await {
1426 Ok(result) => result,
1427 Err(_) => {
1428 if let Some(metrics) = &self.inner().metrics {
1429 metrics.record_worker_close_timeout();
1430 }
1431 Err(SqliteWorkerCloseTimeoutError.into())
1434 }
1435 },
1436 None => {
1437 tokio::pin!(wait_closed);
1438 tokio::select! {
1439 result = &mut wait_closed => result,
1440 _ = tokio::time::sleep(SQLITE_WORKER_CLOSE_TIMEOUT) => {
1441 if let Some(metrics) = &self.inner().metrics {
1442 metrics.record_worker_close_timeout();
1443 }
1444 tracing::warn!(
1445 duration_ms = start.elapsed().as_millis() as u64,
1446 "waiting for SQLite worker to finish strict close",
1447 );
1448 wait_closed.await
1449 }
1450 }
1451 }
1452 }
1453 }
1454
1455 pub async fn wait_for_failure(&self) -> bool {
1456 loop {
1457 let closed = self.inner().closed.notified();
1458 match self.inner().state.load(Ordering::Acquire) {
1459 STATE_DEAD => return true,
1460 STATE_CLOSED => return false,
1461 STATE_RUNNING | STATE_CLOSING => closed.await,
1462 _ => return true,
1463 }
1464 }
1465 }
1466
1467
1468
1469
1470
1471
1472}
1473
1474impl SqliteWorkerInner {
1475 fn mark_closing(&self) -> bool {
1476 self.state
1477 .compare_exchange(
1478 STATE_RUNNING,
1479 STATE_CLOSING,
1480 Ordering::AcqRel,
1481 Ordering::Acquire,
1482 )
1483 .is_ok()
1484 }
1485
1486 fn record_queue_depth(&self, depth: u64) {
1487 if let Some(metrics) = &self.metrics {
1488 metrics.set_worker_queue_depth(depth);
1489 }
1490 }
1491}
1492
1493fn warn_if_direct_overflow_reads() {
1498 let enabled = unsafe { sqlite3_compileoption_used(c"DIRECT_OVERFLOW_READ".as_ptr()) } != 0;
1501 if enabled {
1502 tracing::warn!(
1503 "sqlite was built with SQLITE_DIRECT_OVERFLOW_READ, likely without the repo's cargo config, so reads of large rows bypass the page cache and refetch from depot"
1504 );
1505 }
1506}
1507
1508fn open_database_connection(vfs: &NativeVfsHandle, file_name: &str) -> Result<NativeConnection> {
1509 let connection = open_connection(
1510 vfs.clone(),
1511 file_name,
1512 SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE,
1513 )
1514 .map_err(anyhow::Error::msg)?;
1515 configure_connection_for_database(connection.as_ptr(), vfs, file_name)?;
1516 Ok(connection)
1517}
1518
1519#[allow(clippy::too_many_arguments)]
1520fn run_command(
1521 db: &mut NativeConnection,
1522 vfs: &NativeVfsHandle,
1523 command: SqliteCommand,
1524 max_result_bytes: Option<u64>,
1525 metrics: Option<&dyn SqliteVfsMetrics>,
1526 file_name: &str,
1527 transaction: &mut Option<TransactionTracker>,
1528 pending_uncork: &mut Option<oneshot::Sender<Result<()>>>,
1529 explicit_transaction: &watch::Sender<bool>,
1530) {
1531 let _ = vfs.take_storage_error();
1534 run_command_inner(
1535 db,
1536 vfs,
1537 command,
1538 max_result_bytes,
1539 metrics,
1540 file_name,
1541 transaction,
1542 pending_uncork,
1543 explicit_transaction,
1544 );
1545
1546 if pending_uncork.is_some() && !command_in_tx(db) {
1548 if let Some(reply) = pending_uncork.take() {
1549 let _ = reply.send(db.uncork_commits());
1550 }
1551 }
1552}
1553
1554#[allow(clippy::too_many_arguments)]
1555fn run_command_inner(
1556 db: &mut NativeConnection,
1557 vfs: &NativeVfsHandle,
1558 command: SqliteCommand,
1559 max_result_bytes: Option<u64>,
1560 metrics: Option<&dyn SqliteVfsMetrics>,
1561 file_name: &str,
1562 transaction: &mut Option<TransactionTracker>,
1563 pending_uncork: &mut Option<oneshot::Sender<Result<()>>>,
1564 explicit_transaction: &watch::Sender<bool>,
1565) {
1566 let start = Instant::now();
1567 match command {
1568 SqliteCommand::Execute {
1569 sql,
1570 params,
1571 enqueued_at,
1572 reply,
1573 } => {
1574 if reply.is_closed() {
1575 return;
1576 }
1577 begin_transaction_if_needed(db, transaction);
1578 let in_tx = command_in_tx(db);
1582 let stmt_kind = classify_statement(&sql);
1583 let worker_result = if let Some(enqueued_at) = enqueued_at {
1584 let worker_wait_ns = enqueued_at.elapsed().as_nanos() as u64;
1585 let operation_profile = db.begin_operation_profile();
1586 let execution_start = Instant::now();
1587 let result = execute_single_statement_profiled(
1588 db.as_ptr(),
1589 &sql,
1590 params.as_deref(),
1591 max_result_bytes,
1592 );
1593 let execution_ns = execution_start.elapsed().as_nanos() as u64;
1594 let mut profile = operation_profile.finish();
1595 profile.worker_wait_ns = worker_wait_ns;
1596 profile.sqlite_execution_ns = execution_ns.saturating_sub(profile.storage_ns);
1597 if let Ok((_, query_profile)) = &result {
1598 profile.bind_count = query_profile.bind_count;
1599 profile.bind_logical_bytes = query_profile.bind_logical_bytes;
1600 profile.result_rows = query_profile.result_rows;
1601 profile.result_columns = query_profile.result_columns;
1602 profile.result_logical_bytes = query_profile.result_logical_bytes;
1603 }
1604 record_command_metrics(
1605 metrics,
1606 "execute",
1607 in_tx,
1608 stmt_kind,
1609 &result,
1610 start.elapsed(),
1611 );
1612 SqliteWorkerResult {
1613 result: result.map(|(value, _)| value),
1614 profile,
1615 }
1616 } else {
1617 let result = execute_single_statement(
1618 db.as_ptr(),
1619 &sql,
1620 params.as_deref(),
1621 max_result_bytes,
1622 );
1623 record_command_metrics(
1624 metrics,
1625 "execute",
1626 in_tx,
1627 stmt_kind,
1628 &result,
1629 start.elapsed(),
1630 );
1631 SqliteWorkerResult {
1632 result,
1633 profile: SqliteOperationProfile::default(),
1634 }
1635 };
1636 finalize_transaction_if_complete(db, metrics, file_name, transaction);
1637 publish_explicit_transaction(db, explicit_transaction);
1638 let _ = reply.send(worker_command_result(db, vfs, worker_result));
1639 }
1640 SqliteCommand::Exec {
1641 sql,
1642 enqueued_at,
1643 reply,
1644 } => {
1645 if reply.is_closed() {
1646 return;
1647 }
1648 begin_transaction_if_needed(db, transaction);
1649 let in_tx = command_in_tx(db);
1650 let stmt_kind = classify_statement(&sql);
1651 let worker_result = if let Some(enqueued_at) = enqueued_at {
1652 let worker_wait_ns = enqueued_at.elapsed().as_nanos() as u64;
1653 let operation_profile = db.begin_operation_profile();
1654 let execution_start = Instant::now();
1655 let result = exec_statements_profiled(db.as_ptr(), &sql, max_result_bytes);
1656 let execution_ns = execution_start.elapsed().as_nanos() as u64;
1657 let mut profile = operation_profile.finish();
1658 profile.worker_wait_ns = worker_wait_ns;
1659 profile.sqlite_execution_ns = execution_ns.saturating_sub(profile.storage_ns);
1660 if let Ok((_, query_profile)) = &result {
1661 profile.bind_count = query_profile.bind_count;
1662 profile.bind_logical_bytes = query_profile.bind_logical_bytes;
1663 profile.result_rows = query_profile.result_rows;
1664 profile.result_columns = query_profile.result_columns;
1665 profile.result_logical_bytes = query_profile.result_logical_bytes;
1666 }
1667 record_command_metrics(metrics, "exec", in_tx, stmt_kind, &result, start.elapsed());
1668 SqliteWorkerResult {
1669 result: result.map(|(value, _)| value),
1670 profile,
1671 }
1672 } else {
1673 let result = exec_statements(db.as_ptr(), &sql, max_result_bytes);
1674 record_command_metrics(metrics, "exec", in_tx, stmt_kind, &result, start.elapsed());
1675 SqliteWorkerResult {
1676 result,
1677 profile: SqliteOperationProfile::default(),
1678 }
1679 };
1680 finalize_transaction_if_complete(db, metrics, file_name, transaction);
1681 publish_explicit_transaction(db, explicit_transaction);
1682 let _ = reply.send(worker_command_result(db, vfs, worker_result));
1683 }
1684 SqliteCommand::ExecuteBatch { statements, reply } => {
1685 if reply.is_closed() {
1686 return;
1687 }
1688 begin_transaction_if_needed(db, transaction);
1689 if let Some(tracker) = transaction.as_mut() {
1690 let extra = statements.len().saturating_sub(1) as u64;
1691 tracker.statement_count = tracker.statement_count.saturating_add(extra);
1692 }
1693 let in_tx = command_in_tx(db);
1694 let stmt_kind = statements
1695 .first()
1696 .map_or("other", |(sql, _)| classify_statement(sql));
1697 let mut results = Vec::with_capacity(statements.len());
1698 let mut failure = None;
1699 for (sql, params) in statements {
1700 match execute_single_statement(
1701 db.as_ptr(),
1702 &sql,
1703 params.as_deref(),
1704 max_result_bytes,
1705 ) {
1706 Ok(result) => results.push(result),
1707 Err(error) => {
1708 failure = Some(error);
1709 break;
1710 }
1711 }
1712 }
1713 let result = match failure {
1714 Some(error) => Err(error),
1715 None => Ok(results),
1716 };
1717 record_command_metrics(
1718 metrics,
1719 "execute_batch",
1720 in_tx,
1721 stmt_kind,
1722 &result,
1723 start.elapsed(),
1724 );
1725 finalize_transaction_if_complete(db, metrics, file_name, transaction);
1726 publish_explicit_transaction(db, explicit_transaction);
1727 let worker_result = SqliteWorkerResult {
1728 result,
1729 profile: SqliteOperationProfile::default(),
1730 };
1731 let _ = reply.send(worker_command_result(db, vfs, worker_result));
1732 }
1733 SqliteCommand::Cork { reply } => {
1734 db.cork_commits();
1735 let _ = reply.send(Ok(()));
1736 }
1737 SqliteCommand::Uncork { reply } => {
1738 if !command_in_tx(db) {
1739 let _ = reply.send(db.uncork_commits());
1740 } else if pending_uncork.is_some() {
1741 let _ = reply.send(Err(anyhow::anyhow!(
1742 "an earlier sqlite uncork is still waiting for its transaction"
1743 )));
1744 } else {
1745 *pending_uncork = Some(reply);
1746 }
1747 }
1748
1749
1750 }
1751}
1752
1753fn worker_command_result<T>(
1754 db: &NativeConnection,
1755 vfs: &NativeVfsHandle,
1756 result: SqliteWorkerResult<T>,
1757) -> Result<SqliteWorkerResult<T>> {
1758 if let Some(operation) = vfs.take_indeterminate_result() {
1759 return Err(SqliteCommitIndeterminateError(operation).into());
1760 }
1761
1762 let storage_error = vfs.take_storage_error();
1765 let result = match (storage_error, result) {
1766 (
1767 Some(storage_error),
1768 SqliteWorkerResult {
1769 result: Err(_),
1770 profile,
1771 },
1772 ) => SqliteWorkerResult {
1773 result: Err(storage_error.into()),
1774 profile,
1775 },
1776 (_, result) => result,
1777 };
1778
1779 match (vfs.take_transaction_too_large(), result) {
1782 (
1783 Some(too_large),
1784 SqliteWorkerResult {
1785 result: Err(_),
1786 profile,
1787 },
1788 ) => {
1789 finish_refused_rollback(db);
1790 Ok(SqliteWorkerResult {
1791 result: Err(too_large.into()),
1792 profile,
1793 })
1794 }
1795 (_, result) => Ok(result),
1796 }
1797}
1798
1799fn finish_refused_rollback(db: &NativeConnection) {
1804 if unsafe { sqlite3_get_autocommit(db.as_ptr()) } == 0 {
1805 return;
1806 }
1807 if let Err(error) = sqlite_exec(db.as_ptr(), "PRAGMA main.schema_version;") {
1808 tracing::warn!(
1809 %error,
1810 "sqlite rollback after a refused oversize transaction failed; it retries on the next access"
1811 );
1812 }
1813}
1814
1815fn begin_transaction_if_needed(
1821 db: &mut NativeConnection,
1822 transaction: &mut Option<TransactionTracker>,
1823) {
1824 match transaction {
1825 Some(tracker) => tracker.statement_count = tracker.statement_count.saturating_add(1),
1826 None => {
1827 *transaction = Some(TransactionTracker {
1828 started_at: Instant::now(),
1829 start_round_trips: db.round_trip_counts(),
1830 statement_count: 1,
1831 });
1832 }
1833 }
1834}
1835
1836fn finalize_transaction_if_complete(
1841 db: &mut NativeConnection,
1842 metrics: Option<&dyn SqliteVfsMetrics>,
1843 file_name: &str,
1844 transaction: &mut Option<TransactionTracker>,
1845) {
1846 if command_in_tx(db) {
1849 return;
1850 }
1851 let Some(tracker) = transaction.take() else {
1852 return;
1853 };
1854
1855 let duration = tracker.started_at.elapsed();
1856 let round_trips = db.round_trip_counts().since(tracker.start_round_trips);
1857
1858 if let Some(metrics) = metrics {
1859 metrics.observe_transaction_round_trips(round_trips.get_pages, round_trips.commit);
1860 }
1861
1862 if duration >= SQLITE_SLOW_TRANSACTION_THRESHOLD {
1863 tracing::warn!(
1864 database = %file_name,
1865 duration_ms = duration.as_millis() as u64,
1866 statement_count = tracker.statement_count,
1867 get_pages_round_trips = round_trips.get_pages,
1868 commit_round_trips = round_trips.commit,
1869 total_round_trips = round_trips.total(),
1870 "slow sqlite transaction"
1871 );
1872 } else {
1873 tracing::debug!(
1874 database = %file_name,
1875 duration_ms = duration.as_millis() as u64,
1876 statement_count = tracker.statement_count,
1877 get_pages_round_trips = round_trips.get_pages,
1878 commit_round_trips = round_trips.commit,
1879 total_round_trips = round_trips.total(),
1880 "sqlite transaction complete"
1881 );
1882 }
1883}
1884
1885fn record_command_metrics<T>(
1886 metrics: Option<&dyn SqliteVfsMetrics>,
1887 operation: &'static str,
1888 in_tx: bool,
1889 stmt_kind: &'static str,
1890 result: &Result<T>,
1891 duration: Duration,
1892) {
1893 let Some(metrics) = metrics else {
1894 return;
1895 };
1896 metrics.observe_worker_command_duration(
1897 operation,
1898 in_tx,
1899 stmt_kind,
1900 duration.as_nanos() as u64,
1901 );
1902 if let Err(error) = result {
1903 metrics.record_worker_command_error(operation, worker_error_code(error));
1904 }
1905}
1906
1907fn publish_explicit_transaction(
1910 db: &mut NativeConnection,
1911 explicit_transaction: &watch::Sender<bool>,
1912) {
1913 let open = command_in_tx(db);
1914 explicit_transaction.send_if_modified(|current| {
1915 if *current == open {
1916 return false;
1917 }
1918 *current = open;
1919 true
1920 });
1921}
1922
1923fn command_in_tx(db: &mut NativeConnection) -> bool {
1929 unsafe { sqlite3_get_autocommit(db.as_ptr()) == 0 }
1930}
1931
1932fn classify_statement(sql: &str) -> &'static str {
1938 let trimmed = strip_leading_sql_noise(sql);
1939 let keyword: String = trimmed
1940 .chars()
1941 .take_while(|c| c.is_ascii_alphabetic())
1942 .collect();
1943 match keyword.to_ascii_uppercase().as_str() {
1944 "SELECT" | "WITH" => "select",
1945 "INSERT" | "REPLACE" => "insert",
1946 "UPDATE" => "update",
1947 "DELETE" => "delete",
1948 "BEGIN" => "begin",
1949 "COMMIT" | "END" => "commit",
1950 "ROLLBACK" => "rollback",
1951 "SAVEPOINT" => "savepoint",
1952 "RELEASE" => "release",
1953 "PRAGMA" => "pragma",
1954 "CREATE" => "create",
1955 "DROP" => "drop",
1956 "ALTER" => "alter",
1957 "VACUUM" => "vacuum",
1958 "ANALYZE" => "analyze",
1959 "ATTACH" => "attach",
1960 "DETACH" => "detach",
1961 _ => "other",
1962 }
1963}
1964
1965fn strip_leading_sql_noise(sql: &str) -> &str {
1966 let mut rest = sql.trim_start();
1967 loop {
1968 if let Some(after) = rest.strip_prefix("--") {
1969 rest = match after.find('\n') {
1970 Some(idx) => after[idx + 1..].trim_start(),
1971 None => "",
1972 };
1973 } else if let Some(after) = rest.strip_prefix("/*") {
1974 rest = match after.find("*/") {
1975 Some(idx) => after[idx + 2..].trim_start(),
1976 None => "",
1977 };
1978 } else {
1979 return rest;
1980 }
1981 }
1982}
1983
1984fn fail_command(command: SqliteCommand) {
1985 match command {
1988 SqliteCommand::Execute { reply, .. } => {
1989 let _ = reply.send(Err(sqlite_closing_error()));
1990 }
1991 SqliteCommand::Exec { reply, .. } => {
1992 let _ = reply.send(Err(sqlite_closing_error()));
1993 }
1994 SqliteCommand::ExecuteBatch { reply, .. } => {
1995 let _ = reply.send(Err(sqlite_closing_error()));
1996 }
1997 SqliteCommand::Cork { reply } | SqliteCommand::Uncork { reply } => {
1998 let _ = reply.send(Err(sqlite_closing_error()));
1999 }
2000
2001
2002 }
2003}
2004
2005fn sqlite_closing_error() -> anyhow::Error {
2006 SqliteWorkerClosingError.into()
2007}
2008
2009fn sqlite_worker_dead_error() -> anyhow::Error {
2010 SqliteWorkerDeadError.into()
2011}
2012
2013fn worker_error_code(error: &anyhow::Error) -> &'static str {
2014 if error
2015 .downcast_ref::<SqliteWorkerOverloadedError>()
2016 .is_some()
2017 {
2018 "overloaded"
2019 } else if error.downcast_ref::<SqliteWorkerFatalError>().is_some() {
2020 "fatal"
2021 } else if error.downcast_ref::<SqliteWorkerClosingError>().is_some() {
2022 "closing"
2023 } else if error.downcast_ref::<SqliteWorkerDeadError>().is_some() {
2024 "dead"
2025 } else if error
2026 .downcast_ref::<SqliteWorkerCloseTimeoutError>()
2027 .is_some()
2028 {
2029 "close_timeout"
2030 } else {
2031 "sqlite"
2032 }
2033}
2034
2035#[derive(Debug)]
2036pub struct SqliteWorkerOverloadedError;
2037
2038impl fmt::Display for SqliteWorkerOverloadedError {
2039 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2040 f.write_str("actor.overloaded: sqlite worker command queue is full")
2041 }
2042}
2043
2044impl Error for SqliteWorkerOverloadedError {}
2045
2046#[derive(Debug)]
2047pub struct SqliteWorkerClosingError;
2048
2049impl fmt::Display for SqliteWorkerClosingError {
2050 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2051 f.write_str("sqlite worker is closing")
2052 }
2053}
2054
2055impl Error for SqliteWorkerClosingError {}
2056
2057#[derive(Debug)]
2058pub struct SqliteWorkerDeadError;
2059
2060impl fmt::Display for SqliteWorkerDeadError {
2061 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2062 f.write_str("sqlite worker is closed")
2063 }
2064}
2065
2066impl Error for SqliteWorkerDeadError {}
2067
2068#[derive(Debug)]
2069pub struct SqliteWorkerCloseTimeoutError;
2070
2071impl fmt::Display for SqliteWorkerCloseTimeoutError {
2072 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2073 f.write_str("sqlite worker close timed out")
2074 }
2075}
2076
2077impl Error for SqliteWorkerCloseTimeoutError {}
2078
2079#[derive(Debug, Clone)]
2080pub struct SqliteWorkerFatalError {
2081 message: String,
2082}
2083
2084impl SqliteWorkerFatalError {
2085 pub fn new(message: String) -> Self {
2086 Self { message }
2087 }
2088
2089 pub fn message(&self) -> &str {
2090 &self.message
2091 }
2092}
2093
2094impl fmt::Display for SqliteWorkerFatalError {
2095 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2096 write!(f, "sqlite worker hit fatal storage error: {}", self.message)
2097 }
2098}
2099
2100impl Error for SqliteWorkerFatalError {}
2101
2102fn panic_message(payload: &Box<dyn std::any::Any + Send>) -> String {
2103 if let Some(message) = payload.downcast_ref::<&str>() {
2104 message.to_string()
2105 } else if let Some(message) = payload.downcast_ref::<String>() {
2106 message.clone()
2107 } else {
2108 "unknown panic".to_string()
2109 }
2110}
2111
2112
2113
2114
2115
2116
2117
2118