Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 94 additions & 10 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
7 changes: 4 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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.
Expand Down Expand Up @@ -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.

Expand Down
22 changes: 15 additions & 7 deletions src/admin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand Down
31 changes: 29 additions & 2 deletions src/ingest/flusher.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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))
}
}

Expand All @@ -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 {
Expand Down
Loading
Loading