1#![doc = include_str!("../README.md")]
2
3use std::cell::Cell;
4use std::collections::{hash_map::Entry, HashMap};
5use std::ops::Deref;
6use std::path::{Path, PathBuf};
7use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
8use std::sync::{mpsc, Arc, OnceLock};
9use std::thread::{self, JoinHandle};
10use std::time::{Duration, Instant};
11
12use lru::LruCache;
13use parking_lot::{Condvar, Mutex, MutexGuard};
14use rusqlite::OpenFlags;
15use rusqlite::{params, Connection, OptionalExtension, TransactionBehavior};
16
17use prolly::{
18 BatchOp, Cid, Error, IndexedStore, ManifestStore, ManifestStoreScan, ManifestUpdate,
19 NamedRootManifest, NodePublication, NodeStoreScan, PublicationOrigin, RootCondition,
20 RootManifest, RootWrite, Store, TransactionConflict, TransactionNodeWrite, TransactionUpdate,
21 TransactionalStore,
22};
23
24const DEFAULT_NODE_CACHE_SHARDS: usize = 16;
25static NEXT_SQLITE_READER_SLOT: AtomicUsize = AtomicUsize::new(0);
26
27thread_local! {
28 static SQLITE_READER_SLOT: Cell<usize> = const { Cell::new(usize::MAX) };
29}
30
31struct NodeReadCacheShard {
32 values: LruCache<Vec<u8>, Arc<[u8]>>,
33 retained_bytes: usize,
34 max_bytes: usize,
35}
36
37impl NodeReadCacheShard {
38 fn new(max_bytes: usize) -> Self {
39 Self {
40 values: LruCache::unbounded(),
41 retained_bytes: 0,
42 max_bytes,
43 }
44 }
45
46 fn get(&mut self, key: &[u8]) -> Option<Arc<[u8]>> {
47 self.values.get(key).cloned()
48 }
49
50 fn insert(&mut self, key: &[u8], value: Arc<[u8]>) -> usize {
51 if self.max_bytes == 0 || value.len() > self.max_bytes {
52 return 0;
53 }
54 let mut evictions = 0usize;
55 let previous = self.values.pop(key);
56 self.retained_bytes = self
57 .retained_bytes
58 .saturating_sub(previous.as_ref().map_or(0, |entry| entry.len()));
59 let target_bytes = self.max_bytes.saturating_sub(value.len());
60 while self.retained_bytes > target_bytes {
61 if let Some((_, evicted)) = self.values.pop_lru() {
62 self.retained_bytes = self.retained_bytes.saturating_sub(evicted.len());
63 evictions = evictions.saturating_add(1);
64 } else {
65 break;
66 }
67 }
68 self.retained_bytes = self.retained_bytes.saturating_add(value.len());
69 self.values.put(key.to_vec(), value);
70 evictions
71 }
72
73 fn insert_immutable(&mut self, key: &[u8], value: Arc<[u8]>) -> usize {
74 if self.max_bytes == 0 || value.len() > self.max_bytes {
75 return 0;
76 }
77 if self.get(key).is_some() {
78 return 0;
79 }
80 self.insert(key, value)
81 }
82
83 fn remove(&mut self, key: &[u8]) {
84 if let Some(value) = self.values.pop(key) {
85 self.retained_bytes = self.retained_bytes.saturating_sub(value.len());
86 }
87 }
88}
89
90struct ShardedNodeReadCache {
91 shards: Box<[Mutex<NodeReadCacheShard>]>,
92 metrics: Arc<SqliteMetricsInner>,
93}
94
95impl ShardedNodeReadCache {
96 fn new(max_bytes: usize, shard_count: usize, metrics: Arc<SqliteMetricsInner>) -> Self {
97 let shard_count = if max_bytes == 0 {
98 1
99 } else {
100 shard_count.max(1)
101 };
102 let shard_bytes = max_bytes.div_ceil(shard_count);
103 let shards = (0..shard_count)
104 .map(|_| Mutex::new(NodeReadCacheShard::new(shard_bytes)))
105 .collect::<Vec<_>>()
106 .into_boxed_slice();
107 Self { shards, metrics }
108 }
109
110 fn shard_index(&self, key: &[u8]) -> usize {
111 let hash = key
112 .iter()
113 .take(8)
114 .fold(0xcbf29ce484222325u64, |hash, byte| {
115 (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3)
116 });
117 hash as usize % self.shards.len()
118 }
119
120 fn get(&self, key: &[u8]) -> Option<Arc<[u8]>> {
121 let value = self.shards[self.shard_index(key)].lock().get(key);
122 if value.is_some() {
123 self.metrics.node_cache_hits.fetch_add(1, Ordering::Relaxed);
124 } else {
125 self.metrics
126 .node_cache_misses
127 .fetch_add(1, Ordering::Relaxed);
128 }
129 value
130 }
131
132 fn insert(&self, key: &[u8], value: Arc<[u8]>) {
133 let evictions = self.shards[self.shard_index(key)].lock().insert(key, value);
134 self.metrics
135 .node_cache_evictions
136 .fetch_add(evictions as u64, Ordering::Relaxed);
137 }
138
139 fn insert_immutable(&self, key: &[u8], value: Arc<[u8]>) {
140 let evictions = self.shards[self.shard_index(key)]
141 .lock()
142 .insert_immutable(key, value);
143 self.metrics
144 .node_cache_evictions
145 .fetch_add(evictions as u64, Ordering::Relaxed);
146 }
147
148 fn remove(&self, key: &[u8]) {
149 self.shards[self.shard_index(key)].lock().remove(key);
150 }
151
152 fn retained_bytes(&self) -> usize {
153 self.shards
154 .iter()
155 .map(|shard| shard.lock().retained_bytes)
156 .sum()
157 }
158
159 fn entries(&self) -> usize {
160 self.shards
161 .iter()
162 .map(|shard| shard.lock().values.len())
163 .sum()
164 }
165}
166
167#[derive(Default)]
168struct SqliteMetricsInner {
169 node_cache_hits: AtomicU64,
170 node_cache_misses: AtomicU64,
171 node_cache_evictions: AtomicU64,
172 sql_reads: AtomicU64,
173 write_transactions: AtomicU64,
174 published_nodes: AtomicU64,
175 grouped_publications: AtomicU64,
176 checkpoint_attempts: AtomicU64,
177 checkpoint_busy: AtomicU64,
178 checkpointed_frames: AtomicU64,
179 compressed_input_bytes: AtomicU64,
180 compressed_stored_bytes: AtomicU64,
181}
182
183#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
185pub struct SqliteStoreMetrics {
186 pub node_cache_hits: u64,
188 pub node_cache_misses: u64,
190 pub node_cache_evictions: u64,
192 pub node_cache_entries: usize,
194 pub node_cache_retained_bytes: usize,
196 pub sql_reads: u64,
198 pub write_transactions: u64,
200 pub published_nodes: u64,
202 pub grouped_publications: u64,
204 pub checkpoint_attempts: u64,
206 pub checkpoint_busy: u64,
208 pub checkpointed_frames: u64,
210 pub compressed_input_bytes: u64,
212 pub compressed_stored_bytes: u64,
214 pub sqlite_page_cache_hits: u64,
216 pub sqlite_page_cache_misses: u64,
218 pub sqlite_page_cache_writes: u64,
220 pub sqlite_page_cache_spills: u64,
222}
223
224struct OrderedBatchReadPlan<'a> {
225 unique_keys: Vec<&'a [u8]>,
226 positions: Option<Vec<usize>>,
227}
228
229impl<'a> OrderedBatchReadPlan<'a> {
230 fn new(keys: &[&'a [u8]]) -> Self {
231 let mut unique_indexes = HashMap::with_capacity(keys.len());
232 let mut unique_keys = Vec::with_capacity(keys.len());
233 let mut positions = None;
234 for key in keys {
235 match unique_indexes.entry(*key) {
236 Entry::Occupied(entry) => positions
237 .get_or_insert_with(|| (0..unique_keys.len()).collect::<Vec<_>>())
238 .push(*entry.get()),
239 Entry::Vacant(entry) => {
240 let index = unique_keys.len();
241 unique_keys.push(*key);
242 if let Some(positions) = positions.as_mut() {
243 positions.push(index);
244 }
245 entry.insert(index);
246 }
247 }
248 }
249 Self {
250 unique_keys,
251 positions,
252 }
253 }
254
255 fn unique_keys(&self) -> &[&'a [u8]] {
256 &self.unique_keys
257 }
258
259 fn expand_owned<T: Clone>(&self, values: Vec<Option<T>>) -> Vec<Option<T>> {
260 match &self.positions {
261 Some(positions) => positions
262 .iter()
263 .map(|&index| values[index].clone())
264 .collect(),
265 None => values,
266 }
267 }
268}
269
270fn cid_from_store_key(key: &[u8], context: &str) -> Result<Cid, String> {
271 let bytes: [u8; 32] = key.try_into().map_err(|_| {
272 format!(
273 "{context} key has invalid CID length {}, expected 32",
274 key.len()
275 )
276 })?;
277 Ok(Cid(bytes))
278}
279
280fn sort_cids(cids: &mut [Cid]) {
281 cids.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
282}
283
284fn sort_named_root_manifests(roots: &mut [NamedRootManifest]) {
285 roots.sort_by(|left, right| left.name.cmp(&right.name));
286}
287
288const CREATE_TABLE_SQL: &str = "\
289CREATE TABLE IF NOT EXISTS prolly_nodes (
290 cid BLOB PRIMARY KEY NOT NULL,
291 encoding INTEGER NOT NULL DEFAULT 0,
292 node BLOB NOT NULL
293);";
294
295const CREATE_HINTS_TABLE_SQL: &str = "\
296CREATE TABLE IF NOT EXISTS prolly_hints (
297 namespace BLOB NOT NULL,
298 key BLOB NOT NULL,
299 value BLOB NOT NULL,
300 PRIMARY KEY (namespace, key)
301) WITHOUT ROWID;";
302
303const CREATE_ROOTS_TABLE_SQL: &str = "\
304CREATE TABLE IF NOT EXISTS prolly_roots (
305 name BLOB PRIMARY KEY NOT NULL,
306 manifest BLOB NOT NULL
307) WITHOUT ROWID;";
308
309const SELECT_SQL: &str = "SELECT encoding, node FROM prolly_nodes WHERE cid = ?1";
310const SELECT_NODE_CIDS_SQL: &str = "SELECT cid FROM prolly_nodes ORDER BY cid";
311const UPSERT_SQL: &str = "\
312INSERT INTO prolly_nodes (cid, encoding, node)
313VALUES (?1, ?2, ?3)
314ON CONFLICT(cid) DO UPDATE SET encoding = excluded.encoding, node = excluded.node";
315const INSERT_IMMUTABLE_SQL: &str = "\
316INSERT OR IGNORE INTO prolly_nodes (cid, encoding, node)
317VALUES (?1, ?2, ?3)";
318const DELETE_SQL: &str = "DELETE FROM prolly_nodes WHERE cid = ?1";
319const UPSERT_HINT_SQL: &str = "\
320INSERT INTO prolly_hints (namespace, key, value)
321VALUES (?1, ?2, ?3)
322ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value";
323const SELECT_ROOT_SQL: &str = "SELECT manifest FROM prolly_roots WHERE name = ?1";
324const SELECT_ROOTS_SQL: &str = "SELECT name, manifest FROM prolly_roots ORDER BY name";
325const UPSERT_ROOT_SQL: &str = "\
326INSERT INTO prolly_roots (name, manifest)
327VALUES (?1, ?2)
328ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest";
329const DELETE_ROOT_SQL: &str = "DELETE FROM prolly_roots WHERE name = ?1";
330
331#[derive(Debug, Clone)]
333pub struct SqliteStoreConfig {
334 pub busy_timeout_ms: u64,
336 pub enable_wal: bool,
338 pub synchronous_normal: bool,
340 pub page_size_bytes: u32,
344 pub page_cache_size_bytes: u64,
349 pub wal_autocheckpoint_pages: u32,
354 pub node_read_cache_size_bytes: usize,
360 pub node_compression_min_bytes: usize,
365 pub mmap_size_bytes: u64,
370 pub reader_connections: usize,
374 pub primary_reader_uses_writer: bool,
378 pub node_read_cache_shards: usize,
380 pub max_batch_select_keys: usize,
382 pub background_checkpoints: bool,
385 pub checkpoint_wal_bytes: u64,
387 pub checkpoint_interval_ms: u64,
389 pub journal_size_limit_bytes: u64,
391 pub group_commit_delay_micros: u64,
394 pub group_commit_max_nodes: usize,
396}
397
398impl Default for SqliteStoreConfig {
399 fn default() -> Self {
400 Self {
401 busy_timeout_ms: 5_000,
402 enable_wal: true,
403 synchronous_normal: true,
404 page_size_bytes: 64 * 1024,
405 page_cache_size_bytes: 64 * 1024 * 1024,
406 wal_autocheckpoint_pages: 32 * 1024,
407 node_read_cache_size_bytes: 128 * 1024 * 1024,
408 node_compression_min_bytes: 8 * 1024,
409 mmap_size_bytes: 256 * 1024 * 1024,
410 reader_connections: 4,
411 primary_reader_uses_writer: true,
412 node_read_cache_shards: DEFAULT_NODE_CACHE_SHARDS,
413 max_batch_select_keys: 128,
414 background_checkpoints: true,
415 checkpoint_wal_bytes: 64 * 1024 * 1024,
416 checkpoint_interval_ms: 1_000,
417 journal_size_limit_bytes: 256 * 1024 * 1024,
418 group_commit_delay_micros: 0,
419 group_commit_max_nodes: 16_384,
420 }
421 }
422}
423
424#[derive(Debug)]
426pub struct SqliteStoreError {
427 message: String,
428 source: Option<rusqlite::Error>,
429}
430
431impl SqliteStoreError {
432 pub fn new(message: impl Into<String>) -> Self {
434 Self {
435 message: message.into(),
436 source: None,
437 }
438 }
439
440 pub fn from_sqlite(err: rusqlite::Error, context: impl Into<String>) -> Self {
442 Self {
443 message: format!("{}: {}", context.into(), err),
444 source: Some(err),
445 }
446 }
447}
448
449impl std::fmt::Display for SqliteStoreError {
450 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
451 write!(f, "SQLite error: {}", self.message)
452 }
453}
454
455impl std::error::Error for SqliteStoreError {
456 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
457 self.source
458 .as_ref()
459 .map(|e| e as &(dyn std::error::Error + 'static))
460 }
461}
462
463impl From<rusqlite::Error> for SqliteStoreError {
464 fn from(err: rusqlite::Error) -> Self {
465 Self {
466 message: err.to_string(),
467 source: Some(err),
468 }
469 }
470}
471
472enum ReadConnectionGuard<'a> {
473 Writer(MutexGuard<'a, Connection>),
474 Reader(MutexGuard<'a, Connection>),
475}
476
477impl Deref for ReadConnectionGuard<'_> {
478 type Target = Connection;
479
480 fn deref(&self) -> &Self::Target {
481 match self {
482 Self::Writer(conn) | Self::Reader(conn) => conn,
483 }
484 }
485}
486
487#[derive(Clone, Copy, Debug, Eq, PartialEq)]
489pub enum SqliteCheckpointMode {
490 Passive,
491 Full,
492 Restart,
493 Truncate,
494}
495
496impl SqliteCheckpointMode {
497 const fn sql(self) -> &'static str {
498 match self {
499 Self::Passive => "PRAGMA wal_checkpoint(PASSIVE)",
500 Self::Full => "PRAGMA wal_checkpoint(FULL)",
501 Self::Restart => "PRAGMA wal_checkpoint(RESTART)",
502 Self::Truncate => "PRAGMA wal_checkpoint(TRUNCATE)",
503 }
504 }
505}
506
507#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
509pub struct SqliteCheckpointStats {
510 pub busy: bool,
512 pub log_frames: u64,
514 pub checkpointed_frames: u64,
516}
517
518#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
520pub struct SqliteStorageStats {
521 pub page_size_bytes: u64,
523 pub page_count: u64,
525 pub freelist_pages: u64,
527 pub database_bytes: u64,
529 pub free_bytes: u64,
531 pub wal_bytes: u64,
533}
534
535enum CheckpointSignal {
536 WriteCommitted,
537 Shutdown,
538}
539
540struct CheckpointWorker {
541 sender: mpsc::Sender<CheckpointSignal>,
542 handle: Option<JoinHandle<()>>,
543}
544
545impl CheckpointWorker {
546 fn notify_write(&self) {
547 let _ = self.sender.send(CheckpointSignal::WriteCommitted);
548 }
549}
550
551impl Drop for CheckpointWorker {
552 fn drop(&mut self) {
553 let _ = self.sender.send(CheckpointSignal::Shutdown);
554 if let Some(handle) = self.handle.take() {
555 let _ = handle.join();
556 }
557 }
558}
559
560struct OwnedPublication {
561 entries: Vec<(Vec<u8>, Vec<u8>)>,
562 hint: Option<(Vec<u8>, Vec<u8>, Vec<u8>)>,
563 origin: PublicationOrigin,
564}
565
566struct PublicationCompletion {
567 result: Mutex<Option<Result<(), String>>>,
568 ready: Condvar,
569}
570
571struct PendingPublication {
572 publication: OwnedPublication,
573 completion: Arc<PublicationCompletion>,
574}
575
576#[derive(Default)]
577struct PublicationGroupState {
578 queue: Vec<PendingPublication>,
579 flushing: bool,
580}
581
582struct PublicationGroupCommit {
583 delay: Duration,
584 max_nodes: usize,
585 state: Mutex<PublicationGroupState>,
586}
587
588pub struct SqliteStore {
593 writer: Mutex<Connection>,
594 readers: Box<[Mutex<Connection>]>,
595 primary_reader_thread: OnceLock<thread::ThreadId>,
596 primary_reader_uses_writer: bool,
597 node_read_cache: ShardedNodeReadCache,
598 node_compression_min_bytes: usize,
599 max_batch_select_keys: usize,
600 metrics: Arc<SqliteMetricsInner>,
601 checkpoint_worker: Option<CheckpointWorker>,
602 group_commit: Option<PublicationGroupCommit>,
603 wal_path: Option<PathBuf>,
604}
605
606#[derive(Clone, Copy, Debug, Eq, PartialEq)]
608pub struct SqliteMainFileIdentity {
609 pub device: u64,
611 pub inode: u64,
613 pub length: u64,
615}
616
617impl SqliteStore {
618 pub fn open<P: AsRef<Path>>(path: P) -> Result<Self, SqliteStoreError> {
620 Self::open_with_config(path, SqliteStoreConfig::default())
621 }
622
623 pub fn open_with_config<P: AsRef<Path>>(
625 path: P,
626 config: SqliteStoreConfig,
627 ) -> Result<Self, SqliteStoreError> {
628 let path = path.as_ref().to_path_buf();
629 let conn = Connection::open_with_flags(
630 &path,
631 OpenFlags::SQLITE_OPEN_READ_WRITE
632 | OpenFlags::SQLITE_OPEN_CREATE
633 | OpenFlags::SQLITE_OPEN_NO_MUTEX,
634 )
635 .map_err(|e| {
636 SqliteStoreError::from_sqlite(e, format!("Failed to open database at {path:?}"))
637 })?;
638 Self::from_connection(conn, config.clone())?.attach_file_runtime(path, &config, None)
639 }
640
641 pub fn open_existing<P: AsRef<Path>>(path: P) -> Result<Self, SqliteStoreError> {
647 Self::open_existing_verified(path, |_| Ok(()))
648 }
649
650 pub fn open_existing_with_config<P: AsRef<Path>>(
652 path: P,
653 config: SqliteStoreConfig,
654 ) -> Result<Self, SqliteStoreError> {
655 Self::open_existing_verified_with_config(path, config, |_| Ok(()))
656 }
657
658 #[cfg(unix)]
661 pub fn open_existing_verified<P, F>(path: P, verifier: F) -> Result<Self, SqliteStoreError>
662 where
663 P: AsRef<Path>,
664 F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
665 {
666 Self::open_existing_verified_with_config(path, SqliteStoreConfig::default(), verifier)
667 }
668
669 #[cfg(unix)]
671 pub fn open_existing_verified_with_config<P, F>(
672 path: P,
673 config: SqliteStoreConfig,
674 verifier: F,
675 ) -> Result<Self, SqliteStoreError>
676 where
677 P: AsRef<Path>,
678 F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
679 {
680 let path = path.as_ref().to_path_buf();
681 let conn = Connection::open_with_flags(
682 &path,
683 OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX,
684 )
685 .map_err(|error| {
686 SqliteStoreError::from_sqlite(
687 error,
688 format!("Failed to open existing database at {path:?}"),
689 )
690 })?;
691 let identity = sqlite_main_file_identity(&conn)?;
692 verifier(identity)?;
693 Self::from_existing_connection(conn, config.clone())?.attach_file_runtime(
694 path,
695 &config,
696 Some(identity),
697 )
698 }
699
700 #[cfg(not(unix))]
703 pub fn open_existing_verified<P, F>(_path: P, _verifier: F) -> Result<Self, SqliteStoreError>
704 where
705 P: AsRef<Path>,
706 F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
707 {
708 Err(SqliteStoreError::new(
709 "verified existing SQLite opens are unsupported on this platform",
710 ))
711 }
712
713 #[cfg(not(unix))]
715 pub fn open_existing_verified_with_config<P, F>(
716 _path: P,
717 _config: SqliteStoreConfig,
718 _verifier: F,
719 ) -> Result<Self, SqliteStoreError>
720 where
721 P: AsRef<Path>,
722 F: FnOnce(SqliteMainFileIdentity) -> Result<(), SqliteStoreError>,
723 {
724 Err(SqliteStoreError::new(
725 "verified existing SQLite opens are unsupported on this platform",
726 ))
727 }
728
729 pub fn open_in_memory() -> Result<Self, SqliteStoreError> {
731 let conn = Connection::open_in_memory()
732 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to open in-memory database"))?;
733 Self::from_connection(conn, SqliteStoreConfig::default())
734 }
735
736 fn from_connection(
737 conn: Connection,
738 config: SqliteStoreConfig,
739 ) -> Result<Self, SqliteStoreError> {
740 Self::validate_config(&config)?;
741 conn.pragma_update(None, "page_size", config.page_size_bytes)
744 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set page_size"))?;
745 Self::apply_runtime_config(&conn, &config)?;
746 conn.execute_batch(CREATE_TABLE_SQL)
747 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize schema"))?;
748 ensure_node_encoding_column(&conn)?;
749 conn.execute_batch(CREATE_HINTS_TABLE_SQL)
750 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize hint schema"))?;
751 conn.execute_batch(CREATE_ROOTS_TABLE_SQL)
752 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to initialize root schema"))?;
753 let metrics = Arc::new(SqliteMetricsInner::default());
754 Ok(Self {
755 writer: Mutex::new(conn),
756 readers: Box::new([]),
757 primary_reader_thread: OnceLock::new(),
758 primary_reader_uses_writer: config.primary_reader_uses_writer,
759 node_read_cache: ShardedNodeReadCache::new(
760 config.node_read_cache_size_bytes,
761 config.node_read_cache_shards,
762 metrics.clone(),
763 ),
764 node_compression_min_bytes: config.node_compression_min_bytes,
765 max_batch_select_keys: config.max_batch_select_keys.max(1),
766 metrics,
767 checkpoint_worker: None,
768 group_commit: (config.group_commit_delay_micros > 0).then(|| PublicationGroupCommit {
769 delay: Duration::from_micros(config.group_commit_delay_micros),
770 max_nodes: config.group_commit_max_nodes.max(1),
771 state: Mutex::new(PublicationGroupState::default()),
772 }),
773 wal_path: None,
774 })
775 }
776
777 fn from_existing_connection(
778 conn: Connection,
779 config: SqliteStoreConfig,
780 ) -> Result<Self, SqliteStoreError> {
781 Self::validate_config(&config)?;
782 Self::apply_runtime_config(&conn, &config)?;
783 let metrics = Arc::new(SqliteMetricsInner::default());
784 Ok(Self {
785 writer: Mutex::new(conn),
786 readers: Box::new([]),
787 primary_reader_thread: OnceLock::new(),
788 primary_reader_uses_writer: config.primary_reader_uses_writer,
789 node_read_cache: ShardedNodeReadCache::new(
790 config.node_read_cache_size_bytes,
791 config.node_read_cache_shards,
792 metrics.clone(),
793 ),
794 node_compression_min_bytes: config.node_compression_min_bytes,
795 max_batch_select_keys: config.max_batch_select_keys.max(1),
796 metrics,
797 checkpoint_worker: None,
798 group_commit: (config.group_commit_delay_micros > 0).then(|| PublicationGroupCommit {
799 delay: Duration::from_micros(config.group_commit_delay_micros),
800 max_nodes: config.group_commit_max_nodes.max(1),
801 state: Mutex::new(PublicationGroupState::default()),
802 }),
803 wal_path: None,
804 })
805 }
806
807 fn validate_config(config: &SqliteStoreConfig) -> Result<(), SqliteStoreError> {
808 if !(512..=65_536).contains(&config.page_size_bytes)
809 || !config.page_size_bytes.is_power_of_two()
810 {
811 return Err(SqliteStoreError::new(
812 "page_size_bytes must be a power of two from 512 through 65536",
813 ));
814 }
815 if config.reader_connections > 64 {
816 return Err(SqliteStoreError::new(
817 "reader_connections must not exceed 64",
818 ));
819 }
820 if !(1..=256).contains(&config.node_read_cache_shards) {
821 return Err(SqliteStoreError::new(
822 "node_read_cache_shards must be from 1 through 256",
823 ));
824 }
825 if config.max_batch_select_keys == 0 {
826 return Err(SqliteStoreError::new(
827 "max_batch_select_keys must be greater than zero",
828 ));
829 }
830 if config.group_commit_delay_micros > 0 && config.group_commit_max_nodes == 0 {
831 return Err(SqliteStoreError::new(
832 "group_commit_max_nodes must be greater than zero when group commit is enabled",
833 ));
834 }
835 Ok(())
836 }
837
838 fn attach_file_runtime(
839 mut self,
840 path: PathBuf,
841 config: &SqliteStoreConfig,
842 expected_identity: Option<SqliteMainFileIdentity>,
843 ) -> Result<Self, SqliteStoreError> {
844 let mut readers = Vec::with_capacity(config.reader_connections);
845 for _ in 0..config.reader_connections {
846 let reader = Connection::open_with_flags(
847 &path,
848 OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX,
849 )
850 .map_err(|error| {
851 SqliteStoreError::from_sqlite(error, "Failed to open SQLite read connection")
852 })?;
853 #[cfg(unix)]
854 if let Some(expected) = expected_identity {
855 let actual = sqlite_main_file_identity(&reader)?;
856 if actual.device != expected.device || actual.inode != expected.inode {
857 return Err(SqliteStoreError::new(
858 "SQLite read connection resolved to a different database file",
859 ));
860 }
861 }
862 Self::apply_reader_runtime_config(&reader, config)?;
863 readers.push(Mutex::new(reader));
864 }
865 self.readers = readers.into_boxed_slice();
866 self.wal_path = Some(wal_path_for(&path));
867 if config.enable_wal && config.background_checkpoints {
868 self.writer
869 .lock()
870 .pragma_update(None, "wal_autocheckpoint", 0)
871 .map_err(|error| {
872 SqliteStoreError::from_sqlite(error, "Failed to disable writer autocheckpoint")
873 })?;
874 self.checkpoint_worker =
875 Some(spawn_checkpoint_worker(path, config, self.metrics.clone())?);
876 }
877 Ok(self)
878 }
879
880 fn apply_runtime_config(
881 conn: &Connection,
882 config: &SqliteStoreConfig,
883 ) -> Result<(), SqliteStoreError> {
884 conn.busy_timeout(Duration::from_millis(config.busy_timeout_ms))
885 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set busy timeout"))?;
886
887 if config.enable_wal {
888 conn.pragma_update(None, "journal_mode", "WAL")
889 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to enable WAL mode"))?;
890 }
891 if config.synchronous_normal {
892 conn.pragma_update(None, "synchronous", "NORMAL")
893 .map_err(|e| {
894 SqliteStoreError::from_sqlite(e, "Failed to set synchronous=NORMAL")
895 })?;
896 }
897 conn.pragma_update(None, "temp_store", "MEMORY")
898 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set temp_store=MEMORY"))?;
899 let cache_size_kib = config
900 .page_cache_size_bytes
901 .div_ceil(1024)
902 .min(i64::MAX as u64);
903 let cache_size_kib = -(cache_size_kib as i64);
904 conn.pragma_update(None, "cache_size", cache_size_kib)
905 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set cache_size"))?;
906 conn.pragma_update(None, "wal_autocheckpoint", config.wal_autocheckpoint_pages)
907 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set WAL autocheckpoint"))?;
908 conn.pragma_update(
909 None,
910 "journal_size_limit",
911 config.journal_size_limit_bytes.min(i64::MAX as u64) as i64,
912 )
913 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set journal_size_limit"))?;
914 conn.pragma_update(None, "mmap_size", config.mmap_size_bytes)
915 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set mmap_size"))?;
916 Ok(())
917 }
918
919 fn apply_reader_runtime_config(
920 conn: &Connection,
921 config: &SqliteStoreConfig,
922 ) -> Result<(), SqliteStoreError> {
923 conn.busy_timeout(Duration::from_millis(config.busy_timeout_ms))
924 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set reader busy timeout"))?;
925 conn.pragma_update(None, "query_only", true)
926 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to enable query_only"))?;
927 let reader_cache_bytes = config
928 .page_cache_size_bytes
929 .checked_div(config.reader_connections.max(1) as u64)
930 .unwrap_or(config.page_cache_size_bytes)
931 .max(1024 * 1024);
932 let cache_kib = -(reader_cache_bytes.div_ceil(1024).min(i64::MAX as u64) as i64);
933 conn.pragma_update(None, "cache_size", cache_kib)
934 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set reader cache_size"))?;
935 conn.pragma_update(None, "mmap_size", config.mmap_size_bytes)
936 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to set reader mmap_size"))?;
937 Ok(())
938 }
939
940 fn connection(&self) -> Result<MutexGuard<'_, Connection>, SqliteStoreError> {
941 Ok(self.writer.lock())
942 }
943
944 fn read_connection(&self) -> Result<ReadConnectionGuard<'_>, SqliteStoreError> {
945 if self.readers.is_empty() {
946 return Ok(ReadConnectionGuard::Writer(self.writer.lock()));
947 }
948 if self.primary_reader_uses_writer {
949 let current = thread::current().id();
950 let primary = self.primary_reader_thread.get_or_init(|| current);
951 if *primary == current {
952 return Ok(ReadConnectionGuard::Writer(self.writer.lock()));
953 }
954 }
955 let index = SQLITE_READER_SLOT.with(|slot| {
956 let current = slot.get();
957 if current == usize::MAX {
958 let assigned = NEXT_SQLITE_READER_SLOT.fetch_add(1, Ordering::Relaxed);
959 slot.set(assigned);
960 assigned
961 } else {
962 current
963 }
964 }) % self.readers.len();
965 Ok(ReadConnectionGuard::Reader(self.readers[index].lock()))
966 }
967
968 fn write_committed(&self) {
969 self.metrics
970 .write_transactions
971 .fetch_add(1, Ordering::Relaxed);
972 if let Some(worker) = &self.checkpoint_worker {
973 worker.notify_write();
974 }
975 }
976
977 fn record_compression(&self, input_len: usize, encoding: i64, stored_len: usize) {
978 if encoding == NODE_ENCODING_LZ4 {
979 self.metrics
980 .compressed_input_bytes
981 .fetch_add(input_len as u64, Ordering::Relaxed);
982 self.metrics
983 .compressed_stored_bytes
984 .fetch_add(stored_len as u64, Ordering::Relaxed);
985 }
986 }
987
988 fn own_publication(publication: NodePublication<'_>) -> OwnedPublication {
989 OwnedPublication {
990 entries: publication
991 .entries()
992 .iter()
993 .map(|(key, value)| (key.to_vec(), value.to_vec()))
994 .collect(),
995 hint: publication.hint().map(|hint| {
996 (
997 hint.namespace().to_vec(),
998 hint.key().to_vec(),
999 hint.value().to_vec(),
1000 )
1001 }),
1002 origin: publication.origin(),
1003 }
1004 }
1005
1006 fn publish_owned_batch(
1007 &self,
1008 publications: &[&OwnedPublication],
1009 ) -> Result<(), SqliteStoreError> {
1010 let mut conn = self.connection()?;
1011 let tx = conn
1012 .transaction_with_behavior(TransactionBehavior::Immediate)
1013 .map_err(|error| {
1014 SqliteStoreError::from_sqlite(error, "Failed to start grouped node publication")
1015 })?;
1016 {
1017 let mut insert = tx.prepare_cached(INSERT_IMMUTABLE_SQL).map_err(|error| {
1018 SqliteStoreError::from_sqlite(error, "Failed to prepare grouped node publication")
1019 })?;
1020 let mut compression_scratch = Vec::new();
1021 for publication in publications {
1022 let mut ordered = publication.entries.iter().collect::<Vec<_>>();
1023 ordered.sort_by(|left, right| left.0.cmp(&right.0));
1024 for (key, input) in ordered {
1025 let (encoding, stored) = encode_stored_node_into(
1026 input,
1027 &mut compression_scratch,
1028 self.node_compression_min_bytes,
1029 );
1030 self.record_compression(input.len(), encoding, stored.len());
1031 insert
1032 .execute(params![key, encoding, stored])
1033 .map_err(|error| {
1034 SqliteStoreError::from_sqlite(
1035 error,
1036 "Failed to publish grouped immutable node",
1037 )
1038 })?;
1039 }
1040 }
1041 }
1042 for publication in publications {
1043 if let Some((namespace, key, value)) = &publication.hint {
1044 tx.execute(UPSERT_HINT_SQL, params![namespace, key, value])
1045 .map_err(|error| {
1046 SqliteStoreError::from_sqlite(
1047 error,
1048 "Failed to write grouped publication hint",
1049 )
1050 })?;
1051 }
1052 }
1053 tx.commit().map_err(|error| {
1054 SqliteStoreError::from_sqlite(error, "Failed to commit grouped node publication")
1055 })?;
1056 drop(conn);
1057 self.write_committed();
1058 self.metrics.grouped_publications.fetch_add(
1059 publications.len().saturating_sub(1) as u64,
1060 Ordering::Relaxed,
1061 );
1062 self.metrics.published_nodes.fetch_add(
1063 publications
1064 .iter()
1065 .map(|publication| publication.entries.len() as u64)
1066 .sum(),
1067 Ordering::Relaxed,
1068 );
1069 for publication in publications {
1070 if !matches!(
1071 publication.origin,
1072 PublicationOrigin::TreeBuild | PublicationOrigin::Merge
1073 ) {
1074 for (key, value) in &publication.entries {
1075 self.node_read_cache
1076 .insert_immutable(key, Arc::from(value.as_slice()));
1077 }
1078 }
1079 }
1080 Ok(())
1081 }
1082
1083 fn publish_grouped(&self, publication: NodePublication<'_>) -> Result<(), SqliteStoreError> {
1084 let group = self
1085 .group_commit
1086 .as_ref()
1087 .expect("grouped publication requires configured coordinator");
1088 let completion = Arc::new(PublicationCompletion {
1089 result: Mutex::new(None),
1090 ready: Condvar::new(),
1091 });
1092 let leader = {
1093 let mut state = group.state.lock();
1094 state.queue.push(PendingPublication {
1095 publication: Self::own_publication(publication),
1096 completion: completion.clone(),
1097 });
1098 if state.flushing {
1099 false
1100 } else {
1101 state.flushing = true;
1102 true
1103 }
1104 };
1105
1106 if leader {
1107 thread::sleep(group.delay);
1108 loop {
1109 let pending = {
1110 let mut state = group.state.lock();
1111 if state.queue.is_empty() {
1112 state.flushing = false;
1113 break;
1114 }
1115 let mut nodes = 0usize;
1116 let take = state
1117 .queue
1118 .iter()
1119 .take_while(|pending| {
1120 if nodes > 0
1121 && nodes.saturating_add(pending.publication.entries.len())
1122 > group.max_nodes
1123 {
1124 return false;
1125 }
1126 nodes = nodes.saturating_add(pending.publication.entries.len());
1127 true
1128 })
1129 .count()
1130 .max(1);
1131 state.queue.drain(..take).collect::<Vec<_>>()
1132 };
1133 let publications = pending
1134 .iter()
1135 .map(|pending| &pending.publication)
1136 .collect::<Vec<_>>();
1137 let result = self
1138 .publish_owned_batch(&publications)
1139 .map_err(|error| error.to_string());
1140 for pending in pending {
1141 let mut slot = pending.completion.result.lock();
1142 *slot = Some(result.clone());
1143 pending.completion.ready.notify_one();
1144 }
1145 }
1146 }
1147
1148 let mut result = completion.result.lock();
1149 while result.is_none() {
1150 completion.ready.wait(&mut result);
1151 }
1152 result
1153 .take()
1154 .expect("group publication completion must be populated")
1155 .map_err(SqliteStoreError::new)
1156 }
1157
1158 pub fn reader_connection_count(&self) -> usize {
1160 self.readers.len()
1161 }
1162
1163 pub fn metrics(&self) -> Result<SqliteStoreMetrics, SqliteStoreError> {
1165 let mut pager = SqlitePagerMetrics::default();
1166 {
1167 let writer = self.writer.lock();
1168 pager.add_connection(&writer)?;
1169 }
1170 for reader in &self.readers {
1171 pager.add_connection(&reader.lock())?;
1172 }
1173 Ok(SqliteStoreMetrics {
1174 node_cache_hits: self.metrics.node_cache_hits.load(Ordering::Relaxed),
1175 node_cache_misses: self.metrics.node_cache_misses.load(Ordering::Relaxed),
1176 node_cache_evictions: self.metrics.node_cache_evictions.load(Ordering::Relaxed),
1177 node_cache_entries: self.node_read_cache.entries(),
1178 node_cache_retained_bytes: self.node_read_cache.retained_bytes(),
1179 sql_reads: self.metrics.sql_reads.load(Ordering::Relaxed),
1180 write_transactions: self.metrics.write_transactions.load(Ordering::Relaxed),
1181 published_nodes: self.metrics.published_nodes.load(Ordering::Relaxed),
1182 grouped_publications: self.metrics.grouped_publications.load(Ordering::Relaxed),
1183 checkpoint_attempts: self.metrics.checkpoint_attempts.load(Ordering::Relaxed),
1184 checkpoint_busy: self.metrics.checkpoint_busy.load(Ordering::Relaxed),
1185 checkpointed_frames: self.metrics.checkpointed_frames.load(Ordering::Relaxed),
1186 compressed_input_bytes: self.metrics.compressed_input_bytes.load(Ordering::Relaxed),
1187 compressed_stored_bytes: self.metrics.compressed_stored_bytes.load(Ordering::Relaxed),
1188 sqlite_page_cache_hits: pager.hits,
1189 sqlite_page_cache_misses: pager.misses,
1190 sqlite_page_cache_writes: pager.writes,
1191 sqlite_page_cache_spills: pager.spills,
1192 })
1193 }
1194
1195 pub fn checkpoint(
1197 &self,
1198 mode: SqliteCheckpointMode,
1199 ) -> Result<SqliteCheckpointStats, SqliteStoreError> {
1200 let conn = self.connection()?;
1201 let stats = run_checkpoint(&conn, mode)?;
1202 record_checkpoint_metrics(&self.metrics, stats);
1203 Ok(stats)
1204 }
1205
1206 pub fn storage_stats(&self) -> Result<SqliteStorageStats, SqliteStoreError> {
1208 let conn = self.read_connection()?;
1209 let page_size: u64 = conn
1210 .query_row("PRAGMA page_size", [], |row| row.get(0))
1211 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to read page_size"))?;
1212 let page_count: u64 = conn
1213 .query_row("PRAGMA page_count", [], |row| row.get(0))
1214 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to read page_count"))?;
1215 let freelist_pages: u64 = conn
1216 .query_row("PRAGMA freelist_count", [], |row| row.get(0))
1217 .map_err(|error| {
1218 SqliteStoreError::from_sqlite(error, "Failed to read freelist_count")
1219 })?;
1220 let wal_bytes = self
1221 .wal_path
1222 .as_ref()
1223 .and_then(|path| std::fs::metadata(path).ok())
1224 .map_or(0, |metadata| metadata.len());
1225 Ok(SqliteStorageStats {
1226 page_size_bytes: page_size,
1227 page_count,
1228 freelist_pages,
1229 database_bytes: page_size.saturating_mul(page_count),
1230 free_bytes: page_size.saturating_mul(freelist_pages),
1231 wal_bytes,
1232 })
1233 }
1234
1235 pub fn backup_to<P: AsRef<Path>>(&self, path: P) -> Result<(), SqliteStoreError> {
1237 let conn = self.read_connection()?;
1238 conn.backup(rusqlite::DatabaseName::Main, path, None)
1239 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to back up database"))
1240 }
1241
1242 pub fn compact(&self) -> Result<(), SqliteStoreError> {
1245 let conn = self.connection()?;
1246 conn.execute_batch("VACUUM")
1247 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to compact database"))?;
1248 drop(conn);
1249 self.write_committed();
1250 Ok(())
1251 }
1252
1253 pub fn optimize(&self) -> Result<(), SqliteStoreError> {
1255 let conn = self.connection()?;
1256 conn.execute_batch("PRAGMA optimize")
1257 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to optimize database"))
1258 }
1259
1260 pub fn quick_check(&self) -> Result<(), SqliteStoreError> {
1262 let conn = self.read_connection()?;
1263 let result: String = conn
1264 .query_row("PRAGMA quick_check", [], |row| row.get(0))
1265 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to check database"))?;
1266 if result == "ok" {
1267 Ok(())
1268 } else {
1269 Err(SqliteStoreError::new(format!(
1270 "SQLite quick_check failed: {result}"
1271 )))
1272 }
1273 }
1274}
1275
1276impl IndexedStore for SqliteStore {}
1277
1278#[derive(Default)]
1279struct SqlitePagerMetrics {
1280 hits: u64,
1281 misses: u64,
1282 writes: u64,
1283 spills: u64,
1284}
1285
1286impl SqlitePagerMetrics {
1287 fn add_connection(&mut self, conn: &Connection) -> Result<(), SqliteStoreError> {
1288 self.hits = self.hits.saturating_add(connection_db_status(
1289 conn,
1290 rusqlite::ffi::SQLITE_DBSTATUS_CACHE_HIT,
1291 )?);
1292 self.misses = self.misses.saturating_add(connection_db_status(
1293 conn,
1294 rusqlite::ffi::SQLITE_DBSTATUS_CACHE_MISS,
1295 )?);
1296 self.writes = self.writes.saturating_add(connection_db_status(
1297 conn,
1298 rusqlite::ffi::SQLITE_DBSTATUS_CACHE_WRITE,
1299 )?);
1300 self.spills = self.spills.saturating_add(connection_db_status(
1301 conn,
1302 rusqlite::ffi::SQLITE_DBSTATUS_CACHE_SPILL,
1303 )?);
1304 Ok(())
1305 }
1306}
1307
1308fn connection_db_status(conn: &Connection, operation: i32) -> Result<u64, SqliteStoreError> {
1309 let mut current = 0;
1310 let mut highwater = 0;
1311 let code = unsafe {
1314 rusqlite::ffi::sqlite3_db_status(conn.handle(), operation, &mut current, &mut highwater, 0)
1315 };
1316 if code == rusqlite::ffi::SQLITE_OK {
1317 Ok(current.max(0) as u64)
1318 } else {
1319 Err(SqliteStoreError::new(format!(
1320 "sqlite3_db_status failed with code {code}"
1321 )))
1322 }
1323}
1324
1325fn run_checkpoint(
1326 conn: &Connection,
1327 mode: SqliteCheckpointMode,
1328) -> Result<SqliteCheckpointStats, SqliteStoreError> {
1329 let (busy, log, checkpointed): (i64, i64, i64) = conn
1330 .query_row(mode.sql(), [], |row| {
1331 Ok((row.get(0)?, row.get(1)?, row.get(2)?))
1332 })
1333 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to checkpoint WAL"))?;
1334 Ok(SqliteCheckpointStats {
1335 busy: busy != 0,
1336 log_frames: log.max(0) as u64,
1337 checkpointed_frames: checkpointed.max(0) as u64,
1338 })
1339}
1340
1341fn record_checkpoint_metrics(metrics: &SqliteMetricsInner, stats: SqliteCheckpointStats) {
1342 metrics.checkpoint_attempts.fetch_add(1, Ordering::Relaxed);
1343 if stats.busy {
1344 metrics.checkpoint_busy.fetch_add(1, Ordering::Relaxed);
1345 }
1346 metrics
1347 .checkpointed_frames
1348 .store(stats.checkpointed_frames, Ordering::Relaxed);
1349}
1350
1351fn spawn_checkpoint_worker(
1352 path: PathBuf,
1353 config: &SqliteStoreConfig,
1354 metrics: Arc<SqliteMetricsInner>,
1355) -> Result<CheckpointWorker, SqliteStoreError> {
1356 let conn = Connection::open_with_flags(
1357 &path,
1358 OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_NO_MUTEX,
1359 )
1360 .map_err(|error| {
1361 SqliteStoreError::from_sqlite(error, "Failed to open checkpoint connection")
1362 })?;
1363 conn.busy_timeout(Duration::from_millis(config.busy_timeout_ms))
1364 .map_err(|error| {
1365 SqliteStoreError::from_sqlite(error, "Failed to configure checkpoint timeout")
1366 })?;
1367 conn.pragma_update(None, "wal_autocheckpoint", 0)
1368 .map_err(|error| {
1369 SqliteStoreError::from_sqlite(error, "Failed to disable checkpoint autocheckpoint")
1370 })?;
1371 conn.pragma_update(
1372 None,
1373 "journal_size_limit",
1374 config.journal_size_limit_bytes.min(i64::MAX as u64) as i64,
1375 )
1376 .map_err(|error| {
1377 SqliteStoreError::from_sqlite(error, "Failed to configure checkpoint journal limit")
1378 })?;
1379 let wal_path = wal_path_for(&path);
1380 let threshold = config.checkpoint_wal_bytes;
1381 let interval = Duration::from_millis(config.checkpoint_interval_ms.max(1));
1382 let (sender, receiver) = mpsc::channel();
1383 let handle = thread::Builder::new()
1384 .name("prolly-sqlite-checkpoint".to_string())
1385 .spawn(move || {
1386 let mut writes_pending = false;
1387 let mut next_inspection = Instant::now() + interval;
1388 loop {
1389 let signal = receiver
1390 .recv_timeout(next_inspection.saturating_duration_since(Instant::now()));
1391 if matches!(signal, Ok(CheckpointSignal::Shutdown)) {
1392 if let Ok(stats) = run_checkpoint(&conn, SqliteCheckpointMode::Truncate) {
1393 record_checkpoint_metrics(&metrics, stats);
1394 }
1395 break;
1396 }
1397 if matches!(signal, Ok(CheckpointSignal::WriteCommitted)) {
1398 writes_pending = true;
1399 }
1400 while let Ok(signal) = receiver.try_recv() {
1401 if matches!(signal, CheckpointSignal::Shutdown) {
1402 if let Ok(stats) = run_checkpoint(&conn, SqliteCheckpointMode::Truncate) {
1403 record_checkpoint_metrics(&metrics, stats);
1404 }
1405 return;
1406 }
1407 writes_pending = true;
1408 }
1409 if Instant::now() < next_inspection {
1410 continue;
1411 }
1412 next_inspection = Instant::now() + interval;
1413 if !writes_pending {
1414 continue;
1415 }
1416 let wal_bytes = std::fs::metadata(&wal_path)
1417 .map(|metadata| metadata.len())
1418 .unwrap_or(0);
1419 if wal_bytes >= threshold {
1420 if let Ok(stats) = run_checkpoint(&conn, SqliteCheckpointMode::Passive) {
1421 record_checkpoint_metrics(&metrics, stats);
1422 writes_pending = stats.busy || stats.checkpointed_frames < stats.log_frames;
1423 }
1424 }
1425 }
1426 })
1427 .map_err(|error| {
1428 SqliteStoreError::new(format!("failed to spawn checkpoint worker: {error}"))
1429 })?;
1430 Ok(CheckpointWorker {
1431 sender,
1432 handle: Some(handle),
1433 })
1434}
1435
1436fn wal_path_for(path: &Path) -> PathBuf {
1437 let mut wal_path = path.as_os_str().to_os_string();
1438 wal_path.push("-wal");
1439 PathBuf::from(wal_path)
1440}
1441
1442fn ensure_node_encoding_column(conn: &Connection) -> Result<(), SqliteStoreError> {
1443 let has_encoding = {
1444 let mut stmt = conn
1445 .prepare("PRAGMA table_info(prolly_nodes)")
1446 .map_err(|error| {
1447 SqliteStoreError::from_sqlite(error, "Failed to inspect node schema")
1448 })?;
1449 let columns = stmt
1450 .query_map([], |row| row.get::<_, String>(1))
1451 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to query node schema"))?;
1452 let mut found = false;
1453 for column in columns {
1454 if column.map_err(|error| {
1455 SqliteStoreError::from_sqlite(error, "Failed to read node schema")
1456 })? == "encoding"
1457 {
1458 found = true;
1459 break;
1460 }
1461 }
1462 found
1463 };
1464 if !has_encoding {
1465 conn.execute(
1466 "ALTER TABLE prolly_nodes ADD COLUMN encoding INTEGER NOT NULL DEFAULT 0",
1467 [],
1468 )
1469 .map_err(|error| {
1470 SqliteStoreError::from_sqlite(error, "Failed to migrate node encoding schema")
1471 })?;
1472 }
1473 Ok(())
1474}
1475
1476const NODE_ENCODING_RAW: i64 = 0;
1477const NODE_ENCODING_LZ4: i64 = 1;
1478fn encode_stored_node(node: &[u8], min_compressible_bytes: usize) -> (i64, Vec<u8>) {
1479 let mut scratch = Vec::new();
1480 let (encoding, stored) = encode_stored_node_into(node, &mut scratch, min_compressible_bytes);
1481 (encoding, stored.to_vec())
1482}
1483
1484fn encode_stored_node_into<'a>(
1485 node: &'a [u8],
1486 scratch: &'a mut Vec<u8>,
1487 min_compressible_bytes: usize,
1488) -> (i64, &'a [u8]) {
1489 if node.len() < min_compressible_bytes || node.len() > u32::MAX as usize {
1490 return (NODE_ENCODING_RAW, node);
1491 }
1492 let maximum = 4 + lz4_flex::block::get_maximum_output_size(node.len());
1493 scratch.clear();
1494 scratch.resize(maximum, 0);
1495 scratch[..4].copy_from_slice(&(node.len() as u32).to_le_bytes());
1496 let compressed_len = lz4_flex::block::compress_into(node, &mut scratch[4..])
1497 .expect("maximum LZ4 output size is sufficient");
1498 let stored_len = 4 + compressed_len;
1499 if stored_len >= node.len() {
1500 return (NODE_ENCODING_RAW, node);
1501 }
1502 scratch.truncate(stored_len);
1503 (NODE_ENCODING_LZ4, scratch)
1504}
1505
1506fn decode_stored_node_ref(encoding: i64, node: &[u8]) -> Result<Vec<u8>, SqliteStoreError> {
1507 match encoding {
1508 NODE_ENCODING_RAW => Ok(node.to_vec()),
1509 NODE_ENCODING_LZ4 => lz4_flex::decompress_size_prepended(node)
1510 .map_err(|error| SqliteStoreError::new(format!("failed to decompress node: {error}"))),
1511 other => Err(SqliteStoreError::new(format!(
1512 "unsupported node encoding {other}"
1513 ))),
1514 }
1515}
1516
1517fn select_nodes_ordered_unique(
1518 conn: &Connection,
1519 keys: &[&[u8]],
1520 configured_max_keys: usize,
1521) -> Result<Vec<Option<Vec<u8>>>, SqliteStoreError> {
1522 if keys.is_empty() {
1523 return Ok(Vec::new());
1524 }
1525 if keys.len() == 1 {
1526 let mut stmt = conn.prepare_cached(SELECT_SQL).map_err(|error| {
1527 SqliteStoreError::from_sqlite(error, "Failed to prepare point read")
1528 })?;
1529 let mut rows = stmt
1530 .query(params![keys[0]])
1531 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to read key"))?;
1532 let value = match rows
1533 .next()
1534 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to read key"))?
1535 {
1536 Some(row) => {
1537 let encoding = row.get::<_, i64>(0).map_err(|error| {
1538 SqliteStoreError::from_sqlite(error, "Failed to read node encoding")
1539 })?;
1540 let node_value = row.get_ref(1).map_err(|error| {
1541 SqliteStoreError::from_sqlite(error, "Failed to borrow node bytes")
1542 })?;
1543 let node = node_value.as_blob().map_err(|error| {
1544 SqliteStoreError::new(format!("Failed to borrow node bytes: {error}"))
1545 })?;
1546 Some(decode_stored_node_ref(encoding, node)?)
1547 }
1548 None => None,
1549 };
1550 return Ok(vec![value]);
1551 }
1552
1553 let positions = keys
1554 .iter()
1555 .enumerate()
1556 .map(|(index, key)| (*key, index))
1557 .collect::<HashMap<_, _>>();
1558 let mut values = vec![None; keys.len()];
1559
1560 let sqlite_max_keys = unsafe {
1563 rusqlite::ffi::sqlite3_limit(
1564 conn.handle(),
1565 rusqlite::ffi::SQLITE_LIMIT_VARIABLE_NUMBER,
1566 -1,
1567 )
1568 }
1569 .max(1) as usize;
1570 let max_keys = configured_max_keys.max(1).min(sqlite_max_keys);
1571 for chunk in keys.chunks(max_keys) {
1572 let placeholders = (1..=chunk.len())
1573 .map(|index| format!("?{index}"))
1574 .collect::<Vec<_>>()
1575 .join(",");
1576 let sql =
1577 format!("SELECT cid, encoding, node FROM prolly_nodes WHERE cid IN ({placeholders})");
1578 let mut stmt = conn.prepare_cached(&sql).map_err(|error| {
1579 SqliteStoreError::from_sqlite(error, "Failed to prepare multi-key read")
1580 })?;
1581 let mut rows = stmt
1582 .query(rusqlite::params_from_iter(chunk.iter().copied()))
1583 .map_err(|error| {
1584 SqliteStoreError::from_sqlite(error, "Failed to execute multi-key read")
1585 })?;
1586 while let Some(row) = rows.next().map_err(|error| {
1587 SqliteStoreError::from_sqlite(error, "Failed to read multi-key result")
1588 })? {
1589 let key_value = row.get_ref(0).map_err(|error| {
1590 SqliteStoreError::from_sqlite(error, "Failed to read multi-key result")
1591 })?;
1592 let key = key_value.as_blob().map_err(|error| {
1593 SqliteStoreError::new(format!("Failed to borrow node key: {error}"))
1594 })?;
1595 let encoding = row.get::<_, i64>(1).map_err(|error| {
1596 SqliteStoreError::from_sqlite(error, "Failed to read node encoding")
1597 })?;
1598 let node_value = row.get_ref(2).map_err(|error| {
1599 SqliteStoreError::from_sqlite(error, "Failed to borrow node bytes")
1600 })?;
1601 let node = node_value.as_blob().map_err(|error| {
1602 SqliteStoreError::new(format!("Failed to borrow node bytes: {error}"))
1603 })?;
1604 let index = positions.get(key).copied().ok_or_else(|| {
1605 SqliteStoreError::new("multi-key read returned an unrequested key")
1606 })?;
1607 values[index] = Some(decode_stored_node_ref(encoding, node)?);
1608 }
1609 }
1610
1611 Ok(values)
1612}
1613
1614#[cfg(unix)]
1615pub fn sqlite_main_file_identity(
1617 conn: &Connection,
1618) -> Result<SqliteMainFileIdentity, SqliteStoreError> {
1619 use std::ffi::{c_int, c_void};
1620 use std::fs::File;
1621 use std::os::fd::BorrowedFd;
1622 use std::os::unix::fs::MetadataExt;
1623
1624 #[repr(C)]
1628 struct UnixFilePrefix {
1629 methods: *const rusqlite::ffi::sqlite3_io_methods,
1630 vfs: *mut rusqlite::ffi::sqlite3_vfs,
1631 inode: *mut c_void,
1632 fd: c_int,
1633 }
1634
1635 let mut sqlite_file: *mut rusqlite::ffi::sqlite3_file = std::ptr::null_mut();
1636 let rc = unsafe {
1639 rusqlite::ffi::sqlite3_file_control(
1640 conn.handle(),
1641 c"main".as_ptr(),
1642 rusqlite::ffi::SQLITE_FCNTL_FILE_POINTER,
1643 (&mut sqlite_file as *mut *mut rusqlite::ffi::sqlite3_file).cast(),
1644 )
1645 };
1646 if rc != rusqlite::ffi::SQLITE_OK || sqlite_file.is_null() {
1647 return Err(SqliteStoreError::new(format!(
1648 "SQLite did not expose its main-file handle (code {rc})"
1649 )));
1650 }
1651 let fd = unsafe { (*(sqlite_file.cast::<UnixFilePrefix>())).fd };
1654 let borrowed = unsafe { BorrowedFd::borrow_raw(fd) };
1657 let owned = borrowed.try_clone_to_owned().map_err(|error| {
1658 SqliteStoreError::new(format!(
1659 "failed to duplicate SQLite main-file handle: {error}"
1660 ))
1661 })?;
1662 let metadata = File::from(owned).metadata().map_err(|error| {
1663 SqliteStoreError::new(format!("failed to stat SQLite main-file handle: {error}"))
1664 })?;
1665 Ok(SqliteMainFileIdentity {
1666 device: metadata.dev(),
1667 inode: metadata.ino(),
1668 length: metadata.len(),
1669 })
1670}
1671
1672impl Store for SqliteStore {
1673 type Error = SqliteStoreError;
1674
1675 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1676 if let Some(value) = self.node_read_cache.get(key) {
1677 return Ok(Some(value.as_ref().to_vec()));
1678 }
1679 let conn = self.read_connection()?;
1680 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
1681 let mut stmt = conn
1682 .prepare_cached(SELECT_SQL)
1683 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare point read"))?;
1684 let mut rows = stmt
1685 .query(params![key])
1686 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read key"))?;
1687 let Some(row) = rows
1688 .next()
1689 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read key"))?
1690 else {
1691 return Ok(None);
1692 };
1693 let encoding = row.get::<_, i64>(0).map_err(|error| {
1694 SqliteStoreError::from_sqlite(error, "Failed to read node encoding")
1695 })?;
1696 let node_value = row
1697 .get_ref(1)
1698 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to borrow node bytes"))?;
1699 let node = node_value.as_blob().map_err(|error| {
1700 SqliteStoreError::new(format!("Failed to borrow node bytes: {error}"))
1701 })?;
1702 let decoded: Arc<[u8]> = Arc::from(decode_stored_node_ref(encoding, node)?);
1703 self.node_read_cache.insert(key, decoded.clone());
1704 Ok(Some(decoded.as_ref().to_vec()))
1705 }
1706
1707 fn get_shared(&self, key: &[u8]) -> Result<Option<Arc<[u8]>>, Self::Error> {
1708 if let Some(value) = self.node_read_cache.get(key) {
1709 return Ok(Some(value));
1710 }
1711 let conn = self.read_connection()?;
1712 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
1713 let mut stmt = conn
1714 .prepare_cached(SELECT_SQL)
1715 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare shared point read"))?;
1716 let mut rows = stmt
1717 .query(params![key])
1718 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read shared key"))?;
1719 let Some(row) = rows
1720 .next()
1721 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read shared key"))?
1722 else {
1723 return Ok(None);
1724 };
1725 let encoding = row.get::<_, i64>(0).map_err(|error| {
1726 SqliteStoreError::from_sqlite(error, "Failed to read node encoding")
1727 })?;
1728 let node_value = row
1729 .get_ref(1)
1730 .map_err(|error| SqliteStoreError::from_sqlite(error, "Failed to borrow node bytes"))?;
1731 let node = node_value.as_blob().map_err(|error| {
1732 SqliteStoreError::new(format!("Failed to borrow node bytes: {error}"))
1733 })?;
1734 let decoded: Arc<[u8]> = Arc::from(decode_stored_node_ref(encoding, node)?);
1735 self.node_read_cache.insert(key, decoded.clone());
1736 Ok(Some(decoded))
1737 }
1738
1739 fn has_native_shared_reads(&self) -> bool {
1740 true
1741 }
1742
1743 fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1744 let conn = self.connection()?;
1745 let (encoding, stored) = encode_stored_node(value, self.node_compression_min_bytes);
1746 self.record_compression(value.len(), encoding, stored.len());
1747 conn.execute(UPSERT_SQL, params![key, encoding, stored])
1748 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write key"))?;
1749 drop(conn);
1750 self.write_committed();
1751 self.node_read_cache.insert(key, Arc::from(value));
1752 Ok(())
1753 }
1754
1755 fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
1756 let conn = self.connection()?;
1757 conn.execute(DELETE_SQL, params![key])
1758 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to delete key"))?;
1759 drop(conn);
1760 self.write_committed();
1761 self.node_read_cache.remove(key);
1762 Ok(())
1763 }
1764
1765 fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
1766 let mut conn = self.connection()?;
1767 let tx = conn
1768 .transaction_with_behavior(TransactionBehavior::Immediate)
1769 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
1770
1771 {
1772 let mut upsert = tx
1773 .prepare_cached(UPSERT_SQL)
1774 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare batch write"))?;
1775 let mut delete = tx
1776 .prepare_cached(DELETE_SQL)
1777 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare batch delete"))?;
1778 let mut compression_scratch = Vec::new();
1779
1780 for op in ops {
1781 match op {
1782 BatchOp::Upsert { key, value } => {
1783 let input_len = value.len();
1784 let (encoding, stored) = encode_stored_node_into(
1785 value,
1786 &mut compression_scratch,
1787 self.node_compression_min_bytes,
1788 );
1789 self.record_compression(input_len, encoding, stored.len());
1790 upsert
1791 .execute(params![key, encoding, stored])
1792 .map_err(|e| {
1793 SqliteStoreError::from_sqlite(e, "Failed to write key in batch")
1794 })?;
1795 }
1796 BatchOp::Delete { key } => {
1797 delete.execute(params![key]).map_err(|e| {
1798 SqliteStoreError::from_sqlite(e, "Failed to delete key in batch")
1799 })?;
1800 }
1801 }
1802 }
1803 }
1804
1805 tx.commit()
1806 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))?;
1807 drop(conn);
1808 self.write_committed();
1809 for op in ops {
1810 match op {
1811 BatchOp::Upsert { key, .. } | BatchOp::Delete { key } => {
1812 self.node_read_cache.remove(key)
1813 }
1814 }
1815 }
1816 Ok(())
1817 }
1818
1819 fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
1820 let plan = OrderedBatchReadPlan::new(keys);
1821 let values = self.batch_get_shared_ordered_unique(plan.unique_keys())?;
1822 let mut results = HashMap::with_capacity(plan.unique_keys().len());
1823 for (key, value) in plan.unique_keys().iter().zip(values) {
1824 if let Some(value) = value {
1825 results.insert(key.to_vec(), value.as_ref().to_vec());
1826 }
1827 }
1828
1829 Ok(results)
1830 }
1831
1832 fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1833 let plan = OrderedBatchReadPlan::new(keys);
1834 let unique_values = self
1835 .batch_get_shared_ordered_unique(plan.unique_keys())?
1836 .into_iter()
1837 .map(|value| value.map(|value| value.as_ref().to_vec()))
1838 .collect();
1839 Ok(plan.expand_owned(unique_values))
1840 }
1841
1842 fn batch_get_ordered_unique(
1843 &self,
1844 keys: &[&[u8]],
1845 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
1846 self.batch_get_shared_ordered_unique(keys).map(|values| {
1847 values
1848 .into_iter()
1849 .map(|value| value.map(|value| value.as_ref().to_vec()))
1850 .collect()
1851 })
1852 }
1853
1854 fn batch_get_shared_ordered_unique(
1855 &self,
1856 keys: &[&[u8]],
1857 ) -> Result<Vec<Option<Arc<[u8]>>>, Self::Error> {
1858 let mut values = vec![None; keys.len()];
1859 let mut missing = Vec::new();
1860 for (position, key) in keys.iter().enumerate() {
1861 match self.node_read_cache.get(key) {
1862 Some(value) => values[position] = Some(value),
1863 None => missing.push((position, *key)),
1864 }
1865 }
1866 if missing.is_empty() {
1867 return Ok(values);
1868 }
1869 let missing_keys = missing.iter().map(|(_, key)| *key).collect::<Vec<_>>();
1870 let conn = self.read_connection()?;
1871 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
1872 let loaded = select_nodes_ordered_unique(&conn, &missing_keys, self.max_batch_select_keys)?;
1873 drop(conn);
1874 for ((position, key), loaded) in missing.into_iter().zip(loaded) {
1875 if let Some(loaded) = loaded {
1876 let loaded: Arc<[u8]> = Arc::from(loaded);
1877 self.node_read_cache.insert(key, loaded.clone());
1878 values[position] = Some(loaded);
1879 }
1880 }
1881 Ok(values)
1882 }
1883
1884 fn prefers_batch_reads(&self) -> bool {
1885 true
1886 }
1887
1888 fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
1889 let mut conn = self.connection()?;
1890 let tx = conn
1891 .transaction_with_behavior(TransactionBehavior::Immediate)
1892 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
1893
1894 {
1895 let mut stmt = tx.prepare_cached(UPSERT_SQL).map_err(|e| {
1896 SqliteStoreError::from_sqlite(e, "Failed to prepare batch_put write")
1897 })?;
1898 let mut ordered = entries.iter().collect::<Vec<_>>();
1899 ordered.sort_by(|left, right| left.0.cmp(right.0));
1900 let mut compression_scratch = Vec::new();
1901 for &&(key, value) in &ordered {
1902 let input_len = value.len();
1903 let (encoding, stored) = encode_stored_node_into(
1904 value,
1905 &mut compression_scratch,
1906 self.node_compression_min_bytes,
1907 );
1908 self.record_compression(input_len, encoding, stored.len());
1909 stmt.execute(params![key, encoding, stored]).map_err(|e| {
1910 SqliteStoreError::from_sqlite(e, "Failed to write key in batch_put")
1911 })?;
1912 }
1913 }
1914
1915 tx.commit()
1916 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))?;
1917 drop(conn);
1918 self.write_committed();
1919 for &(key, _) in entries {
1920 self.node_read_cache.remove(key);
1921 }
1922 Ok(())
1923 }
1924
1925 fn supports_hints(&self) -> bool {
1926 true
1927 }
1928
1929 fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
1930 let conn = self.read_connection()?;
1931 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
1932 conn.query_row(
1933 "SELECT value FROM prolly_hints WHERE namespace = ?1 AND key = ?2",
1934 params![namespace, key],
1935 |row| row.get(0),
1936 )
1937 .optional()
1938 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read hint"))
1939 }
1940
1941 fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
1942 let conn = self.connection()?;
1943 conn.execute(
1944 "\
1945 INSERT INTO prolly_hints (namespace, key, value) \
1946 VALUES (?1, ?2, ?3) \
1947 ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
1948 params![namespace, key, value],
1949 )
1950 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write hint"))?;
1951 drop(conn);
1952 self.write_committed();
1953 Ok(())
1954 }
1955
1956 fn batch_put_with_hint(
1957 &self,
1958 entries: &[(&[u8], &[u8])],
1959 namespace: &[u8],
1960 key: &[u8],
1961 value: &[u8],
1962 ) -> Result<(), Self::Error> {
1963 let mut conn = self.connection()?;
1964 let tx = conn
1965 .transaction_with_behavior(TransactionBehavior::Immediate)
1966 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start transaction"))?;
1967
1968 {
1969 let mut upsert_node = tx.prepare_cached(UPSERT_SQL).map_err(|e| {
1970 SqliteStoreError::from_sqlite(e, "Failed to prepare batch_put write")
1971 })?;
1972 let mut ordered = entries.iter().collect::<Vec<_>>();
1973 ordered.sort_by(|left, right| left.0.cmp(right.0));
1974 let mut compression_scratch = Vec::new();
1975 for &&(key, value) in &ordered {
1976 let input_len = value.len();
1977 let (encoding, stored) = encode_stored_node_into(
1978 value,
1979 &mut compression_scratch,
1980 self.node_compression_min_bytes,
1981 );
1982 self.record_compression(input_len, encoding, stored.len());
1983 upsert_node
1984 .execute(params![key, encoding, stored])
1985 .map_err(|e| {
1986 SqliteStoreError::from_sqlite(e, "Failed to write key in batch_put")
1987 })?;
1988 }
1989 }
1990
1991 tx.execute(
1992 "\
1993 INSERT INTO prolly_hints (namespace, key, value) \
1994 VALUES (?1, ?2, ?3) \
1995 ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
1996 params![namespace, key, value],
1997 )
1998 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write hint in batch_put"))?;
1999
2000 tx.commit()
2001 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit transaction"))?;
2002 drop(conn);
2003 self.write_committed();
2004 for &(key, _) in entries {
2005 self.node_read_cache.remove(key);
2006 }
2007 Ok(())
2008 }
2009
2010 fn publish_nodes(&self, publication: NodePublication<'_>) -> Result<(), Self::Error> {
2011 if self.group_commit.is_some() {
2012 return self.publish_grouped(publication);
2013 }
2014 let mut conn = self.connection()?;
2015 let tx = conn
2016 .transaction_with_behavior(TransactionBehavior::Immediate)
2017 .map_err(|error| {
2018 SqliteStoreError::from_sqlite(error, "Failed to start node publication")
2019 })?;
2020 {
2021 let mut insert = tx.prepare_cached(INSERT_IMMUTABLE_SQL).map_err(|error| {
2025 SqliteStoreError::from_sqlite(error, "Failed to prepare node publication")
2026 })?;
2027 let mut ordered = publication.entries().iter().collect::<Vec<_>>();
2028 ordered.sort_by(|left, right| left.0.cmp(right.0));
2029 let mut compression_scratch = Vec::new();
2030 for &&(key, value) in &ordered {
2031 let input_len = value.len();
2032 let (encoding, stored) = encode_stored_node_into(
2033 value,
2034 &mut compression_scratch,
2035 self.node_compression_min_bytes,
2036 );
2037 self.record_compression(input_len, encoding, stored.len());
2038 insert
2039 .execute(params![key, encoding, stored])
2040 .map_err(|error| {
2041 SqliteStoreError::from_sqlite(error, "Failed to publish immutable node")
2042 })?;
2043 }
2044 }
2045 if let Some(hint) = publication.hint() {
2046 tx.execute(
2047 UPSERT_HINT_SQL,
2048 params![hint.namespace(), hint.key(), hint.value()],
2049 )
2050 .map_err(|error| {
2051 SqliteStoreError::from_sqlite(error, "Failed to write publication hint")
2052 })?;
2053 }
2054 tx.commit().map_err(|error| {
2055 SqliteStoreError::from_sqlite(error, "Failed to commit node publication")
2056 })?;
2057 drop(conn);
2058 self.write_committed();
2059 self.metrics
2060 .published_nodes
2061 .fetch_add(publication.entries().len() as u64, Ordering::Relaxed);
2062 if !matches!(
2066 publication.origin(),
2067 PublicationOrigin::TreeBuild | PublicationOrigin::Merge
2068 ) {
2069 for &(key, value) in publication.entries() {
2070 self.node_read_cache.insert_immutable(key, Arc::from(value));
2071 }
2072 }
2073 Ok(())
2074 }
2075}
2076
2077impl NodeStoreScan for SqliteStore {
2078 type Error = SqliteStoreError;
2079
2080 fn list_node_cids(&self) -> Result<Vec<Cid>, Self::Error> {
2081 let conn = self.read_connection()?;
2082 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
2083 let mut stmt = conn
2084 .prepare_cached(SELECT_NODE_CIDS_SQL)
2085 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to prepare node CID listing"))?;
2086 let rows = stmt
2087 .query_map([], |row| row.get::<_, Vec<u8>>(0))
2088 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to list node CIDs"))?;
2089
2090 let mut cids = Vec::new();
2091 for row in rows {
2092 let key = row
2093 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read listed node CID"))?;
2094 cids.push(cid_from_store_key(&key, "SQLite node").map_err(SqliteStoreError::new)?);
2095 }
2096 sort_cids(&mut cids);
2097 Ok(cids)
2098 }
2099}
2100
2101impl ManifestStore for SqliteStore {
2102 type Error = SqliteStoreError;
2103
2104 fn get_root(&self, name: &[u8]) -> Result<Option<RootManifest>, Self::Error> {
2105 let conn = self.read_connection()?;
2106 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
2107 let bytes = conn
2108 .query_row(SELECT_ROOT_SQL, params![name], |row| row.get(0))
2109 .optional()
2110 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read root manifest"))?;
2111 decode_root_manifest(bytes)
2112 }
2113
2114 fn put_root(&self, name: &[u8], manifest: &RootManifest) -> Result<(), Self::Error> {
2115 let conn = self.connection()?;
2116 let bytes = encode_root_manifest(manifest)?;
2117 conn.execute(UPSERT_ROOT_SQL, params![name, bytes])
2118 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to write root manifest"))?;
2119 drop(conn);
2120 self.write_committed();
2121 Ok(())
2122 }
2123
2124 fn delete_root(&self, name: &[u8]) -> Result<(), Self::Error> {
2125 let conn = self.connection()?;
2126 conn.execute(DELETE_ROOT_SQL, params![name])
2127 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to delete root manifest"))?;
2128 drop(conn);
2129 self.write_committed();
2130 Ok(())
2131 }
2132
2133 fn compare_and_swap_root(
2134 &self,
2135 name: &[u8],
2136 expected: Option<&RootManifest>,
2137 new: Option<&RootManifest>,
2138 ) -> Result<ManifestUpdate, Self::Error> {
2139 let expected_bytes = expected.map(encode_root_manifest).transpose()?;
2140 let new_bytes = new.map(encode_root_manifest).transpose()?;
2141
2142 let mut conn = self.connection()?;
2143 let tx = conn
2144 .transaction_with_behavior(TransactionBehavior::Immediate)
2145 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to start root transaction"))?;
2146
2147 let current_bytes = tx
2148 .query_row(SELECT_ROOT_SQL, params![name], |row| row.get(0))
2149 .optional()
2150 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to read root manifest"))?;
2151
2152 if current_bytes.as_deref() != expected_bytes.as_deref() {
2153 return Ok(ManifestUpdate::Conflict {
2154 current: decode_root_manifest(current_bytes)?,
2155 });
2156 }
2157
2158 match new_bytes {
2159 Some(bytes) => {
2160 tx.execute(UPSERT_ROOT_SQL, params![name, bytes])
2161 .map_err(|e| {
2162 SqliteStoreError::from_sqlite(e, "Failed to write root manifest")
2163 })?;
2164 }
2165 None => {
2166 tx.execute(DELETE_ROOT_SQL, params![name]).map_err(|e| {
2167 SqliteStoreError::from_sqlite(e, "Failed to delete root manifest")
2168 })?;
2169 }
2170 }
2171
2172 tx.commit()
2173 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to commit root transaction"))?;
2174 drop(conn);
2175 self.write_committed();
2176 Ok(ManifestUpdate::Applied)
2177 }
2178}
2179
2180impl ManifestStoreScan for SqliteStore {
2181 fn list_roots(&self) -> Result<Vec<NamedRootManifest>, Self::Error> {
2182 let conn = self.read_connection()?;
2183 self.metrics.sql_reads.fetch_add(1, Ordering::Relaxed);
2184 let mut stmt = conn.prepare_cached(SELECT_ROOTS_SQL).map_err(|e| {
2185 SqliteStoreError::from_sqlite(e, "Failed to prepare root manifest listing")
2186 })?;
2187 let rows = stmt
2188 .query_map([], |row| {
2189 Ok((row.get::<_, Vec<u8>>(0)?, row.get::<_, Vec<u8>>(1)?))
2190 })
2191 .map_err(|e| SqliteStoreError::from_sqlite(e, "Failed to list root manifests"))?;
2192
2193 let mut roots = Vec::new();
2194 for row in rows {
2195 let (name, bytes) = row.map_err(|e| {
2196 SqliteStoreError::from_sqlite(e, "Failed to read listed root manifest")
2197 })?;
2198 let manifest = RootManifest::from_bytes(&bytes)
2199 .map_err(|err| SqliteStoreError::new(err.to_string()))?;
2200 roots.push(NamedRootManifest::new(name, manifest));
2201 }
2202 sort_named_root_manifests(&mut roots);
2203 Ok(roots)
2204 }
2205}
2206
2207impl TransactionalStore for SqliteStore {
2208 fn supports_transactions(&self) -> bool {
2209 true
2210 }
2211
2212 fn commit_transaction(
2213 &self,
2214 node_writes: &[TransactionNodeWrite],
2215 root_conditions: &[RootCondition],
2216 root_writes: &[RootWrite],
2217 ) -> Result<TransactionUpdate, Error> {
2218 let mut conn = self
2219 .connection()
2220 .map_err(|err| Error::Store(Box::new(err)))?;
2221 let tx = conn
2222 .transaction_with_behavior(TransactionBehavior::Immediate)
2223 .map_err(|err| {
2224 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2225 err,
2226 "Failed to start transaction commit",
2227 )))
2228 })?;
2229
2230 for condition in root_conditions {
2231 let current_bytes = tx
2232 .query_row(SELECT_ROOT_SQL, params![condition.name], |row| row.get(0))
2233 .optional()
2234 .map_err(|err| {
2235 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2236 err,
2237 "Failed to read root manifest during transaction commit",
2238 )))
2239 })?;
2240 let current =
2241 decode_root_manifest(current_bytes).map_err(|err| Error::Store(Box::new(err)))?;
2242 if current != condition.expected {
2243 return Ok(TransactionUpdate::Conflict(Box::new(
2244 TransactionConflict::new(
2245 condition.name.clone(),
2246 condition.expected.clone(),
2247 current,
2248 ),
2249 )));
2250 }
2251 }
2252
2253 {
2254 let mut upsert_node = tx.prepare_cached(UPSERT_SQL).map_err(|err| {
2255 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2256 err,
2257 "Failed to prepare transaction node write",
2258 )))
2259 })?;
2260 let mut delete_node = tx.prepare_cached(DELETE_SQL).map_err(|err| {
2261 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2262 err,
2263 "Failed to prepare transaction node delete",
2264 )))
2265 })?;
2266 for write in node_writes {
2267 match write {
2268 TransactionNodeWrite::Upsert { key, value } => {
2269 let (encoding, stored) =
2270 encode_stored_node(value, self.node_compression_min_bytes);
2271 self.record_compression(value.len(), encoding, stored.len());
2272 upsert_node
2273 .execute(params![key, encoding, stored])
2274 .map_err(|err| {
2275 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2276 err,
2277 "Failed to write node during transaction commit",
2278 )))
2279 })?;
2280 }
2281 TransactionNodeWrite::Delete { key } => {
2282 delete_node.execute(params![key]).map_err(|err| {
2283 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2284 err,
2285 "Failed to delete node during transaction commit",
2286 )))
2287 })?;
2288 }
2289 }
2290 }
2291 }
2292
2293 {
2294 let mut upsert_root = tx.prepare_cached(UPSERT_ROOT_SQL).map_err(|err| {
2295 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2296 err,
2297 "Failed to prepare transaction root write",
2298 )))
2299 })?;
2300 let mut delete_root = tx.prepare_cached(DELETE_ROOT_SQL).map_err(|err| {
2301 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2302 err,
2303 "Failed to prepare transaction root delete",
2304 )))
2305 })?;
2306 for write in root_writes {
2307 match write {
2308 RootWrite::Put { name, manifest } => {
2309 let bytes = encode_root_manifest(manifest)
2310 .map_err(|err| Error::Store(Box::new(err)))?;
2311 upsert_root.execute(params![name, bytes]).map_err(|err| {
2312 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2313 err,
2314 "Failed to write root during transaction commit",
2315 )))
2316 })?;
2317 }
2318 RootWrite::Delete { name } => {
2319 delete_root.execute(params![name]).map_err(|err| {
2320 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2321 err,
2322 "Failed to delete root during transaction commit",
2323 )))
2324 })?;
2325 }
2326 }
2327 }
2328 }
2329
2330 tx.commit().map_err(|err| {
2331 Error::Store(Box::new(SqliteStoreError::from_sqlite(
2332 err,
2333 "Failed to commit transaction",
2334 )))
2335 })?;
2336 drop(conn);
2337 self.write_committed();
2338 for write in node_writes {
2339 match write {
2340 TransactionNodeWrite::Upsert { key, .. } | TransactionNodeWrite::Delete { key } => {
2341 self.node_read_cache.remove(key)
2342 }
2343 }
2344 }
2345 Ok(TransactionUpdate::Applied {
2346 nodes_written: node_writes.len(),
2347 roots_written: root_writes.len(),
2348 })
2349 }
2350}
2351
2352fn encode_root_manifest(manifest: &RootManifest) -> Result<Vec<u8>, SqliteStoreError> {
2353 manifest
2354 .to_bytes()
2355 .map_err(|e| SqliteStoreError::new(format!("failed to encode root manifest: {e}")))
2356}
2357
2358fn decode_root_manifest(bytes: Option<Vec<u8>>) -> Result<Option<RootManifest>, SqliteStoreError> {
2359 bytes
2360 .as_deref()
2361 .map(RootManifest::from_bytes)
2362 .transpose()
2363 .map_err(|e| SqliteStoreError::new(format!("failed to decode root manifest: {e}")))
2364}
2365
2366#[cfg(test)]
2367mod tests {
2368 use super::*;
2369
2370 #[test]
2371 fn indexed_map_cas_coordinates_separate_handles() {
2372 let directory = tempfile::tempdir().unwrap();
2373 let path = directory.path().join("coordinated.db");
2374 let config = SqliteStoreConfig {
2375 synchronous_normal: false,
2376 background_checkpoints: false,
2377 ..SqliteStoreConfig::default()
2378 };
2379 let first = SqliteStore::open_with_config(&path, config.clone()).unwrap();
2380 let second = SqliteStore::open_existing_with_config(&path, config).unwrap();
2381 let manifest =
2382 RootManifest::new(Some(Cid::from_bytes(b"one")), prolly::Config::default());
2383 assert!(first
2384 .compare_and_swap_root(b"indexed", None, Some(&manifest))
2385 .unwrap()
2386 .is_applied());
2387 assert!(matches!(
2388 second
2389 .compare_and_swap_root(b"indexed", None, Some(&manifest))
2390 .unwrap(),
2391 ManifestUpdate::Conflict { .. }
2392 ));
2393 }
2394
2395 #[test]
2396 fn indexed_store_passes_shared_contract() {
2397 let directory = tempfile::tempdir().unwrap();
2398 let config = SqliteStoreConfig {
2399 synchronous_normal: false,
2400 background_checkpoints: false,
2401 ..SqliteStoreConfig::default()
2402 };
2403 let store =
2404 SqliteStore::open_with_config(directory.path().join("indexed.db"), config).unwrap();
2405 prolly_store_test::assert_indexed_store(store);
2406 }
2407
2408 #[test]
2409 fn sqlite_store_put_get_delete() {
2410 let store = SqliteStore::open_in_memory().unwrap();
2411
2412 store.put(b"key", b"value").unwrap();
2413 assert_eq!(store.get(b"key").unwrap(), Some(b"value".to_vec()));
2414
2415 store.delete(b"key").unwrap();
2416 assert_eq!(store.get(b"key").unwrap(), None);
2417 }
2418
2419 #[test]
2420 fn sqlite_store_reuses_shared_node_reads() {
2421 let store = SqliteStore::open_in_memory().unwrap();
2422 store.put(b"key", b"value").unwrap();
2423
2424 let first = store.get_shared(b"key").unwrap().unwrap();
2425 let second = store.get_shared(b"key").unwrap().unwrap();
2426
2427 assert!(Arc::ptr_eq(&first, &second));
2428 }
2429
2430 #[test]
2431 fn sqlite_store_applies_publication_pragmas_and_rowid_schema() {
2432 let store = SqliteStore::open_in_memory().unwrap();
2433 let conn = store.connection().unwrap();
2434 let page_size: u32 = conn
2435 .query_row("PRAGMA page_size", [], |row| row.get(0))
2436 .unwrap();
2437 let cache_size: i64 = conn
2438 .query_row("PRAGMA cache_size", [], |row| row.get(0))
2439 .unwrap();
2440 let wal_autocheckpoint: u32 = conn
2441 .query_row("PRAGMA wal_autocheckpoint", [], |row| row.get(0))
2442 .unwrap();
2443 let schema: String = conn
2444 .query_row(
2445 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'prolly_nodes'",
2446 [],
2447 |row| row.get(0),
2448 )
2449 .unwrap();
2450
2451 assert_eq!(page_size, 64 * 1024);
2452 assert_eq!(cache_size, -64 * 1024);
2453 assert_eq!(wal_autocheckpoint, 32 * 1024);
2454 assert!(!schema.to_ascii_uppercase().contains("WITHOUT ROWID"));
2455 }
2456
2457 #[test]
2458 fn sqlite_store_rejects_invalid_page_size() {
2459 let config = SqliteStoreConfig {
2460 page_size_bytes: 1_000,
2461 ..SqliteStoreConfig::default()
2462 };
2463 let error = SqliteStore::from_connection(Connection::open_in_memory().unwrap(), config)
2464 .err()
2465 .expect("invalid page size must fail");
2466 assert!(error.to_string().contains("page_size_bytes"));
2467 }
2468
2469 #[test]
2470 fn sqlite_store_batch_is_order_preserving_for_reads() {
2471 let store = SqliteStore::open_in_memory().unwrap();
2472 let ops = vec![
2473 BatchOp::Upsert {
2474 key: b"a",
2475 value: b"1",
2476 },
2477 BatchOp::Upsert {
2478 key: b"b",
2479 value: b"2",
2480 },
2481 BatchOp::Upsert {
2482 key: b"c",
2483 value: b"3",
2484 },
2485 ];
2486
2487 store.batch(&ops).unwrap();
2488
2489 let keys: Vec<&[u8]> = vec![b"c", b"missing", b"a", b"c", b"missing", b"b"];
2490 assert_eq!(
2491 store.batch_get_ordered(&keys).unwrap(),
2492 vec![
2493 Some(b"3".to_vec()),
2494 None,
2495 Some(b"1".to_vec()),
2496 Some(b"3".to_vec()),
2497 None,
2498 Some(b"2".to_vec())
2499 ]
2500 );
2501 }
2502
2503 #[test]
2504 fn sqlite_store_batch_put_updates_existing_keys() {
2505 let store = SqliteStore::open_in_memory().unwrap();
2506
2507 store.put(b"a", b"old").unwrap();
2508 store
2509 .batch_put(&[(b"a".as_slice(), b"new".as_slice()), (b"b", b"2")])
2510 .unwrap();
2511
2512 assert_eq!(store.get(b"a").unwrap(), Some(b"new".to_vec()));
2513 assert_eq!(store.get(b"b").unwrap(), Some(b"2".to_vec()));
2514 }
2515
2516 #[test]
2517 fn sqlite_store_persists_hints_separately_from_nodes() {
2518 let store = SqliteStore::open_in_memory().unwrap();
2519
2520 store.put_hint(b"rightmost", b"root", b"hint-v1").unwrap();
2521 assert_eq!(
2522 store.get_hint(b"rightmost", b"root").unwrap(),
2523 Some(b"hint-v1".to_vec())
2524 );
2525 assert_eq!(store.get_hint(b"rightmost", b"missing").unwrap(), None);
2526 assert_eq!(store.get(b"root").unwrap(), None);
2527
2528 store.put_hint(b"rightmost", b"root", b"hint-v2").unwrap();
2529 assert_eq!(
2530 store.get_hint(b"rightmost", b"root").unwrap(),
2531 Some(b"hint-v2".to_vec())
2532 );
2533 }
2534
2535 #[test]
2536 fn sqlite_store_compresses_repetitive_nodes_transparently() {
2537 let store = SqliteStore::open_in_memory().unwrap();
2538 let node = vec![b'x'; 16 * 1024];
2539
2540 store.put(b"compressed", &node).unwrap();
2541
2542 let conn = store.connection().unwrap();
2543 let (encoding, stored_bytes): (i64, usize) = conn
2544 .query_row(
2545 "SELECT encoding, length(node) FROM prolly_nodes WHERE cid = ?1",
2546 params![b"compressed"],
2547 |row| Ok((row.get(0)?, row.get(1)?)),
2548 )
2549 .unwrap();
2550 drop(conn);
2551 assert_eq!(encoding, NODE_ENCODING_LZ4);
2552 assert!(stored_bytes < node.len());
2553 assert_eq!(store.get(b"compressed").unwrap(), Some(node));
2554 }
2555
2556 #[test]
2557 fn sqlite_store_honors_node_compression_threshold() {
2558 let config = SqliteStoreConfig {
2559 node_compression_min_bytes: 32 * 1024,
2560 ..SqliteStoreConfig::default()
2561 };
2562 let store =
2563 SqliteStore::from_connection(Connection::open_in_memory().unwrap(), config).unwrap();
2564 let node = vec![b'x'; 16 * 1024];
2565
2566 store.put(b"raw", &node).unwrap();
2567
2568 let conn = store.connection().unwrap();
2569 let encoding: i64 = conn
2570 .query_row(
2571 "SELECT encoding FROM prolly_nodes WHERE cid = ?1",
2572 params![b"raw"],
2573 |row| row.get(0),
2574 )
2575 .unwrap();
2576 assert_eq!(encoding, NODE_ENCODING_RAW);
2577 }
2578
2579 #[test]
2580 fn sqlite_store_migrates_legacy_raw_node_tables() {
2581 let conn = Connection::open_in_memory().unwrap();
2582 conn.execute_batch(
2583 "CREATE TABLE prolly_nodes (
2584 cid BLOB PRIMARY KEY NOT NULL,
2585 node BLOB NOT NULL
2586 ) WITHOUT ROWID;
2587 INSERT INTO prolly_nodes (cid, node) VALUES (x'6c6567616379', x'726177');",
2588 )
2589 .unwrap();
2590
2591 let store = SqliteStore::from_connection(conn, SqliteStoreConfig::default()).unwrap();
2592
2593 assert_eq!(store.get(b"legacy").unwrap(), Some(b"raw".to_vec()));
2594 let conn = store.connection().unwrap();
2595 let encoding: i64 = conn
2596 .query_row(
2597 "SELECT encoding FROM prolly_nodes WHERE cid = ?1",
2598 params![b"legacy"],
2599 |row| row.get(0),
2600 )
2601 .unwrap();
2602 assert_eq!(encoding, NODE_ENCODING_RAW);
2603 }
2604
2605 #[test]
2606 fn file_store_opens_reader_pool_and_disables_writer_autocheckpoint() {
2607 let directory = tempfile::tempdir().unwrap();
2608 let config = SqliteStoreConfig {
2609 reader_connections: 3,
2610 background_checkpoints: true,
2611 ..SqliteStoreConfig::default()
2612 };
2613 let store =
2614 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap();
2615
2616 assert_eq!(store.reader_connection_count(), 3);
2617 let writer = store.connection().unwrap();
2618 let autocheckpoint: u32 = writer
2619 .query_row("PRAGMA wal_autocheckpoint", [], |row| row.get(0))
2620 .unwrap();
2621 assert_eq!(autocheckpoint, 0);
2622 }
2623
2624 #[test]
2625 fn primary_reader_reuses_writer_and_other_threads_use_the_pool() {
2626 let directory = tempfile::tempdir().unwrap();
2627 let config = SqliteStoreConfig {
2628 background_checkpoints: false,
2629 reader_connections: 2,
2630 ..SqliteStoreConfig::default()
2631 };
2632 let store = Arc::new(
2633 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap(),
2634 );
2635
2636 assert!(matches!(
2637 store.read_connection().unwrap(),
2638 ReadConnectionGuard::Writer(_)
2639 ));
2640 let other = {
2641 let store = store.clone();
2642 thread::spawn(move || {
2643 matches!(
2644 store.read_connection().unwrap(),
2645 ReadConnectionGuard::Reader(_)
2646 )
2647 })
2648 };
2649 assert!(other.join().unwrap());
2650 }
2651
2652 #[test]
2653 fn reader_pool_serves_reads_while_the_writer_commits() {
2654 let directory = tempfile::tempdir().unwrap();
2655 let config = SqliteStoreConfig {
2656 background_checkpoints: false,
2657 node_read_cache_size_bytes: 0,
2658 reader_connections: 4,
2659 ..SqliteStoreConfig::default()
2660 };
2661 let store = Arc::new(
2662 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap(),
2663 );
2664 store.put(b"stable", b"value").unwrap();
2665 let barrier = Arc::new(std::sync::Barrier::new(5));
2666 let readers = (0..4)
2667 .map(|_| {
2668 let store = store.clone();
2669 let barrier = barrier.clone();
2670 thread::spawn(move || {
2671 barrier.wait();
2672 for _ in 0..100 {
2673 assert_eq!(store.get(b"stable").unwrap(), Some(b"value".to_vec()));
2674 }
2675 })
2676 })
2677 .collect::<Vec<_>>();
2678 barrier.wait();
2679 for index in 0u32..100 {
2680 store
2681 .put(&index.to_be_bytes(), &index.wrapping_add(1).to_be_bytes())
2682 .unwrap();
2683 }
2684 for reader in readers {
2685 reader.join().unwrap();
2686 }
2687
2688 assert!(store.metrics().unwrap().sql_reads >= 400);
2689 }
2690
2691 #[test]
2692 fn background_checkpoint_worker_processes_committed_wal_frames() {
2693 let directory = tempfile::tempdir().unwrap();
2694 let config = SqliteStoreConfig {
2695 background_checkpoints: true,
2696 checkpoint_interval_ms: 5,
2697 checkpoint_wal_bytes: 1,
2698 reader_connections: 1,
2699 ..SqliteStoreConfig::default()
2700 };
2701 let store =
2702 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap();
2703 store.put(b"key", b"value").unwrap();
2704
2705 let deadline = Instant::now() + Duration::from_secs(2);
2706 while store.metrics().unwrap().checkpoint_attempts == 0 && Instant::now() < deadline {
2707 thread::sleep(Duration::from_millis(5));
2708 }
2709
2710 let metrics = store.metrics().unwrap();
2711 assert!(metrics.checkpoint_attempts >= 1);
2712 assert!(metrics.checkpointed_frames >= 1);
2713 }
2714
2715 #[test]
2716 fn adaptive_batch_reads_cross_the_legacy_256_key_boundary() {
2717 let directory = tempfile::tempdir().unwrap();
2718 let config = SqliteStoreConfig {
2719 background_checkpoints: false,
2720 max_batch_select_keys: 512,
2721 node_read_cache_size_bytes: 0,
2722 reader_connections: 2,
2723 ..SqliteStoreConfig::default()
2724 };
2725 let store =
2726 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap();
2727 let keys = (0u32..600)
2728 .map(|index| index.to_be_bytes().to_vec())
2729 .collect::<Vec<_>>();
2730 let values = (0u32..600)
2731 .map(|index| format!("value-{index}").into_bytes())
2732 .collect::<Vec<_>>();
2733 let entries = keys
2734 .iter()
2735 .zip(&values)
2736 .map(|(key, value)| (key.as_slice(), value.as_slice()))
2737 .collect::<Vec<_>>();
2738 store.batch_put(&entries).unwrap();
2739 let key_refs = keys.iter().map(Vec::as_slice).collect::<Vec<_>>();
2740
2741 let observed = store.batch_get_ordered_unique(&key_refs).unwrap();
2742
2743 assert_eq!(observed.len(), values.len());
2744 assert!(observed
2745 .iter()
2746 .zip(values)
2747 .all(|(observed, expected)| observed.as_deref() == Some(expected.as_slice())));
2748 }
2749
2750 #[test]
2751 fn sharded_cache_eviction_and_sqlite_metrics_are_reported() {
2752 let config = SqliteStoreConfig {
2753 node_read_cache_size_bytes: 16,
2754 node_read_cache_shards: 1,
2755 ..SqliteStoreConfig::default()
2756 };
2757 let store =
2758 SqliteStore::from_connection(Connection::open_in_memory().unwrap(), config).unwrap();
2759 store.put(b"a", b"123456789012").unwrap();
2760 store.put(b"b", b"abcdefghijkl").unwrap();
2761 assert_eq!(store.get(b"a").unwrap(), Some(b"123456789012".to_vec()));
2762
2763 let metrics = store.metrics().unwrap();
2764 assert!(metrics.node_cache_evictions >= 1);
2765 assert!(metrics.node_cache_misses >= 1);
2766 assert!(metrics.sql_reads >= 1);
2767 assert!(metrics.sqlite_page_cache_hits + metrics.sqlite_page_cache_misses > 0);
2768 }
2769
2770 #[test]
2771 fn group_commit_combines_concurrent_publications() {
2772 let directory = tempfile::tempdir().unwrap();
2773 let config = SqliteStoreConfig {
2774 background_checkpoints: false,
2775 group_commit_delay_micros: 20_000,
2776 group_commit_max_nodes: 64,
2777 reader_connections: 2,
2778 ..SqliteStoreConfig::default()
2779 };
2780 let store = Arc::new(
2781 SqliteStore::open_with_config(directory.path().join("store.db"), config).unwrap(),
2782 );
2783 let barrier = Arc::new(std::sync::Barrier::new(8));
2784 let threads = (0u8..8)
2785 .map(|index| {
2786 let store = store.clone();
2787 let barrier = barrier.clone();
2788 thread::spawn(move || {
2789 let key = vec![index; 32];
2790 let value = vec![index.wrapping_add(1); 128];
2791 let entries = [(key.as_slice(), value.as_slice())];
2792 barrier.wait();
2793 store
2794 .publish_nodes(NodePublication::new(
2795 &entries,
2796 PublicationOrigin::PointUpsert,
2797 ))
2798 .unwrap();
2799 })
2800 })
2801 .collect::<Vec<_>>();
2802 for thread in threads {
2803 thread.join().unwrap();
2804 }
2805
2806 let metrics = store.metrics().unwrap();
2807 assert_eq!(metrics.published_nodes, 8);
2808 assert!(metrics.grouped_publications >= 1);
2809 assert!(metrics.write_transactions < 8);
2810 }
2811
2812 #[test]
2813 fn checkpoint_backup_compaction_and_quick_check_are_available() {
2814 let directory = tempfile::tempdir().unwrap();
2815 let database = directory.path().join("store.db");
2816 let backup = directory.path().join("backup.db");
2817 let config = SqliteStoreConfig {
2818 background_checkpoints: false,
2819 reader_connections: 2,
2820 ..SqliteStoreConfig::default()
2821 };
2822 let store = SqliteStore::open_with_config(&database, config).unwrap();
2823 store.put(b"key", b"value").unwrap();
2824
2825 let checkpoint = store.checkpoint(SqliteCheckpointMode::Passive).unwrap();
2826 assert!(checkpoint.checkpointed_frames <= checkpoint.log_frames);
2827 store.quick_check().unwrap();
2828 store.backup_to(&backup).unwrap();
2829 store.compact().unwrap();
2830 let stats = store.storage_stats().unwrap();
2831 assert!(stats.page_count > 0);
2832
2833 let backup_store = SqliteStore::open_existing(&backup).unwrap();
2834 assert_eq!(backup_store.get(b"key").unwrap(), Some(b"value".to_vec()));
2835 }
2836}