diff --git a/storages/sqlite-storage/examples/per_connection_memory.rs b/storages/sqlite-storage/examples/per_connection_memory.rs new file mode 100644 index 000000000..39330f2fc --- /dev/null +++ b/storages/sqlite-storage/examples/per_connection_memory.rs @@ -0,0 +1,308 @@ +//! What one extra SQLite connection costs, measured rather than estimated. +//! +//! A process that holds N WhatsApp sessions against the *same* database file +//! opens N stores, and every constructor builds its own r2d2 pool — so N +//! connections. This harness prices that: it seeds one database, then opens N +//! sessions two ways and reports the resident-set delta per session. +//! +//! ```text +//! cargo run -p whatsapp-rust-sqlite-storage --release \ +//! --example per_connection_memory -- [warm] +//! cargo run -p whatsapp-rust-sqlite-storage --release \ +//! --example per_connection_memory -- writes +//! cargo run -p whatsapp-rust-sqlite-storage --release \ +//! --example per_connection_memory -- compile-options +//! ``` +//! +//! * `pools` — one `SqliteStore::new_for_device` per session (today's shape). +//! * `handles` — one store, then `share_for_device` per session (one pool). +//! * `warm` — every session scans the seeded table first, filling its page +//! cache to the `cache_kib` cap. Without it each session only does the small +//! reads an idle session does, which is the realistic steady state. +//! * `writes` — the other side of the trade: every session writes at once, +//! both ways, reporting total time and the spread between the fastest and +//! slowest session (i.e. whether anyone starves). +//! +//! Resident set, not a Rust allocator counter: SQLite's page cache and +//! lookaside are `sqlite3_malloc` allocations from the bundled C library, which +//! a `GlobalAlloc` wrapper never sees. RSS is coarse (page-granular, and +//! includes whatever the allocator declines to return), so it is read after a +//! settle and divided across enough sessions for the per-session figure to +//! outweigh the noise. Linux only. + +// The numbers *are* this binary's output; there is no logger to route them +// through, and a measurement harness whose result lands in a log filter would +// be worse than useless. +#![allow(clippy::print_stdout)] + +use std::time::Duration; + +use diesel::prelude::*; +use wacore::store::traits::SignalStore as _; +use whatsapp_rust_sqlite_storage::{SqliteStore, SqliteStoreConfig}; + +/// Rows of ~1 KiB each: enough database that a full scan can fill a 512 KiB +/// page cache several times over, so the cap is what bounds a warm connection. +const SEED_ROWS: usize = 4_000; +const ROW_BYTES: usize = 1_024; + +fn rss_bytes() -> u64 { + // smaps_rollup rather than statm: it reports `Rss:` already in kB, so the + // reading does not depend on knowing the kernel's page size (which is not + // always 4 KiB, and which std cannot report without libc). + let rollup = + std::fs::read_to_string("/proc/self/smaps_rollup").expect("/proc/self/smaps_rollup"); + let kib: u64 = rollup + .lines() + .find_map(|line| line.strip_prefix("Rss:")) + .and_then(|value| value.split_whitespace().next()) + .and_then(|n| n.parse().ok()) + .expect("Rss: line"); + kib * 1024 +} + +fn db_err(e: diesel::result::Error) -> wacore::store::error::StoreError { + wacore::store::error::StoreError::Database(Box::new(e)) +} + +/// Seed a database large enough that page caches have something to hold. +async fn seed(url: &str) { + let store = SqliteStore::new_for_device(url, 1).await.expect("open"); + let record = vec![0x5au8; ROW_BYTES]; + for chunk in (0..SEED_ROWS).collect::>().chunks(200) { + let batch: Vec<_> = chunk + .iter() + .map(|i| { + ( + format!("seed.{i}:0").into(), + bytes::Bytes::from(record.clone()), + ) + }) + .collect(); + store.put_sessions_batch(&batch).await.expect("seed write"); + } +} + +/// The reads an idle session actually does on connect: a couple of point +/// lookups. Also forces r2d2 to open the connection, which is the allocation +/// this harness is pricing. +async fn touch(store: &SqliteStore) { + store.get_session("seed.0:0").await.expect("point read"); + store.get_session("seed.1:0").await.expect("point read"); +} + +/// A full table scan, which pulls pages in until the cache cap stops it. +async fn scan(store: &SqliteStore) { + #[derive(QueryableByName)] + struct Count { + #[diesel(sql_type = diesel::sql_types::BigInt)] + n: i64, + } + store + .shared() + .read(|conn| { + diesel::sql_query("SELECT count(record) AS n FROM sessions") + .get_result::(conn) + .map(|c| c.n) + .map_err(db_err) + }) + .await + .expect("scan"); +} + +/// Every session issues `WRITES` writes at once. Returns the wall clock for +/// the whole burst and each session's own duration, so a mode that finishes +/// quickly by letting one session hog the lock is visible as a wide spread. +async fn write_burst(stores: Vec) -> (Duration, Duration, Duration) { + const WRITES: usize = 200; + let started = wacore::time::Instant::now(); + let mut tasks = Vec::new(); + for (n, store) in stores.into_iter().enumerate() { + tasks.push(tokio::spawn(async move { + let session_started = wacore::time::Instant::now(); + for i in 0..WRITES { + store + .put_session(&format!("peer.{n}.{i}:0"), &[n as u8; 256]) + .await + .expect("write"); + } + session_started.elapsed() + })); + } + let mut per_session = Vec::new(); + for task in tasks { + per_session.push(task.await.expect("join")); + } + ( + started.elapsed(), + per_session.iter().copied().min().unwrap_or_default(), + per_session.iter().copied().max().unwrap_or_default(), + ) +} + +/// `dir` rather than one database URL: each arm gets a database of its own, so +/// the second is not measured against the pages, WAL and rows the first left +/// behind — which would confound the comparison with run order. +async fn writes(dir: &std::path::Path, sessions: usize, read_pool_size: u32) { + let config = || SqliteStoreConfig { + read_pool_size, + ..Default::default() + }; + let url = |name: &str| dir.join(name).to_string_lossy().into_owned(); + + let shared_url = url("writes_handles.db"); + let base = SqliteStore::with_config_for_device(&shared_url, 1, config()) + .await + .expect("open"); + let mut fleet = vec![base.clone()]; + for device_id in 2..=sessions { + fleet.push(base.share_for_device(device_id as i32)); + } + let (total, fastest, slowest) = write_burst(fleet).await; + println!( + "handles sessions={sessions} read_pool_size={read_pool_size} \ + total={total:?} fastest={fastest:?} slowest={slowest:?}" + ); + + let separate_url = url("writes_pools.db"); + let mut separate = Vec::new(); + for device_id in 1..=sessions { + separate.push( + SqliteStore::with_config_for_device(&separate_url, device_id as i32, config()) + .await + .expect("open"), + ); + } + let (total, fastest, slowest) = write_burst(separate).await; + println!( + "pools sessions={sessions} read_pool_size={read_pool_size} \ + total={total:?} fastest={fastest:?} slowest={slowest:?}" + ); +} + +async fn compile_options(url: &str) { + let store = SqliteStore::new(url).await.expect("open"); + #[derive(QueryableByName)] + struct Opt { + #[diesel(sql_type = diesel::sql_types::Text)] + compile_options: String, + } + let opts: Vec = store + .shared() + .run(|conn| { + diesel::sql_query("PRAGMA compile_options") + .load(conn) + .map_err(db_err) + }) + .await + .expect("compile_options"); + for opt in opts { + println!("{}", opt.compile_options); + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let args: Vec = std::env::args().skip(1).collect(); + let mode = args.first().map(String::as_str).unwrap_or("pools"); + let sessions: usize = args.get(1).and_then(|a| a.parse().ok()).unwrap_or(50); + let cache_kib: u32 = args.get(2).and_then(|a| a.parse().ok()).unwrap_or(512); + let warm = args.iter().any(|a| a == "warm"); + // Zero would divide the per-session figure by nothing, and `handles` would + // still open its base store — a count nobody asked for. + assert!(sessions >= 1, "sessions must be at least 1"); + + // A directory of our own, created exclusively and owner-only: the create + // fails outright if anything already sits at the path (a symlink included), + // and mode 0o700 keeps it that way regardless of umask. Between them, + // nobody else on the machine can pre-place or swap the database, WAL and + // shm files this then writes. + use std::os::unix::fs::DirBuilderExt as _; + let dir = std::env::temp_dir().join(format!("wa_percon_{}", std::process::id())); + std::fs::DirBuilder::new() + .mode(0o700) + .create(&dir) + .expect("exclusive scratch directory"); + // A guard, not a tail cleanup: two of the modes below return early. + struct ScratchDir(std::path::PathBuf); + impl Drop for ScratchDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + let scratch = ScratchDir(dir.clone()); + let url = dir.join("bench.db").to_string_lossy().into_owned(); + + if mode == "compile-options" { + compile_options(&url).await; + return; + } + if mode == "writes" { + // args[2] is the reader-pool size here, not a cache size. + writes(&dir, sessions, cache_kib).await; + return; + } + + seed(&url).await; + let config = || SqliteStoreConfig { + cache_size_kib: cache_kib, + ..Default::default() + }; + + // Settle: the seeding store is dropped, and its connection with it, so the + // baseline is the process without any session on this database. + tokio::time::sleep(Duration::from_millis(200)).await; + let before = rss_bytes(); + + let mut stores: Vec = Vec::with_capacity(sessions); + let mut after_first = before; + for n in 0..sessions { + let device_id = n as i32 + 1; + let store = match (mode, stores.first()) { + ("pools", _) => SqliteStore::with_config_for_device(&url, device_id, config()) + .await + .expect("open"), + // The first handle is a real store; the rest hang off it. + ("handles", None) => SqliteStore::with_config_for_device(&url, device_id, config()) + .await + .expect("open"), + ("handles", Some(base)) => base.share_for_device(device_id), + (other, _) => panic!("unknown mode {other}"), + }; + touch(&store).await; + if warm { + scan(&store).await; + } + stores.push(store); + if n == 0 { + // The *marginal* cost is the number that describes the batch, and + // it is the one that survives a warmed allocator: whatever the + // seeding connection left behind for the first session to reuse is + // in this reading, so subtracting it takes the discount out of the + // per-session figure instead of hiding in it. + tokio::time::sleep(Duration::from_millis(200)).await; + after_first = rss_bytes(); + } + } + + tokio::time::sleep(Duration::from_millis(200)).await; + let after = rss_bytes(); + let delta = after.saturating_sub(before); + let marginal = after.saturating_sub(after_first); + let others = sessions.saturating_sub(1); + println!( + "mode={mode} sessions={sessions} cache_kib={cache_kib} warm={warm} \ + rss_delta={delta}B per_session={:.1}KiB marginal_per_session={}", + delta as f64 / sessions as f64 / 1024.0, + if others == 0 { + // One session is all baseline and no margin; saying "0.0KiB" would + // read as a measurement rather than the absence of one. + "n/a".to_string() + } else { + format!("{:.1}KiB", marginal as f64 / others as f64 / 1024.0) + } + ); + + drop(stores); + drop(scratch); +} diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index f9b2bde4d..6c29b2678 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -183,6 +183,9 @@ pub type ConnectionInitHook = Arc< /// concurrent DB access — it drives both the pool and the internal serialization in /// lockstep — or `cache_size_kib` for a hotter/larger DB; pass a `thread_pool` to control /// r2d2's management threads (e.g. share your own across crates). +/// +/// Sessions that share one database file can go further and share the connection +/// itself: see [`SqliteStore::share_for_device`]. #[derive(Clone)] pub struct SqliteStoreConfig { /// Max concurrent operations: r2d2 `max_size` AND the internal semaphore permits, @@ -212,6 +215,14 @@ pub struct SqliteStoreConfig { /// per-session stores. pub read_pool_size: u32, /// `PRAGMA cache_size`, in KiB per connection. + /// + /// A cap on growth, not a reservation, and not the whole per-connection + /// cost: a connection also carries a 48,000 B lookaside slab that no pragma + /// can shrink (`SQLITE_DBCONFIG_LOOKASIDE` is C-API only, and diesel does + /// not expose the `sqlite3*`). Measured on an idle session, dropping this + /// from 512 to 1 moved resident memory from ~123 to ~92 KiB per connection + /// — so tuning it down does not substitute for holding fewer connections; + /// see [`SqliteStore::share_for_device`]. pub cache_size_kib: u32, /// `PRAGMA mmap_size`, in bytes. `None` (default) leaves mmap off — the /// current behavior. When set, pages are read through a reclaimable, @@ -609,6 +620,76 @@ impl SqliteStore { self.device_id } + /// A store for a *sibling device* in the same database, reusing this + /// store's connections instead of opening more. + /// + /// Every constructor builds its own r2d2 pool, so a process holding N + /// sessions against one database file ends up with N SQLite connections — + /// and a connection costs memory before it reads a single row: a 48,000 B + /// lookaside slab (`SQLITE_DEFAULT_LOOKASIDE` 1200,40, which this build + /// does not override), plus a page cache that grows to + /// [`SqliteStoreConfig::cache_size_kib`]. Measured on an idle session that + /// has only done a couple of point reads, that is ~123 KiB of resident + /// memory per session, and it does not shrink meaningfully with a smaller + /// cache cap: ~92 KiB of it survives `cache_size_kib = 1`. Nothing else + /// about the store is per-session — every query already takes a + /// `device_id` — so sibling sessions on one database only ever needed that + /// field to differ. Same reasoning as [`SqliteStore::shared`], applied to + /// sibling devices instead of sibling crates. + /// + /// The returned store owns clones of the pool handles, so it stays usable + /// for as long as it lives — dropping the store it came from closes + /// nothing. + /// + /// What it does **not** do: + /// + /// - **Create the device row.** It only stamps queries with `device_id`. + /// The row still comes from the usual provisioning path — the same + /// [`create_new_device`](Self::create_new_device) or restore that a store + /// from [`new_for_device`](Self::new_for_device) would need. + /// - **Isolate writes.** Siblings share the write permits, of which + /// [`SqliteStoreConfig::pool_size`] decides the number — so at its + /// default of 1 their writes serialize against each other, and a base + /// store built with a wider pool passes that width on instead. That is + /// the trade, and at the default it is not free: on a burst where every + /// session writes continuously, sharing + /// costs ~2.5x the aggregate write throughput of a pool per session, + /// because a private connection lets one session's queueing overlap + /// another's SQLite work. In exchange the queue is FIFO-fair, where + /// separate connections leave it to SQLite's busy handler and its random + /// backoff (measured: ~2x spread between the fastest and slowest + /// session). So this is for fleets that are mostly idle — the shape + /// sessions actually have — and not for continuously writing ones. + /// [`SqliteStoreConfig::read_pool_size`] widens the *read* side only, + /// and its connections are shared here too. + /// - **Split the resource report.** `resource_report()` describes the + /// *pool*, and siblings share one, so every handle reports the same + /// whole-pool estimate. That is the honest answer for a shared + /// connection — the bytes belong to the pool, not to any one session — + /// but it means summing the report across a fleet of siblings counts + /// those bytes once per sibling. Count them once per pool instead. The + /// saving this method exists for is exactly why there is only one pool + /// left to count. + /// + /// ```no_run + /// # use whatsapp_rust_sqlite_storage::SqliteStore; + /// # async fn run() -> Result<(), Box> { + /// let device_1 = SqliteStore::new_for_device("whatsapp.db", 1).await?; + /// // One pool, one connection, two sessions. + /// let device_2 = device_1.share_for_device(2); + /// # Ok(()) } + /// ``` + pub fn share_for_device(&self, device_id: i32) -> Self { + Self { + pool: self.pool.clone(), + db_semaphore: Arc::clone(&self.db_semaphore), + reads: self.reads.clone(), + snapshot_safe: self.snapshot_safe, + database_path: self.database_path.clone(), + device_id, + } + } + /// Run a **read-only** query on a reader connection, falling back to the /// write queue when none is configured. /// @@ -5718,10 +5799,10 @@ mod read_routing_tests { /// A file-backed store: reader connections need real WAL, which an /// in-memory database has none of. Removed on drop. - struct TempDb(std::path::PathBuf); + pub(super) struct TempDb(std::path::PathBuf); impl TempDb { - fn new(tag: &str) -> Self { + pub(super) fn new(tag: &str) -> Self { use portable_atomic::AtomicU64; use std::sync::atomic::Ordering; static COUNTER: AtomicU64 = AtomicU64::new(0); @@ -5734,7 +5815,7 @@ mod read_routing_tests { Self(path) } - fn url(&self) -> String { + pub(super) fn url(&self) -> String { self.0.to_string_lossy().into_owned() } } @@ -6598,3 +6679,207 @@ impl SqliteStore { ); } } + +#[cfg(test)] +mod share_for_device_tests { + use super::read_routing_tests::TempDb; + use super::*; + use std::time::Duration; + use wacore::time::Instant; + + /// How many sibling sessions the concurrency tests run. Small enough to + /// stay quick, large enough that a serialized queue is visible. + const SESSIONS: usize = 8; + const WRITES_PER_SESSION: usize = 25; + + async fn base_store(db: &TempDb) -> SqliteStore { + SqliteStore::new_for_device(&db.url(), 1) + .await + .expect("store opens") + } + + /// Sibling handles are only plumbing: the `device_id` is what separates + /// their rows, exactly as it does for two independently-opened stores. + #[tokio::test] + async fn siblings_share_the_database_but_not_each_other_s_rows() { + let db = TempDb::new("share_isolation"); + let device_1 = base_store(&db).await; + let device_2 = device_1.share_for_device(2); + assert_eq!(device_2.device_id(), 2); + + device_1 + .put_session("alice.1:0", b"device-1-record") + .await + .expect("write through the first handle"); + + // Same file: the sibling can read the row by asking for the other + // device explicitly. + assert_eq!( + device_2 + .get_session_for_device("alice.1:0", 1) + .await + .expect("read"), + Some(b"device-1-record".to_vec()), + "both handles must be looking at the same database" + ); + // Its own device scope, however, is empty. + assert_eq!( + device_2.get_session("alice.1:0").await.expect("read"), + None, + "a sibling device must not see another device's session" + ); + + // And a write through the sibling lands in its own scope only. + device_2 + .put_session("alice.1:0", b"device-2-record") + .await + .expect("write through the sibling handle"); + assert_eq!( + device_1.get_session("alice.1:0").await.expect("read"), + Some(Bytes::from_static(b"device-1-record")), + "the sibling's write must not clobber the first device's row" + ); + } + + /// The whole point of the method, asserted the only way that proves it: + /// by counting connections. A fleet of handles opens one; a fleet of + /// stores opens one each. + #[tokio::test] + async fn a_fleet_of_handles_opens_one_connection() { + let db = TempDb::new("share_conn_count"); + let base = base_store(&db).await; + let mut fleet = vec![base.clone()]; + for device_id in 2..=SESSIONS as i32 { + fleet.push(base.share_for_device(device_id)); + } + // r2d2 opens connections lazily, so make every handle actually use one. + for store in &fleet { + store.get_session("probe").await.expect("read"); + } + let shared_connections: u32 = fleet + .iter() + .map(|store| store.pool.state().connections) + .max() + .expect("non-empty fleet"); + assert_eq!( + shared_connections, 1, + "sibling handles must reuse the one pooled connection" + ); + // Same semaphore, so they also share the write queue — the trade-off + // the doc comment describes, asserted rather than assumed. + assert!( + fleet + .iter() + .all(|store| Arc::ptr_eq(&store.db_semaphore, &base.db_semaphore)), + "handles must share the write permit, not just the pool" + ); + + // The baseline this replaces: one store per session, one connection each. + let db = TempDb::new("share_conn_count_baseline"); + let mut separate = Vec::new(); + for device_id in 1..=SESSIONS as i32 { + let store = SqliteStore::new_for_device(&db.url(), device_id) + .await + .expect("store opens"); + store.get_session("probe").await.expect("read"); + separate.push(store); + } + let total: u32 = separate + .iter() + .map(|store| store.pool.state().connections) + .sum(); + assert_eq!( + total, SESSIONS as u32, + "one store per session is one connection per session" + ); + } + + /// Every session writes at once; returns wall-clock for the whole burst + /// and each session's own completion time. + async fn write_burst(stores: Vec) -> (Duration, Vec) { + let started = Instant::now(); + let mut tasks = Vec::new(); + for (n, store) in stores.into_iter().enumerate() { + tasks.push(tokio::spawn(async move { + let session_started = Instant::now(); + for i in 0..WRITES_PER_SESSION { + store + .put_session(&format!("peer.{n}.{i}:0"), &[n as u8; 256]) + .await + .expect("write must not fail under contention"); + } + session_started.elapsed() + })); + } + let mut per_session = Vec::new(); + for task in tasks { + per_session.push(task.await.expect("join")); + } + (started.elapsed(), per_session) + } + + /// Sharing a pool means sharing its write permits, and at the default + /// `pool_size` there is exactly one — so sibling sessions serialize on + /// writes. That is the cost of the memory saving and the reason + /// `share_for_device` is not the default shape; it is measured here rather + /// than argued about. `pool_size` is set explicitly, because the claim + /// holds for that value and not for a wider pool. + /// + /// The assertions are the two properties that must hold on any machine: + /// no write fails, and no session starves. The timings are printed for the + /// record; asserting on wall-clock would only buy a flaky test. + #[tokio::test] + async fn concurrent_writes_serialize_across_siblings_at_the_default_pool_size() { + let db = TempDb::new("share_write_contention"); + let base = SqliteStore::with_config_for_device( + &db.url(), + 1, + SqliteStoreConfig { + pool_size: 1, + ..Default::default() + }, + ) + .await + .expect("store opens"); + let mut fleet = vec![base.clone()]; + for device_id in 2..=SESSIONS as i32 { + fleet.push(base.share_for_device(device_id)); + } + let (shared_total, shared_sessions) = write_burst(fleet).await; + + let db = TempDb::new("share_write_contention_baseline"); + let mut separate = Vec::new(); + for device_id in 1..=SESSIONS as i32 { + separate.push( + SqliteStore::new_for_device(&db.url(), device_id) + .await + .expect("store opens"), + ); + } + let (separate_total, separate_sessions) = write_burst(separate).await; + + let summarize = |label: &str, total: Duration, sessions: &[Duration]| { + let slowest = sessions.iter().max().copied().unwrap_or_default(); + let fastest = sessions.iter().min().copied().unwrap_or_default(); + println!( + "{label}: {SESSIONS} sessions x {WRITES_PER_SESSION} writes in {total:?} \ + (session fastest {fastest:?}, slowest {slowest:?})" + ); + }; + summarize("shared pool", shared_total, &shared_sessions); + summarize("pool per session", separate_total, &separate_sessions); + + // Starvation check: a FIFO permit hands every session its turn, so the + // slowest cannot be an order of magnitude behind the fastest. A pool + // per session leans on SQLite's busy handler instead, which backs off + // randomly and offers no such guarantee — so only the shared side is + // asserted. + let fastest = shared_sessions.iter().min().copied().unwrap_or_default(); + let slowest = shared_sessions.iter().max().copied().unwrap_or_default(); + assert!( + slowest < fastest * 10 + Duration::from_secs(1), + "no sibling may starve on the shared write queue: \ + fastest {fastest:?}, slowest {slowest:?}" + ); + } +}