1use std::{
37 alloc::{self, Layout},
38 cell::UnsafeCell,
39 collections::HashMap,
40 ffi::{c_int, c_uint, c_void},
41 ptr::{self, NonNull},
42 sync::{
43 Arc, OnceLock,
44 atomic::{AtomicPtr, AtomicU64, AtomicUsize, Ordering},
45 },
46};
47
48use libsqlite3_sys::{
49 SQLITE_CONFIG_PCACHE2, SQLITE_OK, sqlite3_config, sqlite3_pcache, sqlite3_pcache_methods2,
50 sqlite3_pcache_page,
51};
52use lru::LruCache;
53use parking_lot::Mutex;
54
55use crate::worker::SqlitePoolMetrics;
56
57const SHARDS: usize = 16;
59
60static PROCESS_PAGER_CACHE: AtomicPtr<PagerCache> = AtomicPtr::new(ptr::null_mut());
65
66static INSTALLED: OnceLock<bool> = OnceLock::new();
68
69pub fn install(capacity_bytes: u64, metrics: Option<Arc<dyn SqlitePoolMetrics>>) -> bool {
78 let installed = *INSTALLED.get_or_init(|| {
79 if capacity_bytes == 0 {
80 return false;
81 }
82 let cache: &'static PagerCache = Box::leak(Box::new(PagerCache::new(capacity_bytes)));
86 PROCESS_PAGER_CACHE.store(cache as *const PagerCache as *mut PagerCache, Ordering::Release);
87 let methods = sqlite3_pcache_methods2 {
88 iVersion: 1,
89 pArg: cache as *const PagerCache as *mut c_void,
90 xInit: Some(x_init),
91 xShutdown: Some(x_shutdown),
92 xCreate: Some(x_create),
93 xCachesize: Some(x_cachesize),
94 xPagecount: Some(x_pagecount),
95 xFetch: Some(x_fetch),
96 xUnpin: Some(x_unpin),
97 xRekey: Some(x_rekey),
98 xTruncate: Some(x_truncate),
99 xDestroy: Some(x_destroy),
100 xShrink: Some(x_shrink),
101 };
102 let rc = unsafe { sqlite3_config(SQLITE_CONFIG_PCACHE2, &methods as *const _) };
105 if rc != SQLITE_OK {
106 PROCESS_PAGER_CACHE.store(ptr::null_mut(), Ordering::Release);
107 tracing::warn!(
108 rc,
109 "sqlite was initialized before the process-wide page cache could be installed; connections keep per-connection page caches"
110 );
111 return false;
112 }
113 true
114 });
115 if let (true, Some(metrics)) = (installed, metrics) {
116 attach_metrics(metrics);
117 }
118 installed
119}
120
121pub fn installed() -> Option<&'static PagerCache> {
123 let cache = PROCESS_PAGER_CACHE.load(Ordering::Acquire);
124 unsafe { cache.as_ref() }
126}
127
128pub fn attach_metrics(metrics: Arc<dyn SqlitePoolMetrics>) {
132 if let Some(cache) = installed()
133 && cache.metrics.set(metrics).is_ok()
134 {
135 cache.report_capacity();
136 cache.report();
137 }
138}
139
140pub struct PagerCache {
142 shards: Box<[Mutex<Shard>]>,
143 capacity_bytes: u64,
144 purgeable_bytes: AtomicU64,
146 unbudgeted_bytes: AtomicU64,
148 pages: AtomicU64,
149 refusals: AtomicU64,
151 evictions: AtomicU64,
153 next_cache: AtomicUsize,
154 evict_cursor: AtomicUsize,
155 metrics: OnceLock<Arc<dyn SqlitePoolMetrics>>,
156}
157
158struct Shard {
161 lru: LruCache<(usize, u32), ()>,
162}
163
164struct CacheHandle {
166 cache: &'static PagerCache,
167 shard: usize,
168 state: UnsafeCell<CacheState>,
170}
171
172struct CacheState {
173 layout: Layout,
175 buf_offset: usize,
177 slot_bytes: u64,
179 page_size: usize,
180 purgeable: bool,
181 pages: HashMap<u32, NonNull<Slot>>,
182 max_key: u32,
184}
185
186#[repr(C)]
189struct Slot {
190 page: sqlite3_pcache_page,
191 key: u32,
192 pinned: bool,
193}
194
195type Freed = Vec<(NonNull<Slot>, Layout)>;
197
198impl PagerCache {
199 fn new(capacity_bytes: u64) -> Self {
200 Self {
201 shards: (0..SHARDS)
202 .map(|_| {
203 Mutex::new(Shard {
204 lru: LruCache::unbounded(),
205 })
206 })
207 .collect(),
208 capacity_bytes,
209 purgeable_bytes: AtomicU64::new(0),
210 unbudgeted_bytes: AtomicU64::new(0),
211 pages: AtomicU64::new(0),
212 refusals: AtomicU64::new(0),
213 evictions: AtomicU64::new(0),
214 next_cache: AtomicUsize::new(0),
215 evict_cursor: AtomicUsize::new(0),
216 metrics: OnceLock::new(),
217 }
218 }
219
220 pub fn bytes(&self) -> u64 {
222 self.purgeable_bytes.load(Ordering::Relaxed)
223 }
224
225 pub fn unbudgeted_bytes(&self) -> u64 {
227 self.unbudgeted_bytes.load(Ordering::Relaxed)
228 }
229
230 pub fn pages(&self) -> u64 {
232 self.pages.load(Ordering::Relaxed)
233 }
234
235 pub fn capacity_bytes(&self) -> u64 {
237 self.capacity_bytes
238 }
239
240 pub fn refusals(&self) -> u64 {
242 self.refusals.load(Ordering::Relaxed)
243 }
244
245 pub fn evictions(&self) -> u64 {
247 self.evictions.load(Ordering::Relaxed)
248 }
249
250 fn report_capacity(&self) {
251 if let Some(metrics) = self.metrics.get() {
252 metrics.set_pager_cache_capacity_bytes(self.capacity_bytes);
253 }
254 }
255
256 fn report(&self) {
258 if let Some(metrics) = self.metrics.get() {
259 metrics.set_pager_cache_bytes(self.bytes());
260 metrics.set_pager_cache_unbudgeted_bytes(self.unbudgeted_bytes());
261 metrics.set_pager_cache_pages(self.pages());
262 }
263 }
264
265 fn release(&self, freed: Freed) {
267 if freed.is_empty() {
268 return;
269 }
270 for (slot, layout) in freed {
271 unsafe { alloc::dealloc(slot.as_ptr() as *mut u8, layout) };
274 }
275 self.report();
276 }
277
278 fn account(&self, state: &CacheState, added: bool) {
280 let counter = if state.purgeable {
281 &self.purgeable_bytes
282 } else {
283 &self.unbudgeted_bytes
284 };
285 if added {
286 counter.fetch_add(state.slot_bytes, Ordering::Relaxed);
287 self.pages.fetch_add(1, Ordering::Relaxed);
288 } else {
289 counter.fetch_sub(state.slot_bytes, Ordering::Relaxed);
290 self.pages.fetch_sub(1, Ordering::Relaxed);
291 }
292 }
293
294 fn reserve(&self, bytes: u64) -> bool {
297 loop {
298 let current = self.purgeable_bytes.load(Ordering::Relaxed);
299 if current.saturating_add(bytes) <= self.capacity_bytes {
300 if self
301 .purgeable_bytes
302 .compare_exchange_weak(
303 current,
304 current + bytes,
305 Ordering::Relaxed,
306 Ordering::Relaxed,
307 )
308 .is_ok()
309 {
310 return true;
311 }
312 continue;
313 }
314 if !self.evict_one() {
315 return false;
316 }
317 }
318 }
319
320 fn shed_overshoot(&self) {
322 while self.bytes() > self.capacity_bytes {
323 if !self.evict_one() {
324 return;
325 }
326 }
327 }
328
329 fn evict_one(&self) -> bool {
333 for _ in 0..SHARDS {
334 let shard_index = self.evict_cursor.fetch_add(1, Ordering::Relaxed) % SHARDS;
335 let mut shard = self.shards[shard_index].lock();
336 let Some(((handle_addr, key), ())) = shard.lru.pop_lru() else {
337 continue;
338 };
339 let state = unsafe { &mut *(*(handle_addr as *const CacheHandle)).state.get() };
343 let freed = state.pages.remove(&key).map(|slot| (slot, state.layout));
344 if freed.is_some() {
345 self.account(state, false);
346 self.evictions.fetch_add(1, Ordering::Relaxed);
347 }
348 drop(shard);
349 self.release(freed.into_iter().collect());
350 return true;
351 }
352 false
353 }
354}
355
356fn take_page(
358 cache: &PagerCache,
359 shard: &mut Shard,
360 handle_addr: usize,
361 state: &mut CacheState,
362 key: u32,
363 freed: &mut Freed,
364) {
365 if let Some(slot) = state.pages.remove(&key) {
366 shard.lru.pop(&(handle_addr, key));
367 cache.account(state, false);
368 freed.push((slot, state.layout));
369 }
370}
371
372fn allocate(
375 layout: Layout,
376 buf_offset: usize,
377 page_size: usize,
378 key: u32,
379) -> Option<NonNull<Slot>> {
380 let base = NonNull::new(unsafe { alloc::alloc_zeroed(layout) })?;
382 let slot = base.cast::<Slot>();
383 unsafe {
386 let buf = base.as_ptr().add(buf_offset);
387 ptr::write(
388 slot.as_ptr(),
389 Slot {
390 page: sqlite3_pcache_page {
391 pBuf: buf as *mut c_void,
392 pExtra: buf.add(page_size) as *mut c_void,
393 },
394 key,
395 pinned: true,
396 },
397 );
398 }
399 Some(slot)
400}
401
402unsafe fn handle<'a>(cache: *mut sqlite3_pcache) -> &'a CacheHandle {
407 unsafe { &*(cache as *const CacheHandle) }
408}
409
410unsafe extern "C" fn x_init(_arg: *mut c_void) -> c_int {
411 SQLITE_OK
412}
413
414unsafe extern "C" fn x_shutdown(_arg: *mut c_void) {}
415
416unsafe extern "C" fn x_create(
417 page_size: c_int,
418 extra_size: c_int,
419 purgeable: c_int,
420) -> *mut sqlite3_pcache {
421 match installed() {
423 Some(cache) => create(cache, page_size, extra_size, purgeable),
424 None => ptr::null_mut(),
425 }
426}
427
428fn create(
431 cache: &'static PagerCache,
432 page_size: c_int,
433 extra_size: c_int,
434 purgeable: c_int,
435) -> *mut sqlite3_pcache {
436 let page_size = page_size.max(0) as usize;
437 let extra_size = extra_size.max(0) as usize;
438 let Ok(area) = Layout::from_size_align(page_size + extra_size, 8) else {
439 return ptr::null_mut();
440 };
441 let Ok((layout, buf_offset)) = Layout::new::<Slot>().extend(area) else {
442 return ptr::null_mut();
443 };
444 let layout = layout.pad_to_align();
445 let shard = cache.next_cache.fetch_add(1, Ordering::Relaxed) % SHARDS;
446 let handle = Box::new(CacheHandle {
447 cache,
448 shard,
449 state: UnsafeCell::new(CacheState {
450 layout,
451 buf_offset,
452 slot_bytes: layout.size() as u64,
453 page_size,
454 purgeable: purgeable != 0,
455 pages: HashMap::new(),
456 max_key: 0,
457 }),
458 });
459 Box::into_raw(handle) as *mut sqlite3_pcache
460}
461
462unsafe extern "C" fn x_cachesize(_cache: *mut sqlite3_pcache, _pages: c_int) {}
466
467unsafe extern "C" fn x_pagecount(cache: *mut sqlite3_pcache) -> c_int {
468 let handle = unsafe { handle(cache) };
470 let _shard = handle.cache.shards[handle.shard].lock();
471 let state = unsafe { &*handle.state.get() };
473 state.pages.len().min(c_int::MAX as usize) as c_int
474}
475
476unsafe extern "C" fn x_fetch(
477 cache: *mut sqlite3_pcache,
478 key: c_uint,
479 create_flag: c_int,
480) -> *mut sqlite3_pcache_page {
481 let handle = unsafe { handle(cache) };
483 let pager = handle.cache;
484 let handle_addr = handle as *const CacheHandle as usize;
485
486 let (layout, buf_offset, page_size, slot_bytes, purgeable) = {
488 let mut shard = pager.shards[handle.shard].lock();
489 let state = unsafe { &mut *handle.state.get() };
491 if let Some(slot) = state.pages.get(&key).copied() {
492 let slot_ref = unsafe { &mut *slot.as_ptr() };
494 if !slot_ref.pinned {
495 slot_ref.pinned = true;
496 shard.lru.pop(&(handle_addr, key));
497 }
498 return slot.as_ptr() as *mut sqlite3_pcache_page;
499 }
500 if create_flag == 0 {
501 return ptr::null_mut();
502 }
503 (
504 state.layout,
505 state.buf_offset,
506 state.page_size,
507 state.slot_bytes,
508 state.purgeable,
509 )
510 };
511
512 if purgeable && !pager.reserve(slot_bytes) {
516 if create_flag == 1 {
518 pager.refusals.fetch_add(1, Ordering::Relaxed);
519 return ptr::null_mut();
520 }
521 pager
522 .purgeable_bytes
523 .fetch_add(slot_bytes, Ordering::Relaxed);
524 }
525 let Some(slot) = allocate(layout, buf_offset, page_size, key) else {
526 if purgeable {
527 pager
528 .purgeable_bytes
529 .fetch_sub(slot_bytes, Ordering::Relaxed);
530 }
531 return ptr::null_mut();
532 };
533
534 {
535 let mut shard = pager.shards[handle.shard].lock();
536 let state = unsafe { &mut *handle.state.get() };
538 if let Some(existing) = state.pages.get(&key).copied() {
539 let existing_ref = unsafe { &mut *existing.as_ptr() };
543 if !existing_ref.pinned {
544 existing_ref.pinned = true;
545 shard.lru.pop(&(handle_addr, key));
546 }
547 drop(shard);
548 if purgeable {
549 pager
550 .purgeable_bytes
551 .fetch_sub(slot_bytes, Ordering::Relaxed);
552 }
553 unsafe { alloc::dealloc(slot.as_ptr() as *mut u8, layout) };
555 return existing.as_ptr() as *mut sqlite3_pcache_page;
556 }
557 state.pages.insert(key, slot);
558 state.max_key = state.max_key.max(key);
559 }
560 if !purgeable {
561 pager
562 .unbudgeted_bytes
563 .fetch_add(slot_bytes, Ordering::Relaxed);
564 }
565 pager.pages.fetch_add(1, Ordering::Relaxed);
566 pager.report();
567 slot.as_ptr() as *mut sqlite3_pcache_page
568}
569
570unsafe extern "C" fn x_unpin(
571 cache: *mut sqlite3_pcache,
572 page: *mut sqlite3_pcache_page,
573 discard: c_int,
574) {
575 let handle = unsafe { handle(cache) };
577 let pager = handle.cache;
578 let handle_addr = handle as *const CacheHandle as usize;
579 let mut freed = Freed::new();
580 {
581 let mut shard = pager.shards[handle.shard].lock();
582 let state = unsafe { &mut *handle.state.get() };
584 let slot = unsafe { &mut *(page as *mut Slot) };
586 if discard != 0 {
587 take_page(pager, &mut shard, handle_addr, state, slot.key, &mut freed);
588 } else {
589 slot.pinned = false;
590 if state.purgeable {
591 shard.lru.put((handle_addr, slot.key), ());
592 }
593 }
594 }
595 pager.release(freed);
596 pager.shed_overshoot();
598}
599
600unsafe extern "C" fn x_rekey(
601 cache: *mut sqlite3_pcache,
602 _page: *mut sqlite3_pcache_page,
603 old_key: c_uint,
604 new_key: c_uint,
605) {
606 if old_key == new_key {
607 return;
608 }
609 let handle = unsafe { handle(cache) };
611 let pager = handle.cache;
612 let handle_addr = handle as *const CacheHandle as usize;
613 let mut freed = Freed::new();
614 {
615 let mut shard = pager.shards[handle.shard].lock();
616 let state = unsafe { &mut *handle.state.get() };
618 take_page(pager, &mut shard, handle_addr, state, new_key, &mut freed);
620 if let Some(slot) = state.pages.remove(&old_key) {
621 let pinned = unsafe {
623 let slot = &mut *slot.as_ptr();
624 slot.key = new_key;
625 slot.pinned
626 };
627 state.pages.insert(new_key, slot);
628 state.max_key = state.max_key.max(new_key);
629 if !pinned && shard.lru.pop(&(handle_addr, old_key)).is_some() {
630 shard.lru.put((handle_addr, new_key), ());
631 }
632 }
633 }
634 pager.release(freed);
635}
636
637unsafe extern "C" fn x_truncate(cache: *mut sqlite3_pcache, limit: c_uint) {
638 let handle = unsafe { handle(cache) };
640 let pager = handle.cache;
641 let handle_addr = handle as *const CacheHandle as usize;
642 let mut freed = Freed::new();
643 {
644 let mut shard = pager.shards[handle.shard].lock();
645 let state = unsafe { &mut *handle.state.get() };
647 if limit > state.max_key {
648 return;
649 }
650 let keys: Vec<u32> = state
651 .pages
652 .keys()
653 .copied()
654 .filter(|key| *key >= limit)
655 .collect();
656 for key in keys {
658 take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
659 }
660 state.max_key = limit.saturating_sub(1);
661 }
662 pager.release(freed);
663}
664
665unsafe extern "C" fn x_destroy(cache: *mut sqlite3_pcache) {
666 let handle_ptr = cache as *mut CacheHandle;
669 let handle = unsafe { &*handle_ptr };
670 let pager = handle.cache;
671 let handle_addr = handle_ptr as usize;
672 let mut freed = Freed::new();
673 {
674 let mut shard = pager.shards[handle.shard].lock();
675 let state = unsafe { &mut *handle.state.get() };
677 let keys: Vec<u32> = state.pages.keys().copied().collect();
678 for key in keys {
679 take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
680 }
681 }
682 pager.release(freed);
683 drop(unsafe { Box::from_raw(handle_ptr) });
685}
686
687unsafe extern "C" fn x_shrink(cache: *mut sqlite3_pcache) {
688 let handle = unsafe { handle(cache) };
690 let pager = handle.cache;
691 let handle_addr = handle as *const CacheHandle as usize;
692 let mut freed = Freed::new();
693 {
694 let mut shard = pager.shards[handle.shard].lock();
695 let state = unsafe { &mut *handle.state.get() };
697 if !state.purgeable {
698 return;
699 }
700 let unpinned: Vec<u32> = state
701 .pages
702 .iter()
703 .filter(|(_, slot)| !unsafe { slot.as_ref() }.pinned)
705 .map(|(key, _)| *key)
706 .collect();
707 for key in unpinned {
708 take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
709 }
710 }
711 pager.release(freed);
712}
713
714