diff --git a/Cargo.lock b/Cargo.lock index e3fb3126..8f6bee9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4027,6 +4027,16 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "reclaim-measure" +version = "0.0.0" +dependencies = [ + "anyhow", + "borsh", + "rocksdb", + "sqd-storage", +] + [[package]] name = "recursive" version = "0.1.1" diff --git a/Cargo.toml b/Cargo.toml index ac4a1fae..59935afe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ members = [ "crates/primitives", "crates/query", "crates/query-example", + "crates/reclaim-measure", "crates/storage" ] resolver = "2" diff --git a/Dockerfile b/Dockerfile index bcf8bc6d..f22f7007 100644 --- a/Dockerfile +++ b/Dockerfile @@ -13,12 +13,15 @@ ADD crates crates FROM builder AS hotblocks-builder -RUN cargo build -p sqd-hotblocks --release +RUN cargo build -p sqd-hotblocks -p reclaim-measure --release FROM rust AS hotblocks WORKDIR /app COPY --from=hotblocks-builder /app/target/release/sqd-hotblocks . +# Read-only probe: opens the live db as a RocksDB secondary instance and reports how much +# --startup-disk-reclaim would free. Run it in the pod before turning that flag on. +COPY --from=hotblocks-builder /app/target/release/reclaim-measure . ENTRYPOINT ["/app/sqd-hotblocks"] diff --git a/crates/hotblocks/src/cli.rs b/crates/hotblocks/src/cli.rs index 1e14ff72..1db864ae 100644 --- a/crates/hotblocks/src/cli.rs +++ b/crates/hotblocks/src/cli.rs @@ -61,6 +61,29 @@ pub struct CLI { #[arg(long, value_name = "N", default_value = "10")] pub rocksdb_keep_log_file_num: usize, + /// Concurrent RocksDB background flush + compaction jobs. + /// Defaults to the core count, clamped to 2..=8. + #[arg(long, value_name = "N")] + pub rocksdb_max_background_jobs: Option, + + /// Rewrite every table SST older than this, collecting dead data that never made a file + /// tombstone-dense enough for the deletion collector. Lower means faster reclaim and + /// proportionally more write amplification. 0 disables it, leaving RocksDB's 30-day + /// `ttl` as the only backstop. + #[arg(long, value_name = "SECS", default_value_t = sqd_storage::db::DEFAULT_PERIODIC_COMPACTION_SECS)] + pub rocksdb_periodic_compaction_secs: u64, + + /// Reclaim dead disk at startup: purge orphaned dirty-table markers, then unlink whole + /// SST files below the live-table watermark. Off by default -- the unlink ignores + /// snapshots, so it is only safe before any query exists. Use the `reclaim-measure` + /// binary (shipped in this image) to size the win first. + /// + /// Note that this runs *after* the database opens, and opening replays the WAL and + /// flushes it to L0. A volume at literally zero free bytes needs a few hundred MB + /// cleared by hand before the process can get far enough to reclaim anything. + #[arg(long)] + pub startup_disk_reclaim: bool, + /// Known client IDs for metrics labeling. Client IDs not in this list /// will be reported as "unknown" to prevent metrics cardinality abuse. #[arg(long = "known-client", value_name = "ID")] @@ -80,12 +103,19 @@ impl CLI { pub async fn build_app(&self) -> anyhow::Result { let datasets = DatasetConfig::read_config_file(&self.datasets).context("failed to read datasets config")?; - let db = DatabaseSettings::default() + let mut settings = DatabaseSettings::default() .with_data_cache_size(self.data_cache_size) .with_rocksdb_stats(self.rocksdb_stats) .with_direct_io(!self.rocksdb_disable_direct_io) .with_max_log_file_size(self.rocksdb_max_log_file_size) .with_keep_log_file_num(self.rocksdb_keep_log_file_num) + .with_periodic_compaction_secs(self.rocksdb_periodic_compaction_secs); + + if let Some(jobs) = self.rocksdb_max_background_jobs { + settings = settings.with_max_background_jobs(jobs); + } + + let db = settings .open(&self.database_dir) .map(Arc::new) .context("failed to open rocksdb database")?; @@ -101,7 +131,9 @@ impl CLI { .filter_map(|(id, cfg)| (cfg.retention_strategy == RetentionConfig::Api).then_some(*id)) .collect(); - let data_service = DataService::start(db.clone(), datasets).await.map(Arc::new)?; + let data_service = DataService::start(db.clone(), datasets, self.startup_disk_reclaim) + .await + .map(Arc::new)?; let query_service = { let mut builder = QueryService::builder(db.clone()); diff --git a/crates/hotblocks/src/data_service.rs b/crates/hotblocks/src/data_service.rs index a95256c1..5cf33fdc 100644 --- a/crates/hotblocks/src/data_service.rs +++ b/crates/hotblocks/src/data_service.rs @@ -7,7 +7,7 @@ use anyhow::{Context, anyhow}; use futures::{FutureExt, StreamExt, TryStreamExt}; use sqd_data_client::reqwest::ReqwestDataClient; use sqd_storage::db::DatasetId; -use tracing::{error, info}; +use tracing::{debug, error, info, warn}; use crate::{ dataset_config::{DatasetConfig, RetentionConfig}, @@ -23,14 +23,24 @@ pub struct DataService { } impl DataService { - pub async fn start(db: DBRef, datasets: BTreeMap) -> anyhow::Result { - let all_datasets = db.get_all_datasets()?; - for dataset in all_datasets { - if !datasets.contains_key(&dataset.id) { - info!("deleting unconfigured dataset {}", dataset.id); - if let Err(err) = db.delete_dataset(dataset.id) { - error!("failed to delete dataset {}: {}", dataset.id, err); - } + pub async fn start( + db: DBRef, + datasets: BTreeMap, + disk_reclaim: bool + ) -> anyhow::Result { + let unconfigured: Vec = db + .get_all_datasets()? + .into_iter() + .filter(|ds| !datasets.contains_key(&ds.id)) + .map(|ds| ds.id) + .collect(); + + // Must run before any controller spawns -- see `startup_disk_recovery`. + { + let db = db.clone(); + let recovery = tokio::task::spawn_blocking(move || startup_disk_recovery(&db, &unconfigured, disk_reclaim)); + if let Err(err) = recovery.await { + error!(error =? err, "startup disk recovery panicked"); } } @@ -80,3 +90,54 @@ impl DataService { .ok_or(UnknownDataset { dataset_id }) } } + +/// Startup-only disk recovery; must run before any ingest or query exists (the file +/// unlink ignores snapshots, and the orphan purge treats every dirty marker as an orphan +/// from a dead build). +/// +/// `disk_reclaim` gates both reclaim steps (`--startup-disk-reclaim`, off by default); +/// deleting unconfigured datasets always runs. Ordering matters on a near-full disk: +/// 1. an unlink pass first -- it needs no scratch space, so it frees below-watermark space +/// where every write below would fail with ENOSPC; +/// 2. bookkeeping writes that lift the reclaim watermark: purge orphan dirty markers, +/// delete unconfigured datasets; +/// 3. a second unlink pass to free whatever step 2 unpinned. +/// +/// Every step is best-effort: a failure leaves the watermark pinned until a later startup +/// succeeds, but never blocks startup. +/// +/// Does not rescue a volume at literally zero free bytes: the database has already opened +/// by the time we get here, and opening replays the WAL and flushes it to L0. +fn startup_disk_recovery(db: &DBRef, unconfigured: &[DatasetId], disk_reclaim: bool) { + if disk_reclaim { + if let Err(err) = db.reclaim_disk_space() { + error!(error =? err, "startup disk reclaim (first pass) failed"); + } + + match db.purge_orphan_dirty_tables() { + Ok(0) => {} + Ok(n) => info!("purged {n} orphan dirty table(s) left by an interrupted build"), + Err(err) => warn!(error =? err, "failed to purge orphan dirty tables") + } + } else { + // FUTURE: ungate the orphan purge once the rollout measurement is done. It is safe + // without the unlink, and while gated an interrupted build leaks its data for good; + // it shares the flag only so `reclaim-measure` can still contrast the two watermarks. + info!("startup disk reclaim is off; enable with --startup-disk-reclaim"); + } + + for dataset_id in unconfigured { + info!("deleting unconfigured dataset {dataset_id}"); + if let Err(err) = db.delete_dataset(*dataset_id) { + error!("failed to delete dataset {dataset_id}: {err}; its chunks keep pinning the reclaim watermark"); + } + } + + if disk_reclaim { + if let Err(err) = db.reclaim_disk_space() { + error!(error =? err, "startup disk reclaim (second pass) failed"); + } + } + + debug!("startup disk recovery complete"); +} diff --git a/crates/hotblocks/src/main.rs b/crates/hotblocks/src/main.rs index 409ffecd..ccfdd8e7 100644 --- a/crates/hotblocks/src/main.rs +++ b/crates/hotblocks/src/main.rs @@ -36,6 +36,10 @@ fn main() -> anyhow::Result<()> { .block_on(async { let app = args.build_app().await?; + // NB: startup disk recovery (orphan purge + file unlink, gated by + // --startup-disk-reclaim) already ran inside `build_app` -> `DataService::start`, + // before any controller spawned. + tokio::spawn(db_cleanup_task(app.db.clone())); let api = build_api(app); @@ -89,24 +93,46 @@ async fn shutdown_signal() { } } +const CLEANUP_INTERVAL: Duration = Duration::from_secs(10); +/// Longer pause after a failed tick, so a persistent error (e.g. full disk) doesn't +/// busy-loop failing writes. +const CLEANUP_ERROR_BACKOFF: Duration = Duration::from_secs(30); + +/// Routine Phase-1 cleanup: point-delete deleted tables' data every tick so compaction +/// can reclaim their space -- in normal operation the entire runtime reclaim path. +/// +/// The physical file unlink (`Database::reclaim_disk_space`) is NOT run here: it ignores +/// snapshots and could break an in-flight query, so it runs only at startup. FUTURE: +/// trigger it here under disk pressure too, accepting that risk. #[instrument(name = "db_cleanup", skip_all)] async fn db_cleanup_task(db: DBRef) { - tokio::time::sleep(Duration::from_secs(10)).await; + tokio::time::sleep(CLEANUP_INTERVAL).await; loop { - debug!("db cleanup started"); let db = db.clone(); let result = tokio::task::spawn_blocking(move || db.cleanup()).await; - match result { - Ok(Ok(deleted)) => { - if deleted > 0 { - debug!("purged {} tables", deleted) - } else { - debug!("nothing to purge, pausing cleanup for 10 seconds"); - tokio::time::sleep(Duration::from_secs(10)).await; + + let failed = match result { + Ok(Ok(purged)) => { + if purged > 0 { + debug!("cleanup: logically purged {purged} table(s)"); } + false } - Ok(Err(err)) => error!(error =? err, "database cleanup task failed"), - Err(_) => error!("database cleanup task panicked") - } + Ok(Err(err)) => { + error!(error =? err, "database cleanup failed"); + true + } + Err(_) => { + error!("database cleanup task panicked"); + true + } + }; + + tokio::time::sleep(if failed { + CLEANUP_ERROR_BACKOFF + } else { + CLEANUP_INTERVAL + }) + .await; } } diff --git a/crates/reclaim-measure/Cargo.toml b/crates/reclaim-measure/Cargo.toml new file mode 100644 index 00000000..4d8d4294 --- /dev/null +++ b/crates/reclaim-measure/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "reclaim-measure" +version = "0.0.0" +edition = "2021" +publish = false + +[[bin]] +name = "reclaim-measure" +path = "src/main.rs" + +[dependencies] +sqd-storage = { path = "../storage" } +rocksdb = { version = "0.24.0", features = ["jemalloc"] } +borsh = { workspace = true } +anyhow = { workspace = true } + +[lints] +workspace = true diff --git a/crates/reclaim-measure/src/main.rs b/crates/reclaim-measure/src/main.rs new file mode 100644 index 00000000..480b271f --- /dev/null +++ b/crates/reclaim-measure/src/main.rs @@ -0,0 +1,192 @@ +//! Read-only measurement of reclaimable disk in a hotblocks RocksDB. +//! +//! Opens the live DB as a RocksDB *secondary instance* (no lock, does not disturb the +//! running primary) and reports how much `--startup-disk-reclaim` would file-unlink. +//! +//! Sizes come from `live_files()` metadata -- no SST opens. The rules deciding what is +//! unlinkable live in `sqd_storage::db::reclaim`, shared with the database itself. +//! +//! Usage: reclaim-measure [secondary_scratch_dir] + +use std::collections::BTreeSet; + +use anyhow::{Context, Result}; +use rocksdb::{Options, DB}; +use sqd_storage::db::{ + reclaim::{reclaim_upper_bound, sst_is_unlinkable, watermark}, + Chunk, TableId, CF_CHUNKS, CF_DATASETS, CF_DELETED_TABLES, CF_DIRTY_TABLES, CF_TABLES +}; + +/// Scale to the unit that keeps the number readable -- a prod volume reports GiB, a test db +/// would otherwise report `0.0`. +fn human(bytes: u64) -> String { + const UNITS: [(u64, &str); 3] = [(1 << 30, "GiB"), (1 << 20, "MiB"), (1 << 10, "KiB")]; + for (scale, unit) in UNITS { + if bytes >= scale { + return format!("{:.1} {unit}", bytes as f64 / scale as f64); + } + } + format!("{bytes} B") +} + +/// First 16 bytes of an SST boundary key = the owning `TableId`. +fn id_of(key: &Option>) -> Option { + key.as_ref() + .filter(|k| k.len() >= 16) + .and_then(|k| TableId::try_from_key(&k[..16])) +} + +#[derive(PartialEq)] +enum Cat { + Live, + Dirty, + Dead +} + +fn main() -> Result<()> { + let mut args = std::env::args().skip(1); + let primary = args.next().unwrap_or_else(|| "/run/db".to_string()); + let secondary = args.next().unwrap_or_else(|| "/tmp/reclaim-measure-sec".to_string()); + + let mut opts = Options::default(); + // RocksDB warns against anything else on a secondary: the primary deletes obsolete + // files freely, and only an eagerly-opened table cache keeps them readable here. + opts.set_max_open_files(-1); + + let cfs = [CF_DATASETS, CF_CHUNKS, CF_TABLES, CF_DIRTY_TABLES, CF_DELETED_TABLES]; + eprintln!("[*] opening secondary on {primary} ..."); + let db = DB::open_cf_as_secondary(&opts, &primary, &secondary, cfs) + .with_context(|| format!("open secondary instance on {primary}"))?; + let _ = db.try_catch_up_with_primary(); + + let cf_chunks = db.cf_handle(CF_CHUNKS).context("no CHUNKS cf")?; + let cf_dirty = db.cf_handle(CF_DIRTY_TABLES).context("no DIRTY_TABLES cf")?; + let cf_deleted = db.cf_handle(CF_DELETED_TABLES).context("no DELETED_TABLES cf")?; + + // Live table ids = referenced by a committed chunk. + eprintln!("[*] scanning CHUNKS ..."); + let mut live: BTreeSet = BTreeSet::new(); + let mut chunks = 0u64; + { + let mut it = db.raw_iterator_cf(&cf_chunks); + it.seek_to_first(); + while it.valid() { + let chunk: Chunk = borsh::from_slice(it.value().context("missing chunk value")?).context("decode chunk")?; + for id in chunk.tables().values() { + live.insert(*id); + } + chunks += 1; + it.next(); + } + it.status()?; + } + eprintln!("[*] {chunks} chunks, {} live tables", live.len()); + + // Dirty markers = built but not committed (orphans pin the watermark). + let mut dirty: BTreeSet = BTreeSet::new(); + { + let mut it = db.raw_iterator_cf(&cf_dirty); + it.seek_to_first(); + while it.valid() { + if let Some(id) = TableId::try_from_key(it.key().context("dirty key")?) { + dirty.insert(id); + } + it.next(); + } + it.status()?; + } + + // Pending logical deletes awaiting Phase-1 cleanup. + let mut pending_deletes = 0u64; + { + let mut it = db.raw_iterator_cf(&cf_deleted); + it.seek_to_first(); + while it.valid() { + pending_deletes += 1; + it.next(); + } + it.status()?; + } + + // The watermark today, and the one a startup orphan purge would leave behind. + let wm_now = watermark(live.iter().copied(), dirty.iter().copied()); + let bound_now = reclaim_upper_bound(wm_now); + let bound_after_purge = reclaim_upper_bound(watermark(live.iter().copied(), None)); + + let cat = |id: &TableId| -> Cat { + if live.contains(id) { + Cat::Live + } else if dirty.contains(id) { + Cat::Dirty + } else { + Cat::Dead + } + }; + + eprintln!("[*] reading live SST file metadata ..."); + let files = db.live_files().context("live_files")?; + + let (mut total, mut live_b, mut dirty_b, mut dead_b, mut mixed_b) = (0u64, 0u64, 0u64, 0u64, 0u64); + let (mut below_now, mut below_purge, mut stuck_in_l0) = (0u64, 0u64, 0u64); + let mut n_files = 0u64; + + for f in &files { + if f.column_family_name != CF_TABLES { + continue; + } + let sz = f.size as u64; + total += sz; + n_files += 1; + + // A file whose two ends are dead but which spans a live table in between still + // counts dead, so `dead_b` is an upper bound. + match (id_of(&f.start_key), id_of(&f.end_key)) { + (Some(a), Some(b)) => match (cat(&a), cat(&b)) { + (Cat::Live, Cat::Live) => live_b += sz, + (Cat::Dirty, Cat::Dirty) => dirty_b += sz, + (Cat::Dead, Cat::Dead) => dead_b += sz, + _ => mixed_b += sz + }, + _ => mixed_b += sz + } + + let end = f.end_key.as_deref(); + if sst_is_unlinkable(f.level, end, &bound_now) { + below_now += sz; + } + if sst_is_unlinkable(f.level, end, &bound_after_purge) { + below_purge += sz; + } + // Dead bytes the unlink cannot touch, because DeleteFilesInRange skips level 0. + if f.level == 0 && sst_is_unlinkable(1, end, &bound_after_purge) { + stuck_in_l0 += sz; + } + } + + println!("== reclaim-measure: {primary} =="); + println!("chunks : {chunks}"); + println!("live tables : {}", live.len()); + println!("dirty markers (orphans) : {}", dirty.len()); + println!("pending deletes : {pending_deletes}"); + println!("TABLES sst files : {n_files}"); + match wm_now { + Some(w) => println!("watermark (live+dirty) : {w}"), + None => println!("watermark (live+dirty) : NONE (no live tables)") + } + println!(); + println!("CF_TABLES total : {:>12}", human(total)); + println!(" live (chunk-ref'd) : {:>12}", human(live_b)); + println!(" dirty (orphan builds) : {:>12}", human(dirty_b)); + println!(" DEAD (upper bound) : {:>12}", human(dead_b)); + println!(" mixed (boundary SSTs) : {:>12}", human(mixed_b)); + println!(); + println!("file-unlinkable at startup (whole SSTs below wm, above L0):"); + println!(" below wm now : {:>12}", human(below_now)); + println!(" below wm after purge : {:>12}", human(below_purge)); + println!( + " below wm but in L0 : {:>12} (unlink skips L0; compaction only)", + human(stuck_in_l0) + ); + + Ok(()) +} diff --git a/crates/storage/src/db/db.rs b/crates/storage/src/db/db.rs index 34e0d71c..07736365 100644 --- a/crates/storage/src/db/db.rs +++ b/crates/storage/src/db/db.rs @@ -12,15 +12,21 @@ use super::{ use crate::db::{ ops::{perform_dataset_compaction, CompactionStatus}, read::datasets::list_all_datasets, - write::{ops::deleted_deleted_tables, table_builder::TableBuilder, tx::Tx}, + write::{ops as cleanup_ops, table_builder::TableBuilder, tx::Tx}, Chunk, DatasetUpdate }; -pub(super) const CF_DATASETS: Name = "DATASETS"; -pub(super) const CF_CHUNKS: Name = "CHUNKS"; -pub(super) const CF_TABLES: Name = "TABLES"; -pub(super) const CF_DIRTY_TABLES: Name = "DIRTY_TABLES"; -pub(super) const CF_DELETED_TABLES: Name = "DELETED_TABLES"; +// Public so out-of-process readers (`reclaim-measure`) don't copy the strings. +pub const CF_DATASETS: Name = "DATASETS"; +pub const CF_CHUNKS: Name = "CHUNKS"; +pub const CF_TABLES: Name = "TABLES"; +pub const CF_DIRTY_TABLES: Name = "DIRTY_TABLES"; +pub const CF_DELETED_TABLES: Name = "DELETED_TABLES"; + +/// Whole-file rewrite cadence for `CF_TABLES`. RocksDB leaves `periodic_compaction_seconds` +/// disabled for leveled compaction without a compaction filter, so the effective baseline +/// is the separate 30-day `ttl` default. 7 days buys ~4x that rewrite rate. +pub const DEFAULT_PERIODIC_COMPACTION_SECS: u64 = 7 * 24 * 60 * 60; pub(super) type RocksDB = rocksdb::OptimisticTransactionDB; pub(super) type RocksTransaction<'a> = rocksdb::Transaction<'a, RocksDB>; @@ -36,7 +42,19 @@ pub struct DatabaseSettings { direct_io: bool, cache_index_and_filter_blocks: bool, max_log_file_size: usize, - keep_log_file_num: usize + keep_log_file_num: usize, + auto_compactions: bool, + max_background_jobs: usize, + periodic_compaction_secs: u64 +} + +/// RocksDB's default of 2 could not keep up with ingest during the NET-819 incident, but a +/// fixed 8 would just thrash a 2-core node. +fn default_max_background_jobs() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(2) + .clamp(2, 8) } impl Default for DatabaseSettings { @@ -48,7 +66,10 @@ impl Default for DatabaseSettings { direct_io: false, cache_index_and_filter_blocks: false, max_log_file_size: 10, - keep_log_file_num: 10 + keep_log_file_num: 10, + auto_compactions: true, + max_background_jobs: default_max_background_jobs(), + periodic_compaction_secs: DEFAULT_PERIODIC_COMPACTION_SECS } } } @@ -91,6 +112,28 @@ impl DatabaseSettings { self } + /// Enable/disable RocksDB background auto-compaction of the table data. + /// Defaults to `true`; off mainly for tests needing deterministic compaction. + pub fn with_auto_compactions(mut self, yes: bool) -> Self { + self.auto_compactions = yes; + self + } + + /// Concurrent RocksDB background flush + compaction jobs. Defaults to the core count, + /// clamped to `2..=8`. + pub fn with_max_background_jobs(mut self, jobs: usize) -> Self { + self.max_background_jobs = jobs.max(1); + self + } + + /// Rewrite every `CF_TABLES` SST older than this, collecting dead data the deletion + /// collector never triggers on. `0` disables it, leaving the 30-day `ttl` as the only + /// backstop. See [`DEFAULT_PERIODIC_COMPACTION_SECS`]. + pub fn with_periodic_compaction_secs(mut self, secs: u64) -> Self { + self.periodic_compaction_secs = secs; + self + } + fn db_options(&self) -> RocksOptions { let mut options = RocksOptions::default(); options.create_if_missing(true); @@ -99,6 +142,11 @@ impl DatabaseSettings { // Bound info log (LOG, LOG.old.*) growth options.set_max_log_file_size(self.max_log_file_size * 1024 * 1024); options.set_keep_log_file_num(self.keep_log_file_num); + // Keep compaction ahead of ingest + deletion tombstones: the default 2 jobs could + // not keep up during the NET-819 incident, and routine reclaim now relies entirely + // on compaction (the file unlink is startup-only). + options.set_max_background_jobs(self.max_background_jobs as i32); + options.set_max_subcompactions((self.max_background_jobs as u32 / 2).max(1)); if self.with_rocksdb_stats { options.enable_statistics(); } @@ -141,6 +189,19 @@ impl DatabaseSettings { let mut options = RocksOptions::default(); options.set_block_based_table_factory(&block_based_table_factory); options.set_compression_type(rocksdb::DBCompressionType::Lz4); + // A table purge point-deletes all of its keys, so its SSTs turn tombstone-dense and + // the deletion collector compacts them out; periodic compaction backstops files + // that never cross the density threshold, rewriting live data along with them. + // Thresholds are provisional. + options.add_compact_on_deletion_collector_factory(128 * 1024, 64 * 1024, 0.5); + if self.periodic_compaction_secs > 0 { + options.set_periodic_compaction_seconds(self.periodic_compaction_secs); + } + // Bound space amplification (default since RocksDB 8.4; pinned -- load-bearing here). + options.set_level_compaction_dynamic_level_bytes(true); + if !self.auto_compactions { + options.set_disable_auto_compactions(true); + } options } @@ -287,8 +348,41 @@ impl Database { Ok(()) } + /// Phase 1 -- logically purge deleted tables (snapshot-safe point deletes). + /// Returns the number of tables logically deleted by this call. pub fn cleanup(&self) -> anyhow::Result { - deleted_deleted_tables(&self.db) + cleanup_ops::logical_cleanup(&self.db) + } + + /// Physically reclaim disk by unlinking whole SST files below the live watermark + /// (min live `TableId`). Needs no scratch space, unlike compaction, which must write + /// its merged output before dropping the inputs. IGNORES snapshots, so it is safe only + /// where no live reader exists -- today only at STARTUP. + /// See [`cleanup_ops::reclaim_disk_space`]. + pub fn reclaim_disk_space(&self) -> anyhow::Result<()> { + cleanup_ops::reclaim_disk_space(&self.db) + } + + /// Crash recovery: purge `DIRTY_TABLES` markers left by builds that died before + /// commit -- an orphan pins the reclaim watermark forever. MUST run before any ingest + /// starts: it treats every dirty marker as an orphan. Returns orphans purged. + pub fn purge_orphan_dirty_tables(&self) -> anyhow::Result { + cleanup_ops::purge_orphan_dirty_tables(&self.db) + } + + /// Flush `CF_TABLES`'s memtable to SST files (e.g. before a reclaim, so freshly + /// written data is unlinkable). Bookkeeping column families are not flushed. + pub fn flush_tables(&self) -> anyhow::Result<()> { + self.db.flush_cf(self.db.cf_handle(CF_TABLES).unwrap())?; + Ok(()) + } + + /// Force a full compaction of the table-data column family. Test support: production + /// relies on background compaction. Rewrites files, so it needs scratch space -- the + /// very thing that deadlocks on the full disk [`Database::reclaim_disk_space`] exists for. + pub fn compact_tables(&self) { + let cf = self.db.cf_handle(CF_TABLES).unwrap(); + self.db.compact_range_cf(cf, None::<&[u8]>, None::<&[u8]>); } pub fn get_statistics(&self) -> Option { diff --git a/crates/storage/src/db/mod.rs b/crates/storage/src/db/mod.rs index bb4007cf..90e06547 100644 --- a/crates/storage/src/db/mod.rs +++ b/crates/storage/src/db/mod.rs @@ -2,6 +2,7 @@ mod data; mod db; pub mod ops; mod read; +pub mod reclaim; mod rocks; mod table_id; mod write; diff --git a/crates/storage/src/db/reclaim.rs b/crates/storage/src/db/reclaim.rs new file mode 100644 index 00000000..be068b23 --- /dev/null +++ b/crates/storage/src/db/reclaim.rs @@ -0,0 +1,95 @@ +//! Rules shared by [`crate::db::Database::reclaim_disk_space`] and the out-of-process +//! `reclaim-measure` probe, so the number the probe prints matches what the flag frees. +//! (They cannot share a scan: one runs on a secondary `rocksdb::DB`, the other on an +//! `OptimisticTransactionDB`, and `rocksdb::DBInner` is not exported.) + +use crate::db::table_id::TableId; + +/// Lower bound for `DeleteFilesInRange`. Keys are `table_id (16B) ++ tag ++ ..`. +pub const RECLAIM_LOWER_BOUND: [u8; 16] = [0u8; 16]; + +/// Upper bound for `DeleteFilesInRange`. +/// +/// A bare 16-byte id is exclusive in practice even though the C binding passes +/// `include_end = true` (keeping files with `largest_user_key <= end`): `TableKeyFactory` +/// always appends a tag byte, so no key is ever exactly 16 bytes. With nothing live, +/// 17 x `0xFF` sorts above every key. +pub fn reclaim_upper_bound(watermark: Option) -> Vec { + match watermark { + Some(id) => id.as_ref().to_vec(), + None => vec![0xFFu8; 17] + } +} + +/// Smallest table id still live: referenced by a committed chunk, or pending in +/// `CF_DIRTY_TABLES`. Taken across all datasets, which absorbs UUIDv7 clock skew. +pub fn watermark(live: impl IntoIterator, dirty: impl IntoIterator) -> Option { + live.into_iter().chain(dirty).min() +} + +/// Whether `DeleteFilesInRange(.., RECLAIM_LOWER_BOUND, upper_bound)` unlinks this SST. +/// +/// Mirrors `DBImpl::DeleteFilesInRanges`: only files lying entirely inside the range are +/// dropped, and its loop starts at level 1, so L0 is skipped outright. RocksDB also skips +/// files it is compacting, invisible from here -- this is an upper bound. +pub fn sst_is_unlinkable(level: i32, end_key: Option<&[u8]>, upper_bound: &[u8]) -> bool { + level > 0 && end_key.is_some_and(|end| end <= upper_bound) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn id(byte: u8) -> TableId { + TableId::try_from_key(&[byte; 16]).unwrap() + } + + #[test] + fn upper_bound_excludes_the_watermark_tables_own_keys() { + let wm = id(7); + let bound = reclaim_upper_bound(Some(wm)); + + // The watermark table's own keys sort above the bound, so its files survive. + let mut smallest_key = wm.as_ref().to_vec(); + smallest_key.push(0); + assert!(smallest_key.as_slice() > bound.as_slice()); + assert!(RECLAIM_LOWER_BOUND.as_slice() < bound.as_slice()); + } + + #[test] + fn no_live_table_means_every_key_is_below_the_bound() { + let bound = reclaim_upper_bound(None); + let mut largest_key = [0xFFu8; 16].to_vec(); + largest_key.push(0xFF); // above any real tag byte + assert!(largest_key.as_slice() <= bound.as_slice()); + } + + #[test] + fn watermark_is_the_min_over_live_and_dirty() { + let (a, b) = (id(1), id(2)); + assert_eq!(watermark([b], [a]), Some(a)); + assert_eq!(watermark([a], [b]), Some(a)); + assert_eq!(watermark(None, Some(b)), Some(b)); + assert_eq!(watermark(None::, None), None); + } + + #[test] + fn level_zero_files_are_never_unlinkable() { + let bound = reclaim_upper_bound(None); + assert!(!sst_is_unlinkable(0, Some(b"anything"), &bound)); + assert!(sst_is_unlinkable(1, Some(b"anything"), &bound)); + assert!(!sst_is_unlinkable(1, None, &bound)); + } + + #[test] + fn files_reaching_above_the_bound_are_kept() { + let wm = id(7); + let bound = reclaim_upper_bound(Some(wm)); + + let mut reaches_into_wm = wm.as_ref().to_vec(); + reaches_into_wm.push(0); + assert!(!sst_is_unlinkable(3, Some(&reaches_into_wm), &bound)); + + assert!(sst_is_unlinkable(3, Some(&RECLAIM_LOWER_BOUND), &bound)); + } +} diff --git a/crates/storage/src/db/table_id.rs b/crates/storage/src/db/table_id.rs index dfca7357..d88fb1ae 100644 --- a/crates/storage/src/db/table_id.rs +++ b/crates/storage/src/db/table_id.rs @@ -19,10 +19,12 @@ impl TableId { Self { uuid: Uuid::now_v7() } } - pub fn from_slice(bytes: &[u8]) -> Self { - Self { - uuid: Uuid::from_slice(bytes).unwrap() - } + /// Decode a raw column-family key, or `None` if it is not exactly a 16-byte UUID, so a + /// corrupt bookkeeping key is skipped instead of wedging the scan with a panic. + /// + /// Not named `try_from_slice`: that would shadow the `BorshDeserialize` method. + pub fn try_from_key(bytes: &[u8]) -> Option { + Uuid::from_slice(bytes).ok().map(|uuid| Self { uuid }) } } @@ -31,3 +33,19 @@ impl Display for TableId { self.uuid.fmt(f) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn try_from_key_rejects_non_16_byte_keys() { + let id = TableId::new(); + assert_eq!(TableId::try_from_key(id.as_ref()), Some(id)); + + // Anything not exactly 16 bytes decodes to None instead of panicking. + assert_eq!(TableId::try_from_key(&[]), None); + assert_eq!(TableId::try_from_key(&[0u8; 15]), None); + assert_eq!(TableId::try_from_key(&[0u8; 17]), None); + } +} diff --git a/crates/storage/src/db/write/ops.rs b/crates/storage/src/db/write/ops.rs index 3e0ce35e..2dbc6cd3 100644 --- a/crates/storage/src/db/write/ops.rs +++ b/crates/storage/src/db/write/ops.rs @@ -1,56 +1,170 @@ use crate::{ - db::db::{RocksDB, RocksWriteBatch, CF_DELETED_TABLES, CF_DIRTY_TABLES, CF_TABLES}, - kv::KvReadCursor, + db::{ + db::{RocksDB, RocksWriteBatch, CF_CHUNKS, CF_DELETED_TABLES, CF_DIRTY_TABLES, CF_TABLES}, + reclaim::{reclaim_upper_bound, watermark, RECLAIM_LOWER_BOUND}, + table_id::TableId, + Chunk + }, table::key::TableKeyFactory }; -pub fn deleted_deleted_tables(db: &RocksDB) -> anyhow::Result { - let mut deleted = 0; - let cf_deleted_tables = db.cf_handle(CF_DELETED_TABLES).unwrap(); - let mut it = db.raw_iterator_cf(cf_deleted_tables); - for_each_key(&mut it, |key| { - deleted += 1; - delete_table(db, key) - })?; - Ok(deleted) +/// Phase 1 -- logical, snapshot-safe purge of deleted tables. +/// +/// Point-deletes the `CF_TABLES` keys of every table in `CF_DELETED_TABLES`, then drops +/// its bookkeeping entry. The deletes are MVCC-versioned, so in-flight queries are +/// unaffected; the space is freed later by compaction. Idempotent. Returns tables purged. +pub(crate) fn logical_cleanup(db: &RocksDB) -> anyhow::Result { + // Collect first, then mutate: writing mid-iteration disturbs the cursor. + let (pending, malformed) = scan_table_ids(db, CF_DELETED_TABLES)?; + drop_malformed_keys(db, CF_DELETED_TABLES, &malformed)?; + + for id in &pending { + purge_table(db, id)?; + } + + Ok(pending.len()) } -fn delete_table(db: &RocksDB, table_id: &[u8]) -> anyhow::Result<()> { - let mut key1 = TableKeyFactory::new(table_id); - let mut key2 = TableKeyFactory::new(table_id); - let start = key1.start(); - let end = key2.end(); +/// Scan `cf`, splitting keys into well-formed `TableId`s and malformed leftovers. +/// Malformed keys are returned rather than skipped, so callers can delete them via +/// [`drop_malformed_keys`] -- left in place, every pass would rescan them forever. +fn scan_table_ids(db: &RocksDB, cf: &str) -> anyhow::Result<(Vec, Vec>)> { + let cf = db.cf_handle(cf).unwrap(); + let mut ids = Vec::new(); + let mut malformed = Vec::new(); - let cf_tables = db.cf_handle(CF_TABLES).unwrap(); + let mut it = db.raw_iterator_cf(cf); + it.seek_to_first(); + while it.valid() { + let key = it.key().unwrap(); + match TableId::try_from_key(key) { + Some(id) => ids.push(id), + None => malformed.push(key.to_vec()) + } + it.next(); + } + it.status()?; + + Ok((ids, malformed)) +} + +/// Delete malformed bookkeeping keys outright -- no other code path can remove them. +fn drop_malformed_keys(db: &RocksDB, cf: &str, keys: &[Vec]) -> anyhow::Result<()> { + if keys.is_empty() { + return Ok(()); + } + let cf = db.cf_handle(cf).unwrap(); let mut batch = RocksWriteBatch::default(); - let mut cursor = db.raw_iterator_cf(cf_tables); + for key in keys { + batch.delete_cf(cf, key); + } + db.write(batch)?; + Ok(()) +} - list_keys(&mut cursor, start, end, |key| batch.delete_cf(cf_tables, key))?; +/// Point-delete all of `id`'s `CF_TABLES` data and drop its bookkeeping entries. One write +/// batch per table, so a large purge never holds every table's deletes in memory at once. +/// +/// A table is normally in only one bookkeeping CF, but clearing both costs one tombstone +/// and guarantees a purge never leaves a stale marker pinning the watermark. +/// +/// Point deletes rather than a single `delete_range` tombstone: range deletions are not +/// officially supported on the transactional `OptimisticTransactionDB`. +fn purge_table(db: &RocksDB, id: &TableId) -> anyhow::Result<()> { + let cf_tables = db.cf_handle(CF_TABLES).unwrap(); + let cf_deleted = db.cf_handle(CF_DELETED_TABLES).unwrap(); + let cf_dirty = db.cf_handle(CF_DIRTY_TABLES).unwrap(); - let cf_dirty_tables = db.cf_handle(CF_DIRTY_TABLES).unwrap(); - batch.delete_cf(cf_dirty_tables, table_id); + let mut start = TableKeyFactory::new(id); + let mut end = TableKeyFactory::new(id); + let start_key = start.start(); + let end_key = end.end(); - let cf_deleted_tables = db.cf_handle(CF_DELETED_TABLES).unwrap(); - batch.delete_cf(cf_deleted_tables, table_id); + let mut batch = RocksWriteBatch::default(); + let mut it = db.raw_iterator_cf(cf_tables); + it.seek(start_key); + while it.valid() { + let key = it.key().unwrap(); + if key >= end_key { + break; + } + batch.delete_cf(cf_tables, key); + it.next(); + } + it.status()?; + batch.delete_cf(cf_deleted, id); + batch.delete_cf(cf_dirty, id); db.write(batch)?; + Ok(()) } -fn list_keys(cursor: &mut impl KvReadCursor, from: &[u8], to: &[u8], mut cb: impl FnMut(&[u8])) -> anyhow::Result<()> { - cursor.seek(from)?; - while cursor.is_valid() && cursor.key() < to { - cb(cursor.key()); - cursor.next()?; - } +/// Physically reclaim disk by unlinking whole `CF_TABLES` SST files below the live +/// watermark ([`min_live_table_id`]). Table ids are time-ordered UUIDv7s, so dead tables +/// form a contiguous low range; boundary, level-0 and above-watermark garbage is left to +/// compaction. +/// +/// Needs no scratch space, where compaction must write its merged output before dropping +/// the inputs and so deadlocks on a full disk. Not literally write-free: RocksDB appends a +/// `VersionEdit` to the MANIFEST, a few hundred bytes against gigabytes freed. +/// +/// SAFETY: the unlink IGNORES snapshots -- it can break an in-flight query reading a +/// just-deleted table below the watermark. So it runs only at STARTUP, before any +/// controller/query exists. FUTURE: trigger under runtime disk pressure too, accepting that +/// read risk. That, not this, is the answer to a genuinely full disk -- getting here at all +/// requires the database to have opened, and opening replays the WAL and flushes it to L0. +pub(crate) fn reclaim_disk_space(db: &RocksDB) -> anyhow::Result<()> { + let cf_tables = db.cf_handle(CF_TABLES).unwrap(); + let hi = reclaim_upper_bound(min_live_table_id(db)?); + db.delete_file_in_range_cf(cf_tables, RECLAIM_LOWER_BOUND.as_slice(), hi.as_slice())?; Ok(()) } -fn for_each_key(cursor: &mut impl KvReadCursor, mut cb: impl FnMut(&[u8]) -> anyhow::Result<()>) -> anyhow::Result<()> { - cursor.seek_first()?; - while cursor.is_valid() { - cb(cursor.key())?; - cursor.next()?; +/// Crash recovery -- purge orphaned `CF_DIRTY_TABLES` markers. +/// +/// A dirty marker is written when a build starts and removed when its chunk commits. +/// One left by a build that died before commit counts as live in [`min_live_table_id`], +/// pinning [`reclaim_disk_space`]'s watermark forever. We drop the marker and +/// point-delete its unreferenced data. +/// +/// MUST run only with no build in flight (e.g. startup before ingest): it treats EVERY +/// dirty marker as an orphan, so it would tombstone a live build's data. Returns orphans +/// purged. +pub(crate) fn purge_orphan_dirty_tables(db: &RocksDB) -> anyhow::Result { + // Collect first, mutate after: writing while iterating disturbs the cursor. + let (orphans, malformed) = scan_table_ids(db, CF_DIRTY_TABLES)?; + drop_malformed_keys(db, CF_DIRTY_TABLES, &malformed)?; + + for id in &orphans { + purge_table(db, id)?; } - Ok(()) + + Ok(orphans.len()) +} + +/// See [`crate::db::reclaim::watermark`]. +/// +/// An undecodable `CF_CHUNKS` value aborts with `Err` rather than being skipped, since +/// skipping could lift the watermark over live data. The watermark is global, so one bad +/// chunk anywhere disables the unlink for all datasets -- the safe failure mode. +fn min_live_table_id(db: &RocksDB) -> anyhow::Result> { + let mut min: Option = None; + { + let cf_chunks = db.cf_handle(CF_CHUNKS).unwrap(); + let mut it = db.raw_iterator_cf(cf_chunks); + it.seek_to_first(); + while it.valid() { + let chunk: Chunk = borsh::from_slice(it.value().unwrap())?; + min = watermark(min, chunk.tables().values().copied()); + it.next(); + } + it.status()?; + } + + // Malformed dirty keys pin no data, and the startup purge deletes them anyway. This + // path must not write -- it serves reclaim on a disk with no room for writes. + let (dirty, _) = scan_table_ids(db, CF_DIRTY_TABLES)?; + + Ok(watermark(min, dirty)) } diff --git a/crates/storage/src/db/write/storage.rs b/crates/storage/src/db/write/storage.rs index f7cd89b6..644c2a86 100644 --- a/crates/storage/src/db/write/storage.rs +++ b/crates/storage/src/db/write/storage.rs @@ -25,6 +25,8 @@ impl<'a> TableStorage<'a> { pub fn mark_table_dirty(&mut self, table_id: TableId) { let cf_dirty = self.db.cf_handle(CF_DIRTY_TABLES).unwrap(); + // Value unused; the key marks "table built, chunk not yet committed". Removed on + // commit; an orphan is cleared by `ops::purge_orphan_dirty_tables` at startup. self.write_batch.put_cf(cf_dirty, table_id, []) } diff --git a/crates/storage/src/db/write/tx.rs b/crates/storage/src/db/write/tx.rs index b7ba820f..4fca2e80 100644 --- a/crates/storage/src/db/write/tx.rs +++ b/crates/storage/src/db/write/tx.rs @@ -126,6 +126,8 @@ impl<'a> Tx<'a> { } pub fn delete_table(&self, table_id: &TableId) -> anyhow::Result<()> { + // Value unused; the key's presence is the signal. `ops::logical_cleanup` + // point-deletes the table's data and drops this entry. self.transaction .put_cf(self.cf_handle(CF_DELETED_TABLES), table_id, [])?; Ok(()) diff --git a/crates/storage/tests/cleanup_reclaim.rs b/crates/storage/tests/cleanup_reclaim.rs new file mode 100644 index 00000000..edb2a98e --- /dev/null +++ b/crates/storage/tests/cleanup_reclaim.rs @@ -0,0 +1,267 @@ +//! Tests for table cleanup: +//! * Phase 1 ([`Database::cleanup`]) -- logical, snapshot-safe point deletes. +//! * Physical reclaim ([`Database::reclaim_disk_space`]) -- SST-file unlink below the +//! live watermark. It ignores snapshots, so it runs only where there are no live +//! readers (startup in production); fully decoupled from Phase 1. + +mod mock_db; +#[allow(dead_code)] // shared with the other integration tests; this one uses a subset +mod utils; + +use mock_db::{MockDB, Table}; + +/// Phase 1 is invisible to readers that already hold a snapshot: a query in flight when +/// its chunk is deleted keeps reading every row, while new snapshots see nothing. +#[test] +fn logical_delete_is_snapshot_safe() { + let mut db = MockDB::new(); + let t = db.commit_table(50); + + // Snapshot taken BEFORE the deletion -- models an in-flight query. + let reader = db.snapshot(); + + db.delete(&t); + assert_eq!(db.cleanup(), 1, "one table logically deleted"); + + // A fresh snapshot no longer sees the chunk... + assert!(!db.has_visible_chunk()); + // ...but the pre-deletion snapshot still reads every row. + assert_eq!(db.read(&reader, &t), t.rows); + + assert_eq!(db.cleanup(), 0, "re-running Phase 1 is a no-op"); +} + +/// `DeleteFilesInRange` starts its level loop at 1, so dead data still sitting in L0 is not +/// unlinkable however far below the watermark it is. `reclaim-measure` reports those bytes +/// separately for exactly this reason. +#[test] +fn reclaim_skips_level_zero_files() { + let mut db = MockDB::new(); + let t = db.commit_table(2000); + db.flush(); // memtable -> L0, and auto-compaction is off, so it stays there + let before = db.sst_size(); + assert!(before > 0, "expected flushed SST data"); + + // No live table remains, so the watermark is unbounded and every file is below it. + db.delete(&t); + db.reclaim(); + assert_eq!(db.sst_size(), before, "an L0 file is never unlinked"); + + // The identical call frees the same data once compaction has moved it off L0. + db.compact_to_bottom(); + db.reclaim(); + let after = db.sst_size(); + assert!( + after * 4 < before, + "expected reclaim below L0: before={before} after={after}" + ); +} + +/// Physical reclaim unlinks dead SST files below the watermark. With no live table left, +/// the watermark is unbounded and every dead file is dropped. +#[test] +fn reclaim_unlinks_dead_sst_files() { + let mut db = MockDB::new(); + let tables: Vec = (0..3).map(|_| db.commit_table(1000)).collect(); + db.compact_to_bottom(); + let before = db.sst_size(); + assert!(before > 0, "expected flushed SST data"); + + for t in &tables { + db.delete(t); + } + assert_eq!(db.cleanup(), 3); + // Flush the Phase-1 tombstones out of the memtable. Auto-compaction is off, so + // nothing but the unlink can free space here. + db.flush(); + + db.reclaim(); + let after = db.sst_size(); + assert!( + after * 4 < before, + "expected physical reclaim: before={before} after={after}" + ); +} + +/// The watermark is the min live `TableId` over ALL datasets, so a single old *live* +/// table pins it low and dead tables with larger ids are not file-reclaimable until that +/// live table is gone. Live data is never unlinked; the known limitation is heterogeneous +/// retention. +#[test] +fn live_table_pins_reclaim_watermark() { + let mut db = MockDB::new(); + // Creation order == id order, so `older` is the smaller id. + let older = db.commit_table(1000); + let newer = db.commit_table(1000); + db.compact_to_bottom(); + let before = db.sst_size(); + assert!(before > 0); + + // Delete the NEWER table; the older one stays live and pins the watermark. + db.delete(&newer); + db.cleanup(); + db.flush(); + + // Dead `newer` sits above the watermark pinned by live `older`, so its files survive + // and `older` stays readable. Auto-compaction is off, so only the unlink can shrink + // CF_TABLES -- nothing was unlinked iff the size did not drop. + db.reclaim(); + assert!( + db.sst_size() >= before, + "a dead table above the watermark must NOT be unlinked" + ); + assert_eq!(db.read(&db.snapshot(), &older), older.rows); + + // Remove the older table too: the watermark lifts and everything is reclaimable. + db.delete(&older); + db.cleanup(); + db.flush(); + db.reclaim(); + let after = db.sst_size(); + assert!( + after * 4 < before, + "expected reclaim once the watermark lifts: before={before} after={after}" + ); +} + +/// Both steps are idempotent -- crash safety relies on it. +#[test] +fn cleanup_and_reclaim_are_idempotent() { + let mut db = MockDB::new(); + let t = db.commit_table(1000); + db.delete(&t); + + assert_eq!(db.cleanup(), 1); + assert_eq!(db.cleanup(), 0); + db.flush(); + db.reclaim(); + db.reclaim(); // a second run must not panic or error + assert_eq!(db.cleanup(), 0); +} + +/// End-to-end: `delete_dataset` runs Phase 1 synchronously, and the +/// subsequent startup reclaim (no readers) frees the disk. +#[test] +fn delete_dataset_then_reclaim_frees_space() { + let mut db = MockDB::new(); + let _t0 = db.commit_table(1000); + let _t1 = db.commit_table(1000); + db.compact_to_bottom(); + let before = db.sst_size(); + assert!(before > 0); + + db.delete_dataset(); + assert!(db.has_no_datasets()); + db.flush(); // flush Phase-1 tombstones so the unlink can drop the files + + db.reclaim(); + let after = db.sst_size(); + assert!(after * 4 < before, "expected reclaim: before={before} after={after}"); +} + +/// The trade-off behind running [`Database::reclaim_disk_space`] only at startup: the +/// unlink ignores snapshots, so reclaiming under a live pre-deletion snapshot pulls its +/// files out and the read fails loudly, rather than returning wrong rows. Cache disabled +/// so the data is genuinely gone, not served warm. +#[test] +fn reclaim_breaks_a_live_pre_deletion_reader() { + let mut db = MockDB::uncached(); + let t = db.commit_table(3000); + db.compact_to_bottom(); + + // In-flight query: snapshot taken BEFORE the deletion. + let reader = db.snapshot(); + + db.delete(&t); + assert_eq!(db.cleanup(), 1); + db.flush(); + + // Baseline: the reader still reads every row while its files are present. + assert_eq!(db.read(&reader, &t), t.rows); + + // Reclaim while that snapshot is STILL live. No live table remains, so the watermark + // is unbounded and the reader's files are unlinked from under it. + db.reclaim(); + + assert!( + db.try_read(&reader, &t).is_err(), + "reading a table whose files were unlinked under a live snapshot must fail, not return wrong rows" + ); +} + +/// An orphaned `DIRTY_TABLES` marker -- left when a build dies before committing its +/// chunk -- counts as a live table, pinning disk reclaim for every later table. Startup +/// recovery drops it so the watermark lifts. +#[test] +fn orphan_dirty_marker_unpinned_by_purge() { + let mut db = MockDB::new(); + // Orphan created first, so it has the smaller id and pins the watermark low. + let orphan = db.orphan_table(1000); + let live = db.commit_table(1000); + assert!( + orphan.id < live.id, + "orphan must be the smaller id to pin the watermark" + ); + + db.compact_to_bottom(); + let before = db.sst_size(); + assert!(before > 0); + + db.delete(&live); + assert_eq!(db.cleanup(), 1); + db.flush(); + + // The orphan pins the watermark at its own id, so the dead (higher-id) table cannot + // be unlinked. + db.reclaim(); + assert!( + db.sst_size() >= before, + "orphan pins the watermark, the dead table's files survive" + ); + + // Startup recovery removes the orphan marker (and tombstones its data). + assert_eq!(db.purge_orphans(), 1, "one orphan marker purged"); + db.flush(); + + // The watermark lifts, so the dead table is now reclaimable. + db.reclaim(); + let after = db.sst_size(); + assert!( + after * 4 < before, + "expected physical reclaim once the orphan no longer pins the watermark: before={before} after={after}" + ); + + assert_eq!(db.purge_orphans(), 0, "purge is idempotent"); +} + +/// Phase 1 and physical reclaim are decoupled: reclaim unlinks dead files purely by the +/// watermark, with no dependence on a preceding `cleanup()` -- mirroring the startup path, +/// which runs before any Phase 1. It must also leave the bookkeeping Phase 1 still owes. +#[test] +fn reclaim_is_independent_of_phase1() { + let mut db = MockDB::new(); + let t = db.commit_table(2000); + db.compact_to_bottom(); + let before = db.sst_size(); + assert!(before > 0); + + // Logical-delete but do NOT run Phase 1, so no point deletes are issued. + db.delete(&t); + + // No live table remains, so the watermark is unbounded and the dead table's files are + // unlinked with no preceding tombstone. + db.reclaim(); + let after = db.sst_size(); + assert!( + after * 4 < before, + "reclaim frees dead files without a preceding Phase 1: before={before} after={after}" + ); + + // Reclaim left the bookkeeping entry untouched, so a later Phase 1 still finds it. + assert_eq!( + db.cleanup(), + 1, + "the deletion record survived reclaim and is handled by Phase 1" + ); + assert_eq!(db.cleanup(), 0); +} diff --git a/crates/storage/tests/mock_db/mod.rs b/crates/storage/tests/mock_db/mod.rs new file mode 100644 index 00000000..b93167ef --- /dev/null +++ b/crates/storage/tests/mock_db/mod.rs @@ -0,0 +1,199 @@ +//! [`MockDB`] -- a scratch temp-dir `Database` for driving the table cleanup lifecycle in +//! tests: logical purge (Phase 1, point deletes) + physical SST-file unlink below the live +//! watermark. It hides the fiddly setup (auto-compaction off, small write buffers, UUIDv7 +//! id ordering, flush/compact to the bottom level). + +use std::{sync::Arc, time::Duration}; + +use arrow::datatypes::{DataType, Schema}; +use sqd_storage::{ + db::{Chunk, Database, DatabaseSettings, DatasetId, DatasetKind, ReadSnapshot, TableId}, + table::write::{use_small_buffers, RestoreBufferSizesGuard} +}; +use tempfile::TempDir; + +use crate::utils::{make_irregular_block, make_schema, read_chunk}; + +/// Generous upper bound on the total number of rows a single test builds across +/// all its tables (block numbers double as indices into the shared column data). +const CAPACITY: usize = 20_000; + +fn two_u16_columns(n: usize) -> Vec> { + let col: Vec = (0..n).map(|i| (i % 60_000) as u16).collect(); + vec![col.clone(), col] +} + +/// A table [`MockDB`] built: enough to delete it, read it back, and reason about its +/// watermark position (`id`). +pub struct Table { + pub chunk: Chunk, + pub id: TableId, + pub rows: Vec<(u32, u32)> +} + +/// A scratch `Database` wired for deterministic cleanup/reclaim tests: auto-compaction off +/// (only explicit `compact`/unlink move data) and small write buffers (so modest row counts +/// still produce real SST files). +/// +/// Invariant: tables are created with strictly increasing ids in call order, so "created +/// first" == "smaller id". Tests rely on this to place an orphan/older table below a live +/// one without reading raw ids. +pub struct MockDB { + db: Database, + ds: DatasetId, + schema: Arc, + static_data: Vec>, + next_block: usize, + last_id: Option, + _small_buffers: RestoreBufferSizesGuard, + // Keeps the temp dir alive: flush/reclaim create new WAL/SST files, which + // fails once the directory is cleaned up. + _dir: TempDir +} + +impl MockDB { + pub fn new() -> Self { + Self::open( + DatabaseSettings::default() + .with_rocksdb_stats(true) + .with_auto_compactions(false) + ) + } + + /// Like [`MockDB::new`] but with the block cache disabled, so reads must hit the SST + /// files: once unlinked, the data is genuinely gone, never served from cache. + pub fn uncached() -> Self { + Self::open( + DatabaseSettings::default() + .with_rocksdb_stats(true) + .with_auto_compactions(false) + .with_data_cache_size(0) + .with_chunk_cache_size(0) + ) + } + + fn open(settings: DatabaseSettings) -> Self { + let dir = tempfile::tempdir().unwrap(); + let db = settings.open(dir.path()).unwrap(); + let ds = DatasetId::from_str("solana"); + db.create_dataset(ds, DatasetKind::from_str("solana")).unwrap(); + Self { + db, + ds, + schema: make_schema(DataType::UInt32, DataType::UInt32, false), + static_data: two_u16_columns(CAPACITY), + next_block: 0, + last_id: None, + _small_buffers: use_small_buffers(), + _dir: dir + } + } + + /// Build (and `finish`) a table of `rows` rows over the next free block range. + /// `finish` already persists the dirty marker and the table data; the caller decides + /// whether to commit the chunk. + fn build_table(&mut self, rows: usize) -> (Chunk, Vec<(u32, u32)>, TableId) { + // Advance the UUIDv7 timestamp, keeping creation order == id order. + std::thread::sleep(Duration::from_millis(2)); + + let start = self.next_block; + let end = start + rows; + assert!(end <= self.static_data[0].len(), "harness data capacity exceeded"); + self.next_block = end; + + let (chunk, rows_data) = + make_irregular_block(&self.static_data, start, end, Arc::clone(&self.schema), &self.db); + let id = chunk.tables().get("block").copied().unwrap(); + if let Some(last) = self.last_id { + assert!(id > last, "harness invariant: ids must increase with creation order"); + } + self.last_id = Some(id); + (chunk, rows_data, id) + } + + /// Build a committed table: inserting its chunk removes the dirty marker. + pub fn commit_table(&mut self, rows: usize) -> Table { + let (chunk, rows, id) = self.build_table(rows); + self.db.insert_chunk(self.ds, &chunk).unwrap(); + Table { chunk, id, rows } + } + + /// Build a table but never commit its chunk, so nothing removes its dirty marker -- + /// exactly the orphan a crash-before-commit leaves behind. + pub fn orphan_table(&mut self, rows: usize) -> Table { + let (chunk, rows, id) = self.build_table(rows); + Table { chunk, id, rows } + } + + pub fn snapshot(&self) -> ReadSnapshot<'_> { + self.db.snapshot() + } + + pub fn read(&self, snapshot: &ReadSnapshot, table: &Table) -> Vec<(u32, u32)> { + read_chunk(snapshot, table.chunk.clone()) + } + + /// Like [`MockDB::read`] but propagates errors instead of unwrapping. + pub fn try_read(&self, snapshot: &ReadSnapshot, table: &Table) -> anyhow::Result<()> { + let chunk_reader = snapshot.create_chunk_reader(table.chunk.clone()); + let table_reader = chunk_reader.get_table_reader("block")?; + table_reader.read_column(0, None)?; + table_reader.read_column(1, None)?; + Ok(()) + } + + /// Whether a fresh snapshot still sees any committed chunk in the dataset. + pub fn has_visible_chunk(&self) -> bool { + self.db.snapshot().get_first_chunk(self.ds).unwrap().is_some() + } + + pub fn delete(&self, table: &Table) { + self.db + .update_dataset(self.ds, |tx| tx.delete_chunk(&table.chunk)) + .unwrap(); + } + + /// Delete the whole dataset, which runs Phase 1 synchronously. + pub fn delete_dataset(&self) { + self.db.delete_dataset(self.ds).unwrap(); + } + + /// Whether the database has no datasets left. + pub fn has_no_datasets(&self) -> bool { + self.db.get_all_datasets().unwrap().is_empty() + } + + pub fn cleanup(&self) -> usize { + self.db.cleanup().unwrap() + } + + /// Physically unlink dead SST files below the live watermark. The caller is responsible + /// for there being no live pre-deletion reader (as at startup in production). + pub fn reclaim(&self) { + self.db.reclaim_disk_space().unwrap() + } + + pub fn purge_orphans(&self) -> usize { + self.db.purge_orphan_dirty_tables().unwrap() + } + + pub fn flush(&self) { + self.db.flush_tables().unwrap(); + } + + /// Land the table data in the bottom level, so a later reclaim is a real file unlink + /// (`DeleteFilesInRange` skips L0). In production the dead data lives there too. + pub fn compact_to_bottom(&self) { + self.db.flush_tables().unwrap(); + self.db.compact_tables(); + } + + /// Total size of all SST files in the table-data column family, in bytes. + pub fn sst_size(&self) -> u64 { + self.db + .get_property("TABLES", "rocksdb.total-sst-files-size") + .unwrap() + .and_then(|s| s.parse().ok()) + .unwrap_or(0) + } +}