Skip to main content

sqlite_embedded/
optimization_flags.rs

1//! Process page budgets and I/O concurrency. Prefetch policy is owned by the server.
2
3use std::{env, sync::OnceLock};
4
5pub const VFS_PAGE_CACHE_MODE_ENV: &str = "RIVETKIT_SQLITE_OPT_VFS_PAGE_CACHE_MODE";
6pub const VFS_PAGE_CACHE_CAPACITY_PAGES_ENV: &str =
7	"RIVETKIT_SQLITE_OPT_VFS_PAGE_CACHE_CAPACITY_PAGES";
8pub const VFS_STAGING_CACHE_TTL_MS_ENV: &str = "RIVETKIT_SQLITE_OPT_VFS_STAGING_CACHE_TTL_MS";
9pub const PAGER_CACHE_SIZE_KIB_ENV: &str = "RIVETKIT_SQLITE_OPT_PAGER_CACHE_SIZE_KIB";
10pub const PAGER_PROCESS_CACHE_KIB_ENV: &str = "RIVETKIT_SQLITE_OPT_PAGER_PROCESS_CACHE_KIB";
11pub const IDLE_RELEASE_MS_ENV: &str = "RIVETKIT_SQLITE_OPT_IDLE_RELEASE_MS";
12pub const PAGE_RANGE_CONCURRENCY_ENV: &str = "RIVETKIT_SQLITE_OPT_PAGE_RANGE_CONCURRENCY";
13pub const DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES: u64 = 131_072;
14pub const MAX_VFS_PAGE_CACHE_CAPACITY_PAGES: u64 = 4_194_304;
15pub const DEFAULT_VFS_STAGING_CACHE_TTL_MS: u64 = 30_000;
16pub const MAX_VFS_STAGING_CACHE_TTL_MS: u64 = 300_000;
17pub const DEFAULT_PAGER_CACHE_SIZE_KIB: u64 = 8 * 1024;
18pub const MAX_PAGER_CACHE_SIZE_KIB: u64 = 256 * 1024;
19pub const DEFAULT_PAGER_PROCESS_CACHE_KIB: u64 = 256 * 1024;
20pub const MAX_PAGER_PROCESS_CACHE_KIB: u64 = 16 * 1024 * 1024;
21pub const DEFAULT_IDLE_RELEASE_MS: u64 = 30_000;
22pub const MAX_IDLE_RELEASE_MS: u64 = 3_600_000;
23pub const DEFAULT_PAGE_RANGE_CONCURRENCY: usize = 16;
24pub const MAX_PAGE_RANGE_CONCURRENCY: usize = 64;
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub enum SqliteVfsPageCacheMode {
28	Off,
29	Target,
30	Startup,
31	Prefetch,
32	All,
33}
34
35impl SqliteVfsPageCacheMode {
36	pub fn caches_any_pages(self) -> bool {
37		!matches!(self, Self::Off)
38	}
39
40	pub fn caches_target_pages(self) -> bool {
41		matches!(
42			self,
43			Self::Target | Self::Startup | Self::Prefetch | Self::All
44		)
45	}
46
47	pub fn caches_prefetched_pages(self) -> bool {
48		matches!(self, Self::Prefetch | Self::All)
49	}
50
51	pub fn caches_startup_preloaded_pages(self) -> bool {
52		matches!(self, Self::Startup | Self::All)
53	}
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub struct SqliteOptimizationFlags {
58	pub vfs_page_cache_mode: SqliteVfsPageCacheMode,
59	pub vfs_page_buffer_process_pages: u64,
60	pub vfs_staging_cache_ttl_ms: u64,
61	pub pager_cache_size_kib: u64,
62	pub pager_process_cache_kib: u64,
63	pub idle_release_ms: u64,
64	pub page_range_concurrency: usize,
65}
66impl Default for SqliteOptimizationFlags {
67	fn default() -> Self {
68		Self::from_env_reader(|_| None)
69	}
70}
71impl SqliteOptimizationFlags {
72	fn from_process_env() -> Self {
73		Self::from_env_reader(|key| env::var(key).ok())
74	}
75	pub fn from_env_reader(mut read_env: impl FnMut(&str) -> Option<String>) -> Self {
76		Self {
77			vfs_page_cache_mode: vfs_page_cache_mode_by_default(
78				read_env(VFS_PAGE_CACHE_MODE_ENV).as_deref(),
79				SqliteVfsPageCacheMode::All,
80			),
81			vfs_page_buffer_process_pages: u64_bounded_by_default(
82				read_env(VFS_PAGE_CACHE_CAPACITY_PAGES_ENV).as_deref(),
83				DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES,
84				MAX_VFS_PAGE_CACHE_CAPACITY_PAGES,
85			),
86			vfs_staging_cache_ttl_ms: u64_bounded_by_default(
87				read_env(VFS_STAGING_CACHE_TTL_MS_ENV).as_deref(),
88				DEFAULT_VFS_STAGING_CACHE_TTL_MS,
89				MAX_VFS_STAGING_CACHE_TTL_MS,
90			),
91			pager_cache_size_kib: u64_bounded_by_default(
92				read_env(PAGER_CACHE_SIZE_KIB_ENV).as_deref(),
93				DEFAULT_PAGER_CACHE_SIZE_KIB,
94				MAX_PAGER_CACHE_SIZE_KIB,
95			),
96			pager_process_cache_kib: u64_bounded_by_default(
97				read_env(PAGER_PROCESS_CACHE_KIB_ENV).as_deref(),
98				DEFAULT_PAGER_PROCESS_CACHE_KIB,
99				MAX_PAGER_PROCESS_CACHE_KIB,
100			),
101			idle_release_ms: u64_bounded_by_default(
102				read_env(IDLE_RELEASE_MS_ENV).as_deref(),
103				DEFAULT_IDLE_RELEASE_MS,
104				MAX_IDLE_RELEASE_MS,
105			),
106			page_range_concurrency: usize_bounded_by_default(
107				read_env(PAGE_RANGE_CONCURRENCY_ENV).as_deref(),
108				DEFAULT_PAGE_RANGE_CONCURRENCY,
109				MAX_PAGE_RANGE_CONCURRENCY,
110			)
111			.max(1),
112		}
113	}
114}
115pub fn sqlite_optimization_flags() -> &'static SqliteOptimizationFlags {
116	static FLAGS: OnceLock<SqliteOptimizationFlags> = OnceLock::new();
117	FLAGS.get_or_init(SqliteOptimizationFlags::from_process_env)
118}
119
120fn usize_bounded_by_default(value: Option<&str>, default: usize, max: usize) -> usize {
121	value
122		.and_then(|value| value.trim().parse::<usize>().ok())
123		.unwrap_or(default)
124		.min(max)
125}
126
127fn u64_bounded_by_default(value: Option<&str>, default: u64, max: u64) -> u64 {
128	value
129		.and_then(|value| value.trim().parse::<u64>().ok())
130		.unwrap_or(default)
131		.min(max)
132}
133
134fn vfs_page_cache_mode_by_default(
135	value: Option<&str>,
136	default: SqliteVfsPageCacheMode,
137) -> SqliteVfsPageCacheMode {
138	match value.map(|value| value.trim().to_ascii_lowercase()) {
139		Some(value)
140			if matches!(
141				value.as_str(),
142				"off" | "0" | "false" | "no" | "disabled" | "disable"
143			) =>
144		{
145			SqliteVfsPageCacheMode::Off
146		}
147		Some(value) if value == "target" => SqliteVfsPageCacheMode::Target,
148		Some(value) if value == "startup" => SqliteVfsPageCacheMode::Startup,
149		Some(value) if value == "prefetch" => SqliteVfsPageCacheMode::Prefetch,
150		Some(value) if value == "all" || value == "on" || value == "true" || value == "1" => {
151			SqliteVfsPageCacheMode::All
152		}
153		_ => default,
154	}
155}
156
157