1use std::{env, sync::OnceLock};
4
5pub const READ_AHEAD_MODE_ENV: &str = "RIVETKIT_SQLITE_OPT_READ_AHEAD_MODE";
6pub const RECENT_PAGE_HINTS_ENV: &str = "RIVETKIT_SQLITE_OPT_RECENT_PAGE_HINTS";
7pub const PRELOAD_HINT_FLUSH_ENV: &str = "RIVETKIT_SQLITE_OPT_PRELOAD_HINT_FLUSH";
8pub const STARTUP_PRELOAD_MAX_BYTES_ENV: &str = "RIVETKIT_SQLITE_OPT_STARTUP_PRELOAD_MAX_BYTES";
9pub const STARTUP_PRELOAD_FIRST_PAGES_ENV: &str = "RIVETKIT_SQLITE_OPT_STARTUP_PRELOAD_FIRST_PAGES";
10pub const STARTUP_PRELOAD_FIRST_PAGE_COUNT_ENV: &str =
11 "RIVETKIT_SQLITE_OPT_STARTUP_PRELOAD_FIRST_PAGE_COUNT";
12pub const PRELOAD_HINTS_ON_OPEN_ENV: &str = "RIVETKIT_SQLITE_OPT_PRELOAD_HINTS_ON_OPEN";
13pub const PRELOAD_HINT_HOT_PAGES_ENV: &str = "RIVETKIT_SQLITE_OPT_PRELOAD_HINT_HOT_PAGES";
14pub const PRELOAD_HINT_EARLY_PAGES_ENV: &str = "RIVETKIT_SQLITE_OPT_PRELOAD_HINT_EARLY_PAGES";
15pub const PRELOAD_HINT_SCAN_RANGES_ENV: &str = "RIVETKIT_SQLITE_OPT_PRELOAD_HINT_SCAN_RANGES";
16pub const DEDUP_GET_PAGES_META_ENV: &str = "RIVETKIT_SQLITE_OPT_DEDUP_GET_PAGES_META";
17pub const CACHE_GET_PAGES_VALIDATION_ENV: &str = "RIVETKIT_SQLITE_OPT_CACHE_GET_PAGES_VALIDATION";
18pub const RANGE_READS_ENV: &str = "RIVETKIT_SQLITE_OPT_RANGE_READS";
19pub const BATCH_CHUNK_READS_ENV: &str = "RIVETKIT_SQLITE_OPT_BATCH_CHUNK_READS";
20pub const DECODED_LTX_CACHE_ENV: &str = "RIVETKIT_SQLITE_OPT_DECODED_LTX_CACHE";
21pub const VFS_PAGE_CACHE_MODE_ENV: &str = "RIVETKIT_SQLITE_OPT_VFS_PAGE_CACHE_MODE";
22pub const VFS_PAGE_CACHE_CAPACITY_PAGES_ENV: &str =
23 "RIVETKIT_SQLITE_OPT_VFS_PAGE_CACHE_CAPACITY_PAGES";
24pub const VFS_PROTECTED_CACHE_PAGES_ENV: &str = "RIVETKIT_SQLITE_OPT_VFS_PROTECTED_CACHE_PAGES";
25pub const VFS_STAGING_CACHE_TTL_MS_ENV: &str = "RIVETKIT_SQLITE_OPT_VFS_STAGING_CACHE_TTL_MS";
26pub const PAGER_CACHE_SIZE_KIB_ENV: &str = "RIVETKIT_SQLITE_OPT_PAGER_CACHE_SIZE_KIB";
27
28pub const DEFAULT_STARTUP_PRELOAD_MAX_BYTES: usize = 2 * 1024 * 1024;
29pub const MAX_STARTUP_PRELOAD_MAX_BYTES: usize = 64 * 1024 * 1024;
30pub const DEFAULT_STARTUP_PRELOAD_FIRST_PAGE_COUNT: u32 = 128;
31pub const MAX_STARTUP_PRELOAD_FIRST_PAGE_COUNT: u32 = 16_384;
32pub const DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES: u64 = 50_000;
33pub const MAX_VFS_PAGE_CACHE_CAPACITY_PAGES: u64 = 500_000;
34pub const DEFAULT_VFS_PROTECTED_CACHE_PAGES: usize = 512;
35pub const MAX_VFS_PROTECTED_CACHE_PAGES: usize = 8_192;
36pub const DEFAULT_VFS_STAGING_CACHE_TTL_MS: u64 = 30_000;
37pub const MAX_VFS_STAGING_CACHE_TTL_MS: u64 = 300_000;
38pub const DEFAULT_PAGER_CACHE_SIZE_KIB: u64 = 8 * 1024;
39pub const MAX_PAGER_CACHE_SIZE_KIB: u64 = 256 * 1024;
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum SqliteReadAheadMode {
43 Off,
44 Bounded,
45 Adaptive,
46}
47
48#[derive(Debug, Clone, Copy, PartialEq, Eq)]
49pub enum SqliteVfsPageCacheMode {
50 Off,
51 Target,
52 Startup,
53 Prefetch,
54 All,
55}
56
57impl SqliteReadAheadMode {
58 pub fn uses_bounded_prefetch(self) -> bool {
59 matches!(self, Self::Bounded | Self::Adaptive)
60 }
61
62 pub fn uses_adaptive_prefetch(self) -> bool {
63 matches!(self, Self::Adaptive)
64 }
65}
66
67impl SqliteVfsPageCacheMode {
68 pub fn caches_any_pages(self) -> bool {
69 !matches!(self, Self::Off)
70 }
71
72 pub fn caches_target_pages(self) -> bool {
73 matches!(
74 self,
75 Self::Target | Self::Startup | Self::Prefetch | Self::All
76 )
77 }
78
79 pub fn caches_prefetched_pages(self) -> bool {
80 matches!(self, Self::Prefetch | Self::All)
81 }
82
83 pub fn caches_startup_preloaded_pages(self) -> bool {
84 matches!(self, Self::Startup | Self::All)
85 }
86}
87
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub struct SqliteOptimizationFlags {
90 pub read_ahead_mode: SqliteReadAheadMode,
91 pub read_ahead: bool,
92 pub adaptive_read_ahead: bool,
93 pub recent_page_hints: bool,
94 pub cache_hit_predictor_training: bool,
95 pub preload_hint_flush: bool,
96 pub startup_preload_max_bytes: usize,
97 pub startup_preload_first_pages: bool,
98 pub startup_preload_first_page_count: u32,
99 pub preload_hints_on_open: bool,
100 pub preload_hint_hot_pages: bool,
101 pub preload_hint_early_pages: bool,
102 pub preload_hint_scan_ranges: bool,
103 pub dedup_get_pages_meta: bool,
104 pub cache_get_pages_validation: bool,
105 pub range_reads: bool,
106 pub batch_chunk_reads: bool,
107 pub decoded_ltx_cache: bool,
108 pub vfs_page_cache_mode: SqliteVfsPageCacheMode,
109 pub vfs_page_cache_capacity_pages: u64,
110 pub vfs_protected_cache_pages: usize,
111 pub vfs_staging_cache_ttl_ms: u64,
112 pub pager_cache_size_kib: u64,
113}
114
115impl Default for SqliteOptimizationFlags {
116 fn default() -> Self {
117 Self {
118 read_ahead_mode: SqliteReadAheadMode::Adaptive,
119 read_ahead: true,
120 adaptive_read_ahead: true,
121 recent_page_hints: true,
122 cache_hit_predictor_training: true,
123 preload_hint_flush: true,
124 startup_preload_max_bytes: DEFAULT_STARTUP_PRELOAD_MAX_BYTES,
125 startup_preload_first_pages: true,
126 startup_preload_first_page_count: DEFAULT_STARTUP_PRELOAD_FIRST_PAGE_COUNT,
127 preload_hints_on_open: true,
128 preload_hint_hot_pages: true,
129 preload_hint_early_pages: true,
130 preload_hint_scan_ranges: true,
131 dedup_get_pages_meta: true,
132 cache_get_pages_validation: true,
133 range_reads: true,
134 batch_chunk_reads: true,
135 decoded_ltx_cache: true,
136 vfs_page_cache_mode: SqliteVfsPageCacheMode::All,
137 vfs_page_cache_capacity_pages: DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES,
138 vfs_protected_cache_pages: DEFAULT_VFS_PROTECTED_CACHE_PAGES,
139 vfs_staging_cache_ttl_ms: DEFAULT_VFS_STAGING_CACHE_TTL_MS,
140 pager_cache_size_kib: DEFAULT_PAGER_CACHE_SIZE_KIB,
141 }
142 }
143}
144
145impl SqliteOptimizationFlags {
146 fn from_process_env() -> Self {
147 Self::from_env_reader(|key| env::var(key).ok())
148 }
149
150 pub fn from_env_reader(mut read_env: impl FnMut(&str) -> Option<String>) -> Self {
151 let read_ahead_mode = read_ahead_mode_by_default(
152 read_env(READ_AHEAD_MODE_ENV).as_deref(),
153 SqliteReadAheadMode::Adaptive,
154 );
155 let recent_page_hints = enabled_by_default(read_env(RECENT_PAGE_HINTS_ENV).as_deref());
156 Self {
157 read_ahead_mode,
158 read_ahead: read_ahead_mode.uses_bounded_prefetch(),
159 adaptive_read_ahead: read_ahead_mode.uses_adaptive_prefetch(),
160 recent_page_hints,
161 cache_hit_predictor_training: recent_page_hints,
162 preload_hint_flush: enabled_by_default(read_env(PRELOAD_HINT_FLUSH_ENV).as_deref()),
163 startup_preload_max_bytes: usize_bounded_by_default(
164 read_env(STARTUP_PRELOAD_MAX_BYTES_ENV).as_deref(),
165 DEFAULT_STARTUP_PRELOAD_MAX_BYTES,
166 MAX_STARTUP_PRELOAD_MAX_BYTES,
167 ),
168 startup_preload_first_pages: enabled_by_default(
169 read_env(STARTUP_PRELOAD_FIRST_PAGES_ENV).as_deref(),
170 ),
171 startup_preload_first_page_count: u32_bounded_by_default(
172 read_env(STARTUP_PRELOAD_FIRST_PAGE_COUNT_ENV).as_deref(),
173 DEFAULT_STARTUP_PRELOAD_FIRST_PAGE_COUNT,
174 MAX_STARTUP_PRELOAD_FIRST_PAGE_COUNT,
175 ),
176 preload_hints_on_open: enabled_by_default(
177 read_env(PRELOAD_HINTS_ON_OPEN_ENV).as_deref(),
178 ),
179 preload_hint_hot_pages: enabled_by_default(
180 read_env(PRELOAD_HINT_HOT_PAGES_ENV).as_deref(),
181 ),
182 preload_hint_early_pages: enabled_by_default(
183 read_env(PRELOAD_HINT_EARLY_PAGES_ENV).as_deref(),
184 ),
185 preload_hint_scan_ranges: enabled_by_default(
186 read_env(PRELOAD_HINT_SCAN_RANGES_ENV).as_deref(),
187 ),
188 dedup_get_pages_meta: enabled_by_default(read_env(DEDUP_GET_PAGES_META_ENV).as_deref()),
189 cache_get_pages_validation: enabled_by_default(
190 read_env(CACHE_GET_PAGES_VALIDATION_ENV).as_deref(),
191 ),
192 range_reads: enabled_by_default(read_env(RANGE_READS_ENV).as_deref()),
193 batch_chunk_reads: enabled_by_default(read_env(BATCH_CHUNK_READS_ENV).as_deref()),
194 decoded_ltx_cache: enabled_by_default(read_env(DECODED_LTX_CACHE_ENV).as_deref()),
195 vfs_page_cache_mode: vfs_page_cache_mode_by_default(
196 read_env(VFS_PAGE_CACHE_MODE_ENV).as_deref(),
197 SqliteVfsPageCacheMode::All,
198 ),
199 vfs_page_cache_capacity_pages: u64_bounded_by_default(
200 read_env(VFS_PAGE_CACHE_CAPACITY_PAGES_ENV).as_deref(),
201 DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES,
202 MAX_VFS_PAGE_CACHE_CAPACITY_PAGES,
203 ),
204 vfs_protected_cache_pages: usize_bounded_by_default(
205 read_env(VFS_PROTECTED_CACHE_PAGES_ENV).as_deref(),
206 DEFAULT_VFS_PROTECTED_CACHE_PAGES,
207 MAX_VFS_PROTECTED_CACHE_PAGES,
208 ),
209 vfs_staging_cache_ttl_ms: u64_bounded_by_default(
210 read_env(VFS_STAGING_CACHE_TTL_MS_ENV).as_deref(),
211 DEFAULT_VFS_STAGING_CACHE_TTL_MS,
212 MAX_VFS_STAGING_CACHE_TTL_MS,
213 ),
214 pager_cache_size_kib: u64_bounded_by_default(
215 read_env(PAGER_CACHE_SIZE_KIB_ENV).as_deref(),
216 DEFAULT_PAGER_CACHE_SIZE_KIB,
217 MAX_PAGER_CACHE_SIZE_KIB,
218 ),
219 }
220 }
221}
222
223pub fn sqlite_optimization_flags() -> &'static SqliteOptimizationFlags {
224 static FLAGS: OnceLock<SqliteOptimizationFlags> = OnceLock::new();
225 FLAGS.get_or_init(SqliteOptimizationFlags::from_process_env)
226}
227
228fn enabled_by_default(value: Option<&str>) -> bool {
229 match value.map(|value| value.trim().to_ascii_lowercase()) {
230 Some(value)
231 if matches!(
232 value.as_str(),
233 "0" | "false" | "off" | "no" | "disabled" | "disable"
234 ) =>
235 {
236 false
237 }
238 _ => true,
239 }
240}
241
242fn usize_bounded_by_default(value: Option<&str>, default: usize, max: usize) -> usize {
243 value
244 .and_then(|value| value.trim().parse::<usize>().ok())
245 .unwrap_or(default)
246 .min(max)
247}
248
249fn u64_bounded_by_default(value: Option<&str>, default: u64, max: u64) -> u64 {
250 value
251 .and_then(|value| value.trim().parse::<u64>().ok())
252 .unwrap_or(default)
253 .min(max)
254}
255
256fn u32_bounded_by_default(value: Option<&str>, default: u32, max: u32) -> u32 {
257 value
258 .and_then(|value| value.trim().parse::<u32>().ok())
259 .unwrap_or(default)
260 .min(max)
261}
262
263fn read_ahead_mode_by_default(
264 value: Option<&str>,
265 default: SqliteReadAheadMode,
266) -> SqliteReadAheadMode {
267 match value.map(|value| value.trim().to_ascii_lowercase()) {
268 Some(value)
269 if matches!(
270 value.as_str(),
271 "off" | "0" | "false" | "no" | "disabled" | "disable"
272 ) =>
273 {
274 SqliteReadAheadMode::Off
275 }
276 Some(value) if value == "bounded" => SqliteReadAheadMode::Bounded,
277 Some(value) if value == "adaptive" || value == "on" || value == "true" || value == "1" => {
278 SqliteReadAheadMode::Adaptive
279 }
280 _ => default,
281 }
282}
283
284fn vfs_page_cache_mode_by_default(
285 value: Option<&str>,
286 default: SqliteVfsPageCacheMode,
287) -> SqliteVfsPageCacheMode {
288 match value.map(|value| value.trim().to_ascii_lowercase()) {
289 Some(value)
290 if matches!(
291 value.as_str(),
292 "off" | "0" | "false" | "no" | "disabled" | "disable"
293 ) =>
294 {
295 SqliteVfsPageCacheMode::Off
296 }
297 Some(value) if value == "target" => SqliteVfsPageCacheMode::Target,
298 Some(value) if value == "startup" => SqliteVfsPageCacheMode::Startup,
299 Some(value) if value == "prefetch" => SqliteVfsPageCacheMode::Prefetch,
300 Some(value) if value == "all" || value == "on" || value == "true" || value == "1" => {
301 SqliteVfsPageCacheMode::All
302 }
303 _ => default,
304 }
305}
306
307#[cfg(test)]
308mod tests {
309 use super::*;
310
311 #[test]
312 fn flags_default_enabled_and_explicitly_disableable() {
313 let flags = SqliteOptimizationFlags::from_env_reader(|key| match key {
314 READ_AHEAD_MODE_ENV => Some("off".to_string()),
315 RECENT_PAGE_HINTS_ENV => Some("0".to_string()),
316 PRELOAD_HINT_FLUSH_ENV => Some("false".to_string()),
317 STARTUP_PRELOAD_MAX_BYTES_ENV => Some("0".to_string()),
318 STARTUP_PRELOAD_FIRST_PAGES_ENV => Some("false".to_string()),
319 STARTUP_PRELOAD_FIRST_PAGE_COUNT_ENV => Some("0".to_string()),
320 PRELOAD_HINTS_ON_OPEN_ENV => Some("false".to_string()),
321 PRELOAD_HINT_HOT_PAGES_ENV => Some("false".to_string()),
322 PRELOAD_HINT_EARLY_PAGES_ENV => Some("false".to_string()),
323 PRELOAD_HINT_SCAN_RANGES_ENV => Some("disabled".to_string()),
324 CACHE_GET_PAGES_VALIDATION_ENV => Some("off".to_string()),
325 RANGE_READS_ENV => Some("false".to_string()),
326 BATCH_CHUNK_READS_ENV => Some("no".to_string()),
327 DECODED_LTX_CACHE_ENV => Some("disable".to_string()),
328 VFS_PAGE_CACHE_MODE_ENV => Some("off".to_string()),
329 VFS_PAGE_CACHE_CAPACITY_PAGES_ENV => Some("0".to_string()),
330 VFS_PROTECTED_CACHE_PAGES_ENV => Some("0".to_string()),
331 VFS_STAGING_CACHE_TTL_MS_ENV => Some("0".to_string()),
332 _ => None,
333 });
334
335 assert_eq!(flags.read_ahead_mode, SqliteReadAheadMode::Off);
336 assert!(!flags.recent_page_hints);
337 assert!(!flags.preload_hint_flush);
338 assert_eq!(flags.startup_preload_max_bytes, 0);
339 assert!(!flags.startup_preload_first_pages);
340 assert_eq!(flags.startup_preload_first_page_count, 0);
341 assert!(!flags.preload_hints_on_open);
342 assert!(!flags.preload_hint_hot_pages);
343 assert!(!flags.preload_hint_early_pages);
344 assert!(!flags.preload_hint_scan_ranges);
345 assert!(!flags.cache_get_pages_validation);
346 assert!(!flags.range_reads);
347 assert!(!flags.batch_chunk_reads);
348 assert!(!flags.decoded_ltx_cache);
349 assert_eq!(flags.vfs_page_cache_mode, SqliteVfsPageCacheMode::Off);
350 assert_eq!(flags.vfs_page_cache_capacity_pages, 0);
351 assert_eq!(flags.vfs_protected_cache_pages, 0);
352 assert_eq!(flags.vfs_staging_cache_ttl_ms, 0);
353 }
354
355 #[test]
356 fn preload_numeric_config_defaults_and_clamps() {
357 let invalid = SqliteOptimizationFlags::from_env_reader(|key| match key {
358 STARTUP_PRELOAD_MAX_BYTES_ENV => Some("not-a-number".to_string()),
359 STARTUP_PRELOAD_FIRST_PAGE_COUNT_ENV => Some("nope".to_string()),
360 VFS_PAGE_CACHE_CAPACITY_PAGES_ENV => Some("invalid".to_string()),
361 VFS_PROTECTED_CACHE_PAGES_ENV => Some("invalid".to_string()),
362 VFS_STAGING_CACHE_TTL_MS_ENV => Some("invalid".to_string()),
363 _ => None,
364 });
365 assert_eq!(
366 invalid.startup_preload_max_bytes,
367 DEFAULT_STARTUP_PRELOAD_MAX_BYTES
368 );
369 assert_eq!(
370 invalid.startup_preload_first_page_count,
371 DEFAULT_STARTUP_PRELOAD_FIRST_PAGE_COUNT
372 );
373 assert_eq!(
374 invalid.vfs_page_cache_capacity_pages,
375 DEFAULT_VFS_PAGE_CACHE_CAPACITY_PAGES
376 );
377 assert_eq!(
378 invalid.vfs_protected_cache_pages,
379 DEFAULT_VFS_PROTECTED_CACHE_PAGES
380 );
381 assert_eq!(
382 invalid.vfs_staging_cache_ttl_ms,
383 DEFAULT_VFS_STAGING_CACHE_TTL_MS
384 );
385
386 let clamped = SqliteOptimizationFlags::from_env_reader(|key| match key {
387 STARTUP_PRELOAD_MAX_BYTES_ENV => Some((MAX_STARTUP_PRELOAD_MAX_BYTES + 1).to_string()),
388 STARTUP_PRELOAD_FIRST_PAGE_COUNT_ENV => {
389 Some((MAX_STARTUP_PRELOAD_FIRST_PAGE_COUNT + 1).to_string())
390 }
391 VFS_PAGE_CACHE_CAPACITY_PAGES_ENV => {
392 Some((MAX_VFS_PAGE_CACHE_CAPACITY_PAGES + 1).to_string())
393 }
394 VFS_PROTECTED_CACHE_PAGES_ENV => Some((MAX_VFS_PROTECTED_CACHE_PAGES + 1).to_string()),
395 VFS_STAGING_CACHE_TTL_MS_ENV => Some((MAX_VFS_STAGING_CACHE_TTL_MS + 1).to_string()),
396 PAGER_CACHE_SIZE_KIB_ENV => Some((MAX_PAGER_CACHE_SIZE_KIB + 1).to_string()),
397 _ => None,
398 });
399 assert_eq!(
400 clamped.startup_preload_max_bytes,
401 MAX_STARTUP_PRELOAD_MAX_BYTES
402 );
403 assert_eq!(
404 clamped.startup_preload_first_page_count,
405 MAX_STARTUP_PRELOAD_FIRST_PAGE_COUNT
406 );
407 assert_eq!(
408 clamped.vfs_page_cache_capacity_pages,
409 MAX_VFS_PAGE_CACHE_CAPACITY_PAGES
410 );
411 assert_eq!(
412 clamped.vfs_protected_cache_pages,
413 MAX_VFS_PROTECTED_CACHE_PAGES
414 );
415 assert_eq!(
416 clamped.vfs_staging_cache_ttl_ms,
417 MAX_VFS_STAGING_CACHE_TTL_MS
418 );
419 assert_eq!(clamped.pager_cache_size_kib, MAX_PAGER_CACHE_SIZE_KIB);
420 }
421}