diff --git a/.mise/tasks/sqlite-compress/test b/.mise/tasks/sqlite-compress/test index 4df439e..12f2954 100755 --- a/.mise/tasks/sqlite-compress/test +++ b/.mise/tasks/sqlite-compress/test @@ -1,3 +1,3 @@ #!/usr/bin/env bash #MISE description="Run tests" -cargo test -p sqlite-compress +cargo test -p sqlite-compress -- --test-threads=1 diff --git a/sqlite-compress/src/cache/connection.rs b/sqlite-compress/src/cache/connection.rs index 91c4cc2..15dfee6 100644 --- a/sqlite-compress/src/cache/connection.rs +++ b/sqlite-compress/src/cache/connection.rs @@ -1,11 +1,9 @@ use std::collections::HashMap; -use zstd::dict::{DecoderDictionary, EncoderDictionary}; - -use crate::{ColumnKey, DictId, cache::lru::DictLru}; -/// Per-connection cache of prepared zstd dictionaries and the current dict id per column. +use crate::{ColumnKey, DictId}; +/// Process-wide map of the current dict id per column (shared across all connections via [`crate::Cache`]). /// -/// `encoders` and `decoders` are independent LRUs keyed by [`DictKey`] (schema + dict id). +/// Prepared encoders/decoders live in the separate process-wide pool, not here. /// /// `current_ids` maps a [`ColumnKey`] to the dict id that `compress` should use. It is /// filled by `warm_cache` (latest row in `__compress_dicts`) and `insert_into_caches` @@ -13,16 +11,12 @@ use crate::{ColumnKey, DictId, cache::lru::DictLru}; /// A missing id is treated as raw zstd (`DictId` 0). pub(crate) struct DictCache { current_ids: HashMap, - pub(crate) encoders: DictLru>, - pub(crate) decoders: DictLru>, } impl DictCache { pub(crate) fn new() -> Self { Self { current_ids: HashMap::new(), - encoders: DictLru::new(), - decoders: DictLru::new(), } } diff --git a/sqlite-compress/src/cache/lru.rs b/sqlite-compress/src/cache/lru.rs deleted file mode 100644 index 560e542..0000000 --- a/sqlite-compress/src/cache/lru.rs +++ /dev/null @@ -1,124 +0,0 @@ -use std::{ - collections::{HashMap, VecDeque, hash_map::Entry}, - sync::Arc, -}; - -use crate::dict::DictKey; - -const MAX_CACHED_DICTS: usize = 5; - -/// Bounds the memory used by prepared dictionaries for one SQLite connection. -/// -/// Access through [`DictLru::get`] refreshes recency because the dictionary was -/// used for a codec operation. [`DictLru::peek`] deliberately does not, which -/// lets the registry check for an entry without changing eviction order. -pub(crate) struct DictLru { - map: HashMap>, - order: VecDeque, -} - -impl DictLru { - pub(crate) fn new() -> Self { - Self { - map: HashMap::with_capacity(MAX_CACHED_DICTS), - order: VecDeque::with_capacity(MAX_CACHED_DICTS), - } - } - - pub(crate) fn get(&mut self, key: &DictKey) -> Option> { - let value = self.map.get(key)?.clone(); - self.touch(key); - Some(value) - } - - pub(crate) fn peek(&self, key: &DictKey) -> Option> { - self.map.get(key).cloned() - } - - pub(crate) fn insert(&mut self, key: DictKey, value: Arc) { - if let Entry::Occupied(mut e) = self.map.entry(key.clone()) { - e.insert(value); - self.touch(&key); - return; - } - - while self.map.len() >= MAX_CACHED_DICTS { - if let Some(evicted) = self.order.pop_front() { - self.map.remove(&evicted); - } - } - self.map.insert(key.clone(), value); - self.order.push_back(key); - } - - fn touch(&mut self, key: &DictKey) { - self.order.retain(|x| x != key); - self.order.push_back(key.clone()); - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::dict::DictId; - - fn key(schema: &str, n: u32) -> DictKey { - DictKey::new(schema, DictId::new(n)) - } - - #[test] - fn evicts_oldest_after_five() { - let mut cache = DictLru::new(); - for n in 1..=5 { - cache.insert(key("main", n), Arc::new(n)); - } - cache.insert(key("main", 6), Arc::new(6)); - assert!(cache.peek(&key("main", 1)).is_none()); - assert!(cache.peek(&key("main", 2)).is_some()); - assert!(cache.peek(&key("main", 6)).is_some()); - } - - #[test] - fn get_promotes_to_most_recent() { - let mut cache = DictLru::new(); - for n in 1..=5 { - cache.insert(key("main", n), Arc::new(n)); - } - assert!(cache.get(&key("main", 1)).is_some()); - cache.insert(key("main", 6), Arc::new(6)); - assert!(cache.peek(&key("main", 1)).is_some()); - assert!(cache.peek(&key("main", 2)).is_none()); - } - - #[test] - fn peek_does_not_promote() { - let mut cache = DictLru::new(); - for n in 1..=5 { - cache.insert(key("main", n), Arc::new(n)); - } - assert!(cache.peek(&key("main", 1)).is_some()); - cache.insert(key("main", 6), Arc::new(6)); - assert!(cache.peek(&key("main", 1)).is_none()); - } - - #[test] - fn reinsert_updates_value_and_promotes() { - let mut cache = DictLru::new(); - for n in 1..=5 { - cache.insert(key("main", n), Arc::new(n)); - } - cache.insert(key("main", 1), Arc::new(100)); - cache.insert(key("main", 6), Arc::new(6)); - assert_eq!(*cache.peek(&key("main", 1)).unwrap(), 100); - assert!(cache.peek(&key("main", 2)).is_none()); - } - - #[test] - fn same_id_different_schema_are_distinct() { - let mut cache = DictLru::new(); - cache.insert(key("raw", 1), Arc::new(10)); - cache.insert(key("archive", 1), Arc::new(20)); - assert_eq!(*cache.peek(&key("raw", 1)).unwrap(), 10); - assert_eq!(*cache.peek(&key("archive", 1)).unwrap(), 20); - } -} diff --git a/sqlite-compress/src/cache/mod.rs b/sqlite-compress/src/cache/mod.rs index 220fabd..c34c645 100644 --- a/sqlite-compress/src/cache/mod.rs +++ b/sqlite-compress/src/cache/mod.rs @@ -1,5 +1,5 @@ mod connection; -mod lru; +mod pool; mod prepared; pub(crate) use connection::DictCache; diff --git a/sqlite-compress/src/cache/pool.rs b/sqlite-compress/src/cache/pool.rs new file mode 100644 index 0000000..9c4dd57 --- /dev/null +++ b/sqlite-compress/src/cache/pool.rs @@ -0,0 +1,92 @@ +use std::{ + collections::{HashMap, hash_map::Entry}, + sync::{Arc, LazyLock}, +}; + +use parking_lot::Mutex; +use zstd::dict::{DecoderDictionary, EncoderDictionary}; + +use crate::{dict::DictKey, functions::Level}; + +const MAX_PREPARED_DICTS: usize = 64; + +type PreparedMap = HashMap>; + +static ENCODERS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::with_capacity(MAX_PREPARED_DICTS))); + +static DECODERS: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::with_capacity(MAX_PREPARED_DICTS))); + +pub(crate) fn get_encoder( + key: &DictKey, + level: Level, + load: impl FnOnce() -> Result, crate::DictError>, +) -> Result>, crate::DictError> { + if let Some(dict) = ENCODERS.lock().get(key).cloned() { + return Ok(dict); + } + + let raw = load()?; + + let mut cache = ENCODERS.lock(); + if let Some(dict) = cache.get(key).cloned() { + return Ok(dict); + } + + let encoder = Arc::new(EncoderDictionary::copy(&raw, level.get())); + insert_capped(&mut cache, key.clone(), encoder.clone()); + Ok(encoder) +} + +pub(crate) fn get_decoder( + key: &DictKey, + load: impl FnOnce() -> Result, crate::DictError>, +) -> Result>, crate::DictError> { + if let Some(dict) = DECODERS.lock().get(key).cloned() { + return Ok(dict); + } + + let raw = load()?; + + let mut cache = DECODERS.lock(); + if let Some(dict) = cache.get(key).cloned() { + return Ok(dict); + } + + let decoder = Arc::new(DecoderDictionary::copy(&raw)); + insert_capped(&mut cache, key.clone(), decoder.clone()); + Ok(decoder) +} + +fn insert_capped(map: &mut PreparedMap, key: DictKey, value: Arc) { + let at_cap = map.len() >= MAX_PREPARED_DICTS; + + match map.entry(key) { + Entry::Occupied(mut e) => { + e.insert(value); + } + Entry::Vacant(e) => { + if at_cap { + let key = e.into_key(); + map.clear(); + map.insert(key, value); + } else { + e.insert(value); + } + } + } +} + +pub(crate) fn insert_prepared(key: DictKey, dictionary: &[u8], level: Level) { + insert_capped( + &mut ENCODERS.lock(), + key.clone(), + Arc::new(EncoderDictionary::copy(dictionary, level.get())), + ); + insert_capped( + &mut DECODERS.lock(), + key, + Arc::new(DecoderDictionary::copy(dictionary)), + ); +} diff --git a/sqlite-compress/src/cache/prepared.rs b/sqlite-compress/src/cache/prepared.rs index 65b0eac..89337b6 100644 --- a/sqlite-compress/src/cache/prepared.rs +++ b/sqlite-compress/src/cache/prepared.rs @@ -4,6 +4,7 @@ use zstd::dict::{DecoderDictionary, EncoderDictionary}; use crate::{ Cache, ColumnKey, DictError, DictId, + cache::pool, dict::{DictKey, get_raw_dict, read_current_id}, functions::Level, state::ExtensionState, @@ -18,22 +19,7 @@ pub(crate) fn get_encoder_in_cache( ) -> Result>, DictError> { let key = DictKey::new(schema, dict_id); - if let Some(dict) = state.cache.lock().encoders.peek(&key) { - return Ok(dict); - } - - let raw = get_raw_dict(&key, state)?; - - let mut cache = state.cache.lock(); - - if let Some(dict) = cache.encoders.get(&key) { - return Ok(dict); - } - - let encoder = Arc::new(EncoderDictionary::copy(&raw, level.get())); - cache.encoders.insert(key, encoder.clone()); - - Ok(encoder) + pool::get_encoder(&key, level, || get_raw_dict(&key, state)) } /// Return the prepared zstd decoder for `(schema, dict_id)`, loading it from `__compress_dicts` and caching it on a miss. @@ -44,22 +30,7 @@ pub(crate) fn get_decoder_in_cache( ) -> Result>, DictError> { let key = DictKey::new(schema, dict_id); - if let Some(dict) = state.cache.lock().decoders.peek(&key) { - return Ok(dict); - } - - let raw = get_raw_dict(&key, state)?; - - let mut cache = state.cache.lock(); - - if let Some(dict) = cache.decoders.get(&key) { - return Ok(dict); - } - - let decoder = Arc::new(DecoderDictionary::copy(&raw)); - cache.decoders.insert(key, decoder.clone()); - - Ok(decoder) + pool::get_decoder(&key, || get_raw_dict(&key, state)) } /// Load this column's current dict id, encoder, and decoder. Call it as early as possible so later `compress` and `decompress` hit the cache. @@ -91,20 +62,8 @@ pub(crate) fn insert_into_caches( dictionary: &[u8], level: Level, ) { - let mut cache = cache.lock(); - - cache.set_current_id(column.clone(), dict_id); - - let key = DictKey::new(column.schema(), dict_id); - - cache.encoders.insert( - key.clone(), - Arc::new(EncoderDictionary::copy(dictionary, level.get())), - ); - - cache - .decoders - .insert(key, Arc::new(DecoderDictionary::copy(dictionary))); + cache.lock().set_current_id(column.clone(), dict_id); + pool::insert_prepared(DictKey::new(column.schema(), dict_id), dictionary, level); } #[cfg(test)] diff --git a/sqlite-compress/src/state.rs b/sqlite-compress/src/state.rs index ea55f82..434981c 100644 --- a/sqlite-compress/src/state.rs +++ b/sqlite-compress/src/state.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use libsqlite3_sys::sqlite3; use parking_lot::Mutex; @@ -6,11 +6,17 @@ use sqlite_ffi::{Connection, Context, SqliteError}; use crate::cache::DictCache; -/// Per-connection state for the extension: the SQLite connection and its dictionary cache. +// Shared cache for all connections. +static SHARED: LazyLock = LazyLock::new(|| Cache(Arc::new(Mutex::new(DictCache::new())))); + +/// Per-connection handle: this SQLite connection plus a clone of the process-wide dictionary cache. /// -/// `sqlite3_compress_init` calls [`ExtensionState::attach`], which stores an empty [`Cache`] in the connection's client data. -/// `setup`, `train_all`, `compress`, and `decompress` read that same cache back with [`ExtensionState::from_db`] or [`ExtensionState::from_context`]. -/// The cache is an `Arc`, so every lookup shares the encoders, decoders, and current dict ids for this connection. +/// `sqlite3_compress_init` calls [`ExtensionState::attach`], which stores a clone of the shared +/// [`Cache`] in the connection's client data (not a fresh empty cache). +/// `setup`, `train_all`, `compress`, and `decompress` read it back with [`ExtensionState::from_db`] +/// or [`ExtensionState::from_context`]. +/// All connections therefore share the same `current_ids`; prepared encoders/decoders live in the +/// separate process-wide pool. pub struct ExtensionState { pub(crate) connection: Connection, pub(crate) cache: Cache, @@ -35,9 +41,7 @@ impl ExtensionState { const NAME: &str = "sqlite-compress"; pub fn attach(connection: &Connection) -> Result<(), SqliteError> { - let cache = Cache(Arc::new(Mutex::new(DictCache::new()))); - - connection.set_client_data(Self::NAME, cache) + connection.set_client_data(Self::NAME, SHARED.clone()) } pub fn from_db(db: *mut sqlite3) -> Result { diff --git a/sqlite-compress/tests/common.rs b/sqlite-compress/tests/common.rs index 510a9fe..d149929 100644 --- a/sqlite-compress/tests/common.rs +++ b/sqlite-compress/tests/common.rs @@ -1,4 +1,4 @@ -use std::{ptr, sync::Once}; +use std::{ffi::CString, ptr, sync::Once}; use libsqlite3_sys::{SQLITE_OK, sqlite3, sqlite3_auto_extension, sqlite3_close, sqlite3_open}; use sqlite_compress::{ExtensionState, SetupConnection, sqlite3_compress_init}; @@ -29,6 +29,24 @@ impl TestDb { Self { db } } + #[allow(dead_code)] + pub fn open_path(path: &str) -> Self { + static REGISTER: Once = Once::new(); + REGISTER.call_once(|| unsafe { + #[allow(clippy::missing_transmute_annotations)] + sqlite3_auto_extension(Some(std::mem::transmute( + sqlite3_compress_init as *const (), + ))); + }); + + let c_path = CString::new(path).unwrap(); + let mut db = ptr::null_mut(); + let rc = unsafe { sqlite3_open(c_path.as_ptr(), &mut db) }; + + assert_eq!(rc, SQLITE_OK); + Self { db } + } + pub fn state(&self) -> ExtensionState { ExtensionState::from_db(self.db).unwrap() } diff --git a/sqlite-compress/tests/dict_train.rs b/sqlite-compress/tests/dict_train.rs index 145cad9..1cfbeb7 100644 --- a/sqlite-compress/tests/dict_train.rs +++ b/sqlite-compress/tests/dict_train.rs @@ -8,6 +8,8 @@ mod common; use common::DEFAULT_MAX_SAMPLES; use sqlite_ffi::{SqlValue, first_value}; +use crate::common::TestDb; + pub fn compress_blob(state: &ExtensionState, data: &[u8]) -> Vec { let rows = state .as_ref() @@ -635,6 +637,64 @@ fn compress_uses_current_id_in_header() { assert_eq!(header_dict_id(&blob2), id2); } +#[test] +fn second_connection_compress_uses_trained_dict_id() { + let sample_count = 64; + + let path = std::env::temp_dir().join(format!( + "sqlite-compress-shared-cache-{}.db", + std::process::id() + )); + let _ = std::fs::remove_file(&path); + let path_str = path.to_str().unwrap(); + + let db_a = TestDb::open_path(path_str); + let db_b = TestDb::open_path(path_str); + + let state_a = db_a.state(); + let state_b = db_b.state(); + + let attach = format!( + "ATTACH DATABASE '{}' AS raw; + CREATE TABLE IF NOT EXISTS raw.requests_raw ( + id INTEGER PRIMARY KEY, + data BLOB + );", + path_str.replace('\'', "''") + ); + + state_a.as_ref().batch_execute(&attach).unwrap(); + + state_b + .as_ref() + .batch_execute(&format!( + "ATTACH DATABASE '{}' AS raw;", + path_str.replace('\'', "''") + )) + .unwrap(); + + insert_samples(&state_a, sample_count); + + let config = sample_config(sample_count); + setup(&db_a, &config).unwrap(); + + let dict_ids = train_all(&db_a, &config, 1024, TrainOptions::default()).unwrap(); + let trained_id = dict_ids[0].get(); + assert_ne!(trained_id, 0); + + let payload = b"GET /api/users/42 HTTP/1.1\r\nHost: example.com\r\n\r\n"; + let blob = compress_blob(&state_b, payload); + + assert_eq!( + header_dict_id(&blob), + trained_id, + "conn B must use the dict id trained on conn A" + ); + assert_eq!(decompress_blob(&state_b, &blob), payload); + + let _ = std::fs::remove_file(&path); +} + #[test] fn train_by_column_reports_progress() { let sample_count = 64;