use std::{
alloc::{self, Layout},
cell::UnsafeCell,
collections::HashMap,
ffi::{c_int, c_uint, c_void},
ptr::{self, NonNull},
sync::{
Arc, OnceLock,
atomic::{AtomicPtr, AtomicU64, AtomicUsize, Ordering},
},
};
use libsqlite3_sys::{
SQLITE_CONFIG_PCACHE2, SQLITE_OK, sqlite3_config, sqlite3_pcache, sqlite3_pcache_methods2,
sqlite3_pcache_page,
};
use lru::LruCache;
use parking_lot::Mutex;
use crate::worker::SqlitePoolMetrics;
const SHARDS: usize = 16;
static PROCESS_PAGER_CACHE: AtomicPtr<PagerCache> = AtomicPtr::new(ptr::null_mut());
static INSTALLED: OnceLock<bool> = OnceLock::new();
pub fn install(capacity_bytes: u64, metrics: Option<Arc<dyn SqlitePoolMetrics>>) -> bool {
let installed = *INSTALLED.get_or_init(|| {
if capacity_bytes == 0 {
return false;
}
let cache: &'static PagerCache = Box::leak(Box::new(PagerCache::new(capacity_bytes)));
PROCESS_PAGER_CACHE.store(cache as *const PagerCache as *mut PagerCache, Ordering::Release);
let methods = sqlite3_pcache_methods2 {
iVersion: 1,
pArg: cache as *const PagerCache as *mut c_void,
xInit: Some(x_init),
xShutdown: Some(x_shutdown),
xCreate: Some(x_create),
xCachesize: Some(x_cachesize),
xPagecount: Some(x_pagecount),
xFetch: Some(x_fetch),
xUnpin: Some(x_unpin),
xRekey: Some(x_rekey),
xTruncate: Some(x_truncate),
xDestroy: Some(x_destroy),
xShrink: Some(x_shrink),
};
let rc = unsafe { sqlite3_config(SQLITE_CONFIG_PCACHE2, &methods as *const _) };
if rc != SQLITE_OK {
PROCESS_PAGER_CACHE.store(ptr::null_mut(), Ordering::Release);
tracing::warn!(
rc,
"sqlite was initialized before the process-wide page cache could be installed; connections keep per-connection page caches"
);
return false;
}
true
});
if let (true, Some(metrics)) = (installed, metrics) {
attach_metrics(metrics);
}
installed
}
pub fn installed() -> Option<&'static PagerCache> {
let cache = PROCESS_PAGER_CACHE.load(Ordering::Acquire);
unsafe { cache.as_ref() }
}
pub fn attach_metrics(metrics: Arc<dyn SqlitePoolMetrics>) {
if let Some(cache) = installed()
&& cache.metrics.set(metrics).is_ok()
{
cache.report_capacity();
cache.report();
}
}
pub struct PagerCache {
shards: Box<[Mutex<Shard>]>,
capacity_bytes: u64,
purgeable_bytes: AtomicU64,
unbudgeted_bytes: AtomicU64,
pages: AtomicU64,
refusals: AtomicU64,
evictions: AtomicU64,
next_cache: AtomicUsize,
evict_cursor: AtomicUsize,
metrics: OnceLock<Arc<dyn SqlitePoolMetrics>>,
}
struct Shard {
lru: LruCache<(usize, u32), ()>,
}
struct CacheHandle {
cache: &'static PagerCache,
shard: usize,
state: UnsafeCell<CacheState>,
}
struct CacheState {
layout: Layout,
buf_offset: usize,
slot_bytes: u64,
page_size: usize,
purgeable: bool,
pages: HashMap<u32, NonNull<Slot>>,
max_key: u32,
}
#[repr(C)]
struct Slot {
page: sqlite3_pcache_page,
key: u32,
pinned: bool,
}
type Freed = Vec<(NonNull<Slot>, Layout)>;
impl PagerCache {
fn new(capacity_bytes: u64) -> Self {
Self {
shards: (0..SHARDS)
.map(|_| {
Mutex::new(Shard {
lru: LruCache::unbounded(),
})
})
.collect(),
capacity_bytes,
purgeable_bytes: AtomicU64::new(0),
unbudgeted_bytes: AtomicU64::new(0),
pages: AtomicU64::new(0),
refusals: AtomicU64::new(0),
evictions: AtomicU64::new(0),
next_cache: AtomicUsize::new(0),
evict_cursor: AtomicUsize::new(0),
metrics: OnceLock::new(),
}
}
pub fn bytes(&self) -> u64 {
self.purgeable_bytes.load(Ordering::Relaxed)
}
pub fn unbudgeted_bytes(&self) -> u64 {
self.unbudgeted_bytes.load(Ordering::Relaxed)
}
pub fn pages(&self) -> u64 {
self.pages.load(Ordering::Relaxed)
}
pub fn capacity_bytes(&self) -> u64 {
self.capacity_bytes
}
pub fn refusals(&self) -> u64 {
self.refusals.load(Ordering::Relaxed)
}
pub fn evictions(&self) -> u64 {
self.evictions.load(Ordering::Relaxed)
}
fn report_capacity(&self) {
if let Some(metrics) = self.metrics.get() {
metrics.set_pager_cache_capacity_bytes(self.capacity_bytes);
}
}
fn report(&self) {
if let Some(metrics) = self.metrics.get() {
metrics.set_pager_cache_bytes(self.bytes());
metrics.set_pager_cache_unbudgeted_bytes(self.unbudgeted_bytes());
metrics.set_pager_cache_pages(self.pages());
}
}
fn release(&self, freed: Freed) {
if freed.is_empty() {
return;
}
for (slot, layout) in freed {
unsafe { alloc::dealloc(slot.as_ptr() as *mut u8, layout) };
}
self.report();
}
fn account(&self, state: &CacheState, added: bool) {
let counter = if state.purgeable {
&self.purgeable_bytes
} else {
&self.unbudgeted_bytes
};
if added {
counter.fetch_add(state.slot_bytes, Ordering::Relaxed);
self.pages.fetch_add(1, Ordering::Relaxed);
} else {
counter.fetch_sub(state.slot_bytes, Ordering::Relaxed);
self.pages.fetch_sub(1, Ordering::Relaxed);
}
}
fn reserve(&self, bytes: u64) -> bool {
loop {
let current = self.purgeable_bytes.load(Ordering::Relaxed);
if current.saturating_add(bytes) <= self.capacity_bytes {
if self
.purgeable_bytes
.compare_exchange_weak(
current,
current + bytes,
Ordering::Relaxed,
Ordering::Relaxed,
)
.is_ok()
{
return true;
}
continue;
}
if !self.evict_one() {
return false;
}
}
}
fn shed_overshoot(&self) {
while self.bytes() > self.capacity_bytes {
if !self.evict_one() {
return;
}
}
}
fn evict_one(&self) -> bool {
for _ in 0..SHARDS {
let shard_index = self.evict_cursor.fetch_add(1, Ordering::Relaxed) % SHARDS;
let mut shard = self.shards[shard_index].lock();
let Some(((handle_addr, key), ())) = shard.lru.pop_lru() else {
continue;
};
let state = unsafe { &mut *(*(handle_addr as *const CacheHandle)).state.get() };
let freed = state.pages.remove(&key).map(|slot| (slot, state.layout));
if freed.is_some() {
self.account(state, false);
self.evictions.fetch_add(1, Ordering::Relaxed);
}
drop(shard);
self.release(freed.into_iter().collect());
return true;
}
false
}
}
fn take_page(
cache: &PagerCache,
shard: &mut Shard,
handle_addr: usize,
state: &mut CacheState,
key: u32,
freed: &mut Freed,
) {
if let Some(slot) = state.pages.remove(&key) {
shard.lru.pop(&(handle_addr, key));
cache.account(state, false);
freed.push((slot, state.layout));
}
}
fn allocate(
layout: Layout,
buf_offset: usize,
page_size: usize,
key: u32,
) -> Option<NonNull<Slot>> {
let base = NonNull::new(unsafe { alloc::alloc_zeroed(layout) })?;
let slot = base.cast::<Slot>();
unsafe {
let buf = base.as_ptr().add(buf_offset);
ptr::write(
slot.as_ptr(),
Slot {
page: sqlite3_pcache_page {
pBuf: buf as *mut c_void,
pExtra: buf.add(page_size) as *mut c_void,
},
key,
pinned: true,
},
);
}
Some(slot)
}
unsafe fn handle<'a>(cache: *mut sqlite3_pcache) -> &'a CacheHandle {
unsafe { &*(cache as *const CacheHandle) }
}
unsafe extern "C" fn x_init(_arg: *mut c_void) -> c_int {
SQLITE_OK
}
unsafe extern "C" fn x_shutdown(_arg: *mut c_void) {}
unsafe extern "C" fn x_create(
page_size: c_int,
extra_size: c_int,
purgeable: c_int,
) -> *mut sqlite3_pcache {
match installed() {
Some(cache) => create(cache, page_size, extra_size, purgeable),
None => ptr::null_mut(),
}
}
fn create(
cache: &'static PagerCache,
page_size: c_int,
extra_size: c_int,
purgeable: c_int,
) -> *mut sqlite3_pcache {
let page_size = page_size.max(0) as usize;
let extra_size = extra_size.max(0) as usize;
let Ok(area) = Layout::from_size_align(page_size + extra_size, 8) else {
return ptr::null_mut();
};
let Ok((layout, buf_offset)) = Layout::new::<Slot>().extend(area) else {
return ptr::null_mut();
};
let layout = layout.pad_to_align();
let shard = cache.next_cache.fetch_add(1, Ordering::Relaxed) % SHARDS;
let handle = Box::new(CacheHandle {
cache,
shard,
state: UnsafeCell::new(CacheState {
layout,
buf_offset,
slot_bytes: layout.size() as u64,
page_size,
purgeable: purgeable != 0,
pages: HashMap::new(),
max_key: 0,
}),
});
Box::into_raw(handle) as *mut sqlite3_pcache
}
unsafe extern "C" fn x_cachesize(_cache: *mut sqlite3_pcache, _pages: c_int) {}
unsafe extern "C" fn x_pagecount(cache: *mut sqlite3_pcache) -> c_int {
let handle = unsafe { handle(cache) };
let _shard = handle.cache.shards[handle.shard].lock();
let state = unsafe { &*handle.state.get() };
state.pages.len().min(c_int::MAX as usize) as c_int
}
unsafe extern "C" fn x_fetch(
cache: *mut sqlite3_pcache,
key: c_uint,
create_flag: c_int,
) -> *mut sqlite3_pcache_page {
let handle = unsafe { handle(cache) };
let pager = handle.cache;
let handle_addr = handle as *const CacheHandle as usize;
let (layout, buf_offset, page_size, slot_bytes, purgeable) = {
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
if let Some(slot) = state.pages.get(&key).copied() {
let slot_ref = unsafe { &mut *slot.as_ptr() };
if !slot_ref.pinned {
slot_ref.pinned = true;
shard.lru.pop(&(handle_addr, key));
}
return slot.as_ptr() as *mut sqlite3_pcache_page;
}
if create_flag == 0 {
return ptr::null_mut();
}
(
state.layout,
state.buf_offset,
state.page_size,
state.slot_bytes,
state.purgeable,
)
};
if purgeable && !pager.reserve(slot_bytes) {
if create_flag == 1 {
pager.refusals.fetch_add(1, Ordering::Relaxed);
return ptr::null_mut();
}
pager
.purgeable_bytes
.fetch_add(slot_bytes, Ordering::Relaxed);
}
let Some(slot) = allocate(layout, buf_offset, page_size, key) else {
if purgeable {
pager
.purgeable_bytes
.fetch_sub(slot_bytes, Ordering::Relaxed);
}
return ptr::null_mut();
};
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
if let Some(existing) = state.pages.get(&key).copied() {
let existing_ref = unsafe { &mut *existing.as_ptr() };
if !existing_ref.pinned {
existing_ref.pinned = true;
shard.lru.pop(&(handle_addr, key));
}
drop(shard);
if purgeable {
pager
.purgeable_bytes
.fetch_sub(slot_bytes, Ordering::Relaxed);
}
unsafe { alloc::dealloc(slot.as_ptr() as *mut u8, layout) };
return existing.as_ptr() as *mut sqlite3_pcache_page;
}
state.pages.insert(key, slot);
state.max_key = state.max_key.max(key);
}
if !purgeable {
pager
.unbudgeted_bytes
.fetch_add(slot_bytes, Ordering::Relaxed);
}
pager.pages.fetch_add(1, Ordering::Relaxed);
pager.report();
slot.as_ptr() as *mut sqlite3_pcache_page
}
unsafe extern "C" fn x_unpin(
cache: *mut sqlite3_pcache,
page: *mut sqlite3_pcache_page,
discard: c_int,
) {
let handle = unsafe { handle(cache) };
let pager = handle.cache;
let handle_addr = handle as *const CacheHandle as usize;
let mut freed = Freed::new();
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
let slot = unsafe { &mut *(page as *mut Slot) };
if discard != 0 {
take_page(pager, &mut shard, handle_addr, state, slot.key, &mut freed);
} else {
slot.pinned = false;
if state.purgeable {
shard.lru.put((handle_addr, slot.key), ());
}
}
}
pager.release(freed);
pager.shed_overshoot();
}
unsafe extern "C" fn x_rekey(
cache: *mut sqlite3_pcache,
_page: *mut sqlite3_pcache_page,
old_key: c_uint,
new_key: c_uint,
) {
if old_key == new_key {
return;
}
let handle = unsafe { handle(cache) };
let pager = handle.cache;
let handle_addr = handle as *const CacheHandle as usize;
let mut freed = Freed::new();
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
take_page(pager, &mut shard, handle_addr, state, new_key, &mut freed);
if let Some(slot) = state.pages.remove(&old_key) {
let pinned = unsafe {
let slot = &mut *slot.as_ptr();
slot.key = new_key;
slot.pinned
};
state.pages.insert(new_key, slot);
state.max_key = state.max_key.max(new_key);
if !pinned && shard.lru.pop(&(handle_addr, old_key)).is_some() {
shard.lru.put((handle_addr, new_key), ());
}
}
}
pager.release(freed);
}
unsafe extern "C" fn x_truncate(cache: *mut sqlite3_pcache, limit: c_uint) {
let handle = unsafe { handle(cache) };
let pager = handle.cache;
let handle_addr = handle as *const CacheHandle as usize;
let mut freed = Freed::new();
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
if limit > state.max_key {
return;
}
let keys: Vec<u32> = state
.pages
.keys()
.copied()
.filter(|key| *key >= limit)
.collect();
for key in keys {
take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
}
state.max_key = limit.saturating_sub(1);
}
pager.release(freed);
}
unsafe extern "C" fn x_destroy(cache: *mut sqlite3_pcache) {
let handle_ptr = cache as *mut CacheHandle;
let handle = unsafe { &*handle_ptr };
let pager = handle.cache;
let handle_addr = handle_ptr as usize;
let mut freed = Freed::new();
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
let keys: Vec<u32> = state.pages.keys().copied().collect();
for key in keys {
take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
}
}
pager.release(freed);
drop(unsafe { Box::from_raw(handle_ptr) });
}
unsafe extern "C" fn x_shrink(cache: *mut sqlite3_pcache) {
let handle = unsafe { handle(cache) };
let pager = handle.cache;
let handle_addr = handle as *const CacheHandle as usize;
let mut freed = Freed::new();
{
let mut shard = pager.shards[handle.shard].lock();
let state = unsafe { &mut *handle.state.get() };
if !state.purgeable {
return;
}
let unpinned: Vec<u32> = state
.pages
.iter()
.filter(|(_, slot)| !unsafe { slot.as_ref() }.pinned)
.map(|(key, _)| *key)
.collect();
for key in unpinned {
take_page(pager, &mut shard, handle_addr, state, key, &mut freed);
}
}
pager.release(freed);
}