Merge pull request #19 from pbudzik/feat/polish-bundle-2

Polish bundle: file locking, sort-on-flush, RLE for kind

Author: pbudzik

Cargo.lock

@@ -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"

Cargo.toml

@@ -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<T>` | `event_id`, `kind`, `correction_ref`, `dimensions` |
+| `Plain` (0) | bincode-serialized `Vec<T>` | `event_id`, `correction_ref`, `dimensions` |
 | `Dictionary` (1) | bincode `(Vec<String>, Vec<u32>)` — unique values + index per row | `account_id`, `product_id`, `meter_id`, `model_id`, `source`, `unit`, `subscription_id` |
 | `Delta` (2) | bincode `Vec<i64>` 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.
 

README.md

@@ -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<AppState> {
+/// 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,15 +43,16 @@ pub fn open_state_for_admin(config: Config) -> anyhow::Result<AppState> {
     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)),
         wal: Mutex::new(wal),
         memtable: Mutex::new(Memtable::new()),
         manifest: RwLock::new(manifest),
         flush_sender,
-    }))
+    });
+    Ok((state, lock))
 }
 
 /// `usagedb check [--deep]` — print manifest summary; with --deep, also

src/admin.rs

@@ -153,19 +153,28 @@ 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));
 
         let mut writer = match RawSegmentWriter::new(path.clone()) {
             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 {