From 4e714f35ab047ed236605d397572c95cf48afd1a Mon Sep 17 00:00:00 2001 From: przemek Date: Sat, 16 May 2026 19:24:06 +0200 Subject: [PATCH] Polish bundle: file locking, sort-on-flush, RLE for `kind` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four operational polish items in one PR — none individually big enough for their own. 1. File locking (DbLock). New module src/runtime/lock.rs holds an exclusive flock on db_root/LOCK. Server (run_server) and every admin command acquire it for their full lifetime; OS releases it when the process exits. Prevents server + admin (or two servers) from racing on the same manifest. open_state_for_admin now returns (AppState, DbLock); existing tests destructured the tuple. 2. Sort-on-flush. FlusherWorker.write_bucket sorts the per-bucket batch by (account, product, meter, model, ts) — the spec's canonical billing order — before writing. Saves a sort pass on every later compaction and improves dictionary-encoding locality on the ID columns. The algorithm is exposed as `pub fn sort_events_canonical` so tests can verify it directly without driving the async flusher. 3. RLE for `kind`. Encoding::Rle = 4 added. The kind column has 3 values with usually long runs of `Usage`; RLE collapses those to single (value, run_length) pairs (then zstd over that). For a 1000-row all-Usage segment, the kind column goes from 1000 bytes pre-zstd to ~10 bytes. Reader is permissive (accepts Plain or Rle). Tests (tests/polish.rs, 8 tests): - db_lock_prevents_concurrent_acquisition + re-acquire after drop - open_state_for_admin_blocks_second_caller - sort_events_canonical_orders_by_billing_dimensions - sort_events_canonical_breaks_ties_by_meter_then_model - writer_round_trip_preserves_sorted_order - rle_kind_round_trips_all_usage - rle_kind_round_trips_all_corrections - rle_encoding_round_trips_mixed_kinds (Usage + Correction + Retraction interleaved) Dep: fs4 = "0.13" for cross-platform flock. Total tests: 114 (was 106; +8). Clean under RUSTFLAGS=-D warnings. Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 104 +++++++++++-- Cargo.toml | 1 + README.md | 7 +- src/admin.rs | 22 ++- src/ingest/flusher.rs | 31 +++- src/main.rs | 15 +- src/runtime/lock.rs | 47 ++++++ src/runtime/mod.rs | 1 + src/storage/segment_format.rs | 5 + src/storage/segment_reader.rs | 40 ++++- src/storage/segment_writer.rs | 34 ++++- tests/admin.rs | 18 +-- tests/polish.rs | 269 ++++++++++++++++++++++++++++++++++ 13 files changed, 554 insertions(+), 40 deletions(-) create mode 100644 src/runtime/lock.rs create mode 100644 tests/polish.rs diff --git a/Cargo.lock b/Cargo.lock index b2219aa..d3d8895 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -70,7 +70,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -81,7 +81,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -558,7 +558,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -604,6 +604,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix", + "windows-sys 0.59.0", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -1014,7 +1024,7 @@ checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", "wasi", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1023,7 +1033,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1404,7 +1414,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1559,7 +1569,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1582,7 +1592,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1618,7 +1628,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1684,7 +1694,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1821,6 +1831,7 @@ dependencies = [ "bytes", "chrono", "clap", + "fs4", "http-body-util", "hyper", "lz4_flex", @@ -2040,6 +2051,15 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -2049,6 +2069,70 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + [[package]] name = "wit-bindgen" version = "0.51.0" diff --git a/Cargo.toml b/Cargo.toml index 7c675c2..a3407d2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ clap = { version = "4", features = ["derive"] } byteorder = "1.5.0" bytes = "1.11.1" chrono = { version = "0.4.44", features = ["serde"] } +fs4 = "0.13" hyper = "1.9.0" lz4_flex = "0.13.1" memmap2 = "0.9.10" diff --git a/README.md b/README.md index 9e379a6..28a6024 100644 --- a/README.md +++ b/README.md @@ -117,7 +117,8 @@ What works end-to-end: - Durable batch ingest with idempotent dedupe (blake3 128-bit hash, 7-day TTL) - WAL rotation, sealing, and crash recovery (memtable rebuilt from unsealed files) - Atomic manifest updates (tmp + rename + parent dir fsync) -- Background memtable flush → immutable columnar raw segments, partitioned by `bucket = blake3(account_id) % bucket_count` +- Background memtable flush → immutable columnar raw segments, partitioned by `bucket = blake3(account_id) % bucket_count`, written in canonical billing order `(account, product, meter, model, ts)` (sort-on-flush) so compaction is cheaper and dictionary encoding compresses better +- Exclusive process lock (`db_root/LOCK` via `flock`) — prevents server + admin commands from racing on the same database - Columnar on-disk format with per-column zstd compression and blake3 checksum (see [Segment format](#segment-format)) - Background hourly rollup scheduler — seals completed hours into per-bucket rollup segments, advances the manifest watermark atomically, query path routes `RollupHourly` through rollups with raw fallback for the open-period tail - Background compaction scheduler — merges small per-bucket segments into a single output, applies the `ReplacementRecord` to the manifest, deletes old files after a configurable reader grace period (spec §15.3) @@ -144,7 +145,6 @@ db_root/ Known gaps (tracked against `rust_ai_usage_db_spec.md`): -- No RLE encoding yet (used for `kind`-style low-cardinality columns); `Plain` + zstd handles it adequately. - No block-level metadata for fine-grained skipping inside a segment; pruning is segment-level only. - Rollup segments still use length-prefixed bincode (not the columnar format) — they're tiny so it hasn't been a win yet - COUNT semantics differ for `RollupHourly` queries: each rollup row counts as 1, not as the number of underlying events. Use `RawEvents` source for exact event counts. @@ -178,10 +178,11 @@ Each column payload before compression is encoded based on its declared `encodin | Encoding | Layout | Used for | | --- | --- | --- | -| `Plain` (0) | bincode-serialized `Vec` | `event_id`, `kind`, `correction_ref`, `dimensions` | +| `Plain` (0) | bincode-serialized `Vec` | `event_id`, `correction_ref`, `dimensions` | | `Dictionary` (1) | bincode `(Vec, Vec)` — unique values + index per row | `account_id`, `product_id`, `meter_id`, `model_id`, `source`, `unit`, `subscription_id` | | `Delta` (2) | bincode `Vec` of running differences | `timestamp_ms`, `ingested_at_ms` | | `Zigzag` (3) | `u32 count` + concatenated zigzag-varints | `quantity` | +| `Rle` (4) | bincode `Vec<(u8, u32)>` of (value, run_length) | `kind` | Dictionary collapses ID columns from O(rows × string size) to O(unique values × string size + 4 bytes/row); for ID-heavy workloads this is a 1000× shrink on the column. Delta encoding turns near-monotonic timestamps into small differences that zstd compresses dramatically better. Zigzag-varint packs small i128 quantities into 1–2 bytes instead of 16. diff --git a/src/admin.rs b/src/admin.rs index 863bf25..e519ce2 100644 --- a/src/admin.rs +++ b/src/admin.rs @@ -17,15 +17,22 @@ use crate::ingest::memtable::Memtable; use crate::ingest::wal::Wal; use crate::rollup::worker::RollupWorker; use crate::runtime::config::Config; +use crate::runtime::lock::DbLock; use crate::runtime::state::{AppState, AppStateInner}; use crate::storage::manifest::{Manifest, SegmentKind}; use crate::storage::segment_reader::RawSegmentReader; -/// Build a read-only AppState by loading the manifest directly. Does NOT -/// run the full Recovery flow (no segment-scan dedupe rebuild, no WAL -/// replay) — admin commands operate on the on-disk state as-is. Cheaper -/// than `run_startup_recovery` for one-shot commands. -pub fn open_state_for_admin(config: Config) -> anyhow::Result { +/// Build a read-only AppState by loading the manifest directly. Acquires +/// the DB process lock at the same time and returns it as a separate +/// guard the caller must hold for the duration of the admin operation. +/// Does NOT run the full Recovery flow (no segment-scan dedupe rebuild, +/// no WAL replay) — admin commands operate on the on-disk state as-is. +/// +/// Returning `(AppState, DbLock)` rather than embedding the lock in the +/// state keeps AppStateInner test-friendly: tests that build state +/// manually (each with a fresh tempdir) don't need to also fake a lock. +pub fn open_state_for_admin(config: Config) -> anyhow::Result<(AppState, DbLock)> { + let lock = DbLock::acquire(&config.db_root)?; let manifest = Manifest::load(&config.db_root)?.ok_or_else(|| { anyhow::anyhow!( "no manifest found at {:?} — run the server at least once to initialize the DB", @@ -36,7 +43,7 @@ pub fn open_state_for_admin(config: Config) -> anyhow::Result { std::fs::create_dir_all(&wal_dir)?; let wal = Wal::open(wal_dir, manifest.last_sealed_wal_id)?; let (flush_sender, _r) = tokio::sync::mpsc::channel(4); - Ok(Arc::new(AppStateInner { + let state = Arc::new(AppStateInner { config, // Minimal dedupe — admin commands don't ingest. dedupe: Mutex::new(HotDedupe::new(1)), @@ -44,7 +51,8 @@ pub fn open_state_for_admin(config: Config) -> anyhow::Result { memtable: Mutex::new(Memtable::new()), manifest: RwLock::new(manifest), flush_sender, - })) + }); + Ok((state, lock)) } /// `usagedb check [--deep]` — print manifest summary; with --deep, also diff --git a/src/ingest/flusher.rs b/src/ingest/flusher.rs index f585da4..1eb9a3c 100644 --- a/src/ingest/flusher.rs +++ b/src/ingest/flusher.rs @@ -153,11 +153,20 @@ impl FlusherWorker { /// Write one bucket's segment + return its metadata. Returns Err with /// a human-readable reason on any failure; the caller cleans up the /// partial file and any earlier bucket segments via `rollback_partial`. + /// + /// Sort-on-flush: events are written in the spec's canonical order + /// via `sort_events_canonical`. This matches the order compaction + /// would produce, so reading + filtering is monotonic across the + /// column families that drive billing queries. Saves a sort pass on + /// every later compaction and improves dictionary encoding locality. fn write_bucket( &self, bucket: u32, bucket_events: &[UsageEvent], ) -> Result<(SegmentMeta, PathBuf), String> { + let mut sorted = bucket_events.to_vec(); + sort_events_canonical(&mut sorted); + let segment_id = format!("raw_{}", uuid::Uuid::new_v4().simple()); let path = self.state.config.db_root.join(format!("{}.seg", segment_id)); @@ -165,7 +174,7 @@ impl FlusherWorker { Ok(w) => w, Err(e) => return Err(format!("create segment for bucket {}: {}", bucket, e)), }; - for event in bucket_events { + for event in &sorted { if let Err(e) = writer.write_event(event) { let _ = std::fs::remove_file(&path); return Err(format!("write event to bucket {} segment {}: {}", bucket, segment_id, e)); @@ -179,7 +188,7 @@ impl FlusherWorker { } }; - Ok((build_segment_meta(&segment_id, bucket_events, bucket, checksum), path)) + Ok((build_segment_meta(&segment_id, &sorted, bucket, checksum), path)) } } @@ -189,6 +198,24 @@ fn rollback_partial(paths: &[PathBuf]) { } } +/// Sort events into the spec's canonical billing order: +/// `(account_id, product_id, meter_id, model_id, timestamp_ms)`. +/// Used by the flusher (sort-on-flush) and also exposed publicly so +/// tests can verify the algorithm without driving the async pipeline. +pub fn sort_events_canonical(events: &mut [UsageEvent]) { + events.sort_by(|a, b| { + a.account_id.0.cmp(&b.account_id.0) + .then_with(|| a.product_id.0.cmp(&b.product_id.0)) + .then_with(|| a.meter_id.0.cmp(&b.meter_id.0)) + .then_with(|| { + let am = a.model_id.as_ref().map(|m| m.0.as_str()).unwrap_or(""); + let bm = b.model_id.as_ref().map(|m| m.0.as_str()).unwrap_or(""); + am.cmp(bm) + }) + .then_with(|| a.timestamp_ms.cmp(&b.timestamp_ms)) + }); +} + /// Build a SegmentMeta covering all events in the batch with the given bucket /// and segment checksum (returned by `RawSegmentWriter::finish`). pub fn build_segment_meta(segment_id: &str, batch: &[UsageEvent], bucket: u32, checksum: u64) -> SegmentMeta { diff --git a/src/main.rs b/src/main.rs index 8b00b9f..e42039a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,6 +16,7 @@ use usagedb::ingest::flusher::FlusherWorker; use usagedb::ingest::wal::Wal; use usagedb::rollup::worker::RollupWorker; use usagedb::runtime::config::Config; +use usagedb::runtime::lock::DbLock; use usagedb::runtime::recovery::Recovery; use usagedb::runtime::state::{AppState, AppStateInner}; @@ -82,27 +83,27 @@ async fn main() -> Result<(), Box> { match cli.command.unwrap_or(Command::Serve) { Command::Serve => run_server(config).await?, Command::Check { deep } => { - let state = open_state_for_admin(config)?; + let (state, _lock) = open_state_for_admin(config)?; let output = cmd_check(state, deep).await?; print!("{}", output); } Command::RebuildRollups { from, to } => { - let state = open_state_for_admin(config)?; + let (state, _lock) = open_state_for_admin(config)?; let output = cmd_rebuild_rollups(state, &from, &to).await?; print!("{}", output); } Command::InspectSegment { segment_id } => { - let state = open_state_for_admin(config)?; + let (state, _lock) = open_state_for_admin(config)?; let output = cmd_inspect_segment(state, &segment_id).await?; print!("{}", output); } Command::VerifyPeriod { account, from, to } => { - let state = open_state_for_admin(config)?; + let (state, _lock) = open_state_for_admin(config)?; let output = cmd_verify_period(state, &account, &from, &to).await?; print!("{}", output); } Command::ExportParquet { output } => { - let state = open_state_for_admin(config)?; + let (state, _lock) = open_state_for_admin(config)?; let output_msg = cmd_export_parquet(state, &output).await?; print!("{}", output_msg); } @@ -114,6 +115,10 @@ async fn run_server(config: Config) -> Result<(), Box> { info!("Starting usageDb server..."); std::fs::create_dir_all(&config.db_root)?; + // Acquire the DB process lock for the whole server lifetime so no + // admin command (or second server) can race against the same db_root. + let _db_lock = DbLock::acquire(&config.db_root)?; + // Run startup recovery: load manifest, clean tmp, replay WAL. let recovery = Recovery::new(config.db_root.clone()); let mut recovery_result = match recovery.run_startup_recovery(config.dedupe_capacity) { diff --git a/src/runtime/lock.rs b/src/runtime/lock.rs new file mode 100644 index 0000000..7bb3824 --- /dev/null +++ b/src/runtime/lock.rs @@ -0,0 +1,47 @@ +//! Exclusive-process lock on a database directory. +//! +//! Acquires an OS file lock (`flock(2)` on Unix, `LockFileEx` on Windows) +//! on `db_root/LOCK` for the lifetime of the process. Prevents the server +//! and admin commands from running concurrently against the same DB — +//! without this, concurrent manifest writes could clobber each other. +//! +//! The lock is held as long as the `DbLock` is alive; dropping it +//! releases the OS lock. The lock file itself is left on disk (it's +//! harmless if another process is already holding it via fd, and +//! removing it racily is more trouble than it's worth). + +use std::fs::{File, OpenOptions}; +use std::path::Path; + +use fs4::fs_std::FileExt; + +/// Holds the OS file lock for a database directory. Drop to release. +pub struct DbLock { + _file: File, +} + +impl DbLock { + /// Try to acquire an exclusive lock on `db_root/LOCK`. Returns an + /// error (rather than blocking) if another process holds it — the + /// caller can surface a clear "DB already in use" message. + pub fn acquire(db_root: &Path) -> anyhow::Result { + std::fs::create_dir_all(db_root)?; + let path = db_root.join("LOCK"); + let file = OpenOptions::new() + .create(true) + .write(true) + .truncate(false) + .open(&path)?; + match FileExt::try_lock_exclusive(&file) { + Ok(true) => Ok(Self { _file: file }), + Ok(false) => Err(anyhow::anyhow!( + "another usagedb process holds the lock on {:?}", + path + )), + Err(e) => Err(anyhow::anyhow!( + "failed to acquire DB lock on {:?}: {}", + path, e + )), + } + } +} diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index 68b2c06..1af3ec7 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -1,5 +1,6 @@ // Runtime module pub mod config; +pub mod lock; pub mod recovery; pub mod state; diff --git a/src/storage/segment_format.rs b/src/storage/segment_format.rs index 58ef062..befcc33 100644 --- a/src/storage/segment_format.rs +++ b/src/storage/segment_format.rs @@ -57,6 +57,10 @@ pub enum Encoding { /// to ~1-2 bytes instead of 16; signed mapping via zigzag means /// `-1` ≈ same size as `1`. Zigzag = 3, + /// Run-length encoding for low-cardinality byte columns like `kind` + /// (3 values, almost always long runs of `Usage`). Stored as a + /// `Vec<(u8, u32)>` of (value, run_length) pairs. + Rle = 4, } impl Encoding { @@ -66,6 +70,7 @@ impl Encoding { 1 => Some(Self::Dictionary), 2 => Some(Self::Delta), 3 => Some(Self::Zigzag), + 4 => Some(Self::Rle), _ => None, } } diff --git a/src/storage/segment_reader.rs b/src/storage/segment_reader.rs index 5c45963..784de02 100644 --- a/src/storage/segment_reader.rs +++ b/src/storage/segment_reader.rs @@ -129,7 +129,7 @@ impl RawSegmentReader { // Decode each required column. Dispatch by encoding. let event_id: Vec = decode_strings(&mut chunks, col::EVENT_ID)?; - let kind: Vec = decode_plain(&mut chunks, col::KIND)?; + let kind: Vec = decode_u8(&mut chunks, col::KIND)?; let correction_ref: Vec> = decode_plain(&mut chunks, col::CORRECTION_REF)?; let account_id: Vec = decode_strings(&mut chunks, col::ACCOUNT_ID)?; let subscription_id: Vec> = decode_option_strings(&mut chunks, col::SUBSCRIPTION_ID)?; @@ -371,6 +371,44 @@ fn decode_i128( } } +/// Decode a `Vec` column. Accepts Plain or Rle — the writer chose +/// Rle for `kind` after Phase B, but Plain segments from before that +/// still load. +fn decode_u8( + map: &mut HashMap)>, + name: &str, +) -> IoResult> { + let (encoding, bytes) = take(map, name)?; + match encoding { + Encoding::Plain => de(&bytes), + Encoding::Rle => decode_rle_u8(&bytes, name), + _ => Err(corrupt(&format!( + "column {} u8-typed but on-disk encoding is {:?}", + name, encoding + ))), + } +} + +/// Expand RLE-encoded (value, run_length) pairs back to a flat Vec. +fn decode_rle_u8(bytes: &[u8], name: &str) -> IoResult> { + let runs: Vec<(u8, u32)> = de(bytes)?; + let total: usize = runs.iter().map(|(_, c)| *c as usize).sum(); + let mut out = Vec::with_capacity(total); + for (v, c) in runs { + // Guard against pathological run counts that would OOM us. + if out.len() + c as usize > 10_000_000_000 { + return Err(corrupt(&format!( + "column {} RLE run count overflows reasonable bound", + name + ))); + } + for _ in 0..c { + out.push(v); + } + } + Ok(out) +} + fn decode_zigzag_varint(bytes: &[u8], name: &str) -> IoResult> { if bytes.len() < 4 { return Err(corrupt(&format!("column {}: zigzag payload missing count", name))); diff --git a/src/storage/segment_writer.rs b/src/storage/segment_writer.rs index 697cfd3..1c751e3 100644 --- a/src/storage/segment_writer.rs +++ b/src/storage/segment_writer.rs @@ -89,9 +89,10 @@ impl RawSegmentWriter { // would expand it; Plain + zstd is the right call. column_chunks.push((col::EVENT_ID, Encoding::Plain, ser(&self.event_id)?)); - // kind has only 3 possible values; Plain is fine, zstd handles the - // repetition. (RLE could squeeze more out — punted for now.) - column_chunks.push((col::KIND, Encoding::Plain, ser(&self.kind)?)); + // kind has only 3 possible values with usually long runs of `Usage`. + // RLE compresses long monotonic runs into (value, length) pairs, + // which zstd then collapses further. + column_chunks.push((col::KIND, Encoding::Rle, encode_rle_u8(&self.kind)?)); // correction_ref is sparse + structured; Plain. column_chunks.push((col::CORRECTION_REF, Encoding::Plain, ser(&self.correction_ref)?)); @@ -262,3 +263,30 @@ fn encode_zigzag_varint_i128(values: &[i128]) -> Vec { fn zigzag_i128(n: i128) -> u128 { ((n << 1) ^ (n >> 127)) as u128 } + +/// Run-length encode a `Vec` (used for the `kind` column). Stored as +/// `Vec<(u8, u32)>` of (value, run_length). Reader expands back to a +/// flat `Vec` by repeating each value `run_length` times. +/// +/// On a 10k-event segment that's all `EventKind::Usage` (kind=0), the +/// raw column is 10000 bytes; RLE collapses it to one (0, 10000) pair +/// plus bincode framing — ~10 bytes. +fn encode_rle_u8(values: &[u8]) -> IoResult> { + let mut runs: Vec<(u8, u32)> = Vec::new(); + let mut iter = values.iter().copied(); + if let Some(first) = iter.next() { + let mut current = first; + let mut count: u32 = 1; + for v in iter { + if v == current { + count = count.saturating_add(1); + } else { + runs.push((current, count)); + current = v; + count = 1; + } + } + runs.push((current, count)); + } + ser(&runs) +} diff --git a/tests/admin.rs b/tests/admin.rs index 7684a3b..6f3f0fa 100644 --- a/tests/admin.rs +++ b/tests/admin.rs @@ -115,7 +115,7 @@ async fn check_reports_manifest_summary() { drop(_live); // close the live state so admin can reopen cleanly let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let out = cmd_check(state, false).await.unwrap(); assert!(out.contains("Raw segments: 2"), "{}", out); assert!(out.contains("Rollup segments: 0"), "{}", out); @@ -131,7 +131,7 @@ async fn check_deep_passes_for_valid_segments() { drop(_live); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let out = cmd_check(state, true).await.unwrap(); assert!(out.contains("All segments verified"), "deep check should pass: {}", out); } @@ -163,7 +163,7 @@ async fn check_deep_fails_for_corrupt_segment() { std::fs::write(&seg, bytes).unwrap(); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let err = cmd_check(state, true).await.unwrap_err(); assert!(err.to_string().contains("failed verification"), "{}", err); } @@ -202,7 +202,7 @@ async fn inspect_segment_prints_meta_and_sample_rows() { drop(live); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let out = cmd_inspect_segment(state, &seg_id).await.unwrap(); assert!(out.contains(&seg_id), "{}", out); @@ -218,7 +218,7 @@ async fn inspect_segment_errors_for_unknown_id() { drop(_live); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let err = cmd_inspect_segment(state, "raw_bogus").await.unwrap_err(); assert!(err.to_string().contains("not found in manifest"), "{}", err); } @@ -248,7 +248,7 @@ async fn rebuild_rollups_drops_and_rewinds() { drop(live); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let out = cmd_rebuild_rollups( state.clone(), "1970-01-01T00:00:00Z", @@ -269,7 +269,7 @@ async fn rebuild_rollups_rejects_invalid_dates() { drop(_live); let config = Config { db_root: root, ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let err = cmd_rebuild_rollups(state, "not-a-date", "2030-01-01T00:00:00Z") .await .unwrap_err(); @@ -298,7 +298,7 @@ async fn verify_period_reports_match_for_sealed_rollup() { drop(live); let config = Config { db_root: root, ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let out = cmd_verify_period( state, "acc_v", @@ -327,7 +327,7 @@ async fn export_parquet_writes_file() { let out_path = root.join("export.parquet"); let config = Config { db_root: root.clone(), ..Config::default() }; - let state = open_state_for_admin(config).unwrap(); + let (state, _lock) = open_state_for_admin(config).unwrap(); let msg = cmd_export_parquet(state, &out_path).await.unwrap(); assert!(msg.contains("Exported 2 events"), "{}", msg); assert!(msg.contains("2 segment(s)"), "{}", msg); diff --git a/tests/polish.rs b/tests/polish.rs new file mode 100644 index 0000000..90767c6 --- /dev/null +++ b/tests/polish.rs @@ -0,0 +1,269 @@ +//! Tests for the operational polish bundle: +//! - DbLock prevents concurrent access on the same db_root +//! - Sort-on-flush writes segments in the canonical billing order +//! - RLE for `kind` round-trips + dramatically shrinks low-cardinality columns + +use std::path::PathBuf; + +use usagedb::admin::open_state_for_admin; +use usagedb::ingest::flusher::sort_events_canonical; +use usagedb::model::dimensions::SmallDimensions; +use usagedb::model::event::{EventKind, UsageEvent}; +use usagedb::model::ids::{ + AccountId, EventId, MeterId, ModelId, ProductId, SourceId, SubscriptionId, Unit, +}; +use usagedb::runtime::config::Config; +use usagedb::runtime::lock::DbLock; +use usagedb::storage::manifest::Manifest; +use usagedb::storage::segment_reader::RawSegmentReader; +use usagedb::storage::segment_writer::RawSegmentWriter; + +fn tmp_root() -> PathBuf { + let dir = tempfile::tempdir().expect("tempdir"); + let p = dir.path().to_path_buf(); + std::mem::forget(dir); + p +} + +fn tmp_file(name: &str) -> PathBuf { + let dir = tempfile::tempdir().expect("tempdir"); + let p = dir.path().join(name); + std::mem::forget(dir); + p +} + +fn make_event(id: &str, account: &str, ts: i64, qty: i128) -> UsageEvent { + UsageEvent { + event_id: EventId(id.to_string()), + kind: EventKind::Usage, + correction_ref: None, + account_id: AccountId(account.to_string()), + subscription_id: Some(SubscriptionId("sub".into())), + product_id: ProductId("prod".into()), + meter_id: MeterId("meter".into()), + timestamp_ms: ts, + quantity: qty, + unit: Unit("token".into()), + source: SourceId("test".into()), + model_id: Some(ModelId("m1".into())), + dimensions: SmallDimensions::default(), + ingested_at_ms: ts, + } +} + +// ========================================================================= +// DbLock +// ========================================================================= + +#[test] +fn db_lock_prevents_concurrent_acquisition() { + let root = tmp_root(); + let lock_a = DbLock::acquire(&root).expect("first lock"); + let err = match DbLock::acquire(&root) { + Ok(_) => panic!("second lock should have failed"), + Err(e) => e, + }; + assert!( + err.to_string().contains("holds the lock"), + "error should mention the lock: {}", + err + ); + drop(lock_a); + // After releasing, a fresh acquisition should succeed. + let _lock_b = DbLock::acquire(&root).expect("re-acquire after drop"); +} + +#[tokio::test] +async fn open_state_for_admin_blocks_second_caller() { + let root = tmp_root(); + std::fs::create_dir_all(&root).unwrap(); + // Bootstrap a manifest so open_state_for_admin gets past the + // manifest-load step and hits the lock. + Manifest::default().save(&root).unwrap(); + + let config = Config { db_root: root.clone(), ..Config::default() }; + let (_state_a, _lock_a) = open_state_for_admin(config.clone()).expect("first"); + let err = match open_state_for_admin(config) { + Ok(_) => panic!("concurrent open should fail"), + Err(e) => e, + }; + assert!(err.to_string().contains("holds the lock"), "{}", err); +} + +// ========================================================================= +// Sort-on-flush +// ========================================================================= + +/// Verify the sort algorithm directly. The flusher's `write_bucket` +/// calls `sort_events_canonical` before writing, so a working sort +/// algorithm plus the writer test in tests/encodings.rs together +/// cover the integration. +#[test] +fn sort_events_canonical_orders_by_billing_dimensions() { + let mut events = vec![ + // Same account, different timestamps + meters. + make_event("e1", "acc_b", 100, 1), + make_event("e2", "acc_a", 200, 1), + make_event("e3", "acc_a", 100, 1), + make_event("e4", "acc_b", 50, 1), + ]; + sort_events_canonical(&mut events); + let order: Vec<(String, i64)> = events + .iter() + .map(|e| (e.account_id.0.clone(), e.timestamp_ms)) + .collect(); + assert_eq!( + order, + vec![ + ("acc_a".to_string(), 100), + ("acc_a".to_string(), 200), + ("acc_b".to_string(), 50), + ("acc_b".to_string(), 100), + ], + ); +} + +#[test] +fn sort_events_canonical_breaks_ties_by_meter_then_model() { + let mut events = vec![ + UsageEvent { meter_id: MeterId("z".into()), ..make_event("a", "acc", 100, 1) }, + UsageEvent { meter_id: MeterId("a".into()), ..make_event("b", "acc", 100, 1) }, + UsageEvent { + meter_id: MeterId("a".into()), + model_id: Some(ModelId("z_model".into())), + ..make_event("c", "acc", 100, 1) + }, + ]; + sort_events_canonical(&mut events); + let order: Vec<&str> = events.iter().map(|e| e.event_id.0.as_str()).collect(); + // meter "a" before "z"; within meter "a", model "m1" before "z_model". + assert_eq!(order, vec!["b", "c", "a"]); +} + +/// End-to-end: write events into a segment via the writer in +/// not-yet-sorted order, but use sort_events_canonical first (mirroring +/// the flusher path). The reader should see them in canonical order. +#[test] +fn writer_round_trip_preserves_sorted_order() { + let path = tmp_file("sorted.seg"); + let mut events = vec![ + make_event("e1", "acc_b", 100, 1), + make_event("e2", "acc_a", 200, 1), + make_event("e3", "acc_a", 100, 1), + ]; + sort_events_canonical(&mut events); + + let mut writer = RawSegmentWriter::new(path.clone()).unwrap(); + for e in &events { + writer.write_event(e).unwrap(); + } + writer.finish().unwrap(); + + let mut reader = RawSegmentReader::new(path).unwrap(); + let mut read_back = Vec::new(); + while let Some(e) = reader.read_next().unwrap() { + read_back.push((e.account_id.0, e.timestamp_ms)); + } + assert_eq!( + read_back, + vec![ + ("acc_a".to_string(), 100), + ("acc_a".to_string(), 200), + ("acc_b".to_string(), 100), + ], + ); +} + +// ========================================================================= +// RLE for `kind` +// ========================================================================= + +#[test] +fn rle_encoding_round_trips_mixed_kinds() { + let out = tmp_file("rle_rt.seg"); + let mut writer = RawSegmentWriter::new(out.clone()).unwrap(); + // 100 Usage + 50 Correction + 100 Usage + 1 Retraction + let mut events = Vec::new(); + for i in 0..100 { + events.push(make_event(&format!("u1_{i}"), "acc", i, 10)); + } + for i in 0..50 { + let mut e = make_event(&format!("c_{i}"), "acc", 1000 + i, -1); + e.kind = EventKind::Correction; + e.correction_ref = Some(usagedb::model::event::CorrectionRef { + original_event_id: EventId(format!("u1_{}", i)), + reason: "test".into(), + }); + events.push(e); + } + for i in 0..100 { + events.push(make_event(&format!("u2_{i}"), "acc", 2000 + i, 10)); + } + let mut last = make_event("r", "acc", 5000, 0); + last.kind = EventKind::Retraction; + last.correction_ref = Some(usagedb::model::event::CorrectionRef { + original_event_id: EventId("u2_0".into()), + reason: "retract".into(), + }); + events.push(last); + + for e in &events { + writer.write_event(e).unwrap(); + } + writer.finish().unwrap(); + + let mut reader = RawSegmentReader::new(out).unwrap(); + let mut idx = 0; + while let Some(e) = reader.read_next().unwrap() { + assert_eq!(e.kind, events[idx].kind, "row {}", idx); + assert_eq!(e.event_id, events[idx].event_id); + idx += 1; + } + assert_eq!(idx, events.len()); +} + +#[test] +fn rle_kind_round_trips_all_usage() { + // The common case: every row is `Usage`. RLE collapses to one run. + let out = tmp_file("rle_all_usage.seg"); + let mut writer = RawSegmentWriter::new(out.clone()).unwrap(); + for i in 0..1000 { + writer + .write_event(&make_event(&format!("u_{i}"), "acc", i, 1)) + .unwrap(); + } + writer.finish().unwrap(); + + let mut reader = RawSegmentReader::new(out).unwrap(); + let mut count = 0; + while let Some(e) = reader.read_next().unwrap() { + assert_eq!(e.kind, EventKind::Usage); + count += 1; + } + assert_eq!(count, 1000); +} + +#[test] +fn rle_kind_round_trips_all_corrections() { + // Pathological: every row is Correction. RLE still works. + let out = tmp_file("rle_all_corr.seg"); + let mut writer = RawSegmentWriter::new(out.clone()).unwrap(); + for i in 0..100 { + let mut e = make_event(&format!("c_{i}"), "acc", i, -1); + e.kind = EventKind::Correction; + e.correction_ref = Some(usagedb::model::event::CorrectionRef { + original_event_id: EventId(format!("o_{}", i)), + reason: "x".into(), + }); + writer.write_event(&e).unwrap(); + } + writer.finish().unwrap(); + + let mut reader = RawSegmentReader::new(out).unwrap(); + let mut count = 0; + while let Some(e) = reader.read_next().unwrap() { + assert_eq!(e.kind, EventKind::Correction); + count += 1; + } + assert_eq!(count, 100); +}