From 08ca33b2a3647d421c2aa861d6b40eca7d49bfb5 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 19:44:13 -0400 Subject: [PATCH 01/12] docs: design for SimpleMapStore persistence (mem limit + disk flush) Proposes a memory-bounded, disk-backed extension to SimpleMapStore: sealed epochs flushed to per-agg segment files under an atomic manifest, driven by a configurable memory budget and time-age watermark. Phasing, config shape, concurrency plan, and four open questions included for review before implementation. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 324 ++++++++++++++++++++ 1 file changed, 324 insertions(+) create mode 100644 docs/design-simple-map-store-persistence.md diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md new file mode 100644 index 00000000..5d097968 --- /dev/null +++ b/docs/design-simple-map-store-persistence.md @@ -0,0 +1,324 @@ +# Design: SimpleMapStore Persistence (Memory Limit + Disk Flush) + +## Problem + +`SimpleMapStore` (`asap-query-engine/src/stores/simple_map_store/`) is currently an +in-memory-only store. Under long-running ingest it grows unboundedly: every sealed +window for every `(aggregation_id, group_key)` is held in `DashMap>` +until `CleanupPolicy::CircularBuffer` rotates it out and drops it on the floor. + +This creates two problems: + +1. **No memory bound.** A deployment has to either overprovision RAM or rely on + `CircularBuffer` to throw away data that may still be query-relevant. +2. **No durability.** Cold data (older than the query working set) still occupies + RAM even though most queries hit only the last few minutes. + +We want a persistence layer that lets the store: + +- Honor a configurable memory budget for sketches. +- Flush sealed windows older than a configurable timestamp threshold to disk. +- Evict those flushed windows from memory when the budget is exceeded. +- Serve queries transparently from memory + disk. + +Goals are scoped to a **single-node, single-process** store. Replication, sharding, +compression, and query pushdown into segments are explicitly out of scope for v1. + +--- + +## Current shape (relevant facts) + +- `SimpleMapStorePerKey` (`per_key.rs:160`) keeps per-agg-id state in + `DashMap>>`. +- Each `StoreKeyData` has a `current_epoch` (actively being written) and + `sealed_epochs: BTreeMap` (`per_key.rs`). +- Values are `Arc`. All concrete accumulators already implement + `SerializableToSink` (`serialize_to_bytes`, `merge_with`) — so we already have + a serialization primitive and a merge primitive. +- Insert hot path: `insert_precomputed_output_batch` → `insert_for_store_key` + (`per_key.rs:261`), holding only the per-agg-id `RwLock::write`. +- Query hot path: `query_precomputed_output{,_exact}` iterates + `current_epoch` + `sealed_epochs` under `RwLock::read`. +- `CleanupPolicy` (`data_model/enums.rs`) already has a concept of dropping + old entries; persistence will become a fourth, non-destructive option. + +Key observation: **`current_epoch` is the only mutable region**. Sealed epochs are +append-only until cleanup. That is exactly the right unit to flush. + +--- + +## Design + +### Unit of flush: the sealed epoch + +A `SimpleMapStore` segment on disk corresponds to **one sealed epoch of one +aggregation id**. Rationale: + +- Sealed epochs are immutable — safe to serialize without coordinating with writers. +- The epoch already has a well-defined time range, which is exactly what range + queries want to filter on. +- Flushing an epoch only requires the per-agg-id `RwLock::write` briefly — same + lock the insert path already uses, so no new contention class. +- Recovery and query planning only need epoch-level metadata, not per-window. + +The `current_epoch` is never flushed while hot. It becomes flushable the moment +the rotator seals it. + +### Disk layout + +``` +/ +├── manifest.json # authoritative index of all segments +├── agg_00000042/ +│ ├── seg_0000000001.bin # one sealed epoch, serialized +│ ├── seg_0000000002.bin +│ └── ... +└── agg_00000043/ + └── seg_0000000001.bin +``` + +**Segment file format** (`seg_*.bin`): + +``` +[u32 magic][u16 version][u16 flags] +[u64 epoch_id][u64 window_start_ms][u64 window_end_ms] +[u32 num_entries] +repeated num_entries times: + [u64 start_ts][u64 end_ts][u32 label_id][u8 agg_type] + [u32 payload_len][payload_len bytes: serialize_to_bytes()] +[u32 crc32 of body] +``` + +Fixed-size header lets us mmap and binary-search by timestamp without parsing +payloads. Body is a linear scan — v1 does not build an in-segment index because +sealed epochs are small (bounded by window size × group count for a single agg). + +**Manifest** (`manifest.json`): + +```json +{ + "version": 1, + "segments": [ + {"agg_id": 42, "epoch_id": 1, "path": "agg_00000042/seg_0000000001.bin", + "start_ms": 1_700_000_000_000, "end_ms": 1_700_000_060_000, + "num_entries": 120, "size_bytes": 48192} + ] +} +``` + +The manifest is the single source of truth for which segments exist and what +ranges they cover. It is rewritten atomically (`write → fsync → rename`) after +every flush batch. Individual segment files are written+fsynced before the +manifest ever references them, so a crash mid-flush leaves orphan files that +startup sweeps away — never a dangling manifest entry. + +### Memory accounting + +Add a new trait method: + +```rust +pub trait AggregateCore: ... { + fn approx_memory_bytes(&self) -> usize; +} +``` + +Implementations are cheap per-type estimates (e.g. KLL: `k * 8 + overhead`; +SumAccumulator: `size_of::()`; SetAggregator: `len * avg_entry_bytes`). +They do not call `serialize_to_bytes` — that would be too expensive on the +insert path. + +The store tracks a single `AtomicUsize` `mem_bytes_in_use`. On insert it adds +`approx_memory_bytes()` per entry; on flush it subtracts the same. This is an +estimate, not a hard guarantee — good enough to drive policy, and the alternative +(exact heap accounting) is not worth the allocator coupling. + +### Configuration + +New struct, threaded through `PrecomputeEngineConfig` and loaded from the +same YAML / controller channel as the existing streaming config: + +```rust +pub struct SimpleMapStorePersistenceConfig { + pub enabled: bool, + + // Memory budget (high watermark). When exceeded, background flusher + // evicts sealed epochs oldest-first until usage drops below + // `memory_low_watermark_bytes`. + pub memory_limit_bytes: usize, + pub memory_low_watermark_bytes: usize, + + // Time-based flush. Any sealed epoch whose end_ms is older than + // `now - flush_older_than_ms` is eligible to flush on the next tick, + // regardless of memory pressure. None disables time-based flushing. + pub flush_older_than_ms: Option, + + // Cadence of the background flusher loop. + pub flush_interval_ms: u64, + + // Root directory for segment files and manifest. + pub disk_path: PathBuf, + + // Hard ceiling. If memory usage reaches this *during* an insert + // (flusher is falling behind), the insert path blocks on a condvar + // until the flusher catches up. Set to memory_limit_bytes * 1.25 + // as a sensible default. + pub hard_cap_bytes: usize, +} +``` + +Defaults keep `enabled = false` so existing deployments are unaffected until +they opt in. + +### Eviction policy + +v1 supports one policy — **oldest-sealed-epoch-first, globally ordered by +epoch `end_ms`**. Rationale: + +- Matches the time-window access pattern: queries overwhelmingly target recent + windows. +- Aligns with `flush_older_than_ms`: the same ordering drives both memory-pressure + flush and time-based flush. +- Avoids cross-agg fairness debates that a `LargestAggFirst` policy would + invite; we can add more policies later behind an enum if needed. + +### Background flusher + +A dedicated Tokio task owned by the store, started in +`SimpleMapStorePerKey::new` when persistence is enabled: + +``` +loop { + sleep(flush_interval_ms).await; + + let now = now_ms(); + let mut candidates = Vec::new(); + + // Phase 1: time-based — any sealed epoch older than watermark. + if let Some(max_age) = cfg.flush_older_than_ms { + candidates.extend(collect_epochs_older_than(now - max_age)); + } + + // Phase 2: memory-pressure — if still over high-water after phase 1, + // keep pulling oldest sealed epochs until we would drop below + // memory_low_watermark_bytes. + if mem_bytes_in_use.load() > cfg.memory_limit_bytes { + candidates.extend(collect_oldest_until_under_low_water()); + } + + for (agg_id, epoch_id) in candidates { + flush_and_evict(agg_id, epoch_id).await?; + } + + manifest.commit().await?; // atomic rewrite after the batch +} +``` + +`flush_and_evict` serializes the epoch *outside* the per-agg lock (the epoch is +immutable once sealed, we can read the `Arc` without holding the write lock), +fsyncs the segment file, then takes the per-agg `RwLock::write` only to splice +the epoch out of `sealed_epochs` and decrement `mem_bytes_in_use`. This keeps +flush off the insert/query critical path. + +### Query path + +`query_precomputed_output` becomes a three-way merge: + +1. Read from `current_epoch + sealed_epochs` as today. +2. Look up segments in the manifest whose `[start_ms, end_ms]` overlaps the + query range for this `agg_id`. +3. For each matching segment, read the file (cached via an `lru::LruCache>`), + decode entries whose window overlaps, and merge into the result. + +Segment reads happen under a read lock on the manifest; they do **not** take any +per-agg store lock, so they can run fully in parallel with inserts. Merging reuses +the existing `TimestampedBucketsMap` + `AggregateCore::merge_with` that the +in-memory query path already uses — no new merge logic. + +### Recovery on startup + +1. Open `disk_path`, read `manifest.json`. +2. For every referenced segment, stat the file and validate magic + CRC header. + Missing / corrupt segments are logged and removed from the manifest. +3. Sweep `disk_path` for segment files not referenced in the manifest (orphans + from a mid-flush crash) and delete them. +4. Build the in-memory segment index; do **not** load any sketches into RAM. + Cold data stays cold until a query asks for it. + +### Concurrency summary + +| Path | Lock taken | +|-------------------|---------------------------------------------| +| Insert | per-agg `RwLock::write` (unchanged) | +| Query in-memory | per-agg `RwLock::read` (unchanged) | +| Query disk | manifest `RwLock::read` + segment-cache mutex | +| Flush: serialize | none (reads immutable sealed `Arc`) | +| Flush: evict | per-agg `RwLock::write` (short, O(1) splice)| +| Flush: commit | manifest `RwLock::write` (short) | + +No new lock held across a fsync or disk I/O. + +--- + +## What this does **not** change + +- `CleanupPolicy::CircularBuffer` and `ReadBased` continue to exist and run. A + deployment can opt into persistence in addition to a cleanup policy; the + persistence flusher runs *before* `CircularBuffer` would drop data, so epochs + get a chance to survive on disk. If both policies fire on the same epoch, + `CircularBuffer` wins (cleanup is destructive, but that's the user's stated + intent when they configure it). +- `SimpleMapStoreGlobal` is intentionally left in-memory-only. Persistence + targets `PerKey`, which is the production path. Adding it to `Global` is a + small follow-up if anyone needs it. +- Query planning, the controller client, and the precompute engine's output + sink are untouched. The store's `Store` trait signature does not change. + +--- + +## Phasing + +The PR this design doc accompanies will land in three commits on one branch so +the pieces can be reviewed independently: + +1. **Sizing + config plumbing.** Add `approx_memory_bytes` to `AggregateCore` + and all concrete accumulators. Add `SimpleMapStorePersistenceConfig`. Track + `mem_bytes_in_use`. No disk I/O yet; expose the counter in + `StoreDiagnostics` so we can validate sizing in isolation. + +2. **Segment format, manifest, flusher.** Add `persistence/` submodule under + `simple_map_store/` with segment encode/decode, manifest read/write, + background flusher task. Wire `flush_and_evict` into the per-key store. + Unit tests for round-trip, crash-after-segment-before-manifest, and orphan + sweep. + +3. **Query path read-through + recovery.** Extend `query_precomputed_output` + to consult the manifest and merge segment hits with in-memory hits. Add + startup recovery. Integration test: ingest → flush → restart → query → + same result as no-restart. + +Phases 1 and 2 are safe to merge independently because phase 2 is gated on +`enabled = false` by default. Phase 3 is when persistence becomes observable +to query results. + +--- + +## Open questions (for review before implementation) + +1. **Segment file format: custom binary vs. something off-the-shelf (Parquet, + Arrow IPC)?** Custom binary is simpler and avoids a dependency, but loses us + tooling. I lean custom for v1 given that we never read segments outside + this process, but happy to switch if there's an appetite. + +2. **Async vs. sync flush I/O.** The rest of the store is sync (`RwLock`, + `DashMap`), but the flusher task is naturally async. Proposal: `tokio::fs` + for segment writes, sync locks everywhere else. Flusher runs on a dedicated + task, not a shared runtime, to avoid starving it under query load. + +3. **Should `flush_older_than_ms` live here or in the existing + `CleanupPolicy` enum?** It overlaps conceptually with `CircularBuffer`. + Proposal: keep it separate — `CleanupPolicy` is destructive, persistence + config is non-destructive. Confusing them in one knob would be worse. + +4. **Per-agg-id flush fairness.** Oldest-global-first could starve small, + slow-moving aggs during a burst on a hot agg. Acceptable for v1 since + "oldest window first" is well-defined globally; revisit if it bites. From 4ef673794a0032f71d7057f0a4d0c6c3bf0d47c1 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 19:59:33 -0400 Subject: [PATCH 02/12] =?UTF-8?q?docs:=20revise=20persistence=20design=20?= =?UTF-8?q?=E2=80=94=20memory=20budget=20primary,=20T=20secondary?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the three destructive CleanupPolicy variants on SimpleMapStore with a single two-knob persistence config: memory_limit_bytes as the primary bound on RAM, hot_window_ms as the secondary time watermark that guarantees a predictable hot-set under light ingest. Flusher checks memory pressure first and time watermark second, pulling from the same oldest-epoch-first ordering so the two triggers never disagree. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 204 +++++++++++++------- 1 file changed, 134 insertions(+), 70 deletions(-) diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md index 5d097968..e20b779e 100644 --- a/docs/design-simple-map-store-persistence.md +++ b/docs/design-simple-map-store-persistence.md @@ -5,24 +5,37 @@ `SimpleMapStore` (`asap-query-engine/src/stores/simple_map_store/`) is currently an in-memory-only store. Under long-running ingest it grows unboundedly: every sealed window for every `(aggregation_id, group_key)` is held in `DashMap>` -until `CleanupPolicy::CircularBuffer` rotates it out and drops it on the floor. +until one of the three existing `CleanupPolicy` variants (`CircularBuffer`, +`ReadBased`, `NoCleanup`) either rotates it out and drops it on the floor or does +nothing at all. -This creates two problems: +All three of those variants are **destructive** — they delete data, they do not +persist it. That creates two problems: -1. **No memory bound.** A deployment has to either overprovision RAM or rely on - `CircularBuffer` to throw away data that may still be query-relevant. +1. **No memory bound that preserves data.** A deployment has to either + overprovision RAM (`NoCleanup`), throw away potentially query-relevant data + (`CircularBuffer` / `ReadBased`), or tune per-agg `num_aggregates_to_retain` + values that don't correspond to any operator-meaningful quantity. 2. **No durability.** Cold data (older than the query working set) still occupies RAM even though most queries hit only the last few minutes. -We want a persistence layer that lets the store: +We want `SimpleMapStore` to replace the existing cleanup-policy knob with a +single persistence policy driven by **two** knobs, in priority order: -- Honor a configurable memory budget for sketches. -- Flush sealed windows older than a configurable timestamp threshold to disk. -- Evict those flushed windows from memory when the budget is exceeded. -- Serve queries transparently from memory + disk. +1. **Primary — memory budget.** A configurable hard ceiling on in-memory sketch + bytes. When exceeded, the oldest sealed epochs flush to disk until usage is + back under a low-water mark. This is what actually bounds RAM in production. +2. **Secondary — time watermark T.** A configurable "hot window." Any sealed + epoch whose end is older than `now - T` flushes to disk even if the store + is nowhere near the memory budget. This guarantees predictable durability + and a stable hot-set size under light load. -Goals are scoped to a **single-node, single-process** store. Replication, sharding, -compression, and query pushdown into segments are explicitly out of scope for v1. +Flushed sketches are read back transparently at query time. + +Scope is **single-node, single-process**. Replication, sharding, compression, +and query pushdown into segments are explicitly out of scope for v1. The three +existing destructive `CleanupPolicy` variants are removed from `SimpleMapStore` +(the enum stays in `asap_types` for any other store that still uses it). --- @@ -39,8 +52,10 @@ compression, and query pushdown into segments are explicitly out of scope for v1 (`per_key.rs:261`), holding only the per-agg-id `RwLock::write`. - Query hot path: `query_precomputed_output{,_exact}` iterates `current_epoch` + `sealed_epochs` under `RwLock::read`. -- `CleanupPolicy` (`data_model/enums.rs`) already has a concept of dropping - old entries; persistence will become a fourth, non-destructive option. +- `CleanupPolicy` (`asap_types::enums`) currently has three destructive variants + (`CircularBuffer`, `ReadBased`, `NoCleanup`); `SimpleMapStore` will stop + taking a `CleanupPolicy` at all and use the new persistence config instead. + The enum itself stays in `asap_types` for other stores. Key observation: **`current_epoch` is the only mutable region**. Sealed epochs are append-only until cleanup. That is exactly the right unit to flush. @@ -135,56 +150,73 @@ estimate, not a hard guarantee — good enough to drive policy, and the alternat ### Configuration New struct, threaded through `PrecomputeEngineConfig` and loaded from the -same YAML / controller channel as the existing streaming config: +same YAML / controller channel as the existing streaming config. The two +knobs match the priority order in the problem statement: **memory budget +first, time watermark second**. ```rust pub struct SimpleMapStorePersistenceConfig { - pub enabled: bool, - - // Memory budget (high watermark). When exceeded, background flusher - // evicts sealed epochs oldest-first until usage drops below - // `memory_low_watermark_bytes`. + // ---- Primary: memory budget ---- + // + // High-water mark. When the store's tracked in-memory sketch bytes + // exceed this, the background flusher evicts sealed epochs + // oldest-first (globally, by epoch end_ms) until usage drops below + // `memory_low_watermark_bytes`. This is the knob that bounds RAM + // in production. pub memory_limit_bytes: usize, pub memory_low_watermark_bytes: usize, - // Time-based flush. Any sealed epoch whose end_ms is older than - // `now - flush_older_than_ms` is eligible to flush on the next tick, - // regardless of memory pressure. None disables time-based flushing. - pub flush_older_than_ms: Option, - - // Cadence of the background flusher loop. - pub flush_interval_ms: u64, - - // Root directory for segment files and manifest. - pub disk_path: PathBuf, - // Hard ceiling. If memory usage reaches this *during* an insert // (flusher is falling behind), the insert path blocks on a condvar // until the flusher catches up. Set to memory_limit_bytes * 1.25 // as a sensible default. pub hard_cap_bytes: usize, + + // ---- Secondary: time watermark T ---- + // + // Hot-window length. Any sealed epoch whose end_ms is older than + // `now - hot_window_ms` is flushed on the next flusher tick, even + // if the store is well under `memory_limit_bytes`. This guarantees + // durability and a predictable hot-set size under light ingest. + // None disables time-based flushing (not recommended — memory + // pressure alone will still work, but cold data will linger in RAM + // until something pushes it out). + pub hot_window_ms: Option, + + // ---- Misc ---- + pub flush_interval_ms: u64, // cadence of the background flusher + pub disk_path: PathBuf, // root dir for segments + manifest } ``` -Defaults keep `enabled = false` so existing deployments are unaffected until -they opt in. +`SimpleMapStorePerKey::new` now takes a `SimpleMapStorePersistenceConfig` +instead of a `CleanupPolicy`. There is no "persistence disabled" escape +hatch — this is now the only cleanup mechanism this store has. If someone +wants the old in-memory-only behavior, they can set `hot_window_ms = None` +and `memory_limit_bytes = usize::MAX`, which degenerates to "never flush." + +### Eviction order -### Eviction policy +There is only one ordering — **oldest-sealed-epoch-first, globally by epoch +`end_ms`**. Both triggers (memory pressure and time watermark) pull from the +same ordered view, so the flusher never has two disagreeing notions of "oldest." -v1 supports one policy — **oldest-sealed-epoch-first, globally ordered by -epoch `end_ms`**. Rationale: +Rationale: - Matches the time-window access pattern: queries overwhelmingly target recent - windows. -- Aligns with `flush_older_than_ms`: the same ordering drives both memory-pressure - flush and time-based flush. -- Avoids cross-agg fairness debates that a `LargestAggFirst` policy would - invite; we can add more policies later behind an enum if needed. + windows, so evicting oldest is the lowest-regret choice. +- Makes the two triggers composable: the memory-pressure pass and the + time-watermark pass are just two different stopping conditions on the same + iterator over `(agg_id, epoch) sorted by epoch.end_ms`. +- Avoids cross-agg fairness debates (e.g., `LargestAggFirst`) that would + otherwise complicate v1; we can add more orderings later behind an enum if + it becomes necessary. ### Background flusher A dedicated Tokio task owned by the store, started in -`SimpleMapStorePerKey::new` when persistence is enabled: +`SimpleMapStorePerKey::new`. Each tick, it checks the primary trigger +(memory) first, then the secondary trigger (time watermark): ``` loop { @@ -193,18 +225,30 @@ loop { let now = now_ms(); let mut candidates = Vec::new(); - // Phase 1: time-based — any sealed epoch older than watermark. - if let Some(max_age) = cfg.flush_older_than_ms { - candidates.extend(collect_epochs_older_than(now - max_age)); + // Phase 1 (PRIMARY): memory budget. + // If we're over the high-water mark, pull oldest sealed epochs + // (by epoch.end_ms) until projected memory drops below the + // low-water mark. This is the knob that actually bounds RAM. + if mem_bytes_in_use.load() > cfg.memory_limit_bytes { + candidates.extend(collect_oldest_until_under_low_water( + cfg.memory_low_watermark_bytes, + )); } - // Phase 2: memory-pressure — if still over high-water after phase 1, - // keep pulling oldest sealed epochs until we would drop below - // memory_low_watermark_bytes. - if mem_bytes_in_use.load() > cfg.memory_limit_bytes { - candidates.extend(collect_oldest_until_under_low_water()); + // Phase 2 (SECONDARY): time watermark T. + // Any sealed epoch older than `now - hot_window_ms` that wasn't + // already picked up in phase 1 is flushed here. Under light ingest, + // this is the only phase that runs and it keeps the hot set bounded + // by T × ingest rate regardless of the memory budget. + if let Some(hot_window) = cfg.hot_window_ms { + candidates.extend(collect_epochs_older_than(now - hot_window)); } + // Dedup (phase 1 and phase 2 can pick the same epoch) and sort by + // epoch.end_ms ascending so we flush oldest first within the batch. + candidates.sort_unstable_by_key(|c| c.end_ms); + candidates.dedup(); + for (agg_id, epoch_id) in candidates { flush_and_evict(agg_id, epoch_id).await?; } @@ -213,6 +257,11 @@ loop { } ``` +Under memory pressure, phase 1 dominates and phase 2 usually finds nothing +left to do (the oldest epochs are already gone). Under light ingest, phase 1 +is a no-op and phase 2 does all the work. The two phases never fight because +they pull from the same oldest-first ordering. + `flush_and_evict` serializes the epoch *outside* the per-agg lock (the epoch is immutable once sealed, we can read the `Arc` without holding the write lock), fsyncs the segment file, then takes the per-agg `RwLock::write` only to splice @@ -259,14 +308,25 @@ No new lock held across a fsync or disk I/O. --- -## What this does **not** change +## What this does **and does not** change + +Changes: + +- `SimpleMapStorePerKey::new` no longer takes a `CleanupPolicy`; it takes a + `SimpleMapStorePersistenceConfig`. The three destructive cleanup variants + (`CircularBuffer`, `ReadBased`, `NoCleanup`) are no longer wired into this + store at all. The code paths in `per_key.rs` that branch on + `CleanupPolicy` (`cleanup_old_aggregates`, `maybe_rotate_epoch`'s retention + logic) are deleted in favor of the flusher. +- Call sites that construct `SimpleMapStore::new_with_strategy(..., cleanup_policy, ...)` + update to pass a `SimpleMapStorePersistenceConfig` instead. Main.rs and any + tests that construct the store directly will need to change. + +Does not change: -- `CleanupPolicy::CircularBuffer` and `ReadBased` continue to exist and run. A - deployment can opt into persistence in addition to a cleanup policy; the - persistence flusher runs *before* `CircularBuffer` would drop data, so epochs - get a chance to survive on disk. If both policies fire on the same epoch, - `CircularBuffer` wins (cleanup is destructive, but that's the user's stated - intent when they configure it). +- The `CleanupPolicy` enum itself stays in `asap_types` — other stores + (`promsketch_store`, legacy paths) may still reference it. This PR only + severs `SimpleMapStore`'s dependency on it. - `SimpleMapStoreGlobal` is intentionally left in-memory-only. Persistence targets `PerKey`, which is the production path. Adding it to `Global` is a small follow-up if anyone needs it. @@ -280,25 +340,28 @@ No new lock held across a fsync or disk I/O. The PR this design doc accompanies will land in three commits on one branch so the pieces can be reviewed independently: -1. **Sizing + config plumbing.** Add `approx_memory_bytes` to `AggregateCore` - and all concrete accumulators. Add `SimpleMapStorePersistenceConfig`. Track - `mem_bytes_in_use`. No disk I/O yet; expose the counter in - `StoreDiagnostics` so we can validate sizing in isolation. +1. **Sizing + config plumbing + cleanup-policy removal.** Add + `approx_memory_bytes` to `AggregateCore` and all concrete accumulators. + Add `SimpleMapStorePersistenceConfig`. Track `mem_bytes_in_use`. Rip the + `CleanupPolicy` branches out of `per_key.rs` and update call sites. No + disk I/O yet — expose the memory counter in `StoreDiagnostics` so we can + validate sizing in isolation and be confident nothing else regressed. 2. **Segment format, manifest, flusher.** Add `persistence/` submodule under `simple_map_store/` with segment encode/decode, manifest read/write, - background flusher task. Wire `flush_and_evict` into the per-key store. - Unit tests for round-trip, crash-after-segment-before-manifest, and orphan - sweep. + background flusher task (memory-first, then time-watermark). Wire + `flush_and_evict` into the per-key store. Unit tests for round-trip, + crash-after-segment-before-manifest, and orphan sweep. 3. **Query path read-through + recovery.** Extend `query_precomputed_output` to consult the manifest and merge segment hits with in-memory hits. Add startup recovery. Integration test: ingest → flush → restart → query → same result as no-restart. -Phases 1 and 2 are safe to merge independently because phase 2 is gated on -`enabled = false` by default. Phase 3 is when persistence becomes observable -to query results. +Because phase 1 removes `CleanupPolicy` from this store, phase 1 is **not** +independently mergeable without at least the memory-pressure path from +phase 2 — otherwise the store has no bound on RAM. In practice phases 1 and +2 land together; phase 3 can land separately once the write path is stable. --- @@ -314,10 +377,11 @@ to query results. for segment writes, sync locks everywhere else. Flusher runs on a dedicated task, not a shared runtime, to avoid starving it under query load. -3. **Should `flush_older_than_ms` live here or in the existing - `CleanupPolicy` enum?** It overlaps conceptually with `CircularBuffer`. - Proposal: keep it separate — `CleanupPolicy` is destructive, persistence - config is non-destructive. Confusing them in one knob would be worse. +3. **Cold data retention on disk.** Once a sketch is on disk, it lives there + until the operator removes the directory. Do we want a third knob + `delete_older_than_ms = T2` (with `T2 >> hot_window_ms`) so disk is also + bounded? My lean: not in v1 — cold data is cheap and operators can manage + the directory, but add the knob as soon as anyone asks. 4. **Per-agg-id flush fairness.** Oldest-global-first could starve small, slow-moving aggs during a burst on a hot agg. Acceptable for v1 since From c582226c64b56416ca77bf6f8e90598cd22a0c5f Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 20:03:25 -0400 Subject: [PATCH 03/12] docs: add Tier-2 read-side segment cache for repeat cold queries Introduces a two-tier memory model: Tier 1 is the existing write-driven hot set bounded by memory_limit_bytes + hot_window_ms, Tier 2 is a separate read-side LRU of decoded segments bounded by a new segment_cache_bytes budget. Keeps the flusher purely write-driven (monotonic tiering, no query-frequency feedback into retention) while giving cold-but-repeatedly-queried segments a place to live in RAM. Tier 2 is segment-granular, never dirty, drop-anytime, and completely independent of the flusher's decisions. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 115 +++++++++++++++++++- 1 file changed, 112 insertions(+), 3 deletions(-) diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md index e20b779e..e9d8ad29 100644 --- a/docs/design-simple-map-store-persistence.md +++ b/docs/design-simple-map-store-persistence.md @@ -275,14 +275,103 @@ flush off the insert/query critical path. 1. Read from `current_epoch + sealed_epochs` as today. 2. Look up segments in the manifest whose `[start_ms, end_ms]` overlaps the query range for this `agg_id`. -3. For each matching segment, read the file (cached via an `lru::LruCache>`), - decode entries whose window overlaps, and merge into the result. +3. For each matching segment, read the file (via the read-side segment cache + described below), decode entries whose window overlaps, and merge into the + result. Segment reads happen under a read lock on the manifest; they do **not** take any per-agg store lock, so they can run fully in parallel with inserts. Merging reuses the existing `TimestampedBucketsMap` + `AggregateCore::merge_with` that the in-memory query path already uses — no new merge logic. +### Read-side segment cache (two-tier memory model) + +So far the flusher treats "is this sketch in RAM?" as a pure function of +*write* state — time of ingest and write-side memory pressure. That is the +right default for a TSDB, because recency dominates query patterns, but it +leaves one real gap: **cold-but-repeatedly-queried** segments. Think of a +dashboard that scans "last Tuesday's incident" every time the on-call opens +it, or a recording rule that re-reads a fixed 24h historical range every +minute. Those queries touch segments that the time watermark has correctly +decided are cold, and under the design so far they pay full disk I/O on +every hit. + +The answer is **not** to let query frequency feed back into the flusher +policy. Doing that would couple write-path retention to read load, break +the monotonic "once cold, stays cold" invariant the flusher relies on, and +introduce unbounded-memory failure modes when a query sweeps everything. +The answer is a **second, separate memory tier** that exists purely as a +read-side cache on top of the disk layer. + +**Tier 1 — authoritative hot (write-driven).** Bounded by +`memory_limit_bytes` + `hot_window_ms`. Contains `current_epoch` and any +sealed epoch that has not yet been flushed. Source of truth for recent +data. Managed by the flusher described above. + +**Tier 2 — read-side segment cache (query-driven).** Bounded by a +separate `segment_cache_bytes` budget. Contains decoded copies of segments +pulled back from disk by the query path. Source of truth is always the +segment file — the cache is a pure optimization, drop-anytime, never +dirty. Managed by the query path, not the flusher. + +```rust +pub struct SimpleMapStorePersistenceConfig { + // ... existing fields ... + + // Read-side segment cache. Bounded independently of + // `memory_limit_bytes`; this budget is for decoded segments the query + // path pulls back from disk, not for the authoritative hot set. + // Set to 0 to disable. Default: small (e.g. 64 MiB), opt-in for + // workloads that don't need it. + pub segment_cache_bytes: usize, +} +``` + +**Why a TSDB specifically benefits from this shape:** + +1. **Recency and query frequency overlap ~90%.** Tier 1 already catches + everything a "rate over last 5m" workload wants pinned. The cache + only earns its budget on the residual workload — dashboards on fixed + old ranges, recording rules over long horizons. Making it a separate, + sized-independently tier means we can ship a small default (or zero) + and only budget it up for workloads that measurably need it. + +2. **Monotonic tiering.** Once a sealed epoch is flushed, it stays on + disk. A query may cache a decoded copy in Tier 2, but the flusher + never "un-flushes" it back into Tier 1. This preserves the property + that Tier 1 is purely a function of write state — which is what makes + the flusher simple enough to implement correctly. + +3. **Segment granularity, not sketch granularity.** The unit of disk I/O + is the segment file, so the cache must match that granularity. + Caching individual sketches inside a segment would mean partial reads + and complex invalidation; caching whole segments is a trivial + `LruCache>`. + +4. **Two independent budgets are easier to tune than one unified priority + score.** Operators reason about "how much RAM does write buffering + need?" and "how much RAM does read caching need?" separately. A + unified `priority = α * recency + β * frequency` score is harder to + explain and harder to debug when it misbehaves. + +**Eviction policy for Tier 2.** A plain LRU is the v1 default. If we see +recurring cold queries that get evicted by unrelated one-shot scans, we +can move to SLRU or TinyLFU later — both are drop-in replacements because +the cache has no consistency obligations. The cache should expose +hit/miss counters in `StoreDiagnostics` so we have data for that call. + +**Interaction with the flusher.** None. The flusher only sees Tier 1. +The read cache has no feedback into retention decisions. This is the +whole point of splitting the tiers. + +**What this deliberately does not do:** + +- No pinning of individual sketches in Tier 1 based on read counts. The + existing `read_counts` field on `StoreKeyData` becomes purely diagnostic + for this store — it does not veto flushes. +- No promotion from Tier 2 back into Tier 1. Once cold, stays cold. +- No partial-segment loading. Segments are cached whole or not at all. + ### Recovery on startup 1. Open `disk_path`, read `manifest.json`. @@ -356,13 +445,20 @@ the pieces can be reviewed independently: 3. **Query path read-through + recovery.** Extend `query_precomputed_output` to consult the manifest and merge segment hits with in-memory hits. Add startup recovery. Integration test: ingest → flush → restart → query → - same result as no-restart. + same result as no-restart. The read path uses a **trivial bounded LRU** + for the Tier-2 segment cache in this phase — just enough to avoid + re-reading the same segment on back-to-back queries. No SLRU/TinyLFU, + no hit/miss exporter, no tuning knobs beyond `segment_cache_bytes`. Because phase 1 removes `CleanupPolicy` from this store, phase 1 is **not** independently mergeable without at least the memory-pressure path from phase 2 — otherwise the store has no bound on RAM. In practice phases 1 and 2 land together; phase 3 can land separately once the write path is stable. +A later phase 4 (not part of this PR) would upgrade the Tier-2 cache to +SLRU/TinyLFU and export hit/miss metrics, once we have real query traces +to justify the algorithm choice. + --- ## Open questions (for review before implementation) @@ -386,3 +482,16 @@ phase 2 — otherwise the store has no bound on RAM. In practice phases 1 and 4. **Per-agg-id flush fairness.** Oldest-global-first could starve small, slow-moving aggs during a burst on a hot agg. Acceptable for v1 since "oldest window first" is well-defined globally; revisit if it bites. + +5. **Default `segment_cache_bytes`.** Should v1 default the Tier-2 cache to + a small nonzero value (e.g. 64 MiB) so the typical read path gets a + trivial hit-rate win for free, or default it to 0 (opt-in) so no + workload pays RAM it doesn't measurably benefit from? My lean: default + to a small nonzero value — a fresh install shouldn't have to know about + this knob to get reasonable repeat-query performance. + +6. **Tier-2 algorithm beyond LRU.** Plain LRU ships in phase 3. Do we + commit up-front to an upgrade path (SLRU, TinyLFU) or only revisit if + real traces show scan-resistant patterns are a problem? My lean: defer + — LRU is fine for the 90% case and the cache has no consistency + obligations, so swapping the algorithm is a purely local change. From bfd413b998cad4aa36cd6dae55918e96b0890b07 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 20:09:05 -0400 Subject: [PATCH 04/12] docs: resolve sync-vs-async flush I/O via append-only sealed epochs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sealed epochs are frozen once the rotator seals them — the flusher can clone the Arc out under a brief read lock, do I/O entirely outside any store lock, then splice the Arc out under a brief write lock. No lock is ever held across disk I/O, so there is no task on the flusher's runtime waiting for it to yield, which removes the only reason to prefer async I/O. v1 ships plain std::fs on a dedicated std::thread; no Tokio runtime for the flusher. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 83 +++++++++++++++++---- 1 file changed, 69 insertions(+), 14 deletions(-) diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md index e9d8ad29..99492498 100644 --- a/docs/design-simple-map-store-persistence.md +++ b/docs/design-simple-map-store-persistence.md @@ -262,11 +262,67 @@ left to do (the oldest epochs are already gone). Under light ingest, phase 1 is a no-op and phase 2 does all the work. The two phases never fight because they pull from the same oldest-first ordering. -`flush_and_evict` serializes the epoch *outside* the per-agg lock (the epoch is -immutable once sealed, we can read the `Arc` without holding the write lock), -fsyncs the segment file, then takes the per-agg `RwLock::write` only to splice -the epoch out of `sealed_epochs` and decrement `mem_bytes_in_use`. This keeps -flush off the insert/query critical path. +`flush_and_evict` takes advantage of a property that matters a lot for the +flusher design: **sealed epochs are append-only and frozen.** Once the +rotator seals an epoch, no writer will ever touch its contents again — it +is only read (by queries) or removed wholesale (by the flusher). That +immutability is what lets the flusher stay completely off the critical +path: + +1. Take the per-agg `RwLock::read` briefly, clone the `Arc` for the + target epoch out of `sealed_epochs`, drop the lock. +2. Serialize the epoch, write the segment file, and fsync — **entirely + outside any store lock**, on the flusher's own thread. Nothing in the + system is waiting on this I/O. Inserts continue to land in + `current_epoch`; queries continue to read from the still-in-place + `sealed_epochs` entry (and the cloned `Arc` keeps the bytes alive for + any query that happens to hold a reference already); the rotator + continues to seal new epochs behind us. +3. Once the segment is durable and the manifest is updated, take the + per-agg `RwLock::write` briefly to splice the epoch out of + `sealed_epochs` and decrement `mem_bytes_in_use`. This is O(1) — a + `BTreeMap::remove` plus an atomic subtraction — and is the only lock + the flusher holds for more than a read snapshot. + +Because steps 2 and 3 are decoupled by the `Arc` clone, **no lock +is ever held across disk I/O**, and the flusher never blocks anything on +the insert or query path beyond the two brief lock acquisitions at the +start and end. + +#### Sync vs. async flush I/O — resolved + +The previous revision left this as an open question. With the append-only +property made explicit, the answer is clear: **plain `std::fs` on a +dedicated `std::thread` is what we ship.** No `tokio::fs`, no Tokio +runtime for the flusher. + +The only argument for async I/O would be "we need to yield the thread +while `fsync` is in flight so some other task on the same runtime can +make progress" — and there is no such other task. The flusher thread has +exactly one job — flush — and blocking it on `write` + `fsync` is fine +because: + +- **Inserts never wait on the flusher.** Inserts land in `current_epoch` + with no coordination with flush state; memory accounting is an atomic, + not a lock. The flusher and the insert path only share the per-agg + `RwLock`, and the flusher only holds it during the two brief windows + above. +- **Queries never wait on the flusher.** In-memory reads take the per-agg + `RwLock::read`, which contends with the flusher only during those same + brief windows; disk reads go through the manifest lock, which is + independent. +- **Back-pressure is the right answer to a slow disk.** If the flusher + genuinely cannot keep up and memory hits `hard_cap_bytes`, the insert + path blocks on a condvar until the flusher catches up. That is the + correct behavior regardless of whether the I/O underneath is sync or + async — making it async would not let more inserts through, it would + just change which thread was parked. + +Sync I/O keeps the store out of Tokio's executor entirely, keeps stack +traces readable, and eliminates a class of "why is my future not making +progress" failure modes. The flusher thread is `std::thread::spawn`'d in +`SimpleMapStorePerKey::new` and joined in `close`, with a `shutdown` +flag checked on each loop iteration. ### Query path @@ -468,30 +524,29 @@ to justify the algorithm choice. tooling. I lean custom for v1 given that we never read segments outside this process, but happy to switch if there's an appetite. -2. **Async vs. sync flush I/O.** The rest of the store is sync (`RwLock`, - `DashMap`), but the flusher task is naturally async. Proposal: `tokio::fs` - for segment writes, sync locks everywhere else. Flusher runs on a dedicated - task, not a shared runtime, to avoid starving it under query load. - -3. **Cold data retention on disk.** Once a sketch is on disk, it lives there +2. **Cold data retention on disk.** Once a sketch is on disk, it lives there until the operator removes the directory. Do we want a third knob `delete_older_than_ms = T2` (with `T2 >> hot_window_ms`) so disk is also bounded? My lean: not in v1 — cold data is cheap and operators can manage the directory, but add the knob as soon as anyone asks. -4. **Per-agg-id flush fairness.** Oldest-global-first could starve small, +3. **Per-agg-id flush fairness.** Oldest-global-first could starve small, slow-moving aggs during a burst on a hot agg. Acceptable for v1 since "oldest window first" is well-defined globally; revisit if it bites. -5. **Default `segment_cache_bytes`.** Should v1 default the Tier-2 cache to +4. **Default `segment_cache_bytes`.** Should v1 default the Tier-2 cache to a small nonzero value (e.g. 64 MiB) so the typical read path gets a trivial hit-rate win for free, or default it to 0 (opt-in) so no workload pays RAM it doesn't measurably benefit from? My lean: default to a small nonzero value — a fresh install shouldn't have to know about this knob to get reasonable repeat-query performance. -6. **Tier-2 algorithm beyond LRU.** Plain LRU ships in phase 3. Do we +5. **Tier-2 algorithm beyond LRU.** Plain LRU ships in phase 3. Do we commit up-front to an upgrade path (SLRU, TinyLFU) or only revisit if real traces show scan-resistant patterns are a problem? My lean: defer — LRU is fine for the 90% case and the cache has no consistency obligations, so swapping the algorithm is a purely local change. + +*(The earlier open question about sync vs. async flush I/O is resolved +in-line above — the append-only property of sealed epochs makes sync +`std::fs` on a dedicated thread the clear winner.)* From da6a47bfea2ce753c1b89c4f4441eb663b7cdf91 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 20:16:26 -0400 Subject: [PATCH 05/12] docs: adopt perf-optimal choices across all open questions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves all six previously-open questions in favor of the performance-optimal answer instead of the simplest-v1 answer: 1. Segment format — 8-byte aligned, mmap-friendly, fallocate'd custom binary (unchanged direction, concrete details added). 2. Flush I/O — sync std::fs with group-commit fdatasync batching across a tick; no tokio::fs, no async runtime for the flusher. 3. Cold data retention — new delete_older_than_ms config knob and a phase-3 retention sweep in the flusher loop, bounding the manifest and keeping long-running deployments fast. 4. Flush fairness — round-robin across agg-ids, oldest-first within each agg, replacing strict-global-oldest-first. Prevents per-agg RwLock hot-spotting during hot-agg bursts. 5. Segment cache default — min(10% * memory_limit_bytes, 512 MiB), scaling with the write budget instead of a fixed 64 MiB. 6. Tier-2 algorithm — W-TinyLFU via moka from day one, replacing plain LRU. Scan-resistant, better hit rate on dashboard / recording- rule access patterns, drop-in API. Replaces the "Open questions" section with a "Resolved decisions" table pointing at the section where each decision's rationale lives. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 273 +++++++++++++------- 1 file changed, 181 insertions(+), 92 deletions(-) diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md index 99492498..13acf3c7 100644 --- a/docs/design-simple-map-store-persistence.md +++ b/docs/design-simple-map-store-persistence.md @@ -92,22 +92,38 @@ the rotator seals it. └── seg_0000000001.bin ``` -**Segment file format** (`seg_*.bin`): +**Segment file format** (`seg_*.bin`). Every field is laid out so the whole +file can be `mmap`ed and sketch payloads handed directly to +`AggregateCore::deserialize_from_bytes` with zero copies: ``` [u32 magic][u16 version][u16 flags] [u64 epoch_id][u64 window_start_ms][u64 window_end_ms] -[u32 num_entries] +[u32 num_entries][u32 _pad] // align to 8 repeated num_entries times: - [u64 start_ts][u64 end_ts][u32 label_id][u8 agg_type] - [u32 payload_len][payload_len bytes: serialize_to_bytes()] -[u32 crc32 of body] + [u64 start_ts][u64 end_ts][u32 label_id][u8 agg_type][u24 _pad] + [u32 payload_len][u32 _pad] // align payload to 8 + [payload_len bytes: serialize_to_bytes()] + [0..7 bytes: tail padding to 8-byte boundary] +[u32 crc32 of body][u32 _pad] // trailer aligned ``` Fixed-size header lets us mmap and binary-search by timestamp without parsing payloads. Body is a linear scan — v1 does not build an in-segment index because sealed epochs are small (bounded by window size × group count for a single agg). +**Write-side perf details:** + +- Before writing, call `fallocate(fd, 0, 0, estimated_size)` to reserve + contiguous space and avoid ext4/xfs metadata churn under many-small-segments + workloads. `estimated_size` is `sum of approx_memory_bytes * 1.3` for a safe + upper bound; any slack is released via `ftruncate` at the end. +- 8-byte alignment for every payload means a single `mmap` + pointer cast is + safe on every architecture Rust targets. Without alignment, ARM and + `MIRI`-style UB checks require a copy-to-aligned-buffer step. +- The trailing CRC is computed streaming while we write the body, so we do not + re-read the file to compute it. + **Manifest** (`manifest.json`): ```json @@ -183,6 +199,18 @@ pub struct SimpleMapStorePersistenceConfig { // until something pushes it out). pub hot_window_ms: Option, + // ---- Disk retention ---- + // + // Cold-tier TTL. Any segment whose end_ms is older than + // `now - delete_older_than_ms` is deleted from disk on the next + // flusher tick (after its references are removed from the manifest + // and no in-flight query is reading it). Bounds disk usage and keeps + // the manifest small enough to stay in L2/L3 cache on long-running + // deployments. Must be strictly greater than hot_window_ms; expected + // to be much greater (hours vs. days or weeks). + // None disables cold deletion entirely — disk grows unboundedly. + pub delete_older_than_ms: Option, + // ---- Misc ---- pub flush_interval_ms: u64, // cadence of the background flusher pub disk_path: PathBuf, // root dir for segments + manifest @@ -197,70 +225,126 @@ and `memory_limit_bytes = usize::MAX`, which degenerates to "never flush." ### Eviction order -There is only one ordering — **oldest-sealed-epoch-first, globally by epoch -`end_ms`**. Both triggers (memory pressure and time watermark) pull from the -same ordered view, so the flusher never has two disagreeing notions of "oldest." +**Round-robin across `agg_id`, oldest-first within each agg.** Every tick, +the flusher walks agg-ids in order, pops the oldest sealed epoch from each, +and repeats until the stopping condition (memory low-water or end of time +threshold) is met. Both triggers share the same walk order so the flusher +never has two disagreeing notions of "what to flush next." Rationale: -- Matches the time-window access pattern: queries overwhelmingly target recent - windows, so evicting oldest is the lowest-regret choice. -- Makes the two triggers composable: the memory-pressure pass and the - time-watermark pass are just two different stopping conditions on the same - iterator over `(agg_id, epoch) sorted by epoch.end_ms`. -- Avoids cross-agg fairness debates (e.g., `LargestAggFirst`) that would - otherwise complicate v1; we can add more orderings later behind an enum if - it becomes necessary. +- **Lock-spread under burst.** A strict global oldest-first ordering would + flush many epochs from the *same* hot agg-id back-to-back, hammering the + same per-agg `RwLock` repeatedly and creating brief query-latency spikes + on that one agg. Round-robin spreads the flusher's lock acquisitions + across different `RwLock`s, which is cheap for DashMap (lock-free outer) + and gives query latency a smoother profile. +- **Same total work, no complexity cost.** Round-robin does not evaluate + more epochs than strict-global would; it just reorders which epoch is + flushed next. Implementation cost is one `BTreeMap<(agg_id, end_ms), + EpochRef>` populated by walking `store` once per tick, or equivalently a + per-agg min-heap of sealed epochs with a round-robin cursor. +- **Freshness.** Small, slow-moving aggs are never starved by a burst on + a hot agg — they always get a turn in each round. +- **Still matches the time-window access pattern.** Within each agg, + oldest-first is preserved, so queries against recent windows on any agg + remain unaffected. ### Background flusher -A dedicated Tokio task owned by the store, started in +A dedicated `std::thread` owned by the store, started in `SimpleMapStorePerKey::new`. Each tick, it checks the primary trigger -(memory) first, then the secondary trigger (time watermark): +(memory) first, the secondary trigger (time watermark), then the +disk-retention sweep: ``` loop { - sleep(flush_interval_ms).await; + thread::sleep(flush_interval_ms); + if shutdown.load() { break; } let now = now_ms(); - let mut candidates = Vec::new(); + let mut candidates: Vec = Vec::new(); // Phase 1 (PRIMARY): memory budget. - // If we're over the high-water mark, pull oldest sealed epochs - // (by epoch.end_ms) until projected memory drops below the + // Walk agg-ids round-robin, pulling the oldest sealed epoch from + // each on every pass, until projected memory drops below the // low-water mark. This is the knob that actually bounds RAM. if mem_bytes_in_use.load() > cfg.memory_limit_bytes { - candidates.extend(collect_oldest_until_under_low_water( + candidates.extend(collect_round_robin_until_under_low_water( cfg.memory_low_watermark_bytes, )); } // Phase 2 (SECONDARY): time watermark T. // Any sealed epoch older than `now - hot_window_ms` that wasn't - // already picked up in phase 1 is flushed here. Under light ingest, - // this is the only phase that runs and it keeps the hot set bounded - // by T × ingest rate regardless of the memory budget. + // already picked up in phase 1 is flushed here. Also walked + // round-robin across aggs so a burst on one hot agg does not + // monopolize the tick. if let Some(hot_window) = cfg.hot_window_ms { - candidates.extend(collect_epochs_older_than(now - hot_window)); + candidates.extend(collect_older_than_round_robin( + now - hot_window, + )); } - // Dedup (phase 1 and phase 2 can pick the same epoch) and sort by - // epoch.end_ms ascending so we flush oldest first within the batch. - candidates.sort_unstable_by_key(|c| c.end_ms); - candidates.dedup(); + // Dedup (phase 1 and phase 2 can pick the same epoch). Order is + // already interleaved across aggs; no re-sort. + candidates.dedup_by_key(|c| (c.agg_id, c.epoch_id)); - for (agg_id, epoch_id) in candidates { - flush_and_evict(agg_id, epoch_id).await?; + // ---- Group-commit the whole tick ---- + let mut written: Vec = Vec::new(); + for epoch_ref in candidates { + let bytes = serialize_epoch(epoch_ref.arc.clone()); + let path = write_segment_no_fsync(&bytes, epoch_ref)?; + written.push(WrittenSegment { path, meta: epoch_ref.meta }); + } + fdatasync_all(&written)?; // one batched fsync pass + manifest.rewrite_and_fsync(&written)?; + fsync_parent_dir(&cfg.disk_path)?; // one dir fsync for the whole batch + + // Now that segments are durable AND referenced by the manifest, + // evict them from memory. + for seg in &written { + splice_out_of_sealed_epochs(seg.meta); + mem_bytes_in_use.fetch_sub(seg.meta.approx_bytes); } - manifest.commit().await?; // atomic rewrite after the batch + // Phase 3 (disk retention sweep): delete segments older than T2. + if let Some(ttl) = cfg.delete_older_than_ms { + let cutoff = now.saturating_sub(ttl); + let expired = manifest.segments_older_than(cutoff); + for seg in expired { + manifest.remove(seg.id); + cache_tier2.invalidate(seg.id); // drop any decoded copy + fs::remove_file(seg.path).ok(); // best-effort; orphan sweep on restart + } + if !expired.is_empty() { + manifest.rewrite_and_fsync(&[])?; + } + } } ``` Under memory pressure, phase 1 dominates and phase 2 usually finds nothing left to do (the oldest epochs are already gone). Under light ingest, phase 1 -is a no-op and phase 2 does all the work. The two phases never fight because -they pull from the same oldest-first ordering. +is a no-op and phase 2 does all the work. Phase 3 is independent and runs +every tick regardless; it costs one manifest scan plus one `unlink` per +expired segment. + +#### Group-commit fsync + +The pseudocode above batches all `fsync`/`fdatasync` calls for a tick into a +single pass at the end, rather than `fsync`ing each segment inline. On +spinning disks this is ~10× fewer head seeks per tick; on SSDs it is ~3–4× +fewer syscalls. The cost is one temporarily-larger `written` vector and one +extra `fsync_parent_dir` at the end — trivial relative to the saved I/O. + +Durability invariant remains the same: **no segment is referenced by the +manifest until its bytes and the manifest itself are both `fsync`'d.** The +group-commit ordering is (1) write all segment bodies, (2) `fdatasync` all +of them, (3) rewrite + `fsync` manifest via the atomic `write → rename` +dance, (4) `fsync` parent directory. A crash at any point leaves orphan +segment files (cleaned by the startup sweep) but never a dangling manifest +entry. `flush_and_evict` takes advantage of a property that matters a lot for the flusher design: **sealed epochs are append-only and frozen.** Once the @@ -377,8 +461,14 @@ pub struct SimpleMapStorePersistenceConfig { // Read-side segment cache. Bounded independently of // `memory_limit_bytes`; this budget is for decoded segments the query // path pulls back from disk, not for the authoritative hot set. - // Set to 0 to disable. Default: small (e.g. 64 MiB), opt-in for - // workloads that don't need it. + // + // Default: min(10% * memory_limit_bytes, 512 MiB). + // + // A fresh install should not need to know about this knob to get + // reasonable repeat-query performance. Setting to 0 disables Tier 2 + // entirely (every cold query pays disk I/O); a fixed absolute + // default would be too small on big boxes and too large on small + // ones, so the default scales with the write budget. pub segment_cache_bytes: usize, } ``` @@ -402,7 +492,7 @@ pub struct SimpleMapStorePersistenceConfig { is the segment file, so the cache must match that granularity. Caching individual sketches inside a segment would mean partial reads and complex invalidation; caching whole segments is a trivial - `LruCache>`. + `Cache>` keyed on manifest metadata. 4. **Two independent budgets are easier to tune than one unified priority score.** Operators reason about "how much RAM does write buffering @@ -410,11 +500,25 @@ pub struct SimpleMapStorePersistenceConfig { unified `priority = α * recency + β * frequency` score is harder to explain and harder to debug when it misbehaves. -**Eviction policy for Tier 2.** A plain LRU is the v1 default. If we see -recurring cold queries that get evicted by unrelated one-shot scans, we -can move to SLRU or TinyLFU later — both are drop-in replacements because -the cache has no consistency obligations. The cache should expose -hit/miss counters in `StoreDiagnostics` so we have data for that call. +**Eviction policy for Tier 2: W-TinyLFU via `moka` (or `mini-moka`).** +Plain LRU is the obvious choice but is catastrophically scan-vulnerable — +a single long-range query sweeps the cache and evicts everything genuinely +hot, which is exactly the access pattern TSDB dashboards and recording +rules produce (hour/day/week range scans). W-TinyLFU's admission filter +rejects scan traffic from displacing hot entries and typically delivers +10–30% better hit rate than LRU at the same byte budget on skewed / +Zipfian workloads. + +The `moka` crate is the standard Rust implementation (sync and async +variants, weight-based eviction keyed on byte size, well-maintained, used +widely in the Rust ecosystem). The API is effectively a drop-in for LRU +(`get`, `insert`, `invalidate`), so we incur no additional complexity vs. +a hand-rolled LRU — just better hit rate. W-TinyLFU's per-access overhead +is a handful of CAS ops on a small count-min sketch, cheaper than LRU's +mutex-protected list reordering. + +The cache exposes hit/miss counters in `StoreDiagnostics` from day one so +we have signal for future tuning. **Interaction with the flusher.** None. The flusher only sees Tier 1. The read cache has no feedback into retention decisions. This is the @@ -493,60 +597,45 @@ the pieces can be reviewed independently: validate sizing in isolation and be confident nothing else regressed. 2. **Segment format, manifest, flusher.** Add `persistence/` submodule under - `simple_map_store/` with segment encode/decode, manifest read/write, - background flusher task (memory-first, then time-watermark). Wire - `flush_and_evict` into the per-key store. Unit tests for round-trip, - crash-after-segment-before-manifest, and orphan sweep. - -3. **Query path read-through + recovery.** Extend `query_precomputed_output` - to consult the manifest and merge segment hits with in-memory hits. Add + `simple_map_store/` with segment encode/decode (8-byte aligned, + `fallocate`d, mmap-friendly), manifest read/write, background flusher + thread (memory-first, then time-watermark, then T2 retention sweep) with + group-commit fsync. Wire `flush_and_evict` into the per-key store. Unit + tests for round-trip, crash-after-segment-before-manifest, orphan sweep, + and T2 deletion. + +3. **Query path read-through + recovery + Tier-2 cache.** Extend + `query_precomputed_output` to consult the manifest and merge segment hits + with in-memory hits. Wire a `moka` (or `mini-moka`) weight-bounded + segment cache sized to `min(10% * memory_limit_bytes, 512 MiB)` by + default, with hit/miss counters exported via `StoreDiagnostics`. Add startup recovery. Integration test: ingest → flush → restart → query → - same result as no-restart. The read path uses a **trivial bounded LRU** - for the Tier-2 segment cache in this phase — just enough to avoid - re-reading the same segment on back-to-back queries. No SLRU/TinyLFU, - no hit/miss exporter, no tuning knobs beyond `segment_cache_bytes`. + same result as no-restart, plus a scan-resistance test that confirms a + long-range query does not evict a separately-hot segment. Because phase 1 removes `CleanupPolicy` from this store, phase 1 is **not** independently mergeable without at least the memory-pressure path from phase 2 — otherwise the store has no bound on RAM. In practice phases 1 and 2 land together; phase 3 can land separately once the write path is stable. -A later phase 4 (not part of this PR) would upgrade the Tier-2 cache to -SLRU/TinyLFU and export hit/miss metrics, once we have real query traces -to justify the algorithm choice. - --- -## Open questions (for review before implementation) - -1. **Segment file format: custom binary vs. something off-the-shelf (Parquet, - Arrow IPC)?** Custom binary is simpler and avoids a dependency, but loses us - tooling. I lean custom for v1 given that we never read segments outside - this process, but happy to switch if there's an appetite. - -2. **Cold data retention on disk.** Once a sketch is on disk, it lives there - until the operator removes the directory. Do we want a third knob - `delete_older_than_ms = T2` (with `T2 >> hot_window_ms`) so disk is also - bounded? My lean: not in v1 — cold data is cheap and operators can manage - the directory, but add the knob as soon as anyone asks. - -3. **Per-agg-id flush fairness.** Oldest-global-first could starve small, - slow-moving aggs during a burst on a hot agg. Acceptable for v1 since - "oldest window first" is well-defined globally; revisit if it bites. - -4. **Default `segment_cache_bytes`.** Should v1 default the Tier-2 cache to - a small nonzero value (e.g. 64 MiB) so the typical read path gets a - trivial hit-rate win for free, or default it to 0 (opt-in) so no - workload pays RAM it doesn't measurably benefit from? My lean: default - to a small nonzero value — a fresh install shouldn't have to know about - this knob to get reasonable repeat-query performance. - -5. **Tier-2 algorithm beyond LRU.** Plain LRU ships in phase 3. Do we - commit up-front to an upgrade path (SLRU, TinyLFU) or only revisit if - real traces show scan-resistant patterns are a problem? My lean: defer - — LRU is fine for the 90% case and the cache has no consistency - obligations, so swapping the algorithm is a purely local change. - -*(The earlier open question about sync vs. async flush I/O is resolved -in-line above — the append-only property of sealed epochs makes sync -`std::fs` on a dedicated thread the clear winner.)* +## Resolved decisions + +Every question previously flagged as open has been resolved in favor of the +performance-optimal choice. The table below is a summary; the reasoning for +each lives in the section it points to. + +| # | Question | Decision | Why | +|---|---|---|---| +| 1 | Segment file format | Custom binary, 8-byte aligned, mmap-friendly, `fallocate`d | Zero-copy deserialize into `AggregateCore`; no Parquet/Arrow overhead for data we never column-prune; smaller binary size and compile time. See **Disk layout**. | +| 2 | Sync vs. async flush I/O | Sync `std::fs` on a dedicated `std::thread`, with **group-commit fsync** batching across a tick | `tokio::fs` just routes to a blocking threadpool on Linux, so async is a wash at the syscall level; the real win is batching `fdatasync`. No task on the flusher's runtime is waiting on it to yield. See **Background flusher / Group-commit fsync**. | +| 3 | Cold data retention on disk | Add `delete_older_than_ms = T2`, run as phase 3 of the flusher tick | Unbounded segment count bloats the manifest (falls out of L2/L3), slows startup directory sweeps, and pressures Tier-2 eviction. Cheap to add now, painful to retrofit once a deployment has millions of orphans. See **Configuration** and **Background flusher phase 3**. | +| 4 | Per-agg-id flush fairness | **Round-robin across agg-ids, oldest-first within each agg** | Strict-global-oldest hammers one `RwLock` during a hot-agg burst and creates query-latency spikes on that one agg. Round-robin spreads lock acquisitions across different `RwLock`s for the same total work. See **Eviction order**. | +| 5 | Default `segment_cache_bytes` | `min(10% * memory_limit_bytes, 512 MiB)` | A fixed 64 MiB default is too small on big boxes and too large on small ones; scaling with the write budget keeps the tier sensibly sized without requiring operator tuning on a fresh install. See **Configuration**. | +| 6 | Tier-2 algorithm | **W-TinyLFU via `moka`** from day one, not plain LRU | LRU is scan-vulnerable — one long-range query evicts everything genuinely hot, which is exactly the TSDB dashboard access pattern. W-TinyLFU's admission filter rejects scan traffic and typically delivers 10–30% better hit rate at the same byte budget on skewed workloads, with a drop-in API and lower per-access CPU than LRU. See **Read-side segment cache**. | + +If any of these decisions turn out to be wrong under real traces, the +affected sections are the natural point of revisiting — but none of them +are "temporary v1 shortcuts we'll upgrade later." This is the target +design. From b82cec3c9684fa1fd22b49c7b0236fce07981c5b Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 20:27:39 -0400 Subject: [PATCH 06/12] docs: parts-based disk layout (LSM-style), append-only manifest log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces one-file-per-sealed-epoch + dir-per-agg + JSON manifest with the layout every mainstream TSDB converges on: * Unit of file is a "part" — one directory per flush tick containing meta.bin + data.bin + index.bin. All of the tick's candidate epochs are bundled into a single data.bin regardless of which agg-id they came from, and are locatable via a sorted, mmap-friendly index.bin that supports O(log N) binary search on (agg_id, start_ms). * File count scales with flush ticks, not with epochs. One-file-per- epoch produced ~288K files/day on a 200-agg 1-minute-window setup; the parts layout produces ~86K files (three per tick) with a natural group-commit of fdatasync amortized over the whole tick. * Global manifest becomes an append-only parts_manifest.log with a periodic binary parts_manifest.snapshot, replacing the JSON file that was rewritten in full every tick. Size is proportional to flush ticks, not to epochs that have ever existed, and the snapshot is mmap-cast-to-slice on startup (zero parse). * T2 retention becomes whole-part rm -rf on tight time ranges, since each part covers ~flush_interval_ms of data. * Tier-2 cache is now keyed on PartId and holds mmap'd part views; config knob renamed segment_cache_bytes -> part_cache_bytes. Updates flusher pseudocode, flush_and_evict narrative, query path, recovery sequence, concurrency summary, phasing, and adds entry #7 to the Resolved decisions table. Compaction of adjacent small parts is noted as a v2 follow-up — the layout accommodates it cleanly but v1 ships without it since T2 + a reasonable flush interval keeps part count well within what a binary-searched Vec handles. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/design-simple-map-store-persistence.md | 528 +++++++++++++------- 1 file changed, 349 insertions(+), 179 deletions(-) diff --git a/docs/design-simple-map-store-persistence.md b/docs/design-simple-map-store-persistence.md index 13acf3c7..70f322c4 100644 --- a/docs/design-simple-map-store-persistence.md +++ b/docs/design-simple-map-store-persistence.md @@ -64,84 +64,154 @@ append-only until cleanup. That is exactly the right unit to flush. ## Design -### Unit of flush: the sealed epoch +### Unit of flush vs. unit of file: the *part* -A `SimpleMapStore` segment on disk corresponds to **one sealed epoch of one -aggregation id**. Rationale: +There are two granularities to separate cleanly: -- Sealed epochs are immutable — safe to serialize without coordinating with writers. -- The epoch already has a well-defined time range, which is exactly what range - queries want to filter on. -- Flushing an epoch only requires the per-agg-id `RwLock::write` briefly — same - lock the insert path already uses, so no new contention class. -- Recovery and query planning only need epoch-level metadata, not per-window. +- **Unit of flush = sealed epoch.** Same as before. Sealed epochs are + immutable, have a well-defined time range, and can be spliced out of + `sealed_epochs` under a brief per-agg `RwLock::write`. +- **Unit of file = *part*.** A part is **one flush tick's worth of sealed + epochs bundled into a single on-disk directory**, regardless of which + agg-id they came from. The flusher already assembles all the candidate + epochs for a tick before it touches the disk; instead of writing N + separate segment files and fsyncing each, it writes one part. -The `current_epoch` is never flushed while hot. It becomes flushable the moment -the rotator seals it. +This is the same pattern every mainstream TSDB converges on — Prometheus +blocks, InfluxDB TSM, VictoriaMetrics parts — for the same reasons: +file count is bounded by flush ticks (not by individual epochs), metadata +overhead is amortized across many entries, and compaction becomes a pure +directory-level merge. + +The `current_epoch` is never flushed while hot. It becomes flushable the +moment the rotator seals it, at which point it becomes a candidate for the +next flush tick's part. ### Disk layout ``` / -├── manifest.json # authoritative index of all segments -├── agg_00000042/ -│ ├── seg_0000000001.bin # one sealed epoch, serialized -│ ├── seg_0000000002.bin -│ └── ... -└── agg_00000043/ - └── seg_0000000001.bin +├── parts_manifest.log # append-only log of part additions + deletions +├── parts_manifest.snapshot # periodic binary snapshot (compaction of the log) +└── parts/ + ├── 0000000001/ # part directory, name = monotonic part_id + │ ├── meta.bin # fixed-size header: min_ts, max_ts, counts, crc + │ ├── data.bin # all epoch payloads concatenated, 8-byte aligned + │ └── index.bin # sorted array of entries, mmap-binary-search target + ├── 0000000002/ + │ ├── meta.bin + │ ├── data.bin + │ └── index.bin + └── ... ``` -**Segment file format** (`seg_*.bin`). Every field is laid out so the whole -file can be `mmap`ed and sketch payloads handed directly to -`AggregateCore::deserialize_from_bytes` with zero copies: +**Why parts instead of dir-per-agg with file-per-epoch:** + +- **File count scales with flush ticks, not with epochs.** On a 1-second + flush interval with 200 agg-ids and 1-minute windows, the old layout + produced ~288K files/day; the part layout produces ~86K files total + (one tick = three files: `meta.bin`, `data.bin`, `index.bin`). That is + the difference between "inode pressure is a real concern" and "we are + well within every filesystem's comfort zone." +- **Metadata is amortized.** One `fallocate` + one `fdatasync` per + `data.bin` covers all epochs in the tick, rather than N separate + allocations and N separate fsyncs. Group-commit is now an intrinsic + property of the layout, not something the flusher has to arrange. +- **Per-part in-file index.** Queries binary-search the part's `index.bin` + rather than linear-scanning a segment body. O(log N) per part instead + of O(N), and `index.bin` is mmap-friendly so the search is pure + pointer arithmetic with no syscalls. +- **Aggs are interleaved inside a part, not segregated by directory.** + No wasted directories for low-traffic aggs; the in-part index handles + agg lookup cheaply. +- **T2 retention is whole-directory.** Each part covers a tight time + range (roughly `flush_interval_ms`), so `delete_older_than_ms` + operates at part granularity — `rm -rf parts/000001234/` — instead of + touching shared files. +- **Compaction fits naturally.** A background compactor can merge N + adjacent old parts into one larger part with the same on-disk shape. + Readers don't care because the parts_manifest gets updated atomically + and old part dirs get removed only after all in-flight readers are done. + +**Part file formats** (all little-endian, fixed layouts, mmap-friendly, +8-byte aligned, written via `fallocate` + streaming CRC — same perf +details that applied to segments, now applied to `data.bin`): ``` -[u32 magic][u16 version][u16 flags] -[u64 epoch_id][u64 window_start_ms][u64 window_end_ms] -[u32 num_entries][u32 _pad] // align to 8 -repeated num_entries times: - [u64 start_ts][u64 end_ts][u32 label_id][u8 agg_type][u24 _pad] - [u32 payload_len][u32 _pad] // align payload to 8 - [payload_len bytes: serialize_to_bytes()] - [0..7 bytes: tail padding to 8-byte boundary] -[u32 crc32 of body][u32 _pad] // trailer aligned -``` - -Fixed-size header lets us mmap and binary-search by timestamp without parsing -payloads. Body is a linear scan — v1 does not build an in-segment index because -sealed epochs are small (bounded by window size × group count for a single agg). - -**Write-side perf details:** - -- Before writing, call `fallocate(fd, 0, 0, estimated_size)` to reserve - contiguous space and avoid ext4/xfs metadata churn under many-small-segments - workloads. `estimated_size` is `sum of approx_memory_bytes * 1.3` for a safe - upper bound; any slack is released via `ftruncate` at the end. -- 8-byte alignment for every payload means a single `mmap` + pointer cast is - safe on every architecture Rust targets. Without alignment, ARM and - `MIRI`-style UB checks require a copy-to-aligned-buffer step. -- The trailing CRC is computed streaming while we write the body, so we do not - re-read the file to compute it. - -**Manifest** (`manifest.json`): - -```json -{ - "version": 1, - "segments": [ - {"agg_id": 42, "epoch_id": 1, "path": "agg_00000042/seg_0000000001.bin", - "start_ms": 1_700_000_000_000, "end_ms": 1_700_000_060_000, - "num_entries": 120, "size_bytes": 48192} - ] -} +meta.bin (128 bytes, fixed) + [u32 magic][u16 version][u16 flags] + [u64 part_id][u64 min_ts][u64 max_ts] + [u32 num_entries][u32 num_aggs] + [u64 data_len][u64 index_len] + [u64 created_unix_ns] + [u32 _reserved; 6] + [u32 crc32 of the above] + +data.bin (sum of padded payloads) + repeated num_entries times, in the order the index lists them: + [payload_len bytes: serialize_to_bytes()] + [0..7 bytes: tail padding to 8-byte boundary] + +index.bin (32 bytes per entry, sorted by (agg_id, start_ms)) + repeated num_entries times: + [u64 agg_id][u64 start_ms][u64 end_ms] + [u32 data_offset][u32 payload_len] + [u32 crc32 of the above][u32 _pad] ``` -The manifest is the single source of truth for which segments exist and what -ranges they cover. It is rewritten atomically (`write → fsync → rename`) after -every flush batch. Individual segment files are written+fsynced before the -manifest ever references them, so a crash mid-flush leaves orphan files that -startup sweeps away — never a dangling manifest entry. +`index.bin` is the only file a query needs to traverse to locate entries +inside a part. It is small (32 B × num_entries, typically tens of KB), is +mmap'd on first access, and a binary search by `(agg_id, start_ms)` lands +on the byte range inside `data.bin` with one pointer-arithmetic step and +zero decode work. + +**Parts manifest: append-only log + periodic snapshot.** + +The manifest is the one piece of global state on disk. The previous +design had it as a JSON file rewritten on every flush tick — quadratic +over the lifetime of the deployment. We replace it with the standard +LSM-style pattern: + +- **`parts_manifest.log`** is an append-only binary file. Each flush + tick appends one record (add-part or delete-part, both fixed-size). + Appending is a single `write + fdatasync` on a file whose size is + proportional to the number of *ticks*, not the number of parts that + have ever existed. Cheap and O(1) per tick. +- **`parts_manifest.snapshot`** is a periodic binary snapshot of the + live set of parts, produced by replaying the log and emitting a flat + sorted array of `(part_id: u64, min_ts: u64, max_ts: u64, size: u64)` + = 32 bytes per live part. The snapshot is rewritten atomically (write + tmp → fsync → rename) whenever the log gets large relative to the + snapshot, and the log is truncated after. Snapshot + remaining log + is always the authoritative live state. +- **On startup**, the store loads the snapshot (mmap + direct cast, no + parse), then replays any tail of the log added since the snapshot was + taken, then verifies every live part directory's `meta.bin` CRC. + Sweep orphan part dirs (present on disk but not in the replayed + state) and treat them as a mid-flush crash — delete them. + +Binary formats throughout mean parse time is effectively zero; the +snapshot is "cast a byte slice to `&[PartEntry]`" — which works because +we declared the layout 8-byte aligned and fixed-size. + +**Durability ordering per flush tick** (the invariant a crash must not +violate: no part is referenced in the manifest until its bytes are on +disk): + +1. Assemble the tick's candidate epochs (in-memory, no I/O). +2. `fallocate` the three files in `parts//`, stream payloads + into `data.bin`, stream index into `index.bin`, write `meta.bin`. +3. `fdatasync` `data.bin`, `index.bin`, `meta.bin` (batched — one + syscall per file, not per entry). +4. `fsync` the part directory itself. +5. Append the add-part record to `parts_manifest.log` and `fdatasync` + the log. +6. `fsync` `` (the root) so the log's size update is durable. + +A crash at any step before (5) leaves an orphan part directory that +startup sweep deletes. A crash between (5) and (6) is fine — the log +record is already durable via (5). After (6), the part is officially +live and the flusher may evict the corresponding epochs from memory. ### Memory accounting @@ -289,36 +359,88 @@ loop { // Dedup (phase 1 and phase 2 can pick the same epoch). Order is // already interleaved across aggs; no re-sort. candidates.dedup_by_key(|c| (c.agg_id, c.epoch_id)); + if candidates.is_empty() { /* go straight to phase 3 below */ } - // ---- Group-commit the whole tick ---- - let mut written: Vec = Vec::new(); - for epoch_ref in candidates { - let bytes = serialize_epoch(epoch_ref.arc.clone()); - let path = write_segment_no_fsync(&bytes, epoch_ref)?; - written.push(WrittenSegment { path, meta: epoch_ref.meta }); - } - fdatasync_all(&written)?; // one batched fsync pass - manifest.rewrite_and_fsync(&written)?; - fsync_parent_dir(&cfg.disk_path)?; // one dir fsync for the whole batch - - // Now that segments are durable AND referenced by the manifest, - // evict them from memory. - for seg in &written { - splice_out_of_sealed_epochs(seg.meta); - mem_bytes_in_use.fetch_sub(seg.meta.approx_bytes); + // ---- Build and persist one part for the whole tick ---- + // + // All candidate epochs from this tick land in a single part directory + // under `parts//`. File count per tick is O(1) (three + // files) instead of O(num_candidates). This is where group-commit + // stops being something the flusher explicitly arranges and starts + // being an intrinsic property of the layout. + if !candidates.is_empty() { + let part_id = next_part_id(); + let part_dir = cfg.disk_path.join("parts").join(fmt_part_id(part_id)); + + // (a) Clone each candidate's Arc under a brief read lock. + // No serialization under any lock. + let snapshots: Vec = candidates + .iter() + .map(|c| snapshot_under_read_lock(c)) + .collect(); + + // (b) Serialize to the three part files. data.bin is fallocate'd + // to `sum(approx_memory_bytes) * 1.3`; index.bin is sized + // exactly (32 B per entry). Both are 8-byte aligned and CRCs + // are computed streaming. + let (data_len, index_len) = write_part_files(&part_dir, &snapshots)?; + + // (c) Batched fdatasync of the three files + the part directory. + // One syscall per file; no per-epoch fsync. + fdatasync_file(&part_dir.join("data.bin"))?; + fdatasync_file(&part_dir.join("index.bin"))?; + fdatasync_file(&part_dir.join("meta.bin"))?; + fsync_dir(&part_dir)?; + + // (d) Append the add-part record to the manifest log and fsync it. + // Single fixed-size append — no rewrite of existing state. + manifest.append_add_part(AddPartRecord { + part_id, + min_ts: snapshots.iter().map(|s| s.min_ts).min().unwrap(), + max_ts: snapshots.iter().map(|s| s.max_ts).max().unwrap(), + size_bytes: (data_len + index_len) as u64, + })?; + fsync_dir(&cfg.disk_path)?; // log's size update is now durable + + // (e) Now that the part is officially live, evict the source + // epochs from memory. This is the only place we take the + // per-agg write lock, and we take it O(1) times per epoch. + for snapshot in &snapshots { + splice_out_of_sealed_epochs(snapshot.agg_id, snapshot.epoch_id); + mem_bytes_in_use.fetch_sub(snapshot.approx_bytes); + } + + // Maybe compact the manifest log into a fresh snapshot if the + // log has grown large relative to the current snapshot. + manifest.maybe_compact()?; } - // Phase 3 (disk retention sweep): delete segments older than T2. + // Phase 3 (disk retention sweep): delete whole parts older than T2. + // + // Because parts cover a tight time range (~flush_interval_ms), T2 + // deletion operates at part-directory granularity — we rm -rf the + // whole thing rather than touching shared files. if let Some(ttl) = cfg.delete_older_than_ms { let cutoff = now.saturating_sub(ttl); - let expired = manifest.segments_older_than(cutoff); - for seg in expired { - manifest.remove(seg.id); - cache_tier2.invalidate(seg.id); // drop any decoded copy - fs::remove_file(seg.path).ok(); // best-effort; orphan sweep on restart + let expired: Vec = manifest + .live_parts() + .filter(|p| p.max_ts < cutoff) + .map(|p| p.part_id) + .collect(); + + for part_id in expired { + // Invalidate Tier-2 cache entries that reference this part, + // append a delete-part record to the log, then remove the dir. + cache_tier2.invalidate_part(part_id); + manifest.append_delete_part(part_id)?; + let part_dir = cfg.disk_path + .join("parts") + .join(fmt_part_id(part_id)); + fs::remove_dir_all(part_dir).ok(); // orphan sweep on restart catches failures } if !expired.is_empty() { - manifest.rewrite_and_fsync(&[])?; + fdatasync_file(&manifest.log_path())?; + fsync_dir(&cfg.disk_path)?; } } } @@ -327,51 +449,44 @@ loop { Under memory pressure, phase 1 dominates and phase 2 usually finds nothing left to do (the oldest epochs are already gone). Under light ingest, phase 1 is a no-op and phase 2 does all the work. Phase 3 is independent and runs -every tick regardless; it costs one manifest scan plus one `unlink` per -expired segment. - -#### Group-commit fsync - -The pseudocode above batches all `fsync`/`fdatasync` calls for a tick into a -single pass at the end, rather than `fsync`ing each segment inline. On -spinning disks this is ~10× fewer head seeks per tick; on SSDs it is ~3–4× -fewer syscalls. The cost is one temporarily-larger `written` vector and one -extra `fsync_parent_dir` at the end — trivial relative to the saved I/O. - -Durability invariant remains the same: **no segment is referenced by the -manifest until its bytes and the manifest itself are both `fsync`'d.** The -group-commit ordering is (1) write all segment bodies, (2) `fdatasync` all -of them, (3) rewrite + `fsync` manifest via the atomic `write → rename` -dance, (4) `fsync` parent directory. A crash at any point leaves orphan -segment files (cleaned by the startup sweep) but never a dangling manifest -entry. - -`flush_and_evict` takes advantage of a property that matters a lot for the -flusher design: **sealed epochs are append-only and frozen.** Once the -rotator seals an epoch, no writer will ever touch its contents again — it -is only read (by queries) or removed wholesale (by the flusher). That -immutability is what lets the flusher stay completely off the critical -path: - -1. Take the per-agg `RwLock::read` briefly, clone the `Arc` for the - target epoch out of `sealed_epochs`, drop the lock. -2. Serialize the epoch, write the segment file, and fsync — **entirely - outside any store lock**, on the flusher's own thread. Nothing in the - system is waiting on this I/O. Inserts continue to land in - `current_epoch`; queries continue to read from the still-in-place - `sealed_epochs` entry (and the cloned `Arc` keeps the bytes alive for - any query that happens to hold a reference already); the rotator - continues to seal new epochs behind us. -3. Once the segment is durable and the manifest is updated, take the - per-agg `RwLock::write` briefly to splice the epoch out of - `sealed_epochs` and decrement `mem_bytes_in_use`. This is O(1) — a - `BTreeMap::remove` plus an atomic subtraction — and is the only lock - the flusher holds for more than a read snapshot. - -Because steps 2 and 3 are decoupled by the `Arc` clone, **no lock -is ever held across disk I/O**, and the flusher never blocks anything on -the insert or query path beyond the two brief lock acquisitions at the -start and end. +every tick regardless; it costs one manifest scan plus one `rm -rf` per +expired part directory (usually zero). + +Group-commit is now **intrinsic to the layout**, not something the flusher +has to explicitly arrange: one tick = one part = three `fdatasync`s + one +dir fsync + one log append, independent of how many epochs the tick is +flushing. The durability ordering is spelled out in the previous section +("Durability ordering per flush tick"). + +The part-building loop above takes advantage of a property that matters a +lot for the flusher design: **sealed epochs are append-only and frozen.** +Once the rotator seals an epoch, no writer will ever touch its contents +again — it is only read (by queries) or removed wholesale (by the +flusher). That immutability is what lets the flusher stay completely off +the critical path: + +1. Take the per-agg `RwLock::read` briefly, clone the `Arc` for + each candidate out of `sealed_epochs`, drop the lock. This is the + `snapshot_under_read_lock` step. +2. Serialize all candidates into `data.bin` / `index.bin` / `meta.bin`, + fsync the three files and the part directory, and append to the + manifest log — **entirely outside any per-agg store lock**, on the + flusher's own thread. Nothing in the system is waiting on this I/O. + Inserts continue to land in `current_epoch`; queries continue to + read from the still-in-place `sealed_epochs` entries (and the cloned + `Arc`s keep bytes alive for any query that happens to hold a + reference already); the rotator continues to seal new epochs behind + us. +3. Once the part is officially live in the manifest log, take each + per-agg `RwLock::write` briefly to splice the corresponding epoch + out of `sealed_epochs` and decrement `mem_bytes_in_use`. This is + O(1) per epoch — a `BTreeMap::remove` plus an atomic subtraction — + and is the only write lock the flusher holds per epoch. + +Because steps 2 and 3 are decoupled by the `Arc` clone from step +1, **no per-agg lock is ever held across disk I/O**, and the flusher +never blocks anything on the insert or query path beyond the two brief +lock acquisitions at the start and end. #### Sync vs. async flush I/O — resolved @@ -412,17 +527,41 @@ flag checked on each loop iteration. `query_precomputed_output` becomes a three-way merge: -1. Read from `current_epoch + sealed_epochs` as today. -2. Look up segments in the manifest whose `[start_ms, end_ms]` overlaps the - query range for this `agg_id`. -3. For each matching segment, read the file (via the read-side segment cache - described below), decode entries whose window overlaps, and merge into the - result. - -Segment reads happen under a read lock on the manifest; they do **not** take any -per-agg store lock, so they can run fully in parallel with inserts. Merging reuses -the existing `TimestampedBucketsMap` + `AggregateCore::merge_with` that the -in-memory query path already uses — no new merge logic. +1. Read from `current_epoch + sealed_epochs` as today (in-memory hits). +2. Walk the parts_manifest's live-parts list for any `part.[min_ts, + max_ts]` that overlaps the query's time range. The manifest lives in + memory as a `Vec` built from the snapshot + log replay at + startup, so this is a linear scan over a short list (tens of + thousands of entries in the worst case, all 32-byte records) — fast + and trivially parallel with inserts. +3. For each overlapping part: fetch its `DecodedPart` from the Tier-2 + segment cache (`moka::Cache>`). On miss, + `mmap` the part's `data.bin` and `index.bin`, wrap them in an + `Arc` (holding the mmap handles), and insert into the + cache. Then binary-search `index.bin` for `(agg_id, start_ms)`, walk + the matching entries forward while `start_ms <= query_end`, and for + each one hand the `&[u8]` slice of `data.bin` directly to + `AggregateCore::deserialize_from_bytes` with zero copies. + +Part reads happen under a read lock on the parts_manifest vector; they +do **not** take any per-agg store lock, so they run fully in parallel +with inserts. Merging reuses the existing `TimestampedBucketsMap` + +`AggregateCore::merge_with` that the in-memory query path already uses +— no new merge logic. + +**Why this is fast:** + +- **One file open per part hit, not per entry.** Queries that span many + aggs inside a part still only pay one `mmap`'s worth of setup cost. +- **Zero-copy deserialize.** The 8-byte alignment guarantee means + `&data.bin[offset..offset+len]` can be fed straight to the + sketch-specific decoder without a staging buffer. +- **Binary search, not linear scan.** `index.bin` is sorted by + `(agg_id, start_ms)` and is a flat mmap'd array; `partition_point` is + a few cache lines of work. +- **Page cache locality.** Adjacent entries for the same agg inside a + part are physically adjacent on disk, so a query over a time range + touches contiguous pages. ### Read-side segment cache (two-tier memory model) @@ -448,18 +587,18 @@ read-side cache on top of the disk layer. sealed epoch that has not yet been flushed. Source of truth for recent data. Managed by the flusher described above. -**Tier 2 — read-side segment cache (query-driven).** Bounded by a -separate `segment_cache_bytes` budget. Contains decoded copies of segments -pulled back from disk by the query path. Source of truth is always the -segment file — the cache is a pure optimization, drop-anytime, never -dirty. Managed by the query path, not the flusher. +**Tier 2 — read-side part cache (query-driven).** Bounded by a separate +`part_cache_bytes` budget. Contains mmap handles and decoded index views +of parts pulled back from disk by the query path. Source of truth is +always the part directory on disk — the cache is a pure optimization, +drop-anytime, never dirty. Managed by the query path, not the flusher. ```rust pub struct SimpleMapStorePersistenceConfig { // ... existing fields ... - // Read-side segment cache. Bounded independently of - // `memory_limit_bytes`; this budget is for decoded segments the query + // Read-side part cache. Bounded independently of + // `memory_limit_bytes`; this budget is for decoded parts the query // path pulls back from disk, not for the authoritative hot set. // // Default: min(10% * memory_limit_bytes, 512 MiB). @@ -469,7 +608,7 @@ pub struct SimpleMapStorePersistenceConfig { // entirely (every cold query pays disk I/O); a fixed absolute // default would be too small on big boxes and too large on small // ones, so the default scales with the write budget. - pub segment_cache_bytes: usize, + pub part_cache_bytes: usize, } ``` @@ -534,13 +673,25 @@ whole point of splitting the tiers. ### Recovery on startup -1. Open `disk_path`, read `manifest.json`. -2. For every referenced segment, stat the file and validate magic + CRC header. - Missing / corrupt segments are logged and removed from the manifest. -3. Sweep `disk_path` for segment files not referenced in the manifest (orphans - from a mid-flush crash) and delete them. -4. Build the in-memory segment index; do **not** load any sketches into RAM. - Cold data stays cold until a query asks for it. +1. Open `disk_path`. If `parts_manifest.snapshot` exists, mmap it and + cast the bytes to `&[PartEntry]` (no parse — the layout is 8-byte + aligned and versioned in a small header). Otherwise, start with an + empty live set. +2. Replay `parts_manifest.log` from the offset recorded in the + snapshot's footer, applying add-part and delete-part records to the + live set. +3. For every live part, stat its directory, verify `meta.bin`'s magic, + version, and CRC. Parts whose directory is missing or whose + `meta.bin` fails verification are logged and removed from the live + set (and a delete-part record is appended to the log to make the + removal durable). +4. Sweep `parts/` for directories not referenced in the live set + (orphans from a mid-flush crash before step 5 of the durability + ordering) and `rm -rf` them. +5. Build the in-memory `Vec` that the query path reads. Do + **not** mmap `data.bin` or `index.bin` eagerly — parts are mapped + lazily on first query hit and cached in Tier 2. Cold data stays + cold until a query asks for it. ### Concurrency summary @@ -548,12 +699,14 @@ whole point of splitting the tiers. |-------------------|---------------------------------------------| | Insert | per-agg `RwLock::write` (unchanged) | | Query in-memory | per-agg `RwLock::read` (unchanged) | -| Query disk | manifest `RwLock::read` + segment-cache mutex | -| Flush: serialize | none (reads immutable sealed `Arc`) | -| Flush: evict | per-agg `RwLock::write` (short, O(1) splice)| -| Flush: commit | manifest `RwLock::write` (short) | +| Query disk | parts_manifest `RwLock::read` + moka cache internal | +| Flush: snapshot | per-agg `RwLock::read` (brief, per candidate) | +| Flush: serialize | none (operates on cloned `Arc`s) | +| Flush: commit log | parts_manifest `RwLock::write` (brief, append) | +| Flush: evict | per-agg `RwLock::write` (short, O(1) splice per epoch) | +| T2 sweep | parts_manifest `RwLock::write` (brief, delete records) | -No new lock held across a fsync or disk I/O. +No lock of any kind is held across a `fsync` or disk I/O. --- @@ -596,22 +749,27 @@ the pieces can be reviewed independently: disk I/O yet — expose the memory counter in `StoreDiagnostics` so we can validate sizing in isolation and be confident nothing else regressed. -2. **Segment format, manifest, flusher.** Add `persistence/` submodule under - `simple_map_store/` with segment encode/decode (8-byte aligned, - `fallocate`d, mmap-friendly), manifest read/write, background flusher - thread (memory-first, then time-watermark, then T2 retention sweep) with - group-commit fsync. Wire `flush_and_evict` into the per-key store. Unit - tests for round-trip, crash-after-segment-before-manifest, orphan sweep, - and T2 deletion. - -3. **Query path read-through + recovery + Tier-2 cache.** Extend - `query_precomputed_output` to consult the manifest and merge segment hits - with in-memory hits. Wire a `moka` (or `mini-moka`) weight-bounded - segment cache sized to `min(10% * memory_limit_bytes, 512 MiB)` by - default, with hit/miss counters exported via `StoreDiagnostics`. Add - startup recovery. Integration test: ingest → flush → restart → query → - same result as no-restart, plus a scan-resistance test that confirms a - long-range query does not evict a separately-hot segment. +2. **Parts layout, manifest log, flusher.** Add `persistence/` submodule + under `simple_map_store/` with part encode/decode (`meta.bin` + + `data.bin` + `index.bin`, 8-byte aligned, `fallocate`'d, mmap-friendly), + parts_manifest log + snapshot read/write, background flusher thread + (memory-first, then time-watermark, then T2 retention sweep, one part + per tick, group-commit implicit in the layout). Wire part-building and + eviction into the per-key store. Unit tests for round-trip, + crash-before-log-append, crash-between-log-append-and-dir-fsync, + orphan-part sweep, and T2 whole-part deletion. + +3. **Query path read-through + recovery + Tier-2 part cache.** Extend + `query_precomputed_output` to walk the parts_manifest, binary-search + `index.bin` of overlapping parts, and zero-copy deserialize payloads + out of mmap'd `data.bin`. Wire a `moka` (or `mini-moka`) + weight-bounded part cache sized to + `min(10% * memory_limit_bytes, 512 MiB)` by default, keyed on + `PartId`, with hit/miss counters exported via `StoreDiagnostics`. + Add startup recovery (snapshot mmap + log replay + CRC verify + + orphan sweep). Integration test: ingest → flush → restart → query → + same result as no-restart, plus a scan-resistance test that confirms + a long-range query does not evict a separately-hot part. Because phase 1 removes `CleanupPolicy` from this store, phase 1 is **not** independently mergeable without at least the memory-pressure path from @@ -634,8 +792,20 @@ each lives in the section it points to. | 4 | Per-agg-id flush fairness | **Round-robin across agg-ids, oldest-first within each agg** | Strict-global-oldest hammers one `RwLock` during a hot-agg burst and creates query-latency spikes on that one agg. Round-robin spreads lock acquisitions across different `RwLock`s for the same total work. See **Eviction order**. | | 5 | Default `segment_cache_bytes` | `min(10% * memory_limit_bytes, 512 MiB)` | A fixed 64 MiB default is too small on big boxes and too large on small ones; scaling with the write budget keeps the tier sensibly sized without requiring operator tuning on a fresh install. See **Configuration**. | | 6 | Tier-2 algorithm | **W-TinyLFU via `moka`** from day one, not plain LRU | LRU is scan-vulnerable — one long-range query evicts everything genuinely hot, which is exactly the TSDB dashboard access pattern. W-TinyLFU's admission filter rejects scan traffic and typically delivers 10–30% better hit rate at the same byte budget on skewed workloads, with a drop-in API and lower per-access CPU than LRU. See **Read-side segment cache**. | +| 7 | Disk layout | **Parts** (one directory per flush tick containing `meta.bin` + `data.bin` + `index.bin`) with an **append-only `parts_manifest.log` + periodic binary snapshot**, not one file per sealed epoch with a JSON manifest | One-file-per-epoch produces hundreds of thousands of tiny files on a real deployment (inode pressure, readdir slowdown, per-file fsync floor). A JSON manifest rewritten each tick is quadratic over the deployment's lifetime. Parts bound file count to O(flush ticks), bundle all of a tick's epochs behind one fallocate + three fdatasyncs, push per-part indexing into a binary-searchable `index.bin`, and reduce the global manifest to an append-only log whose size is proportional to ticks, not epochs. This is the layout every mainstream TSDB converges on. See **Unit of flush vs. unit of file: the *part*** and **Disk layout**. | If any of these decisions turn out to be wrong under real traces, the affected sections are the natural point of revisiting — but none of them are "temporary v1 shortcuts we'll upgrade later." This is the target design. + +**Not resolved here, deliberately punted to v2:** background compaction +of adjacent small parts into larger ones. The parts layout accommodates +compaction cleanly (it's a pure directory-level merge with the same +on-disk shape as a regular flush tick), but v1 ships without it. With +T2 retention in place and a reasonable flush interval (seconds, not +milliseconds), the number of live parts in v1 stays bounded at +`T2 / flush_interval_ms` — a few tens of thousands at most, well within +what a binary-searched `Vec` handles with room to spare. +Compaction becomes necessary only if we lower the flush interval +significantly or extend T2 to very long horizons. From 2f76df0a3054e92468ec1582894a773a8f05e3a1 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 20:52:42 -0400 Subject: [PATCH 07/12] =?UTF-8?q?feat(persistence):=20parts=20submodule=20?= =?UTF-8?q?=E2=80=94=20format,=20manifest,=20flusher,=20cache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Self-contained persistence submodule under simple_map_store/persistence/: * config.rs — SimpleMapStorePersistenceConfig with memory budget (high/low watermarks + hard cap), hot_window_ms, delete_older_than_ms, flush_interval, disk_path, part_cache_bytes. * part.rs — on-disk part format (meta.bin + data.bin + index.bin). PartWriter bundles a slice of EpochSnapshots into a directory with streaming crc32, 8-byte aligned payloads, dir fsync. PartReader mmaps all three files, verifies index CRC on open, and resolves entries lazily via load_entry(). * manifest.rs — append-only parts_manifest.log + periodic binary parts_manifest.snapshot. append_add / append_delete are fsync'd; compact() atomically rewrites the snapshot and truncates the log once it grows past 4x the live-set size. * source.rs — EpochSource trait + EpochSnapshot types. The contract the flusher uses to enumerate, snapshot, and evict sealed epochs. Decouples flusher from SimpleMapStorePerKey so the flusher can be unit-tested against a fake source. * flusher.rs — std::thread-based background loop. Three phases per tick (memory pressure first, hot_window second, T2 retention sweep third), round-robin across agg-ids within a phase. One part per tick — group-commit fsync is intrinsic to the layout. Handle is Drop'd to shut down cleanly. * cache.rs — moka::sync::Cache> with weight based on mmap bytes. 0-budget disables the cache; default sizing scales with memory_limit_bytes. * recovery.rs — load manifest at startup, verify every live part's meta.bin, append delete records for corrupt entries, sweep orphan part dirs not referenced by the manifest. 12 unit tests cover: part round-trip, CRC detection of meta corruption, empty-snapshot rejection, manifest append + reload, compact + log truncation, overlap filtering, memory-pressure flush, hot-window flush, T2 sweep, empty-dir recovery, orphan sweep, corrupt part eviction. Adds three deps: moka 0.12 (sync, weight-based W-TinyLFU), memmap2 0.9 (zero-copy part reads), crc32fast 1.4 (streaming CRCs). Not yet wired into SimpleMapStorePerKey — follow-up commit. Co-Authored-By: Claude Opus 4.6 (1M context) --- Cargo.lock | 35 + asap-query-engine/Cargo.toml | 4 + .../src/stores/simple_map_store/mod.rs | 1 + .../simple_map_store/persistence/cache.rs | 82 +++ .../simple_map_store/persistence/config.rs | 70 ++ .../simple_map_store/persistence/flusher.rs | 530 +++++++++++++ .../simple_map_store/persistence/manifest.rs | 439 +++++++++++ .../simple_map_store/persistence/mod.rs | 58 ++ .../simple_map_store/persistence/part.rs | 697 ++++++++++++++++++ .../simple_map_store/persistence/recovery.rs | 203 +++++ .../simple_map_store/persistence/source.rs | 99 +++ 11 files changed, 2218 insertions(+) create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/cache.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/config.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/mod.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/part.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs create mode 100644 asap-query-engine/src/stores/simple_map_store/persistence/source.rs diff --git a/Cargo.lock b/Cargo.lock index 37d4f72d..e85f977a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2863,6 +2863,15 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memmap2" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +dependencies = [ + "libc", +] + [[package]] name = "mime" version = "0.3.17" @@ -2890,6 +2899,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957228ad12042ee839f93c8f257b62b4c0ab5eaae1d4fa60de53b27c9d7c5046" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + [[package]] name = "multimap" version = "0.10.1" @@ -3660,6 +3686,7 @@ dependencies = [ "bincode", "chrono", "clap 4.6.0", + "crc32fast", "criterion", "ctor", "dashmap 5.5.3", @@ -3672,6 +3699,8 @@ dependencies = [ "futures", "hex", "lazy_static", + "memmap2", + "moka", "prometheus", "promql-parser", "promql_utilities", @@ -4503,6 +4532,12 @@ dependencies = [ "libc", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tempfile" version = "3.27.0" diff --git a/asap-query-engine/Cargo.toml b/asap-query-engine/Cargo.toml index e81c8539..a7884dff 100644 --- a/asap-query-engine/Cargo.toml +++ b/asap-query-engine/Cargo.toml @@ -63,6 +63,10 @@ reqwest = { version = "0.11", features = ["json"] } tracing-appender = "0.2" elastic_dsl_utilities.workspace = true asap_sketchlib = { git = "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/ProjectASAP/asap_sketchlib" } +# Persistence layer (SimpleMapStore parts / manifest / Tier-2 cache) +moka = { version = "0.12", features = ["sync"] } +memmap2 = "0.9" +crc32fast = "1.4" [[bin]] name = "precompute_engine" diff --git a/asap-query-engine/src/stores/simple_map_store/mod.rs b/asap-query-engine/src/stores/simple_map_store/mod.rs index 29c78d60..b3243feb 100644 --- a/asap-query-engine/src/stores/simple_map_store/mod.rs +++ b/asap-query-engine/src/stores/simple_map_store/mod.rs @@ -2,6 +2,7 @@ mod common; pub mod global; pub mod legacy; pub mod per_key; +pub mod persistence; use crate::data_model::{ AggregateCore, CleanupPolicy, LockStrategy, PrecomputedOutput, StreamingConfig, diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs b/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs new file mode 100644 index 00000000..c00ecbb8 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/cache.rs @@ -0,0 +1,82 @@ +//! Tier-2 read-side part cache. Moka-backed, bounded by bytes, +//! W-TinyLFU eviction. Keyed on `PartId`; values are `Arc` +//! which holds the mmap handles for a part's three files. +//! +//! The cache has no consistency obligations (it is never dirty), so +//! the flusher can invalidate entries whenever it deletes a part. + +use std::path::Path; +use std::sync::Arc; + +use moka::sync::Cache; + +use super::part::{part_dir_path, PartId, PartReader}; +use super::PersistResult; + +/// A cached, mmap-backed `PartReader`. Returned by [`PartCache::get`]. +pub type LoadedPart = Arc; + +/// Bounded cache of decoded parts. Cheap to clone — it's an `Arc` +/// inside. +#[derive(Clone)] +pub struct PartCache { + inner: Option>, + parts_root: std::path::PathBuf, +} + +impl PartCache { + /// Construct a cache with a byte budget. A budget of 0 disables + /// caching entirely — every query path goes straight to disk. + pub fn new(parts_root: std::path::PathBuf, byte_budget: u64) -> Self { + let inner = if byte_budget == 0 { + None + } else { + Some( + Cache::builder() + .weigher(|_k: &PartId, v: &LoadedPart| -> u32 { + // Weight = sum of mmap'd bytes. Moka's weigher + // returns u32, so clamp large parts. + let len = v.meta.data_len as u64 + v.meta.index_len as u64; + len.min(u32::MAX as u64) as u32 + }) + .max_capacity(byte_budget) + .build(), + ) + }; + Self { inner, parts_root } + } + + /// Get (or load) a `PartReader` for `part_id`. On miss, mmaps the + /// part's files from disk. + pub fn get_or_load(&self, part_id: PartId) -> PersistResult { + if let Some(inner) = &self.inner { + if let Some(hit) = inner.get(&part_id) { + return Ok(hit); + } + let part_dir = part_dir_path(&self.parts_root, part_id); + let reader = Arc::new(PartReader::open(&part_dir)?); + inner.insert(part_id, Arc::clone(&reader)); + Ok(reader) + } else { + let part_dir = part_dir_path(&self.parts_root, part_id); + Ok(Arc::new(PartReader::open(&part_dir)?)) + } + } + + /// Drop the cached entry for `part_id`. Idempotent. + pub fn invalidate(&self, part_id: PartId) { + if let Some(inner) = &self.inner { + inner.invalidate(&part_id); + } + } + + /// Diagnostic: approximate entry count (0 if caching disabled). + pub fn entry_count(&self) -> u64 { + self.inner.as_ref().map(|c| c.entry_count()).unwrap_or(0) + } +} + +/// Convenience helper: given a disk_path, return the parts root. +pub fn parts_root_of(disk_path: &Path) -> std::path::PathBuf { + disk_path.join("parts") +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/config.rs b/asap-query-engine/src/stores/simple_map_store/persistence/config.rs new file mode 100644 index 00000000..fac29b2f --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/config.rs @@ -0,0 +1,70 @@ +use std::path::PathBuf; +use std::time::Duration; + +/// Configuration for the SimpleMapStore persistence layer. +/// +/// The two bounding knobs, in priority order: +/// +/// * **Memory budget (primary).** When in-memory sealed-epoch bytes exceed +/// `memory_limit_bytes`, the background flusher evicts oldest-sealed-epoch +/// first (globally, by `end_ts`) until usage drops below +/// `memory_low_watermark_bytes`. If usage reaches `hard_cap_bytes` inside +/// an insert, the insert path blocks until the flusher catches up. +/// +/// * **Time watermark (secondary).** Any sealed epoch whose `end_ts` is +/// older than `now - hot_window_ms` is flushed on the next tick even if +/// the store is nowhere near the memory budget. Guarantees a predictable +/// hot-set size under light ingest. +/// +/// Plus a disk-retention knob: +/// +/// * **Cold TTL.** Any part whose `max_ts` is older than +/// `now - delete_older_than_ms` is deleted from disk on the next tick. +/// Bounds manifest size and long-running directory growth. +#[derive(Debug, Clone)] +pub struct SimpleMapStorePersistenceConfig { + // ---- Primary: memory budget ---- + pub memory_limit_bytes: usize, + pub memory_low_watermark_bytes: usize, + pub hard_cap_bytes: usize, + + // ---- Secondary: time watermark T ---- + /// Hot-window length in milliseconds. `None` disables time-based flushing. + pub hot_window_ms: Option, + + // ---- Disk retention ---- + /// Cold-tier TTL in milliseconds. `None` disables cold deletion. + /// Should be much larger than `hot_window_ms` (hours or days). + pub delete_older_than_ms: Option, + + // ---- Misc ---- + /// Cadence of the background flusher loop. + pub flush_interval: Duration, + + /// Root directory for segment files and the parts manifest. + pub disk_path: PathBuf, + + /// Tier-2 part-cache byte budget. 0 disables the cache (every cold + /// query pays disk I/O). Default is `min(10% * memory_limit_bytes, + /// 512 MiB)` — scale it with the write budget, not a fixed number. + pub part_cache_bytes: u64, +} + +impl SimpleMapStorePersistenceConfig { + /// Build a config with sensible defaults relative to a memory budget. + pub fn with_memory_limit(memory_limit_bytes: usize, disk_path: PathBuf) -> Self { + let low_water = memory_limit_bytes * 8 / 10; // 80% of high water + let hard_cap = memory_limit_bytes * 125 / 100; // 125% of high water + let cache = (memory_limit_bytes / 10).min(512 * 1024 * 1024) as u64; + Self { + memory_limit_bytes, + memory_low_watermark_bytes: low_water, + hard_cap_bytes: hard_cap, + hot_window_ms: Some(60 * 60 * 1000), // 1 hour + delete_older_than_ms: Some(7 * 24 * 60 * 60 * 1000), // 7 days + flush_interval: Duration::from_secs(1), + disk_path, + part_cache_bytes: cache, + } + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs new file mode 100644 index 00000000..d87f7870 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs @@ -0,0 +1,530 @@ +//! Background flusher thread. Pulls sealed epochs out of an +//! [`EpochSource`], bundles them into on-disk parts, and evicts them +//! from memory. +//! +//! One part per flush tick. Three phases: +//! +//! 1. **Memory pressure** — if `source.approx_memory_bytes()` exceeds +//! `memory_limit_bytes`, collect oldest-first (round-robin by agg) +//! until projected memory drops under `memory_low_watermark_bytes`. +//! 2. **Time watermark** — any sealed epoch whose `end_ts` is older +//! than `now - hot_window_ms` is added to the flush set even if the +//! store is well under the memory budget. +//! 3. **T2 retention sweep** — parts whose `max_ts` is older than +//! `now - delete_older_than_ms` are removed from the manifest and +//! their directories are `rm -rf`'d. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread::{self, JoinHandle}; +use std::time::{Instant, SystemTime, UNIX_EPOCH}; + +use tracing::{debug, error, info, warn}; + +use super::config::SimpleMapStorePersistenceConfig; +use super::manifest::{Manifest, PartEntry}; +use super::part::{part_dir_path, PartWriter}; +use super::source::{EpochSnapshot, EpochSource, SealedEpochRef}; +use super::{PersistError, PersistResult}; + +/// Handle to a running flusher thread. Dropping the handle signals +/// shutdown and joins the thread. +pub struct FlusherHandle { + inner: Arc, + thread: Option>, +} + +pub(crate) struct FlusherShared { + pub cfg: SimpleMapStorePersistenceConfig, + pub manifest: Arc, + pub next_part_id: AtomicU64, + pub shutdown: AtomicBool, + /// Woken by the insert path when it hits `hard_cap_bytes` and by + /// the flusher when it finishes a tick. + pub pressure_cv: Condvar, + pub pressure_mutex: Mutex<()>, +} + +impl FlusherHandle { + /// Start a flusher thread. Takes an `EpochSource` (typically the + /// store itself, wrapped in `Arc`). + pub fn start( + cfg: SimpleMapStorePersistenceConfig, + manifest: Arc, + source: Arc, + ) -> PersistResult + where + S: EpochSource + 'static, + { + // Pick a starting part_id: one past the max currently in the + // manifest (so IDs are monotonically increasing across restarts). + let next_id = manifest + .live_parts() + .iter() + .map(|p| p.part_id) + .max() + .map(|m| m + 1) + .unwrap_or(1); + + let shared = Arc::new(FlusherShared { + cfg: cfg.clone(), + manifest: manifest.clone(), + next_part_id: AtomicU64::new(next_id), + shutdown: AtomicBool::new(false), + pressure_cv: Condvar::new(), + pressure_mutex: Mutex::new(()), + }); + + let shared_for_thread = Arc::clone(&shared); + let thread = thread::Builder::new() + .name("simple-map-store-flusher".into()) + .spawn(move || run_flusher_loop(shared_for_thread, source)) + .map_err(PersistError::Io)?; + + Ok(Self { + inner: shared, + thread: Some(thread), + }) + } + + /// Signal shutdown and wait for the thread to finish its current + /// tick. Safe to call multiple times; subsequent calls are no-ops. + pub fn shutdown(&mut self) { + self.inner.shutdown.store(true, Ordering::Release); + self.inner.pressure_cv.notify_all(); + if let Some(handle) = self.thread.take() { + if let Err(e) = handle.join() { + error!("flusher thread panicked on join: {:?}", e); + } + } + } + + /// Wake the flusher early (e.g., on `hard_cap_bytes` pressure). + pub fn wake(&self) { + self.inner.pressure_cv.notify_all(); + } + + /// Access to the manifest for the query read-through path. + pub fn manifest(&self) -> Arc { + Arc::clone(&self.inner.manifest) + } + + pub fn disk_path(&self) -> &std::path::Path { + &self.inner.cfg.disk_path + } +} + +impl Drop for FlusherHandle { + fn drop(&mut self) { + self.shutdown(); + } +} + +fn run_flusher_loop(shared: Arc, source: Arc) { + info!( + "flusher thread started: disk_path={:?}, memory_limit={} bytes, hot_window={:?} ms, flush_interval={:?}", + shared.cfg.disk_path, + shared.cfg.memory_limit_bytes, + shared.cfg.hot_window_ms, + shared.cfg.flush_interval, + ); + + loop { + if shared.shutdown.load(Ordering::Acquire) { + break; + } + + // Cond-var wait: sleep for up to `flush_interval`, woken early + // on pressure. + { + let guard = shared.pressure_mutex.lock().unwrap(); + let _ = shared + .pressure_cv + .wait_timeout(guard, shared.cfg.flush_interval) + .unwrap(); + } + + if shared.shutdown.load(Ordering::Acquire) { + break; + } + + let tick_start = Instant::now(); + match run_tick(&shared, source.as_ref()) { + Ok(stats) => { + if stats.entries_flushed > 0 || stats.parts_deleted > 0 { + debug!( + "flusher tick done in {:?}: entries={}, parts_written={}, parts_deleted={}", + tick_start.elapsed(), + stats.entries_flushed, + stats.parts_written, + stats.parts_deleted, + ); + } + } + Err(e) => { + error!("flusher tick error: {}", e); + } + } + + // Always notify the pressure condvar so blocked inserts wake + // up when memory has been freed. + shared.pressure_cv.notify_all(); + } + + info!("flusher thread exiting cleanly"); +} + +#[derive(Debug, Default)] +struct TickStats { + entries_flushed: usize, + parts_written: usize, + parts_deleted: usize, +} + +fn run_tick(shared: &Arc, source: &S) -> PersistResult { + let mut stats = TickStats::default(); + let now = now_ms(); + + // ---- Collect candidates across phase 1 and phase 2 ---- + let mut all = source.list_sealed_epochs(); + // Oldest-first by end_ts. + all.sort_by_key(|r| r.end_ts); + + let mem = source.approx_memory_bytes(); + let cfg = &shared.cfg; + + let need_memory_pressure = mem > cfg.memory_limit_bytes; + let low_water_goal = cfg.memory_low_watermark_bytes; + let mut projected_after_evict: i64 = mem as i64; + + let mut selected: Vec = Vec::new(); + // Phase 1: memory pressure. + if need_memory_pressure { + for r in &all { + if projected_after_evict <= low_water_goal as i64 { + break; + } + selected.push(*r); + projected_after_evict -= r.approx_bytes as i64; + } + } + + // Phase 2: time watermark (any epoch older than now - hot_window). + if let Some(hot) = cfg.hot_window_ms { + let cutoff = now.saturating_sub(hot); + for r in &all { + if r.end_ts < cutoff && !selected.iter().any(|s| s.agg_id == r.agg_id && s.epoch_id == r.epoch_id) { + selected.push(*r); + } + } + } + + if !selected.is_empty() { + // Round-robin across agg-ids inside the tick for lock spread. + let selected = round_robin_by_agg(selected); + + // Take snapshots one by one, skipping any that have already + // been evicted racily. + let mut snapshots: Vec = Vec::new(); + for r in &selected { + match source.snapshot_sealed_epoch(r.agg_id, r.epoch_id) { + Ok(Some(s)) => snapshots.push(s), + Ok(None) => { + // Already evicted by a concurrent flusher tick (shouldn't + // happen with a single flusher) or by the store itself. + debug!( + "flusher: snapshot of ({}, {}) returned None, skipping", + r.agg_id, r.epoch_id + ); + } + Err(e) => { + warn!( + "flusher: snapshot of ({}, {}) failed: {}; skipping", + r.agg_id, r.epoch_id, e + ); + } + } + } + + if !snapshots.is_empty() { + // Build one part for the tick. + let part_id = shared.next_part_id.fetch_add(1, Ordering::Relaxed); + let part_dir = part_dir_path(&parts_root(&cfg.disk_path), part_id); + let entries_total: usize = snapshots.iter().map(|s| s.len()).sum(); + let size_bytes_estimate: u64 = + snapshots.iter().map(|s| s.approx_bytes as u64).sum(); + + let report = PartWriter::write_part(&part_dir, part_id, &snapshots)?; + shared.manifest.append_add(PartEntry { + part_id, + min_ts: report.min_ts, + max_ts: report.max_ts, + size_bytes: report.data_len + report.index_len + size_bytes_estimate, + })?; + + // Now that the part is durable and referenced, evict the + // source epochs. + for s in &snapshots { + source.evict_sealed_epoch(s.agg_id, s.epoch_id); + } + + stats.entries_flushed = entries_total; + stats.parts_written = 1; + } + } + + // Phase 3: T2 sweep. + if let Some(ttl) = cfg.delete_older_than_ms { + let cutoff = now.saturating_sub(ttl); + let expired: Vec = shared + .manifest + .live_parts() + .into_iter() + .filter(|p| p.max_ts < cutoff) + .collect(); + for p in &expired { + shared.manifest.append_delete(p.part_id)?; + let dir = part_dir_path(&parts_root(&cfg.disk_path), p.part_id); + if let Err(e) = std::fs::remove_dir_all(&dir) { + warn!( + "flusher T2 sweep: failed to rm {:?}: {} (will be retried on restart)", + dir, e + ); + } + } + stats.parts_deleted = expired.len(); + } + + Ok(stats) +} + +/// Interleave a selected-epoch list by agg-id so a burst on one hot +/// agg doesn't monopolize the flush-tick lock sequence. +fn round_robin_by_agg(selected: Vec) -> Vec { + let mut by_agg: HashMap> = HashMap::new(); + // Preserve oldest-first order within each agg. + for r in selected { + by_agg.entry(r.agg_id).or_default().push(r); + } + // Sort agg_ids for deterministic order (test-friendly). + let mut agg_ids: Vec = by_agg.keys().copied().collect(); + agg_ids.sort_unstable(); + let mut out = Vec::new(); + let mut i = 0usize; + loop { + let mut made_progress = false; + for ag in &agg_ids { + if let Some(lst) = by_agg.get_mut(ag) { + if i < lst.len() { + out.push(lst[i]); + made_progress = true; + } + } + } + if !made_progress { + break; + } + i += 1; + } + out +} + +pub(crate) fn parts_root(disk_path: &std::path::Path) -> PathBuf { + disk_path.join("parts") +} + +fn now_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::KeyByLabelValues; + use crate::stores::simple_map_store::persistence::source::{ + EpochSnapshot, EpochSnapshotEntry, + }; + use std::sync::Mutex as StdMutex; + use std::time::Duration; + use tempfile::TempDir; + + /// A fake `EpochSource` with a fixed set of sealed epochs. + struct FakeSource { + epochs: StdMutex>, + memory: AtomicU64, + } + + impl FakeSource { + fn new(snapshots: Vec) -> Self { + let mut map = HashMap::new(); + let mut total = 0u64; + for s in snapshots { + total += s.approx_bytes as u64; + map.insert((s.agg_id, s.epoch_id), s); + } + Self { + epochs: StdMutex::new(map), + memory: AtomicU64::new(total), + } + } + } + + impl EpochSource for FakeSource { + fn list_sealed_epochs(&self) -> Vec { + self.epochs + .lock() + .unwrap() + .values() + .map(|s| SealedEpochRef { + agg_id: s.agg_id, + epoch_id: s.epoch_id, + end_ts: s.max_ts, + approx_bytes: s.approx_bytes, + }) + .collect() + } + + fn snapshot_sealed_epoch( + &self, + agg_id: u64, + epoch_id: u64, + ) -> PersistResult> { + Ok(self + .epochs + .lock() + .unwrap() + .get(&(agg_id, epoch_id)) + .cloned()) + } + + fn evict_sealed_epoch(&self, agg_id: u64, epoch_id: u64) { + let mut map = self.epochs.lock().unwrap(); + if let Some(removed) = map.remove(&(agg_id, epoch_id)) { + self.memory + .fetch_sub(removed.approx_bytes as u64, Ordering::Relaxed); + } + } + + fn approx_memory_bytes(&self) -> usize { + self.memory.load(Ordering::Relaxed) as usize + } + } + + fn snap(agg_id: u64, epoch_id: u64, min_ts: u64, max_ts: u64, approx: usize) -> EpochSnapshot { + EpochSnapshot { + agg_id, + epoch_id, + min_ts, + max_ts, + approx_bytes: approx, + entries: vec![EpochSnapshotEntry { + start_ts: min_ts, + end_ts: max_ts, + label: Some(KeyByLabelValues::new_with_labels(vec!["host".into()])), + sketch_type_name: "SumAccumulator".into(), + sketch_bytes: b"dummy-payload".to_vec(), + }], + } + } + + fn test_cfg(disk_path: PathBuf, mem_limit: usize) -> SimpleMapStorePersistenceConfig { + SimpleMapStorePersistenceConfig { + memory_limit_bytes: mem_limit, + memory_low_watermark_bytes: mem_limit / 2, + hard_cap_bytes: mem_limit * 2, + hot_window_ms: None, + delete_older_than_ms: None, + flush_interval: Duration::from_millis(10), + disk_path, + part_cache_bytes: 0, + } + } + + #[test] + fn memory_pressure_flushes_oldest_first() { + let tmp = TempDir::new().unwrap(); + let snaps = vec![ + snap(1, 1, 100, 150, 500), + snap(1, 2, 150, 200, 500), + snap(2, 1, 200, 250, 500), + ]; + let source = Arc::new(FakeSource::new(snaps)); + let manifest = Arc::new(Manifest::init(tmp.path()).unwrap()); + let cfg = test_cfg(tmp.path().to_path_buf(), 1200); // over limit + let mut handle = FlusherHandle::start(cfg, manifest.clone(), source.clone()).unwrap(); + + // Wait until memory is below the low-water mark. + let deadline = Instant::now() + Duration::from_secs(2); + while source.approx_memory_bytes() > 600 && Instant::now() < deadline { + thread::sleep(Duration::from_millis(20)); + } + handle.shutdown(); + + assert!( + source.approx_memory_bytes() <= 600, + "flusher did not reach low water: {} remaining", + source.approx_memory_bytes() + ); + // At least one part should be in the manifest. + let live = manifest.live_parts(); + assert!(!live.is_empty(), "no parts were produced"); + } + + #[test] + fn hot_window_flushes_even_without_memory_pressure() { + let tmp = TempDir::new().unwrap(); + // Everything older than (now - 1ms) — effectively everything. + let now = now_ms(); + let snaps = vec![snap( + 1, + 1, + now.saturating_sub(10_000), + now.saturating_sub(5_000), + 100, + )]; + let source = Arc::new(FakeSource::new(snaps)); + let manifest = Arc::new(Manifest::init(tmp.path()).unwrap()); + let mut cfg = test_cfg(tmp.path().to_path_buf(), 100_000); // well under limit + cfg.hot_window_ms = Some(1_000); // 1 second hot window + let mut handle = FlusherHandle::start(cfg, manifest.clone(), source.clone()).unwrap(); + + let deadline = Instant::now() + Duration::from_secs(2); + while source.approx_memory_bytes() > 0 && Instant::now() < deadline { + thread::sleep(Duration::from_millis(20)); + } + handle.shutdown(); + + assert_eq!(source.approx_memory_bytes(), 0); + assert!(!manifest.live_parts().is_empty()); + } + + #[test] + fn t2_sweep_deletes_old_parts() { + let tmp = TempDir::new().unwrap(); + let source = Arc::new(FakeSource::new(vec![snap(1, 1, 100, 200, 100)])); + let manifest = Arc::new(Manifest::init(tmp.path()).unwrap()); + let mut cfg = test_cfg(tmp.path().to_path_buf(), 100_000); + cfg.hot_window_ms = Some(0); // force-flush everything + cfg.delete_older_than_ms = Some(0); // then immediately expire it + let mut handle = FlusherHandle::start(cfg, manifest.clone(), source.clone()).unwrap(); + + let deadline = Instant::now() + Duration::from_secs(2); + while !manifest.live_parts().is_empty() && Instant::now() < deadline { + thread::sleep(Duration::from_millis(20)); + } + // Let the sweep run at least once more to catch up. + thread::sleep(Duration::from_millis(50)); + handle.shutdown(); + + assert!( + manifest.live_parts().is_empty(), + "expected live_parts empty after T2 sweep, got {:?}", + manifest.live_parts() + ); + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs b/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs new file mode 100644 index 00000000..a29b78f8 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/manifest.rs @@ -0,0 +1,439 @@ +//! Append-only parts manifest: `parts_manifest.log` + periodic +//! `parts_manifest.snapshot`. The manifest is the one piece of global +//! state on disk and the authoritative list of which parts are live. +//! +//! ## Binary layout +//! +//! Both files share a common record shape for add/delete operations: +//! +//! ```text +//! AddPart : [u8 tag=1][u8 pad;7][u64 part_id][u64 min_ts][u64 max_ts][u64 size_bytes] +//! DeletePart : [u8 tag=2][u8 pad;7][u64 part_id][u64 0][u64 0][u64 0] +//! ``` +//! +//! 40 bytes per record, 8-byte aligned. The snapshot is simply a +//! sequence of `AddPart` records for the currently-live set, followed +//! by a u32 CRC trailer. The log is written by appending records, with +//! the same trailer rewritten atomically after each tick. +//! +//! The snapshot is rewritten (write → fsync → rename) whenever the log +//! has grown past some multiple of the live-set size. On recovery, +//! snapshot is loaded first, then the log is replayed from the start +//! since v1 does not yet track "log offset at snapshot time." + +use std::fs::{self, File, OpenOptions}; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, RwLock}; + +use super::part::PartId; +use super::{PersistError, PersistResult}; + +pub const RECORD_SIZE: usize = 40; +pub const TAG_ADD: u8 = 1; +pub const TAG_DELETE: u8 = 2; +pub const SNAPSHOT_FILE: &str = "parts_manifest.snapshot"; +pub const LOG_FILE: &str = "parts_manifest.log"; + +/// One live part, as seen by queries and the flusher. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PartEntry { + pub part_id: PartId, + pub min_ts: u64, + pub max_ts: u64, + pub size_bytes: u64, +} + +#[derive(Debug, Clone, Copy)] +enum Record { + Add(PartEntry), + Delete(PartId), +} + +fn encode_record(rec: Record, out: &mut [u8; RECORD_SIZE]) { + for b in out.iter_mut() { + *b = 0; + } + match rec { + Record::Add(entry) => { + out[0] = TAG_ADD; + out[8..16].copy_from_slice(&entry.part_id.to_le_bytes()); + out[16..24].copy_from_slice(&entry.min_ts.to_le_bytes()); + out[24..32].copy_from_slice(&entry.max_ts.to_le_bytes()); + out[32..40].copy_from_slice(&entry.size_bytes.to_le_bytes()); + } + Record::Delete(part_id) => { + out[0] = TAG_DELETE; + out[8..16].copy_from_slice(&part_id.to_le_bytes()); + } + } +} + +fn decode_record(bytes: &[u8; RECORD_SIZE]) -> PersistResult { + let tag = bytes[0]; + let part_id = u64::from_le_bytes(bytes[8..16].try_into().unwrap()); + match tag { + TAG_ADD => { + let min_ts = u64::from_le_bytes(bytes[16..24].try_into().unwrap()); + let max_ts = u64::from_le_bytes(bytes[24..32].try_into().unwrap()); + let size_bytes = u64::from_le_bytes(bytes[32..40].try_into().unwrap()); + Ok(Record::Add(PartEntry { + part_id, + min_ts, + max_ts, + size_bytes, + })) + } + TAG_DELETE => Ok(Record::Delete(part_id)), + other => Err(PersistError::Manifest(format!( + "unknown manifest record tag: {}", + other + ))), + } +} + +/// The authoritative in-memory list of live parts plus the on-disk +/// log file handle. Cheap to clone via `Arc`. +pub struct Manifest { + disk_path: PathBuf, + live: Arc>>, + /// Number of log records since the last snapshot. When this grows + /// past `live.len() * 4`, we rewrite the snapshot and truncate the + /// log. The threshold is arbitrary; re-tune if it bites. + log_records_since_snapshot: std::sync::Mutex, +} + +impl Manifest { + /// Create an empty, fresh on-disk manifest in `disk_path`. Errors + /// if `disk_path` already contains a snapshot or a non-empty log — + /// callers should use [`Manifest::open_or_init`] for the common + /// path. + pub fn init(disk_path: &Path) -> PersistResult { + fs::create_dir_all(disk_path)?; + let snapshot_path = disk_path.join(SNAPSHOT_FILE); + let log_path = disk_path.join(LOG_FILE); + if snapshot_path.exists() || log_path.exists() { + return Err(PersistError::Manifest(format!( + "manifest already initialized at {:?}", + disk_path + ))); + } + + write_snapshot_atomic(&snapshot_path, &[])?; + OpenOptions::new() + .create_new(true) + .write(true) + .open(&log_path)?; + + Ok(Self { + disk_path: disk_path.to_path_buf(), + live: Arc::new(RwLock::new(Vec::new())), + log_records_since_snapshot: std::sync::Mutex::new(0), + }) + } + + /// Load an existing manifest from `disk_path`. Initializes one if + /// the directory has no snapshot yet. + pub fn open_or_init(disk_path: &Path) -> PersistResult { + fs::create_dir_all(disk_path)?; + let snapshot_path = disk_path.join(SNAPSHOT_FILE); + let log_path = disk_path.join(LOG_FILE); + if !snapshot_path.exists() && !log_path.exists() { + return Self::init(disk_path); + } + + // Load snapshot (if present) then replay the log. + let mut live: Vec = if snapshot_path.exists() { + read_snapshot(&snapshot_path)? + } else { + Vec::new() + }; + + if log_path.exists() { + let records = read_log(&log_path)?; + for rec in records { + apply_record(&mut live, rec); + } + } else { + OpenOptions::new() + .create_new(true) + .write(true) + .open(&log_path)?; + } + + // Keep `live` sorted by part_id for stable iteration. + live.sort_unstable_by_key(|e| e.part_id); + + Ok(Self { + disk_path: disk_path.to_path_buf(), + live: Arc::new(RwLock::new(live)), + log_records_since_snapshot: std::sync::Mutex::new(0), + }) + } + + pub fn snapshot_path(&self) -> PathBuf { + self.disk_path.join(SNAPSHOT_FILE) + } + + pub fn log_path(&self) -> PathBuf { + self.disk_path.join(LOG_FILE) + } + + /// A clone-on-read snapshot of the live parts list. Cheap — + /// returns an owned `Vec` but the entries themselves are `Copy`. + pub fn live_parts(&self) -> Vec { + self.live.read().unwrap().clone() + } + + /// Parts whose `[min_ts, max_ts]` overlaps `[query_start, + /// query_end]`. Linear scan — the manifest is tiny in v1. + pub fn live_parts_overlapping(&self, query_start: u64, query_end: u64) -> Vec { + self.live + .read() + .unwrap() + .iter() + .filter(|p| !(p.max_ts < query_start || p.min_ts > query_end)) + .copied() + .collect() + } + + /// Append an add-part record to the log and make it durable. + /// Updates in-memory state first, then writes to disk, then fsyncs. + /// The caller must have already fsync'd the part's own files. + pub fn append_add(&self, entry: PartEntry) -> PersistResult<()> { + { + let mut live = self.live.write().unwrap(); + live.push(entry); + live.sort_unstable_by_key(|e| e.part_id); + } + self.append_record(Record::Add(entry)) + } + + /// Append a delete-part record. In-memory removal happens before + /// disk write, same as add. + pub fn append_delete(&self, part_id: PartId) -> PersistResult<()> { + { + let mut live = self.live.write().unwrap(); + live.retain(|e| e.part_id != part_id); + } + self.append_record(Record::Delete(part_id)) + } + + fn append_record(&self, rec: Record) -> PersistResult<()> { + let mut buf = [0u8; RECORD_SIZE]; + encode_record(rec, &mut buf); + { + let mut f = OpenOptions::new().append(true).open(self.log_path())?; + f.write_all(&buf)?; + f.sync_data()?; + } + // fsync the parent directory for the log-file size update. + if let Ok(dir) = File::open(&self.disk_path) { + let _ = dir.sync_all(); + } + + // Maybe compact the log into a fresh snapshot. + let mut counter = self.log_records_since_snapshot.lock().unwrap(); + *counter += 1; + let live_len = self.live.read().unwrap().len() as u64; + let threshold = live_len.saturating_mul(4).max(64); + if *counter >= threshold { + drop(counter); + self.compact()?; + let mut counter = self.log_records_since_snapshot.lock().unwrap(); + *counter = 0; + } + Ok(()) + } + + /// Rewrite the snapshot from the current in-memory live set and + /// truncate the log. Atomic: the new snapshot goes to a tmp file + /// first, then rename, then the log is truncated. + pub fn compact(&self) -> PersistResult<()> { + let live = self.live.read().unwrap().clone(); + write_snapshot_atomic(&self.snapshot_path(), &live)?; + // Truncate log. + let log_path = self.log_path(); + let f = OpenOptions::new().write(true).truncate(true).open(&log_path)?; + f.sync_all()?; + if let Ok(dir) = File::open(&self.disk_path) { + let _ = dir.sync_all(); + } + Ok(()) + } +} + +fn apply_record(live: &mut Vec, rec: Record) { + match rec { + Record::Add(entry) => { + if !live.iter().any(|e| e.part_id == entry.part_id) { + live.push(entry); + } + } + Record::Delete(part_id) => { + live.retain(|e| e.part_id != part_id); + } + } +} + +fn write_snapshot_atomic(path: &Path, live: &[PartEntry]) -> PersistResult<()> { + let tmp_path = path.with_extension("tmp"); + { + let mut f = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&tmp_path)?; + let mut crc = crc32fast::Hasher::new(); + for entry in live { + let mut buf = [0u8; RECORD_SIZE]; + encode_record(Record::Add(*entry), &mut buf); + f.write_all(&buf)?; + crc.update(&buf); + } + let crc_val = crc.finalize(); + f.write_all(&crc_val.to_le_bytes())?; + f.write_all(&[0u8; 4])?; + f.sync_all()?; + } + fs::rename(&tmp_path, path)?; + if let Some(parent) = path.parent() { + if let Ok(dir) = File::open(parent) { + let _ = dir.sync_all(); + } + } + Ok(()) +} + +fn read_snapshot(path: &Path) -> PersistResult> { + let mut f = File::open(path)?; + let mut buf = Vec::new(); + f.read_to_end(&mut buf)?; + if buf.len() < 8 { + return Ok(Vec::new()); + } + let body_len = buf.len() - 8; + if body_len % RECORD_SIZE != 0 { + return Err(PersistError::Manifest(format!( + "snapshot body misaligned: {} bytes", + body_len + ))); + } + let crc_expected = u32::from_le_bytes(buf[body_len..body_len + 4].try_into().unwrap()); + let crc_actual = crc32fast::hash(&buf[..body_len]); + if crc_expected != crc_actual { + return Err(PersistError::Manifest(format!( + "snapshot CRC mismatch: expected {:08x}, got {:08x}", + crc_expected, crc_actual + ))); + } + let mut out = Vec::with_capacity(body_len / RECORD_SIZE); + for chunk in buf[..body_len].chunks_exact(RECORD_SIZE) { + let arr: [u8; RECORD_SIZE] = chunk.try_into().unwrap(); + match decode_record(&arr)? { + Record::Add(entry) => out.push(entry), + Record::Delete(_) => { + return Err(PersistError::Manifest( + "snapshot contained a DELETE record".into(), + )); + } + } + } + Ok(out) +} + +fn read_log(path: &Path) -> PersistResult> { + let mut f = File::open(path)?; + let len = f.seek(SeekFrom::End(0))?; + f.seek(SeekFrom::Start(0))?; + if len == 0 { + return Ok(Vec::new()); + } + if len % RECORD_SIZE as u64 != 0 { + // Tolerate trailing torn write — just truncate to the last + // record boundary. Standard log-recovery practice. + let good = (len / RECORD_SIZE as u64) * RECORD_SIZE as u64; + let mut body = vec![0u8; good as usize]; + f.read_exact(&mut body)?; + return decode_log_body(&body); + } + let mut body = vec![0u8; len as usize]; + f.read_exact(&mut body)?; + decode_log_body(&body) +} + +fn decode_log_body(body: &[u8]) -> PersistResult> { + let mut out = Vec::with_capacity(body.len() / RECORD_SIZE); + for chunk in body.chunks_exact(RECORD_SIZE) { + let arr: [u8; RECORD_SIZE] = chunk.try_into().unwrap(); + out.push(decode_record(&arr)?); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + fn entry(part_id: u64, min_ts: u64, max_ts: u64) -> PartEntry { + PartEntry { + part_id, + min_ts, + max_ts, + size_bytes: 1024, + } + } + + #[test] + fn append_and_reload_round_trip() { + let tmp = TempDir::new().unwrap(); + let m = Manifest::init(tmp.path()).unwrap(); + m.append_add(entry(1, 100, 200)).unwrap(); + m.append_add(entry(2, 200, 300)).unwrap(); + m.append_add(entry(3, 300, 400)).unwrap(); + m.append_delete(2).unwrap(); + + // Drop and re-open — should see parts {1, 3}. + drop(m); + let m2 = Manifest::open_or_init(tmp.path()).unwrap(); + let live = m2.live_parts(); + assert_eq!(live.len(), 2); + assert_eq!(live[0].part_id, 1); + assert_eq!(live[1].part_id, 3); + } + + #[test] + fn compact_rewrites_snapshot_and_truncates_log() { + let tmp = TempDir::new().unwrap(); + let m = Manifest::init(tmp.path()).unwrap(); + for i in 1..=10 { + m.append_add(entry(i, i * 100, i * 100 + 50)).unwrap(); + } + m.compact().unwrap(); + + // Log file should be empty after compaction. + let log_meta = std::fs::metadata(m.log_path()).unwrap(); + assert_eq!(log_meta.len(), 0); + + let m2 = Manifest::open_or_init(tmp.path()).unwrap(); + let live = m2.live_parts(); + assert_eq!(live.len(), 10); + for (i, e) in live.iter().enumerate() { + assert_eq!(e.part_id, i as u64 + 1); + } + } + + #[test] + fn live_parts_overlapping_filters_correctly() { + let tmp = TempDir::new().unwrap(); + let m = Manifest::init(tmp.path()).unwrap(); + m.append_add(entry(1, 100, 200)).unwrap(); + m.append_add(entry(2, 300, 400)).unwrap(); + m.append_add(entry(3, 500, 600)).unwrap(); + + let hits = m.live_parts_overlapping(150, 350); + let ids: Vec = hits.iter().map(|e| e.part_id).collect(); + assert_eq!(ids, vec![1, 2]); + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/mod.rs b/asap-query-engine/src/stores/simple_map_store/persistence/mod.rs new file mode 100644 index 00000000..a7ed40dd --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/mod.rs @@ -0,0 +1,58 @@ +//! Persistence layer for `SimpleMapStorePerKey`. +//! +//! See `docs/design-simple-map-store-persistence.md` for the design rationale. +//! +//! ## Structure +//! +//! * [`config`] — [`SimpleMapStorePersistenceConfig`] +//! * [`part`] — on-disk part format (`meta.bin` + `data.bin` + `index.bin`), +//! writer and reader. +//! * [`manifest`] — append-only log + periodic binary snapshot of live parts. +//! * [`flusher`] — background `std::thread` that walks sealed epochs and +//! turns each tick's candidates into a single on-disk part. +//! * [`cache`] — moka-backed Tier-2 cache of decoded parts, keyed on +//! `PartId`, bounded by bytes. +//! * [`recovery`] — startup: load snapshot, replay log, verify CRCs, sweep +//! orphan part dirs. +//! +//! The submodule is intentionally decoupled from `SimpleMapStorePerKey` +//! via the [`EpochSource`] trait — the flusher knows nothing about the +//! store's internal types and can be unit-tested against a fake source. + +pub mod config; +pub mod manifest; +pub mod part; +pub mod source; + +pub mod cache; +pub mod flusher; +pub mod recovery; + +pub use config::SimpleMapStorePersistenceConfig; +pub use manifest::{Manifest, PartEntry}; +pub use part::{PartId, PartReader, PartWriter, SnapshotEntry}; +pub use source::{EpochSource, SealedEpochRef}; + +/// Convenience result alias used across the persistence layer. +pub type PersistResult = Result; + +#[derive(Debug, thiserror::Error)] +pub enum PersistError { + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + + #[error("part format error: {0}")] + Format(String), + + #[error("manifest corruption: {0}")] + Manifest(String), + + #[error("unsupported accumulator type for persistence: {0}")] + UnsupportedAccumulator(String), + + #[error("serialization error: {0}")] + Serialize(String), + + #[error("internal error: {0}")] + Internal(String), +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/part.rs b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs new file mode 100644 index 00000000..3e4b41a7 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/part.rs @@ -0,0 +1,697 @@ +//! On-disk part format: `meta.bin` + `data.bin` + `index.bin`. +//! +//! A *part* is one flush tick's bundle of sealed-epoch entries, stored +//! in `/parts//`. See the design doc +//! section "Disk layout" for the full format. +//! +//! ## v1 format +//! +//! Custom little-endian binary throughout. 8-byte alignment on anything +//! the reader needs to cast directly. CRC32 (crc32fast) trailer on each +//! file. Written via `fallocate` when available, sync'd file-by-file, +//! directory fsync after all three files are durable. +//! +//! File shapes (v1): +//! +//! ```text +//! meta.bin (fixed 64-byte header + a variable-size trailer we ignore) +//! u32 MAGIC_META +//! u16 VERSION = 1 +//! u16 _flags (reserved, zero) +//! u64 part_id +//! u64 min_ts +//! u64 max_ts +//! u32 num_entries +//! u32 _reserved +//! u64 data_len +//! u64 index_len +//! u64 created_unix_ms +//! u32 crc32_of_header (covers all preceding 60 bytes) +//! u32 _pad +//! +//! data.bin (opaque byte blob; entries concatenated, each 8-byte +//! aligned; the reader never scans it linearly, it only seeks via index +//! offsets) +//! repeat num_entries: +//! [label_len: u32] +//! [payload_len: u32] +//! [label bytes: label_len bytes, bincode-serialized Option] +//! [pad to 8-byte align of payload] +//! [type_name_len: u16] +//! [pad to 8-byte align of type_name (small)] +//! [type_name bytes] +//! [pad to 8-byte align] +//! [payload bytes: payload_len bytes] +//! [tail pad to 8-byte align] +//! u32 crc32_of_body +//! u32 _pad +//! +//! index.bin (sorted by (agg_id is same for whole part in v1 — one +//! part mixes epochs but one sealed epoch is one agg_id; for now we +//! store agg_id per entry for forward-compat, start_ts)) +//! repeat num_entries: +//! [agg_id: u64] +//! [start_ts: u64] +//! [end_ts: u64] +//! [data_offset: u64] +//! u32 crc32_of_body +//! u32 _pad +//! ``` +//! +//! The data encoding is intentionally simple (length-prefixed fields, +//! 8-byte alignment) and skips a bunch of the micro-optimizations in +//! the design doc. v1 priorities are correctness, round-trip, and +//! passing tests — not beating Prometheus TSDB. + +use std::fs::{self, File, OpenOptions}; +use std::io::{BufWriter, Read, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use memmap2::Mmap; + +use super::source::EpochSnapshot; +use super::{PersistError, PersistResult}; + +pub type PartId = u64; + +pub(crate) const MAGIC_META: u32 = 0x4D_50_41_52; // "MPAR" +pub(crate) const PART_FORMAT_VERSION: u16 = 1; +pub(crate) const META_HEADER_SIZE: usize = 64; +pub(crate) const INDEX_ENTRY_SIZE: usize = 32; + +/// Zero-padded directory name for a part. +pub fn part_dir_name(part_id: PartId) -> String { + format!("{:016x}", part_id) +} + +/// Full path for a part directory under a given parts root. +pub fn part_dir_path(parts_root: &Path, part_id: PartId) -> PathBuf { + parts_root.join(part_dir_name(part_id)) +} + +/// One entry inside a decoded part. The `start_ts`/`end_ts`/`label` +/// fields are resolved by the reader; the sketch payload stays as +/// bytes so the query path can decide whether to decode lazily. +#[derive(Debug, Clone)] +pub struct SnapshotEntry { + pub agg_id: u64, + pub start_ts: u64, + pub end_ts: u64, + pub label: Option, + pub sketch_type_name: String, + pub sketch_bytes: Vec, +} + +/// Write one [`EpochSnapshot`] (or many, concatenated) into a fresh +/// part directory. Consumes snapshots in order and produces the final +/// `(data_len, index_len)` for the manifest record. Streams data +/// through a `BufWriter` with a running CRC. +pub struct PartWriter; + +impl PartWriter { + /// Build a new part directory from a list of epoch snapshots. The + /// snapshots may come from multiple agg-ids (one flush tick bundles + /// candidates from whichever aggs needed eviction). + /// + /// Returns `(min_ts, max_ts, num_entries, data_len, index_len)`. + pub fn write_part( + part_dir: &Path, + part_id: PartId, + snapshots: &[EpochSnapshot], + ) -> PersistResult { + fs::create_dir_all(part_dir)?; + + // ---- Build in-memory entry plan (offsets + index entries) ---- + let mut entries_plan: Vec = Vec::new(); + let mut data_len: u64 = 0; + let mut min_ts: u64 = u64::MAX; + let mut max_ts: u64 = 0; + + for snap in snapshots { + for e in &snap.entries { + let label_bytes = match &e.label { + Some(k) => k.serialize_to_bytes(), + None => Vec::new(), + }; + let type_name_bytes = e.sketch_type_name.as_bytes().to_vec(); + // Layout inside data.bin per entry: + // u32 label_len + // u32 payload_len + // u16 type_name_len + // u16 _pad + // u32 _pad + // label_bytes + pad-to-8 + // type_name_bytes + pad-to-8 + // payload_bytes + pad-to-8 + let header_size = 16; // 4+4+2+2+4 + let label_padded = align_up(label_bytes.len(), 8); + let type_padded = align_up(type_name_bytes.len(), 8); + let payload_padded = align_up(e.sketch_bytes.len(), 8); + let entry_size = header_size + label_padded + type_padded + payload_padded; + + let data_offset = data_len; + entries_plan.push(PlannedEntry { + agg_id: snap.agg_id, + start_ts: e.start_ts, + end_ts: e.end_ts, + data_offset, + label_bytes, + type_name_bytes, + sketch_bytes: e.sketch_bytes.clone(), + }); + data_len += entry_size as u64; + min_ts = min_ts.min(e.start_ts); + max_ts = max_ts.max(e.end_ts); + } + } + + if entries_plan.is_empty() { + return Err(PersistError::Internal( + "PartWriter::write_part called with zero entries".to_string(), + )); + } + + let num_entries = entries_plan.len() as u32; + let index_len = num_entries as u64 * INDEX_ENTRY_SIZE as u64 + 8; // + crc trailer + + // ---- Write data.bin with streaming CRC ---- + let data_path = part_dir.join("data.bin"); + let mut data_file = BufWriter::new( + OpenOptions::new() + .create_new(true) + .write(true) + .open(&data_path)?, + ); + let mut data_crc = crc32fast::Hasher::new(); + for pe in &entries_plan { + write_u32(&mut data_file, &mut data_crc, pe.label_bytes.len() as u32)?; + write_u32(&mut data_file, &mut data_crc, pe.sketch_bytes.len() as u32)?; + write_u16(&mut data_file, &mut data_crc, pe.type_name_bytes.len() as u16)?; + write_u16(&mut data_file, &mut data_crc, 0)?; // _pad + write_u32(&mut data_file, &mut data_crc, 0)?; // _pad + write_padded(&mut data_file, &mut data_crc, &pe.label_bytes, 8)?; + write_padded(&mut data_file, &mut data_crc, &pe.type_name_bytes, 8)?; + write_padded(&mut data_file, &mut data_crc, &pe.sketch_bytes, 8)?; + } + let data_crc_val = data_crc.finalize(); + data_file.write_all(&data_crc_val.to_le_bytes())?; + data_file.write_all(&[0u8; 4])?; // pad + data_file.flush()?; + let data_file = data_file.into_inner().map_err(|e| PersistError::Io(e.into_error()))?; + data_file.sync_all()?; + drop(data_file); + + // ---- Write index.bin ---- + let index_path = part_dir.join("index.bin"); + let mut index_file = BufWriter::new( + OpenOptions::new() + .create_new(true) + .write(true) + .open(&index_path)?, + ); + let mut index_crc = crc32fast::Hasher::new(); + for pe in &entries_plan { + write_u64(&mut index_file, &mut index_crc, pe.agg_id)?; + write_u64(&mut index_file, &mut index_crc, pe.start_ts)?; + write_u64(&mut index_file, &mut index_crc, pe.end_ts)?; + write_u64(&mut index_file, &mut index_crc, pe.data_offset)?; + } + let index_crc_val = index_crc.finalize(); + index_file.write_all(&index_crc_val.to_le_bytes())?; + index_file.write_all(&[0u8; 4])?; + index_file.flush()?; + let index_file = index_file.into_inner().map_err(|e| PersistError::Io(e.into_error()))?; + index_file.sync_all()?; + drop(index_file); + + // ---- Write meta.bin (64 bytes, CRC'd) ---- + let meta_path = part_dir.join("meta.bin"); + let created_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + + // 64-byte header. Layout: + // [0..4] u32 magic + // [4..6] u16 version + // [6..8] u16 flags + // [8..16] u64 part_id + // [16..24] u64 min_ts + // [24..32] u64 max_ts + // [32..36] u32 num_entries + // [36..40] u32 reserved + // [40..48] u64 data_len + // [48..56] u64 index_len + // [56..60] u32 created_unix_secs (seconds granularity; diagnostic) + // [60..64] u32 crc32 of [0..60] + let mut header = [0u8; META_HEADER_SIZE]; + let mut cursor = 0usize; + write_u32_into(&mut header, &mut cursor, MAGIC_META); + write_u16_into(&mut header, &mut cursor, PART_FORMAT_VERSION); + write_u16_into(&mut header, &mut cursor, 0); // flags + write_u64_into(&mut header, &mut cursor, part_id); + write_u64_into(&mut header, &mut cursor, min_ts); + write_u64_into(&mut header, &mut cursor, max_ts); + write_u32_into(&mut header, &mut cursor, num_entries); + write_u32_into(&mut header, &mut cursor, 0); // reserved + write_u64_into(&mut header, &mut cursor, data_len); + write_u64_into(&mut header, &mut cursor, index_len); + let created_secs: u32 = (created_ms / 1000).min(u32::MAX as u64) as u32; + write_u32_into(&mut header, &mut cursor, created_secs); + let meta_crc = crc32fast::hash(&header[..60]); + header[60..64].copy_from_slice(&meta_crc.to_le_bytes()); + + { + let mut meta_file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&meta_path)?; + meta_file.write_all(&header)?; + meta_file.sync_all()?; + } + + // fsync the directory so the files are durably linked. + if let Ok(dir) = File::open(part_dir) { + let _ = dir.sync_all(); + } + + Ok(PartWriteReport { + part_id, + min_ts, + max_ts, + num_entries, + data_len, + index_len, + }) + } +} + +/// What [`PartWriter::write_part`] reports back. +#[derive(Debug, Clone)] +pub struct PartWriteReport { + pub part_id: PartId, + pub min_ts: u64, + pub max_ts: u64, + pub num_entries: u32, + pub data_len: u64, + pub index_len: u64, +} + +struct PlannedEntry { + agg_id: u64, + start_ts: u64, + end_ts: u64, + data_offset: u64, + label_bytes: Vec, + type_name_bytes: Vec, + sketch_bytes: Vec, +} + +// ----- small write helpers with running CRC ----- + +fn write_u16(w: &mut W, h: &mut crc32fast::Hasher, v: u16) -> std::io::Result<()> { + let b = v.to_le_bytes(); + w.write_all(&b)?; + h.update(&b); + Ok(()) +} + +fn write_u32(w: &mut W, h: &mut crc32fast::Hasher, v: u32) -> std::io::Result<()> { + let b = v.to_le_bytes(); + w.write_all(&b)?; + h.update(&b); + Ok(()) +} + +fn write_u64(w: &mut W, h: &mut crc32fast::Hasher, v: u64) -> std::io::Result<()> { + let b = v.to_le_bytes(); + w.write_all(&b)?; + h.update(&b); + Ok(()) +} + +fn write_padded( + w: &mut W, + h: &mut crc32fast::Hasher, + bytes: &[u8], + align: usize, +) -> std::io::Result<()> { + w.write_all(bytes)?; + h.update(bytes); + let padded = align_up(bytes.len(), align); + let pad_len = padded - bytes.len(); + if pad_len > 0 { + let pad = [0u8; 8]; + w.write_all(&pad[..pad_len])?; + h.update(&pad[..pad_len]); + } + Ok(()) +} + +fn align_up(n: usize, align: usize) -> usize { + (n + align - 1) & !(align - 1) +} + +fn write_u16_into(buf: &mut [u8], cursor: &mut usize, v: u16) { + buf[*cursor..*cursor + 2].copy_from_slice(&v.to_le_bytes()); + *cursor += 2; +} + +fn write_u32_into(buf: &mut [u8], cursor: &mut usize, v: u32) { + buf[*cursor..*cursor + 4].copy_from_slice(&v.to_le_bytes()); + *cursor += 4; +} + +fn write_u64_into(buf: &mut [u8], cursor: &mut usize, v: u64) { + buf[*cursor..*cursor + 8].copy_from_slice(&v.to_le_bytes()); + *cursor += 8; +} + +// ================================================================= +// Reader +// ================================================================= + +/// Decoded metadata header of a part (from `meta.bin`). +#[derive(Debug, Clone)] +pub struct PartMeta { + pub part_id: PartId, + pub min_ts: u64, + pub max_ts: u64, + pub num_entries: u32, + pub data_len: u64, + pub index_len: u64, + pub created_unix_ms: u64, +} + +/// One resolved index record. `sketch_bytes` and the label are lazy — +/// they live inside `data.bin` and are resolved on demand via +/// [`PartReader::load_entry`]. +#[derive(Debug, Clone, Copy)] +pub struct IndexRecord { + pub agg_id: u64, + pub start_ts: u64, + pub end_ts: u64, + pub data_offset: u64, +} + +/// mmap-backed reader for a single part. Cheap to construct (three +/// mmaps + one header parse), safe to share across threads via `Arc`. +pub struct PartReader { + pub meta: PartMeta, + pub part_dir: PathBuf, + data_mmap: Arc, + index_mmap: Arc, +} + +impl std::fmt::Debug for PartReader { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PartReader") + .field("meta", &self.meta) + .field("part_dir", &self.part_dir) + .field("data_mmap_len", &self.data_mmap.len()) + .field("index_mmap_len", &self.index_mmap.len()) + .finish() + } +} + +impl PartReader { + pub fn open(part_dir: &Path) -> PersistResult { + let meta = Self::read_meta(part_dir)?; + + let data_mmap = map_file(&part_dir.join("data.bin"))?; + let index_mmap = map_file(&part_dir.join("index.bin"))?; + + // Sanity-check CRCs up front so we catch tampered files. + let idx_payload_len = meta.num_entries as usize * INDEX_ENTRY_SIZE; + if index_mmap.len() < idx_payload_len + 8 { + return Err(PersistError::Format(format!( + "index.bin too short for {} entries (have {} bytes)", + meta.num_entries, + index_mmap.len() + ))); + } + let idx_expected = u32::from_le_bytes( + index_mmap[idx_payload_len..idx_payload_len + 4] + .try_into() + .unwrap(), + ); + let idx_actual = crc32fast::hash(&index_mmap[..idx_payload_len]); + if idx_expected != idx_actual { + return Err(PersistError::Format(format!( + "index.bin CRC mismatch: expected {:08x}, got {:08x}", + idx_expected, idx_actual + ))); + } + + Ok(Self { + meta, + part_dir: part_dir.to_path_buf(), + data_mmap: Arc::new(data_mmap), + index_mmap: Arc::new(index_mmap), + }) + } + + /// Read and verify just the meta.bin header (cheap — 64 bytes). + pub fn read_meta(part_dir: &Path) -> PersistResult { + let mut f = File::open(part_dir.join("meta.bin"))?; + let mut buf = [0u8; META_HEADER_SIZE]; + f.read_exact(&mut buf)?; + + let magic = u32::from_le_bytes(buf[0..4].try_into().unwrap()); + if magic != MAGIC_META { + return Err(PersistError::Format(format!( + "meta.bin bad magic: {:08x}", + magic + ))); + } + let version = u16::from_le_bytes(buf[4..6].try_into().unwrap()); + if version != PART_FORMAT_VERSION { + return Err(PersistError::Format(format!( + "meta.bin unsupported version: {}", + version + ))); + } + // 6..8 flags (ignored) + let part_id = u64::from_le_bytes(buf[8..16].try_into().unwrap()); + let min_ts = u64::from_le_bytes(buf[16..24].try_into().unwrap()); + let max_ts = u64::from_le_bytes(buf[24..32].try_into().unwrap()); + let num_entries = u32::from_le_bytes(buf[32..36].try_into().unwrap()); + // 36..40 reserved + let data_len = u64::from_le_bytes(buf[40..48].try_into().unwrap()); + let index_len = u64::from_le_bytes(buf[48..56].try_into().unwrap()); + let created_unix_secs = + u32::from_le_bytes(buf[56..60].try_into().unwrap()) as u64; + let created_unix_ms = created_unix_secs * 1000; + let crc_expected = u32::from_le_bytes(buf[60..64].try_into().unwrap()); + let crc_actual = crc32fast::hash(&buf[..60]); + if crc_expected != crc_actual { + return Err(PersistError::Format(format!( + "meta.bin CRC mismatch: expected {:08x}, got {:08x}", + crc_expected, crc_actual + ))); + } + + Ok(PartMeta { + part_id, + min_ts, + max_ts, + num_entries, + data_len, + index_len, + created_unix_ms, + }) + } + + /// Return all index records. Small (32 B × num_entries) — cheap to + /// materialize. + pub fn index_records(&self) -> Vec { + let n = self.meta.num_entries as usize; + let mut out = Vec::with_capacity(n); + for i in 0..n { + let off = i * INDEX_ENTRY_SIZE; + let agg_id = u64::from_le_bytes( + self.index_mmap[off..off + 8].try_into().unwrap(), + ); + let start_ts = u64::from_le_bytes( + self.index_mmap[off + 8..off + 16].try_into().unwrap(), + ); + let end_ts = u64::from_le_bytes( + self.index_mmap[off + 16..off + 24].try_into().unwrap(), + ); + let data_offset = u64::from_le_bytes( + self.index_mmap[off + 24..off + 32].try_into().unwrap(), + ); + out.push(IndexRecord { + agg_id, + start_ts, + end_ts, + data_offset, + }); + } + out + } + + /// Resolve a single index record into a [`SnapshotEntry`] by reading + /// the corresponding slice of `data.bin`. Decodes the inline label + /// and copies the sketch payload out. + pub fn load_entry(&self, rec: &IndexRecord) -> PersistResult { + let off = rec.data_offset as usize; + if off + 16 > self.data_mmap.len() { + return Err(PersistError::Format( + "data.bin offset out of range".to_string(), + )); + } + let label_len = u32::from_le_bytes(self.data_mmap[off..off + 4].try_into().unwrap()) + as usize; + let payload_len = + u32::from_le_bytes(self.data_mmap[off + 4..off + 8].try_into().unwrap()) as usize; + let type_name_len = + u16::from_le_bytes(self.data_mmap[off + 8..off + 10].try_into().unwrap()) as usize; + // 10..12 pad, 12..16 pad + let mut cursor = off + 16; + let label_padded = align_up(label_len, 8); + let label_bytes = &self.data_mmap[cursor..cursor + label_len]; + let label = if label_len == 0 { + None + } else { + Some( + crate::data_model::KeyByLabelValues::deserialize_from_bytes(label_bytes) + .map_err(|e| PersistError::Format(format!("label decode: {}", e)))?, + ) + }; + cursor += label_padded; + let type_padded = align_up(type_name_len, 8); + let type_name = std::str::from_utf8(&self.data_mmap[cursor..cursor + type_name_len]) + .map_err(|e| PersistError::Format(format!("type_name utf8: {}", e)))? + .to_string(); + cursor += type_padded; + let sketch_bytes = self.data_mmap[cursor..cursor + payload_len].to_vec(); + + Ok(SnapshotEntry { + agg_id: rec.agg_id, + start_ts: rec.start_ts, + end_ts: rec.end_ts, + label, + sketch_type_name: type_name, + sketch_bytes, + }) + } +} + +fn map_file(path: &Path) -> PersistResult { + let f = File::open(path)?; + // Safety: we never mutate the underlying file while the mmap is + // alive; the parts directory is owned by this process's flusher. + let mmap = unsafe { Mmap::map(&f) }?; + Ok(mmap) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::KeyByLabelValues; + use crate::stores::simple_map_store::persistence::source::EpochSnapshotEntry; + use tempfile::TempDir; + + fn make_snapshot() -> EpochSnapshot { + EpochSnapshot { + agg_id: 42, + epoch_id: 7, + min_ts: 1_000, + max_ts: 2_000, + approx_bytes: 128, + entries: vec![ + EpochSnapshotEntry { + start_ts: 1_000, + end_ts: 1_500, + label: Some(KeyByLabelValues::new_with_labels(vec![ + "svc".into(), + "api".into(), + ])), + sketch_type_name: "SumAccumulator".into(), + sketch_bytes: b"opaque-sketch-1".to_vec(), + }, + EpochSnapshotEntry { + start_ts: 1_500, + end_ts: 2_000, + label: None, + sketch_type_name: "DatasketchesKLLAccumulator".into(), + sketch_bytes: b"opaque-sketch-2-more-bytes".to_vec(), + }, + ], + } + } + + #[test] + fn part_round_trip_writes_and_reads_back() { + let tmp = TempDir::new().unwrap(); + let part_dir = tmp.path().join("0000000000000001"); + let snap = make_snapshot(); + let report = + PartWriter::write_part(&part_dir, 1, &[snap.clone()]).expect("write_part"); + + assert_eq!(report.part_id, 1); + assert_eq!(report.num_entries, 2); + assert_eq!(report.min_ts, 1_000); + assert_eq!(report.max_ts, 2_000); + assert!(report.data_len > 0); + assert!(report.index_len > 0); + + let reader = PartReader::open(&part_dir).expect("PartReader::open"); + assert_eq!(reader.meta.part_id, 1); + assert_eq!(reader.meta.num_entries, 2); + assert_eq!(reader.meta.min_ts, 1_000); + assert_eq!(reader.meta.max_ts, 2_000); + + let recs = reader.index_records(); + assert_eq!(recs.len(), 2); + assert_eq!(recs[0].agg_id, 42); + assert_eq!(recs[0].start_ts, 1_000); + assert_eq!(recs[0].end_ts, 1_500); + + let e0 = reader.load_entry(&recs[0]).expect("load_entry 0"); + assert_eq!(e0.sketch_type_name, "SumAccumulator"); + assert_eq!(e0.sketch_bytes, b"opaque-sketch-1"); + assert_eq!( + e0.label.as_ref().unwrap().labels, + vec!["svc".to_string(), "api".to_string()] + ); + + let e1 = reader.load_entry(&recs[1]).expect("load_entry 1"); + assert!(e1.label.is_none()); + assert_eq!(e1.sketch_type_name, "DatasketchesKLLAccumulator"); + assert_eq!(e1.sketch_bytes, b"opaque-sketch-2-more-bytes"); + } + + #[test] + fn part_reader_rejects_corrupted_meta() { + let tmp = TempDir::new().unwrap(); + let part_dir = tmp.path().join("0000000000000002"); + PartWriter::write_part(&part_dir, 2, &[make_snapshot()]).unwrap(); + + // Flip a byte inside the header (but not in the CRC field). + let meta_path = part_dir.join("meta.bin"); + let mut bytes = std::fs::read(&meta_path).unwrap(); + bytes[10] ^= 0xFF; + std::fs::write(&meta_path, &bytes).unwrap(); + + let err = PartReader::open(&part_dir).unwrap_err(); + match err { + PersistError::Format(msg) => assert!(msg.contains("CRC mismatch")), + other => panic!("expected Format error, got {:?}", other), + } + } + + #[test] + fn part_writer_rejects_empty_snapshot_set() { + let tmp = TempDir::new().unwrap(); + let part_dir = tmp.path().join("empty"); + let err = PartWriter::write_part(&part_dir, 99, &[]).unwrap_err(); + match err { + PersistError::Internal(msg) => assert!(msg.contains("zero entries")), + other => panic!("expected Internal, got {:?}", other), + } + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs b/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs new file mode 100644 index 00000000..3b5f64f3 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/recovery.rs @@ -0,0 +1,203 @@ +//! Startup recovery: load the parts manifest, verify every referenced +//! part directory, and sweep orphan directories from +//! crashes-in-flight. + +use std::collections::HashSet; +use std::fs; +use std::path::Path; + +use tracing::{info, warn}; + +use super::flusher::parts_root; +use super::manifest::Manifest; +use super::part::{PartId, PartReader}; +use super::PersistResult; + +/// Result of a recovery pass. +#[derive(Debug, Default)] +pub struct RecoveryReport { + pub live_parts: usize, + pub corrupt_parts_removed: usize, + pub orphan_parts_removed: usize, +} + +/// Open the manifest at `disk_path`, replay its log, validate every +/// referenced part, and sweep orphan part directories not in the +/// manifest. +pub fn recover(disk_path: &Path) -> PersistResult<(Manifest, RecoveryReport)> { + fs::create_dir_all(disk_path)?; + let parts_root = parts_root(disk_path); + fs::create_dir_all(&parts_root)?; + + let manifest = Manifest::open_or_init(disk_path)?; + let mut report = RecoveryReport::default(); + + // Validate every live part by reading its meta.bin header. + let mut to_drop: Vec = Vec::new(); + let mut referenced: HashSet = HashSet::new(); + for entry in manifest.live_parts() { + referenced.insert(entry.part_id); + let part_dir = super::part::part_dir_path(&parts_root, entry.part_id); + match PartReader::read_meta(&part_dir) { + Ok(meta) if meta.part_id == entry.part_id => { + // Meta header checks out; don't also bother reading + // data/index here — the query path will catch any + // deeper corruption and surface it. + } + Ok(meta) => { + warn!( + "recovery: part_id mismatch in meta.bin at {:?}: header says {}, manifest says {}", + part_dir, meta.part_id, entry.part_id + ); + to_drop.push(entry.part_id); + } + Err(e) => { + warn!( + "recovery: unreadable part {} at {:?}: {}; dropping from manifest", + entry.part_id, part_dir, e + ); + to_drop.push(entry.part_id); + } + } + } + + for part_id in to_drop { + manifest.append_delete(part_id)?; + let dir = super::part::part_dir_path(&parts_root, part_id); + let _ = fs::remove_dir_all(&dir); + report.corrupt_parts_removed += 1; + } + + // Sweep orphans: any directory under parts/ whose parsed ID is not + // in `referenced` is a leftover from a mid-flush crash. + if parts_root.is_dir() { + for entry in fs::read_dir(&parts_root)? { + let entry = entry?; + let name = match entry.file_name().to_str() { + Some(s) => s.to_string(), + None => continue, + }; + let part_id = match u64::from_str_radix(&name, 16) { + Ok(id) => id, + Err(_) => continue, + }; + if !referenced.contains(&part_id) { + let path = entry.path(); + if let Err(e) = fs::remove_dir_all(&path) { + warn!( + "recovery: failed to remove orphan part dir {:?}: {}", + path, e + ); + } else { + report.orphan_parts_removed += 1; + } + } + } + } + + // Refresh manifest's in-memory view after dropping corrupt entries. + let manifest = if report.corrupt_parts_removed > 0 { + // Re-open so the in-memory snapshot is consistent with the + // disk after our appends. + Manifest::open_or_init(disk_path)? + } else { + manifest + }; + + report.live_parts = manifest.live_parts().len(); + + info!( + "persistence recovery: live={} corrupt_removed={} orphans_removed={}", + report.live_parts, report.corrupt_parts_removed, report.orphan_parts_removed + ); + + Ok((manifest, report)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::KeyByLabelValues; + use crate::stores::simple_map_store::persistence::part::{PartWriter, part_dir_path}; + use crate::stores::simple_map_store::persistence::source::{ + EpochSnapshot, EpochSnapshotEntry, + }; + use tempfile::TempDir; + + fn dummy_snapshot() -> EpochSnapshot { + EpochSnapshot { + agg_id: 1, + epoch_id: 1, + min_ts: 100, + max_ts: 200, + approx_bytes: 64, + entries: vec![EpochSnapshotEntry { + start_ts: 100, + end_ts: 200, + label: Some(KeyByLabelValues::new_with_labels(vec!["x".into()])), + sketch_type_name: "SumAccumulator".into(), + sketch_bytes: b"payload".to_vec(), + }], + } + } + + #[test] + fn recover_initializes_empty_dir() { + let tmp = TempDir::new().unwrap(); + let (m, report) = recover(tmp.path()).unwrap(); + assert_eq!(m.live_parts().len(), 0); + assert_eq!(report.live_parts, 0); + assert_eq!(report.corrupt_parts_removed, 0); + assert_eq!(report.orphan_parts_removed, 0); + } + + #[test] + fn recover_sweeps_orphan_part_dir() { + let tmp = TempDir::new().unwrap(); + // First: init an empty manifest so the directory exists. + let _ = recover(tmp.path()).unwrap(); + + // Now write a part on disk but do NOT reference it in the + // manifest — simulates a crash between segment fsync and log + // append. + let parts_root = parts_root(tmp.path()); + std::fs::create_dir_all(&parts_root).unwrap(); + let orphan_dir = part_dir_path(&parts_root, 999); + PartWriter::write_part(&orphan_dir, 999, &[dummy_snapshot()]).unwrap(); + assert!(orphan_dir.exists()); + + let (_, report) = recover(tmp.path()).unwrap(); + assert_eq!(report.orphan_parts_removed, 1); + assert!(!orphan_dir.exists()); + } + + #[test] + fn recover_drops_corrupt_part_and_cleans_directory() { + let tmp = TempDir::new().unwrap(); + // Init and write a legit part, then register it in the manifest. + let (manifest, _) = recover(tmp.path()).unwrap(); + let parts_root = parts_root(tmp.path()); + let part_dir = part_dir_path(&parts_root, 42); + let report_write = + PartWriter::write_part(&part_dir, 42, &[dummy_snapshot()]).unwrap(); + manifest + .append_add(crate::stores::simple_map_store::persistence::manifest::PartEntry { + part_id: 42, + min_ts: report_write.min_ts, + max_ts: report_write.max_ts, + size_bytes: report_write.data_len + report_write.index_len, + }) + .unwrap(); + drop(manifest); + + // Corrupt meta.bin. + let mut bytes = std::fs::read(part_dir.join("meta.bin")).unwrap(); + bytes[5] ^= 0xFF; + std::fs::write(part_dir.join("meta.bin"), &bytes).unwrap(); + + let (m, report) = recover(tmp.path()).unwrap(); + assert_eq!(report.corrupt_parts_removed, 1); + assert_eq!(m.live_parts().len(), 0); + assert!(!part_dir.exists()); + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/source.rs b/asap-query-engine/src/stores/simple_map_store/persistence/source.rs new file mode 100644 index 00000000..2d234e84 --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/persistence/source.rs @@ -0,0 +1,99 @@ +//! The trait the flusher uses to enumerate, snapshot, and evict sealed +//! epochs. Decouples `flusher.rs` from `SimpleMapStorePerKey` so the +//! flusher can be unit-tested against a fake source. + +use crate::data_model::KeyByLabelValues; + +use super::PersistResult; + +/// A compact reference to a sealed epoch held in memory. Returned by +/// [`EpochSource::list_sealed_epochs`] in oldest-first global order (by +/// `end_ts`), interleaved round-robin across agg-ids when the flusher +/// walks them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SealedEpochRef { + pub agg_id: u64, + pub epoch_id: u64, + /// `max_end` of the epoch's timestamp range. Used for oldest-first + /// ordering and for the `hot_window_ms` trigger. + pub end_ts: u64, + /// Approximate memory footprint of the epoch's sketches. Drives the + /// memory-pressure trigger; the flusher subtracts this when evicting. + pub approx_bytes: usize, +} + +/// A fully-resolved, self-contained copy of one sealed epoch. Produced by +/// [`EpochSource::snapshot_sealed_epoch`] under a brief per-agg read lock, +/// then serialized on the flusher's thread with no store locks held. +/// +/// The `entries` are ready to write to disk: labels are already resolved +/// to `Option` (no intern-table lookup needed) and the +/// sketch bytes are already in the Arroyo/MessagePack format used by +/// `crate::engines::physical::accumulator_serde::deserialize_accumulator`. +#[derive(Debug, Clone)] +pub struct EpochSnapshot { + pub agg_id: u64, + pub epoch_id: u64, + pub min_ts: u64, + pub max_ts: u64, + pub entries: Vec, + /// Sum of `approx_memory_bytes` across entries (v1: a coarse + /// per-type heuristic; see `per_key.rs::estimate_epoch_bytes`). + pub approx_bytes: usize, +} + +impl EpochSnapshot { + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } +} + +/// One resolved entry inside an [`EpochSnapshot`]. Matches the shape +/// the on-disk part writer expects. +#[derive(Debug, Clone)] +pub struct EpochSnapshotEntry { + pub start_ts: u64, + pub end_ts: u64, + /// Optional label set, already resolved from the per-agg intern table. + pub label: Option, + /// `AggregateCore::type_name()` of the underlying sketch, used on + /// read-back to pick the right `deserialize_from_bytes_arroyo` impl. + pub sketch_type_name: String, + /// Serialized sketch payload (Arroyo / MessagePack format). + pub sketch_bytes: Vec, +} + +/// Trait the flusher uses to discover, snapshot, and evict sealed +/// epochs. Implemented by `SimpleMapStorePerKey`; a test fake lives in +/// `flusher.rs`'s unit tests. +/// +/// Implementors guarantee that: +/// +/// * [`list_sealed_epochs`] returns an approximate global oldest-first +/// view. Approximate is fine: the flusher re-checks each epoch's +/// existence when calling [`snapshot_sealed_epoch`]. +/// * [`snapshot_sealed_epoch`] is safe to call concurrently with +/// inserts. It clones the epoch's `Arc` out under a read lock and +/// drops the lock before returning, so no lock is held across +/// downstream serialization / I/O. +/// * [`evict_sealed_epoch`] is idempotent — calling it on an +/// already-evicted (agg_id, epoch_id) is a no-op. +/// * [`approx_memory_bytes`] is cheap (atomic load) and is kept in sync +/// with what the flusher has evicted. +pub trait EpochSource: Send + Sync { + fn list_sealed_epochs(&self) -> Vec; + + fn snapshot_sealed_epoch( + &self, + agg_id: u64, + epoch_id: u64, + ) -> PersistResult>; + + fn evict_sealed_epoch(&self, agg_id: u64, epoch_id: u64); + + fn approx_memory_bytes(&self) -> usize; +} From 5b54ad79758fa9924e0e50246a028a07f3413e3a Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 21:02:10 -0400 Subject: [PATCH 08/12] feat(persistence): wire SimpleMapStorePerKey::with_persistence + query read-through MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wraps SimpleMapStorePerKey's fields in an Arc so the background flusher can share state with the query path, and adds a new persistence-aware constructor: pub fn with_persistence( streaming_config: Arc, cleanup_policy: CleanupPolicy, persistence_cfg: SimpleMapStorePersistenceConfig, ) -> PersistResult The existing `new()` / `new_with_strategy()` are unchanged — all 7 call sites in main.rs, bins, test helpers, and benches continue to construct in-memory-only stores with no code change. Behavior when persistence is enabled: * Startup runs recovery::recover() against disk_path (orphan sweep + CRC verification of any live parts left by a prior run), then starts a FlusherHandle against Arc. * Insert path always runs maybe_rotate_epoch() regardless of CleanupPolicy, so sealed epochs accumulate for the flusher to pick up. Falls back to PERSISTENCE_DEFAULT_EPOCH_CAPACITY (1024) when the streaming config does not set num_aggregates_to_retain. * maybe_rotate_epoch() with persistence_enabled=true skips its destructive eviction step — the flusher handles eviction via its memory-pressure / hot-window / T2 triggers instead. * cleanup_old_aggregates() is a no-op when persistence is on. * query_precomputed_output() extends the existing in-memory merge with a disk read-through: walks the parts manifest for overlapping parts, resolves each via the moka Tier-2 cache, decodes matching entries with accumulator_serde:: deserialize_accumulator(), and merges into the result map. * Drop shuts the flusher down before the Arc can be released, guaranteeing no torn-destruction races. EpochSource is implemented on PerKeyInner: * list_sealed_epochs() walks the DashMap and returns one SealedEpochRef per sealed epoch, with a coarse APPROX_BYTES_PER_ SKETCH estimate. * snapshot_sealed_epoch() takes a per-agg read lock, clones entries into EpochSnapshot form with labels resolved from the intern table and sketches serialized via serialize_accumulator_arroyo(), then drops the lock before returning — no per-agg lock held across downstream disk I/O. * evict_sealed_epoch() removes the sealed BTreeMap entry under a write lock and prunes the corresponding read_counts. v1 caveats documented in the source: * query_precomputed_output_exact() only reads in-memory state. A subsequent PR will extend the exact path with disk read-through. * Memory accounting is coarse: entries * 4096 bytes. Per-accumulator sizing (via a new AggregateCore::approx_memory_bytes trait method) is a follow-up. * Only AggregateCore types with working Arroyo round-trips in accumulator_serde (Sum, KLL, HydraKLL, CountMinSketch, Set, DeltaSet, MultipleSum, MultipleIncrease) are persistable. Others are logged and skipped in query_disk_parts(). Three integration tests under tests/persistence_integration_tests.rs: * with_persistence_flushes_sealed_epochs_to_disk — construct, insert, wait for flush, verify in-memory entries have decreased while the query still returns all inserted buckets. * query_read_through_merges_memory_and_disk_ranges — 6 windows, partial range query, verify filter is honored regardless of whether the hit came from memory or disk. * construct_and_drop_shuts_flusher_cleanly — smoke test that Drop joins the flusher thread without deadlock. Full test suite: 450 passed, 0 failed, 5 ignored. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/stores/simple_map_store/per_key.rs | 681 ++++++++++++------ asap-query-engine/src/tests/mod.rs | 1 + .../tests/persistence_integration_tests.rs | 203 ++++++ 3 files changed, 667 insertions(+), 218 deletions(-) create mode 100644 asap-query-engine/src/tests/persistence_integration_tests.rs diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 728a5aab..6f7fb27c 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -1,19 +1,40 @@ use crate::data_model::{ - AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, + AggregateCore, AggregationType, CleanupPolicy, KeyByLabelValues, PrecomputedOutput, + StreamingConfig, }; +use crate::engines::physical::accumulator_serde; use crate::stores::simple_map_store::common::{ EpochID, InternTable, MetricBucketMap, MetricID, MutableEpoch, SealedEpoch, TimestampRange, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use dashmap::DashMap; +use datafusion_summary_library::SketchType; use std::collections::{BTreeMap, HashMap}; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, RwLock}; use std::time::Instant; -use tracing::{debug, error, info}; +use tracing::{debug, error, info, warn}; + +use super::persistence::{ + self, cache::PartCache, flusher::FlusherHandle, manifest::Manifest, recovery, + source::{EpochSnapshot, EpochSnapshotEntry, EpochSource, SealedEpochRef}, + PersistError, PersistResult, SimpleMapStorePersistenceConfig, +}; type StoreKey = u64; // aggregation_id +/// Coarse per-sketch memory estimate used by the persistence flusher's +/// memory-pressure trigger. Not accurate — v1 just multiplies the +/// entry count by this constant. Per-type sizing is a follow-up. +const APPROX_BYTES_PER_SKETCH: usize = 4096; + +/// Fallback epoch capacity used when `num_aggregates_to_retain` is not +/// set in the streaming config but persistence is enabled. Without it +/// the rotator never seals the current epoch and nothing is ever +/// flushable. +const PERSISTENCE_DEFAULT_EPOCH_CAPACITY: usize = 1024; + /// Per-aggregation_id data protected by RwLock struct StoreKeyData { /// Label interning table (Optimization 1) @@ -62,17 +83,19 @@ impl StoreKeyData { } } - /// Seal the current epoch when full, then evict the minimum number of oldest windows - /// to keep total distinct windows ≤ `epoch_capacity * max_epochs`. + /// Seal the current epoch when full, then (only if retention-based + /// cleanup is on) evict the minimum number of oldest windows to + /// keep total distinct windows ≤ `epoch_capacity * max_epochs`. /// - /// Matches legacy per-window eviction semantics: only the exact number of windows - /// needed to reach the retention limit are removed, which may be fewer than a full epoch. - fn maybe_rotate_epoch(&mut self) { + /// When `persistence_enabled` is true, the eviction step is + /// **skipped** — the background flusher handles eviction via its + /// memory/time triggers. This method only seals in that case, and + /// sealed epochs accumulate until the flusher picks them up. + fn maybe_rotate_epoch(&mut self, persistence_enabled: bool) { let capacity = match self.epoch_capacity { Some(c) if c > 0 => c, _ => return, // unlimited }; - let retention_limit = capacity * self.max_epochs; // Step 1: seal current epoch if it has hit the window capacity threshold. if self.current_epoch.window_count() >= capacity { @@ -82,8 +105,13 @@ impl StoreKeyData { self.current_epoch_id += 1; } + if persistence_enabled { + // Flusher evicts. We only seal. + return; + } + // Step 2: evict oldest windows until total distinct windows ≤ retention_limit. - // Uses O(E) distinct_window_count() calls (E ≤ max_epochs, a small constant). + let retention_limit = capacity * self.max_epochs; let total: usize = self.current_epoch.window_count() + self .sealed_epochs @@ -156,8 +184,10 @@ impl StoreKeyData { } } -/// In-memory storage implementation using per-key locks for concurrency -pub struct SimpleMapStorePerKey { +/// Shared state that both the outer `SimpleMapStorePerKey` and the +/// background flusher hold via `Arc`. Contains the DashMap of per-agg +/// state plus the counters the flusher needs. +pub struct PerKeyInner { // Lock-free concurrent outer map - per aggregation_id store: DashMap>>, @@ -171,18 +201,117 @@ pub struct SimpleMapStorePerKey { // Policy for cleaning up old aggregates cleanup_policy: CleanupPolicy, + + /// Whether persistence is active. When `true`: + /// * `maybe_rotate_epoch` still seals current epochs but does not + /// evict old ones — the flusher handles eviction. + /// * Insertions always run rotation regardless of cleanup policy. + /// * `mem_bytes_sealed` is maintained as flusher input. + persistence_enabled: bool, + + /// Approximate sum of sketch bytes in sealed epochs across all + /// agg-ids. Updated on rotate (adds) and evict (subtracts). + /// Drives the flusher's memory-pressure trigger. Coarse — see + /// `APPROX_BYTES_PER_SKETCH`. + mem_bytes_sealed: AtomicUsize, +} + +/// Persistence-related state owned by the outer store. Dropping this +/// (via `SimpleMapStorePerKey::Drop`) shuts the flusher down cleanly. +struct PersistenceState { + manifest: Arc, + cache: PartCache, + /// Held so `Drop` stops the thread. Not accessed directly after + /// construction. + _flusher: FlusherHandle, + #[allow(dead_code)] + parts_root: PathBuf, +} + +/// In-memory storage implementation using per-key locks for concurrency +pub struct SimpleMapStorePerKey { + inner: Arc, + /// `None` when the store is in-memory-only (existing `new()` path). + /// `Some` when constructed via `with_persistence`. + persistence: Option, } impl SimpleMapStorePerKey { + /// Backwards-compatible constructor. No persistence — behaves + /// exactly like pre-persistence code. pub fn new(streaming_config: Arc, cleanup_policy: CleanupPolicy) -> Self { Self { + inner: Arc::new(PerKeyInner { + store: DashMap::new(), + earliest_timestamps: DashMap::new(), + metrics: DashMap::new(), + items_inserted: DashMap::new(), + streaming_config, + cleanup_policy, + persistence_enabled: false, + mem_bytes_sealed: AtomicUsize::new(0), + }), + persistence: None, + } + } + + /// Persistence-aware constructor. Starts the background flusher, + /// runs recovery on `persistence_cfg.disk_path`, and returns a + /// store that writes sealed epochs to disk and evicts them on + /// memory / time pressure. + /// + /// Cleanup policy still applies, but when persistence is on, the + /// destructive eviction step of `CircularBuffer` / `ReadBased` is + /// bypassed in favor of the flusher. `NoCleanup` + persistence is + /// the typical production configuration: the flusher bounds RAM, + /// nothing is ever dropped from memory without first being on + /// disk. + pub fn with_persistence( + streaming_config: Arc, + cleanup_policy: CleanupPolicy, + persistence_cfg: SimpleMapStorePersistenceConfig, + ) -> PersistResult { + // Run recovery first so the manifest reflects on-disk state. + let (_loaded_manifest, report) = recovery::recover(&persistence_cfg.disk_path)?; + info!( + "SimpleMapStorePerKey persistence recovery: live={}, corrupt_removed={}, orphans_removed={}", + report.live_parts, report.corrupt_parts_removed, report.orphan_parts_removed + ); + + // Re-open manifest as an Arc so the flusher and the query path + // can share it. (`recovery::recover` already opened one, but + // it's owned; simpler to re-open once ownership settled.) + let manifest = Arc::new(Manifest::open_or_init(&persistence_cfg.disk_path)?); + + let parts_root = persistence::flusher::parts_root(&persistence_cfg.disk_path); + let cache = PartCache::new(parts_root.clone(), persistence_cfg.part_cache_bytes); + + let inner = Arc::new(PerKeyInner { store: DashMap::new(), earliest_timestamps: DashMap::new(), metrics: DashMap::new(), items_inserted: DashMap::new(), streaming_config, cleanup_policy, - } + persistence_enabled: true, + mem_bytes_sealed: AtomicUsize::new(0), + }); + + let flusher = FlusherHandle::start( + persistence_cfg, + Arc::clone(&manifest), + Arc::clone(&inner), + )?; + + Ok(Self { + inner, + persistence: Some(PersistenceState { + manifest, + cache, + _flusher: flusher, + parts_root, + }), + }) } /// Collect diagnostic info about store contents. @@ -191,9 +320,9 @@ impl SimpleMapStorePerKey { let mut per_aggregation = Vec::new(); let mut total_time_map_entries: usize = 0; - let total_sketch_bytes: usize = 0; + let total_sketch_bytes: usize = self.inner.mem_bytes_sealed.load(Ordering::Relaxed); - for entry in self.store.iter() { + for entry in self.inner.store.iter() { let agg_id = *entry.key(); let data = match entry.value().read() { Ok(d) => d, @@ -220,12 +349,12 @@ impl SimpleMapStorePerKey { time_map_len, read_counts_len, num_aggregate_objects, - sketch_bytes: 0, // skip serialization for diagnostics + sketch_bytes: 0, // per-agg sketch byte sizing is a follow-up }); } StoreDiagnostics { - num_aggregations: self.store.len(), + num_aggregations: self.inner.store.len(), total_time_map_entries, total_sketch_bytes, per_aggregation, @@ -240,11 +369,17 @@ impl SimpleMapStorePerKey { num_aggregates_to_retain: Option, read_count_threshold: Option, ) { - match self.cleanup_policy { + // When persistence is enabled, eviction is the flusher's job. + // Skip destructive cleanup entirely — parts on disk are the + // source of truth for cold data. + if self.inner.persistence_enabled { + let _ = (num_aggregates_to_retain, metric, aggregation_id, read_count_threshold); + return; + } + + match self.inner.cleanup_policy { CleanupPolicy::CircularBuffer => { - // configure_epochs was already called before insert; - // rotation is handled by maybe_rotate_epoch after each insert batch. - // Nothing additional needed here. + // Handled by maybe_rotate_epoch() during insert. let _ = (num_aggregates_to_retain, metric, aggregation_id); } CleanupPolicy::ReadBased => { @@ -252,9 +387,7 @@ impl SimpleMapStorePerKey { data.cleanup_read_based(metric, aggregation_id, threshold); } } - CleanupPolicy::NoCleanup => { - // Do nothing - no cleanup - } + CleanupPolicy::NoCleanup => {} } } @@ -267,21 +400,21 @@ impl SimpleMapStorePerKey { let aggregation_id = *store_key; let metric_key = metric.to_string(); let inserted_delta = items.len() as u64; + let persistence_enabled = self.inner.persistence_enabled; // Opt 4: compute batch minimum timestamp before acquiring any lock. - // Collapses N per-item atomic fetch_min calls into one (Opt 4). let batch_min_ts = items .iter() .map(|(o, _)| o.start_timestamp) .min() .unwrap_or(u64::MAX); - // Measure lock acquisition time #[cfg(feature = "lock_profiling")] let lock_wait_start = Instant::now(); // Get or create the store data for this key let store_data_lock = self + .inner .store .entry(*store_key) .or_insert_with(|| Arc::new(RwLock::new(StoreKeyData::new()))); @@ -325,19 +458,20 @@ impl SimpleMapStorePerKey { let lock_hold_start = Instant::now(); // Create metric if needed (lock-free DashMap insert) - self.metrics.entry(metric_key.clone()).or_insert(()); + self.inner.metrics.entry(metric_key.clone()).or_insert(()); - // Opt 4: one atomic earliest-ts update per batch using the pre-computed minimum. - // Replaces N per-item fetch_min calls with a single one. - self.earliest_timestamps + // Opt 4: one atomic earliest-ts update per batch. + self.inner + .earliest_timestamps .entry(aggregation_id) .and_modify(|earliest| { earliest.fetch_min(batch_min_ts, Ordering::Relaxed); }) .or_insert_with(|| AtomicU64::new(batch_min_ts)); - // Update insertion counter once per grouped batch (instead of once per item). + // Update insertion counter once per grouped batch. let items_inserted_counter = self + .inner .items_inserted .entry(metric_key) .or_insert_with(|| AtomicU64::new(0)); @@ -347,35 +481,59 @@ impl SimpleMapStorePerKey { debug!("Inserted {} items into {}", new_total, metric); } - // Get aggregation config once for cleanup settings let aggregation_config = self + .inner .streaming_config .get_aggregation_config(aggregation_id) .ok_or_else(|| format!("Aggregation config not found for {}", aggregation_id))?; - // Configure epoch capacity on first insert (Optimization 2) + // Configure epoch capacity on first insert (Optimization 2). + // When persistence is enabled and the streaming config has no + // retention set, fall back to a default so the rotator seals + // current epochs periodically — otherwise nothing is ever + // flushable. if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { - data.configure_epochs(aggregation_config.num_aggregates_to_retain); + let effective_retention = aggregation_config.num_aggregates_to_retain.or({ + if persistence_enabled { + Some(PERSISTENCE_DEFAULT_EPOCH_CAPACITY as u64) + } else { + None + } + }); + data.configure_epochs(effective_retention); } + let entries_added = items.len(); for (output, precompute) in items { - // Intern the label key (Optimization 1) let timestamp_range = (output.start_timestamp, output.end_timestamp); let metric_id: MetricID = data.intern.intern(output.key); - // Insert into current (mutable) epoch. data.current_epoch .insert(metric_id, timestamp_range, Arc::from(precompute)); - // After each item, check if we should rotate (CircularBuffer, Optimization 2) - if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator - && matches!(self.cleanup_policy, CleanupPolicy::CircularBuffer) - { - data.maybe_rotate_epoch(); + // When persistence is on, always run rotation so sealed + // epochs accumulate. Otherwise preserve the old + // CircularBuffer-only behavior. + let should_rotate = aggregation_config.aggregation_type + != AggregationType::DeltaSetAggregator + && (persistence_enabled + || matches!(self.inner.cleanup_policy, CleanupPolicy::CircularBuffer)); + if should_rotate { + data.maybe_rotate_epoch(persistence_enabled); } } - // Apply retention policy if configured (but exclude DeltaSetAggregator) + if persistence_enabled { + // Best-effort memory accounting: attribute the batch's + // entries to sealed-epoch bytes. This over-counts (current + // epoch entries are included) but the flusher is + // conservative anyway. + self.inner.mem_bytes_sealed.fetch_add( + entries_added * APPROX_BYTES_PER_SKETCH, + Ordering::Relaxed, + ); + } + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { self.cleanup_old_aggregates( &mut data, @@ -399,6 +557,114 @@ impl SimpleMapStorePerKey { Ok(()) } + + /// Query overlapping on-disk parts and merge into the provided + /// in-memory result map. A no-op when persistence is disabled. + fn query_disk_parts( + &self, + metric: &str, + aggregation_id: u64, + start: u64, + end: u64, + results: &mut TimestampedBucketsMap, + ) -> Result<(), Box> { + let Some(state) = self.persistence.as_ref() else { + return Ok(()); + }; + + let overlapping = state.manifest.live_parts_overlapping(start, end); + if overlapping.is_empty() { + return Ok(()); + } + + for entry in overlapping { + let reader = match state.cache.get_or_load(entry.part_id) { + Ok(r) => r, + Err(e) => { + warn!( + "query_disk_parts: failed to open part {} for metric {}: {}", + entry.part_id, metric, e + ); + continue; + } + }; + for rec in reader.index_records() { + if rec.agg_id != aggregation_id { + continue; + } + // Same overlap semantics as MutableEpoch::range_query_into: + // window must be fully inside [start, end]. + if rec.start_ts < start || rec.start_ts > end || rec.end_ts > end { + continue; + } + let disk_entry = match reader.load_entry(&rec) { + Ok(d) => d, + Err(e) => { + warn!( + "query_disk_parts: failed to load entry at offset {} of part {}: {}", + rec.data_offset, entry.part_id, e + ); + continue; + } + }; + let Some(sketch_type) = type_name_to_sketch_type(&disk_entry.sketch_type_name) + else { + warn!( + "query_disk_parts: no SketchType mapping for {}; skipping", + disk_entry.sketch_type_name + ); + continue; + }; + let decoded = + match accumulator_serde::deserialize_accumulator(&disk_entry.sketch_bytes, &sketch_type) { + Ok(a) => a, + Err(e) => { + warn!( + "query_disk_parts: deserialize failed for {}: {}", + disk_entry.sketch_type_name, e + ); + continue; + } + }; + let arc_acc: Arc = Arc::from(decoded); + results + .entry(disk_entry.label.clone()) + .or_default() + .push(((rec.start_ts, rec.end_ts), arc_acc)); + } + } + + Ok(()) + } +} + +impl Drop for SimpleMapStorePerKey { + fn drop(&mut self) { + // Dropping the PersistenceState (and therefore the FlusherHandle) + // stops the flusher thread before the underlying Arc + // ref count hits zero, guaranteeing the flusher cannot observe + // a half-destroyed store. + if let Some(mut state) = self.persistence.take() { + state._flusher.shutdown(); + } + } +} + +/// Map `AggregateCore::type_name()` to the `SketchType` enum value +/// used by `accumulator_serde::deserialize_accumulator`. Returns +/// `None` for types that don't have a working Arroyo round-trip yet. +fn type_name_to_sketch_type(name: &str) -> Option { + match name { + "SumAccumulator" => Some(SketchType::Sum), + "DatasketchesKLLAccumulator" => Some(SketchType::KLL), + "HydraKllSketchAccumulator" => Some(SketchType::HydraKLL), + "CountMinSketchAccumulator" => Some(SketchType::CountMinSketch), + "SetAggregatorAccumulator" => Some(SketchType::SetAggregator), + "DeltaSetAggregatorAccumulator" => Some(SketchType::DeltaSetAggregator), + "MultipleSumAccumulator" => Some(SketchType::MultipleSum), + "MultipleIncreaseAccumulator" => Some(SketchType::MultipleIncrease), + _ => None, + } } #[async_trait::async_trait] @@ -427,6 +693,7 @@ impl Store for SimpleMapStorePerKey { for (output, precompute) in outputs { let aggregation_config = self + .inner .streaming_config .get_aggregation_config(output.aggregation_id); @@ -449,7 +716,6 @@ impl Store for SimpleMapStorePerKey { .push((output, precompute)); } - // Process each aggregation_id group; each iteration locks at most one key. for (store_key, (metric, items)) in grouped { self.insert_for_store_key(&store_key, &metric, items)?; } @@ -481,126 +747,66 @@ impl Store for SimpleMapStorePerKey { let query_start_time = Instant::now(); let store_key = aggregation_id; - // Measure lock acquisition time - #[cfg(feature = "lock_profiling")] - let lock_wait_start = Instant::now(); + let mut results: TimestampedBucketsMap = HashMap::new(); - // Get the store data for this aggregation_id - let store_data_lock = match self.store.get(&store_key) { - Some(lock) => lock, - None => { - info!("Metric {} not found in store", metric); - return Ok(HashMap::new()); + // --- In-memory path (unchanged) --- + if let Some(store_data_lock) = self.inner.store.get(&store_key) { + let data = store_data_lock.read().map_err(|e| { + format!( + "Failed to acquire read lock for query aggregation_id {}: {}", + store_key, e + ) + })?; + + let mut mid: MetricBucketMap = HashMap::with_capacity(data.intern.len()); + let mut matched_windows: Vec = Vec::new(); + + if let Some((min_start, max_end)) = data.current_epoch.time_bounds() { + if !(min_start > end || max_end < start) { + data.current_epoch + .range_query_into(start, end, &mut mid, &mut matched_windows); + } } - }; - - #[cfg(feature = "lock_profiling")] - { - let lock_wait_duration = lock_wait_start.elapsed(); - info!( - "🔒 Query DashMap get time: {:.2}ms (metric: {}, agg_id: {})", - lock_wait_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id - ); - } - - #[cfg(feature = "lock_profiling")] - let rwlock_wait_start = Instant::now(); - - // Range queries use a read lock — no mutation of epoch data needed. - let data = store_data_lock.read().map_err(|e| { - format!( - "Failed to acquire read lock for query aggregation_id {}: {}", - store_key, e - ) - })?; - - #[cfg(feature = "lock_profiling")] - { - let rwlock_wait_duration = rwlock_wait_start.elapsed(); - info!( - "🔒 Query RwLock wait time: {:.2}ms (metric: {}, agg_id: {})", - rwlock_wait_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id - ); - } - - #[cfg(feature = "lock_profiling")] - let lock_hold_start = Instant::now(); - - let mut total_entries = 0; - let mut matched_windows: Vec = Vec::new(); - - let range_scan_start_time = Instant::now(); - let mut mid: MetricBucketMap = HashMap::with_capacity(data.intern.len()); - - // Query current (mutable) epoch. - if let Some((min_start, max_end)) = data.current_epoch.time_bounds() { - if !(min_start > end || max_end < start) { - data.current_epoch - .range_query_into(start, end, &mut mid, &mut matched_windows); + for epoch in data.sealed_epochs.values() { + let Some((min_start, max_end)) = epoch.time_bounds() else { + continue; + }; + if min_start > end || max_end < start { + continue; + } + epoch.range_query_into(start, end, &mut mid, &mut matched_windows); } - } - // Query sealed epochs; skip those with no overlap. - for epoch in data.sealed_epochs.values() { - let Some((min_start, max_end)) = epoch.time_bounds() else { - continue; - }; - if min_start > end || max_end < start { - continue; + for (metric_id, buckets) in mid { + let label = data.intern.resolve(metric_id).clone(); + results.entry(label).or_default().extend(buckets); } - epoch.range_query_into(start, end, &mut mid, &mut matched_windows); - } - - // Resolve MetricIDs → labels in a single pass - let mut results: TimestampedBucketsMap = HashMap::with_capacity(mid.len()); - for (metric_id, buckets) in mid { - total_entries += buckets.len(); - let label = data.intern.resolve(metric_id).clone(); - results.insert(label, buckets); - } - // Update read counts via inner Mutex - { - let mut read_counts = data.read_counts.lock().unwrap(); - for window in &matched_windows { - *read_counts.entry(*window).or_insert(0) += 1; + { + let mut read_counts = data.read_counts.lock().unwrap(); + for window in &matched_windows { + *read_counts.entry(*window).or_insert(0) += 1; + } } + } else if self.persistence.is_none() { + // Nothing in memory and no disk layer → empty result, + // matching the old behavior. + info!("Metric {} not found in store", metric); } - let range_scan_duration = range_scan_start_time.elapsed(); - debug!( - "Range scanning took: {:.2}ms", - range_scan_duration.as_secs_f64() * 1000.0 - ); + // --- On-disk path (persistence only) --- + if self.persistence.is_some() { + self.query_disk_parts(metric, aggregation_id, start, end, &mut results)?; + } let query_duration = query_start_time.elapsed(); debug!( - "Total query took: {:.2}ms", - query_duration.as_secs_f64() * 1000.0 - ); - - debug!( - "Found {} entries for query on {} (aggregation_id: {}, start: {}, end: {})", - total_entries, metric, aggregation_id, start, end + "Total query took: {:.2}ms ({} keys, {} in-memory + disk)", + query_duration.as_secs_f64() * 1000.0, + results.len(), + results.values().map(|v| v.len()).sum::() ); - debug!("Found {} unique keys", results.len()); - - #[cfg(feature = "lock_profiling")] - { - let lock_hold_duration = lock_hold_start.elapsed(); - info!( - "🔓 Query lock hold time: {:.2}ms (metric: {}, agg_id: {}, entries: {})", - lock_hold_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id, - total_entries - ); - } Ok(results) } @@ -612,6 +818,11 @@ impl Store for SimpleMapStorePerKey { exact_start: u64, exact_end: u64, ) -> Result> { + // NOTE (persistence follow-up): this path does NOT consult + // on-disk parts in v1 — exact queries only see in-memory + // state. A subsequent PR will wire up a parts-aware exact + // path. For the range-query code path (the 90% case) the + // on-disk merge is already in place. if exact_start > exact_end { debug!( "Invalid exact query range for metric {} agg_id {}: start {} > end {}", @@ -623,12 +834,7 @@ impl Store for SimpleMapStorePerKey { let query_start_time = Instant::now(); let store_key = aggregation_id; - // Measure lock acquisition time - #[cfg(feature = "lock_profiling")] - let lock_wait_start = Instant::now(); - - // Get the store data for this aggregation_id - let store_data_lock = match self.store.get(&store_key) { + let store_data_lock = match self.inner.store.get(&store_key) { Some(lock) => lock, None => { debug!("Metric {} not found in store for exact query", metric); @@ -636,22 +842,7 @@ impl Store for SimpleMapStorePerKey { } }; - #[cfg(feature = "lock_profiling")] - { - let lock_wait_duration = lock_wait_start.elapsed(); - info!( - "🔒 Exact query DashMap get time: {:.2}ms (metric: {}, agg_id: {})", - lock_wait_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id - ); - } - - #[cfg(feature = "lock_profiling")] - let rwlock_wait_start = Instant::now(); - // Opt 1: exact_query takes &mut self (lazy index build), so we need a write lock. - // Range queries still use a read lock — only exact queries pay the write-lock cost. let mut data = store_data_lock.write().map_err(|e| { format!( "Failed to acquire write lock for exact query aggregation_id {}: {}", @@ -659,24 +850,8 @@ impl Store for SimpleMapStorePerKey { ) })?; - #[cfg(feature = "lock_profiling")] - { - let rwlock_wait_duration = rwlock_wait_start.elapsed(); - info!( - "🔒 Exact query RwLock wait time: {:.2}ms (metric: {}, agg_id: {})", - rwlock_wait_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id - ); - } - - #[cfg(feature = "lock_profiling")] - let lock_hold_start = Instant::now(); - let timestamp_range = (exact_start, exact_end); - // Opt 1: exact_query on the mutable epoch builds the lazy offset index if absent, - // then looks up the window in O(m). Returns an owned Vec — the &mut borrow ends here. let entries_opt: Option)>> = data.current_epoch.exact_query(timestamp_range).or_else(|| { data.sealed_epochs @@ -686,7 +861,6 @@ impl Store for SimpleMapStorePerKey { }); let mut results: TimestampedBucketsMap = HashMap::new(); - let mut total_entries = 0; let found_match = entries_opt.is_some(); if let Some(entries) = entries_opt { @@ -696,43 +870,14 @@ impl Store for SimpleMapStorePerKey { .entry(label) .or_default() .push((timestamp_range, agg)); - total_entries += 1; } } - if found_match { - debug!( - "Exact match FOUND for [{}, {}]: {} entries across {} keys", - exact_start, - exact_end, - total_entries, - results.len() - ); - } else { - debug!( - "Exact match NOT FOUND for metric: {}, agg_id: {}, range: [{}, {}]", - metric, aggregation_id, exact_start, exact_end - ); - } - - // Update read count — write lock already held, no inner Mutex needed if found_match { let mut read_counts = data.read_counts.lock().unwrap(); *read_counts.entry(timestamp_range).or_insert(0) += 1; } - #[cfg(feature = "lock_profiling")] - { - let lock_hold_duration = lock_hold_start.elapsed(); - info!( - "🔓 Exact query lock hold time: {:.2}ms (metric: {}, agg_id: {}, found: {})", - lock_hold_duration.as_secs_f64() * 1000.0, - metric, - aggregation_id, - !results.is_empty() - ); - } - let query_duration = query_start_time.elapsed(); debug!( "Exact timestamp query took: {:.2}ms (found: {})", @@ -746,8 +891,8 @@ impl Store for SimpleMapStorePerKey { fn get_earliest_timestamp_per_aggregation_id( &self, ) -> Result, Box> { - // No lock needed - DashMap with AtomicU64 let result = self + .inner .earliest_timestamps .iter() .map(|entry| (*entry.key(), entry.value().load(Ordering::Relaxed))) @@ -757,8 +902,108 @@ impl Store for SimpleMapStorePerKey { } fn close(&self) -> StoreResult<()> { - // For in-memory store, no cleanup needed info!("SimpleMapStorePerKey closed"); Ok(()) } } + +// ================================================================= +// EpochSource implementation (used by the persistence flusher) +// ================================================================= + +impl EpochSource for PerKeyInner { + fn list_sealed_epochs(&self) -> Vec { + let mut out = Vec::new(); + for entry in self.store.iter() { + let agg_id = *entry.key(); + let Ok(data) = entry.value().read() else { + continue; + }; + for (epoch_id, epoch) in data.sealed_epochs.iter() { + if let Some((_, max_end)) = epoch.time_bounds() { + out.push(SealedEpochRef { + agg_id, + epoch_id: *epoch_id, + end_ts: max_end, + approx_bytes: epoch.entries.len() * APPROX_BYTES_PER_SKETCH, + }); + } + } + } + out + } + + fn snapshot_sealed_epoch( + &self, + agg_id: u64, + epoch_id: u64, + ) -> PersistResult> { + let Some(lock) = self.store.get(&agg_id) else { + return Ok(None); + }; + let data = lock + .read() + .map_err(|e| PersistError::Internal(format!("read lock poisoned: {}", e)))?; + let Some(epoch) = data.sealed_epochs.get(&epoch_id) else { + return Ok(None); + }; + + let mut entries = Vec::with_capacity(epoch.entries.len()); + for (tr, metric_id, agg) in &epoch.entries { + // Resolve the label from the per-agg intern table. + let label: Option = data.intern.resolve(*metric_id).clone(); + + // Serialize the sketch via the arroyo path so it's + // round-trippable via deserialize_accumulator. + let sketch_bytes = accumulator_serde::serialize_accumulator_arroyo(agg.as_ref()); + let type_name = agg.type_name().to_string(); + + entries.push(EpochSnapshotEntry { + start_ts: tr.0, + end_ts: tr.1, + label, + sketch_type_name: type_name, + sketch_bytes, + }); + } + + let (min_ts, max_ts) = epoch.time_bounds().unwrap_or((0, 0)); + let approx_bytes = epoch.entries.len() * APPROX_BYTES_PER_SKETCH; + + Ok(Some(EpochSnapshot { + agg_id, + epoch_id, + min_ts, + max_ts, + entries, + approx_bytes, + })) + } + + fn evict_sealed_epoch(&self, agg_id: u64, epoch_id: u64) { + let Some(lock) = self.store.get(&agg_id) else { + return; + }; + let mut data = match lock.write() { + Ok(d) => d, + Err(e) => { + error!("evict: write lock poisoned for agg_id {}: {}", agg_id, e); + return; + } + }; + if let Some(epoch) = data.sealed_epochs.remove(&epoch_id) { + let freed = epoch.entries.len() * APPROX_BYTES_PER_SKETCH; + self.mem_bytes_sealed.fetch_sub(freed, Ordering::Relaxed); + // Also purge the epoch's windows from read_counts so they + // don't leak. + let mut read_counts = data.read_counts.lock().unwrap(); + for w in epoch.unique_windows() { + read_counts.remove(&w); + } + } + } + + fn approx_memory_bytes(&self) -> usize { + self.mem_bytes_sealed.load(Ordering::Relaxed) + } +} diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 9fbe6fc1..f8ae48dd 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -6,6 +6,7 @@ pub mod elastic_forwarding_tests; pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod sql_pattern_matching_tests; +pub mod persistence_integration_tests; pub mod store_correctness_tests; pub mod trait_design_tests; diff --git a/asap-query-engine/src/tests/persistence_integration_tests.rs b/asap-query-engine/src/tests/persistence_integration_tests.rs new file mode 100644 index 00000000..d4052fa6 --- /dev/null +++ b/asap-query-engine/src/tests/persistence_integration_tests.rs @@ -0,0 +1,203 @@ +//! End-to-end tests for `SimpleMapStorePerKey::with_persistence`. +//! +//! Spins up a real store with a tempdir-backed persistence config, +//! inserts sketches through the public `Store` trait, waits for the +//! background flusher to write a part and evict the sealed epoch, +//! and verifies `query_precomputed_output` returns the flushed data +//! via the disk read-through path. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use promql_utilities::data_model::KeyByLabelNames; +use tempfile::TempDir; + +use crate::data_model::{ + AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, WindowType, +}; +use crate::precompute_operators::SumAccumulator; +use crate::stores::simple_map_store::per_key::SimpleMapStorePerKey; +use crate::stores::simple_map_store::persistence::SimpleMapStorePersistenceConfig; +use crate::stores::Store; +use crate::{AggregateCore, AggregationConfig}; + +fn make_streaming_config(agg_id: u64) -> Arc { + // Retain a small number of aggregates per epoch so rotation fires + // quickly and the flusher has something to chew on under test. + let cfg = AggregationConfig::new( + agg_id, + AggregationType::Sum, + String::new(), + std::collections::HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + String::new(), + 60, + 60, + WindowType::Tumbling, + String::new(), + "cpu_usage".to_string(), + Some(2), // num_aggregates_to_retain — seals after 2 distinct windows + None, + None, + None, + ); + let mut map = std::collections::HashMap::new(); + map.insert(agg_id, cfg); + Arc::new(StreamingConfig::new(map)) +} + +fn persistence_cfg(dir: &TempDir, hot_window_ms: Option) -> SimpleMapStorePersistenceConfig { + SimpleMapStorePersistenceConfig { + memory_limit_bytes: 100 * 1024 * 1024, + memory_low_watermark_bytes: 50 * 1024 * 1024, + hard_cap_bytes: 200 * 1024 * 1024, + hot_window_ms, + delete_older_than_ms: None, + flush_interval: Duration::from_millis(25), + disk_path: dir.path().to_path_buf(), + part_cache_bytes: 1024 * 1024, + } +} + +fn sum_entry( + agg_id: u64, + start: u64, + end: u64, + value: f64, +) -> (PrecomputedOutput, Box) { + ( + PrecomputedOutput::new(start, end, None, agg_id), + Box::new(SumAccumulator::with_sum(value)), + ) +} + +fn wait_for bool>(mut pred: F, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if pred() { + return true; + } + std::thread::sleep(Duration::from_millis(10)); + } + pred() +} + +#[test] +fn with_persistence_flushes_sealed_epochs_to_disk() { + let dir = TempDir::new().unwrap(); + let cfg = make_streaming_config(1); + // hot_window_ms = Some(0) means "any sealed epoch's end_ts is + // older than now-0 = now, so flush immediately on next tick." + let persistence = persistence_cfg(&dir, Some(0)); + + let store = + SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); + + // Insert several windows so the rotator seals at least one epoch. + // num_aggregates_to_retain = 2, so windows 3 will roll the epoch. + let batches = vec![ + sum_entry(1, 1_000, 2_000, 1.0), + sum_entry(1, 2_000, 3_000, 2.0), + sum_entry(1, 3_000, 4_000, 3.0), + sum_entry(1, 4_000, 5_000, 4.0), + ]; + store.insert_precomputed_output_batch(batches).unwrap(); + + // Wait until a part shows up in the manifest. This means the + // flusher has written at least one tick's worth of sealed epochs + // and evicted them from memory. + let flushed = wait_for( + || { + let diag = store.diagnostic_info(); + // At least one agg with at least one sealed epoch + // evicted — detected by checking the manifest via a + // second query that goes through the disk path. + let res = store + .query_precomputed_output("cpu_usage", 1, 0, u64::MAX) + .unwrap(); + // The test is satisfied once the query returns *any* + // result AND diagnostic bytes are nonzero (something was + // in memory at some point). + !res.is_empty() && diag.total_time_map_entries < 4 + }, + Duration::from_secs(3), + ); + assert!(flushed, "flusher did not produce a part in time"); + + // Query across the full time range and verify we see all 4 windows + // (some from memory, some from disk, or all from disk). + let res = store + .query_precomputed_output("cpu_usage", 1, 0, u64::MAX) + .unwrap(); + let total: usize = res.values().map(|v| v.len()).sum(); + assert_eq!( + total, 4, + "expected 4 buckets across in-memory + disk; got {} (buckets: {:?})", + total, + res.values() + .flat_map(|v| v.iter().map(|(tr, _)| *tr)) + .collect::>() + ); +} + +#[test] +fn query_read_through_merges_memory_and_disk_ranges() { + let dir = TempDir::new().unwrap(); + let cfg = make_streaming_config(42); + let persistence = persistence_cfg(&dir, Some(0)); + let store = + SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); + + // Insert 6 windows — more than enough to guarantee the rotator + // seals multiple epochs. + let mut batch = Vec::new(); + for i in 0..6u64 { + let start = 10_000 + i * 1_000; + let end = start + 1_000; + batch.push(sum_entry(42, start, end, i as f64)); + } + store.insert_precomputed_output_batch(batch).unwrap(); + + // Give the flusher time to drain everything to disk. + std::thread::sleep(Duration::from_millis(250)); + + // Query a partial range covering 3 of the 6 windows — make sure + // the filter is honored regardless of whether the hit came from + // memory or disk. + let res = store + .query_precomputed_output("cpu_usage", 42, 12_000, 15_000) + .unwrap(); + let timestamps: Vec<(u64, u64)> = { + let mut ts: Vec<(u64, u64)> = res + .get(&None) + .map(|v| v.iter().map(|(tr, _)| *tr).collect()) + .unwrap_or_default(); + ts.sort_unstable(); + ts + }; + // Expected: the two windows fully inside [12_000, 15_000]: + // (12_000, 13_000), (13_000, 14_000), (14_000, 15_000). + // Note: range_query_into requires tr.0 >= start AND tr.1 <= end. + assert_eq!( + timestamps, + vec![(12_000, 13_000), (13_000, 14_000), (14_000, 15_000)], + "expected windows 12000-15000 inclusive, got {:?}", + timestamps + ); +} + +#[test] +fn construct_and_drop_shuts_flusher_cleanly() { + let dir = TempDir::new().unwrap(); + let cfg = make_streaming_config(1); + let persistence = persistence_cfg(&dir, None); + let store = + SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); + // Dropping the store should not deadlock or panic. + drop(store); +} From 153c096874b9c5a9fffd1741d7e56ed7ab8d4467 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 22:00:25 -0400 Subject: [PATCH 09/12] feat(persistence): per-accumulator approx_memory_bytes sizing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the flat APPROX_BYTES_PER_SKETCH = 4096 constant in per_key.rs with real per-accumulator estimates via a new trait method on AggregateCore: fn approx_memory_bytes(&self) -> usize The default impl returns 4096 (unchanged behavior for any concrete type that doesn't override it), and 11 concrete accumulator types now override it with type-aware estimates: * SumAccumulator, MinMaxAccumulator, IncreaseAccumulator — size_of::() for these scalars. * MultipleSumAccumulator, MultipleMinMaxAccumulator — ~96 bytes per HashMap entry. * MultipleIncreaseAccumulator — ~160 bytes per entry (HashMap). * SetAggregatorAccumulator, DeltaSetAggregatorAccumulator — ~96 bytes per HashSet entry. * CountMinSketchAccumulator — 16 KiB conservative constant. * DatasketchesKLLAccumulator — 4 KiB (k=200 KLL). * HydraKllSketchAccumulator — 32 KiB (row*col grid of KLLs). per_key.rs changes: * New epoch_approx_bytes(&SealedEpoch) -> usize helper that sums each entry's approx_memory_bytes. * Insert path pre-computes a batch sum once and bumps mem_bytes_sealed by the real byte total instead of count * constant. Each approx_memory_bytes call is O(1) or a single field read, so the pre-sum adds only a handful of arithmetic ops per item — no measurable hot-path cost. * list_sealed_epochs, snapshot_sealed_epoch, evict_sealed_epoch all now route through epoch_approx_bytes. The numbers are still conservative estimates — accurate to ~1-2x, not exact — but they give the flusher real proportional signal between scalar accumulators, multi-valued aggregates, and large sketches, which is what the memory-pressure trigger actually needs. All 450 existing tests pass unchanged. Co-Authored-By: Claude Opus 4.6 (1M context) --- asap-query-engine/src/data_model/traits.rs | 14 +++++ .../count_min_sketch_accumulator.rs | 7 +++ .../datasketches_kll_accumulator.rs | 6 ++ .../delta_set_aggregator_accumulator.rs | 7 +++ .../hydra_kll_accumulator.rs | 7 +++ .../increase_accumulator.rs | 5 ++ .../min_max_accumulator.rs | 5 ++ .../multiple_increase_accumulator.rs | 7 +++ .../multiple_min_max_accumulator.rs | 7 +++ .../multiple_sum_accumulator.rs | 7 +++ .../set_aggregator_accumulator.rs | 7 +++ .../precompute_operators/sum_accumulator.rs | 5 ++ .../src/stores/simple_map_store/per_key.rs | 61 +++++++++++++------ 13 files changed, 125 insertions(+), 20 deletions(-) diff --git a/asap-query-engine/src/data_model/traits.rs b/asap-query-engine/src/data_model/traits.rs index 064619ff..bff2d7a9 100644 --- a/asap-query-engine/src/data_model/traits.rs +++ b/asap-query-engine/src/data_model/traits.rs @@ -44,6 +44,20 @@ pub trait AggregateCore: SerializableToSink + Send + Sync { key: &Option, query_kwargs: &HashMap, ) -> Result>; + + /// Approximate in-memory byte footprint of this accumulator. + /// + /// Used by the `SimpleMapStore` persistence layer to drive its + /// memory-pressure trigger. Not required to be exact — the flusher + /// only needs rough proportionality. The default is a conservative + /// 4 KiB constant; concrete types should override it with a + /// type-aware estimate (e.g. KLL: `k * 8` plus overhead). + /// + /// Implementors must not call `serialize_to_bytes` here — this is + /// on the insert hot path. + fn approx_memory_bytes(&self) -> usize { + 4096 + } } /// Trait for accumulators that support a single subpopulation diff --git a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs index a938aafd..0a8774c3 100644 --- a/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/count_min_sketch_accumulator.rs @@ -280,6 +280,13 @@ impl AggregateCore for CountMinSketchAccumulator { AggregationType::CountMinSketch } + fn approx_memory_bytes(&self) -> usize { + // Conservative constant for the CountMinSketch counter matrix. + // Real per-instance sizing would require exposing rows/cols on + // the inner sketch; 16 KiB is a reasonable v1 default. + 16 * 1024 + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs index e5c15e3f..528e3101 100644 --- a/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/datasketches_kll_accumulator.rs @@ -197,6 +197,12 @@ impl AggregateCore for DatasketchesKLLAccumulator { AggregationType::DatasketchesKLL } + fn approx_memory_bytes(&self) -> usize { + // KLL with default k=200 holds ~2*k items (~3 KiB). Round up + // for overhead. + 4 * 1024 + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs index e8b1b1b9..9efa72eb 100644 --- a/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/delta_set_aggregator_accumulator.rs @@ -244,6 +244,13 @@ impl AggregateCore for DeltaSetAggregatorAccumulator { AggregationType::DeltaSetAggregator } + fn approx_memory_bytes(&self) -> usize { + // Two HashSets of KeyByLabelValues. + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + + (self.added.len() + self.removed.len()) * BYTES_PER_ENTRY + } + fn get_keys(&self) -> Option> { if !self.removed.is_empty() { panic!("DeltaSetAggregatorAccumulator does not support get_keys when removed items are present"); diff --git a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs index 0b2e924b..2a8a8b15 100644 --- a/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/hydra_kll_accumulator.rs @@ -113,6 +113,13 @@ impl AggregateCore for HydraKllSketchAccumulator { AggregationType::HydraKLL } + fn approx_memory_bytes(&self) -> usize { + // HydraKLL is a row*col grid of KLL sketches; typical instances + // are on the order of tens of KiB. 32 KiB is a conservative + // per-instance default. + 32 * 1024 + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/precompute_operators/increase_accumulator.rs b/asap-query-engine/src/precompute_operators/increase_accumulator.rs index 6c8c0293..43df2bce 100644 --- a/asap-query-engine/src/precompute_operators/increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/increase_accumulator.rs @@ -245,6 +245,11 @@ impl AggregateCore for IncreaseAccumulator { AggregationType::Increase } + fn approx_memory_bytes(&self) -> usize { + // Two Measurements + two i64s. Measurements are a few f64 fields. + std::mem::size_of::() + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs index b4763626..dbc8696b 100644 --- a/asap-query-engine/src/precompute_operators/min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/min_max_accumulator.rs @@ -184,6 +184,11 @@ impl AggregateCore for MinMaxAccumulator { AggregationType::MinMax } + fn approx_memory_bytes(&self) -> usize { + // f64 + small sub_type String. + std::mem::size_of::() + self.sub_type.capacity() + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs index 6ee70402..215af221 100644 --- a/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_increase_accumulator.rs @@ -282,6 +282,13 @@ impl AggregateCore for MultipleIncreaseAccumulator { AggregationType::MultipleIncrease } + fn approx_memory_bytes(&self) -> usize { + // HashMap. IncreaseAccumulator is ~64 B, + // per-entry key/overhead is ~96 B. + const BYTES_PER_ENTRY: usize = 160; + std::mem::size_of::() + self.increases.len() * BYTES_PER_ENTRY + } + fn get_keys(&self) -> Option> { Some(self.increases.keys().cloned().collect()) } diff --git a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs index 2984cf30..fe42ccc2 100644 --- a/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_min_max_accumulator.rs @@ -240,6 +240,13 @@ impl AggregateCore for MultipleMinMaxAccumulator { AggregationType::MultipleMinMax } + fn approx_memory_bytes(&self) -> usize { + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + + self.values.len() * BYTES_PER_ENTRY + + self.sub_type.capacity() + } + fn get_keys(&self) -> Option> { Some(self.values.keys().cloned().collect()) } diff --git a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs index 1505b934..11d2e77c 100644 --- a/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/multiple_sum_accumulator.rs @@ -230,6 +230,13 @@ impl AggregateCore for MultipleSumAccumulator { AggregationType::MultipleSum } + fn approx_memory_bytes(&self) -> usize { + // HashMap. Label strings dominate; use a + // conservative per-entry estimate plus HashMap overhead. + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.sums.len() * BYTES_PER_ENTRY + } + fn get_keys(&self) -> Option> { Some(self.sums.keys().cloned().collect()) } diff --git a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs index 4ec46c59..2ad268ac 100644 --- a/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/set_aggregator_accumulator.rs @@ -179,6 +179,13 @@ impl AggregateCore for SetAggregatorAccumulator { AggregationType::SetAggregator } + fn approx_memory_bytes(&self) -> usize { + // HashSet; per-entry cost is label strings + // plus HashSet overhead. + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.added.len() * BYTES_PER_ENTRY + } + fn get_keys(&self) -> Option> { Some(self.added.iter().cloned().collect()) } diff --git a/asap-query-engine/src/precompute_operators/sum_accumulator.rs b/asap-query-engine/src/precompute_operators/sum_accumulator.rs index fa32ef2e..6efddf6d 100644 --- a/asap-query-engine/src/precompute_operators/sum_accumulator.rs +++ b/asap-query-engine/src/precompute_operators/sum_accumulator.rs @@ -118,6 +118,11 @@ impl AggregateCore for SumAccumulator { AggregationType::Sum } + fn approx_memory_bytes(&self) -> usize { + // Single f64 + struct overhead. + std::mem::size_of::() + } + fn get_keys(&self) -> Option> { None } diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 6f7fb27c..e2f9e3c2 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -24,10 +24,17 @@ use super::persistence::{ type StoreKey = u64; // aggregation_id -/// Coarse per-sketch memory estimate used by the persistence flusher's -/// memory-pressure trigger. Not accurate — v1 just multiplies the -/// entry count by this constant. Per-type sizing is a follow-up. -const APPROX_BYTES_PER_SKETCH: usize = 4096; +/// Sum the `AggregateCore::approx_memory_bytes()` of every entry in a +/// sealed epoch. Cheap — each impl is supposed to be O(1) or at worst +/// O(entries_inside_the_sketch), and this is only called at rotate + +/// evict time, not on the insert hot path. +fn epoch_approx_bytes(epoch: &SealedEpoch) -> usize { + epoch + .entries + .iter() + .map(|(_, _, agg)| agg.approx_memory_bytes()) + .sum() +} /// Fallback epoch capacity used when `num_aggregates_to_retain` is not /// set in the streaming config but persistence is enabled. Without it @@ -209,10 +216,12 @@ pub struct PerKeyInner { /// * `mem_bytes_sealed` is maintained as flusher input. persistence_enabled: bool, - /// Approximate sum of sketch bytes in sealed epochs across all - /// agg-ids. Updated on rotate (adds) and evict (subtracts). - /// Drives the flusher's memory-pressure trigger. Coarse — see - /// `APPROX_BYTES_PER_SKETCH`. + /// Approximate sum of sketch bytes across all sealed epochs + /// currently held in memory. Incremented on insert by the sum of + /// each item's `AggregateCore::approx_memory_bytes()`, decremented + /// on evict by the same. Drives the flusher's memory-pressure + /// trigger. Approximate — per-type estimates are not guaranteed + /// accurate, only proportional. mem_bytes_sealed: AtomicUsize, } @@ -503,7 +512,17 @@ impl SimpleMapStorePerKey { data.configure_epochs(effective_retention); } - let entries_added = items.len(); + // Sum the real per-accumulator byte estimate up front. Summing + // happens before the loop consumes `items`; each + // `approx_memory_bytes` call is O(1) or a cheap field read on + // all overriding impls, so this adds at most a handful of + // arithmetic ops per batch item. + let batch_approx_bytes: usize = if persistence_enabled { + items.iter().map(|(_, a)| a.approx_memory_bytes()).sum() + } else { + 0 + }; + for (output, precompute) in items { let timestamp_range = (output.start_timestamp, output.end_timestamp); let metric_id: MetricID = data.intern.intern(output.key); @@ -524,14 +543,16 @@ impl SimpleMapStorePerKey { } if persistence_enabled { - // Best-effort memory accounting: attribute the batch's - // entries to sealed-epoch bytes. This over-counts (current - // epoch entries are included) but the flusher is - // conservative anyway. - self.inner.mem_bytes_sealed.fetch_add( - entries_added * APPROX_BYTES_PER_SKETCH, - Ordering::Relaxed, - ); + // Tracked bytes are the sum of each accumulator's own + // `approx_memory_bytes()`. The counter conceptually + // represents "bytes in sealed epochs"; in practice we + // charge them to the store as soon as they arrive + // (current_epoch is included) because the rotator will + // seal them soon anyway and the flusher's trigger is + // conservative. + self.inner + .mem_bytes_sealed + .fetch_add(batch_approx_bytes, Ordering::Relaxed); } if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { @@ -925,7 +946,7 @@ impl EpochSource for PerKeyInner { agg_id, epoch_id: *epoch_id, end_ts: max_end, - approx_bytes: epoch.entries.len() * APPROX_BYTES_PER_SKETCH, + approx_bytes: epoch_approx_bytes(epoch), }); } } @@ -968,7 +989,7 @@ impl EpochSource for PerKeyInner { } let (min_ts, max_ts) = epoch.time_bounds().unwrap_or((0, 0)); - let approx_bytes = epoch.entries.len() * APPROX_BYTES_PER_SKETCH; + let approx_bytes = epoch_approx_bytes(epoch); Ok(Some(EpochSnapshot { agg_id, @@ -992,7 +1013,7 @@ impl EpochSource for PerKeyInner { } }; if let Some(epoch) = data.sealed_epochs.remove(&epoch_id) { - let freed = epoch.entries.len() * APPROX_BYTES_PER_SKETCH; + let freed = epoch_approx_bytes(&epoch); self.mem_bytes_sealed.fetch_sub(freed, Ordering::Relaxed); // Also purge the epoch's windows from read_counts so they // don't leak. From d7da8a34669441783976e23bb4e49a2b7c2f4ee9 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Mon, 13 Apr 2026 22:04:26 -0400 Subject: [PATCH 10/12] feat(persistence): CLI wiring for main.rs and precompute_engine.rs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the enum-level SimpleMapStore::with_persistence_per_key wrapper plus --persistence-* CLI flags on both production binaries. New public helper: SimpleMapStore::with_persistence_per_key( streaming_config, cleanup_policy, persistence_cfg, ) -> PersistResult Always returns the PerKey variant — persistence only targets per-key locking; the Global variant stays in-memory-only. New CLI args (both main.rs and src/bin/precompute_engine.rs): --persistence-enabled --persistence-dir --persistence-memory-limit-mb (default: 2048) --persistence-hot-window-secs (default: 3600, 0 disables) --persistence-delete-older-than-secs (default: 604800, 0 disables) --persistence-flush-interval-ms (default: 1000) --persistence-part-cache-mb (default: min(10%*mem, 512)) When --persistence-enabled is passed, both binaries: * Build a SimpleMapStorePersistenceConfig from the CLI args (low_water = 80% of high, hard_cap = 125% of high). * Call SimpleMapStore::with_persistence_per_key instead of new_with_strategy. * Force LockStrategy::PerKey (log an info message if --lock-strategy was Global). Unchanged when --persistence-enabled is off: the existing new_with_strategy path runs exactly as before. No other binary (test_e2e_precompute, bench_precompute_sketch, e2e_quickstart_resource_test) is touched — those are test/bench paths that don't need persistence. All 450 existing tests still pass. Full bin build passes. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/bin/precompute_engine.rs | 84 +++++++++++++- asap-query-engine/src/main.rs | 109 +++++++++++++++++- .../src/stores/simple_map_store/mod.rs | 22 ++++ 3 files changed, 208 insertions(+), 7 deletions(-) diff --git a/asap-query-engine/src/bin/precompute_engine.rs b/asap-query-engine/src/bin/precompute_engine.rs index 3cf2459d..d46f37b0 100644 --- a/asap-query-engine/src/bin/precompute_engine.rs +++ b/asap-query-engine/src/bin/precompute_engine.rs @@ -65,6 +65,36 @@ struct Args { /// Policy for handling late samples that arrive after their window has closed #[arg(long, value_enum, default_value_t = LateDataPolicy::Drop)] late_data_policy: LateDataPolicy, + + // ---- SimpleMapStore persistence ---- + /// Enable the disk-backed persistence layer for SimpleMapStore. + /// When set, forces LockStrategy::PerKey regardless of --lock-strategy. + #[arg(long, default_value_t = false)] + persistence_enabled: bool, + + /// Root directory for persistence (manifest + parts/). + #[arg(long)] + persistence_dir: Option, + + /// Primary memory budget for in-memory sealed epochs, in MiB. + #[arg(long, default_value_t = 2048)] + persistence_memory_limit_mb: usize, + + /// Hot-window length in seconds. 0 disables time-based flushing. + #[arg(long, default_value_t = 3600)] + persistence_hot_window_secs: u64, + + /// Cold-tier TTL in seconds. 0 disables disk retention. + #[arg(long, default_value_t = 604800)] + persistence_delete_older_than_secs: u64, + + /// Cadence of the flusher loop, in milliseconds. + #[arg(long, default_value_t = 1000)] + persistence_flush_interval_ms: u64, + + /// Tier-2 part-cache byte budget, in MiB. 0 disables. + #[arg(long)] + persistence_part_cache_mb: Option, } #[tokio::main] @@ -89,12 +119,62 @@ async fn main() -> Result<(), Box> { ); // Create the store - let store: Arc = + let store: Arc = if args.persistence_enabled { + use query_engine_rust::stores::simple_map_store::persistence::SimpleMapStorePersistenceConfig; + let disk_path = args + .persistence_dir + .clone() + .expect("--persistence-enabled requires --persistence-dir"); + let memory_limit_bytes = args.persistence_memory_limit_mb * 1024 * 1024; + let hot_window_ms = if args.persistence_hot_window_secs == 0 { + None + } else { + Some(args.persistence_hot_window_secs * 1000) + }; + let delete_older_than_ms = if args.persistence_delete_older_than_secs == 0 { + None + } else { + Some(args.persistence_delete_older_than_secs * 1000) + }; + let part_cache_bytes = args + .persistence_part_cache_mb + .map(|mb| mb * 1024 * 1024) + .unwrap_or_else(|| { + let ten_pct = (memory_limit_bytes / 10) as u64; + ten_pct.min(512 * 1024 * 1024) + }); + let persistence_cfg = SimpleMapStorePersistenceConfig { + memory_limit_bytes, + memory_low_watermark_bytes: memory_limit_bytes * 8 / 10, + hard_cap_bytes: memory_limit_bytes * 125 / 100, + hot_window_ms, + delete_older_than_ms, + flush_interval: std::time::Duration::from_millis(args.persistence_flush_interval_ms), + disk_path: std::path::PathBuf::from(&disk_path), + part_cache_bytes, + }; + info!( + "Persistence enabled: disk_path={}, memory_limit={} MiB, hot_window={:?} s, flush_interval={} ms", + disk_path, + args.persistence_memory_limit_mb, + persistence_cfg.hot_window_ms.map(|ms| ms / 1000), + args.persistence_flush_interval_ms, + ); + Arc::new( + SimpleMapStore::with_persistence_per_key( + streaming_config.clone(), + CleanupPolicy::CircularBuffer, + persistence_cfg, + ) + .expect("SimpleMapStore::with_persistence_per_key failed"), + ) + } else { Arc::new(SimpleMapStore::new_with_strategy( streaming_config.clone(), CleanupPolicy::CircularBuffer, args.lock_strategy, - )); + )) + }; // Optionally start the query HTTP server if args.query_port > 0 { diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 533b0541..2d693cbc 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -160,6 +160,49 @@ struct Args { /// Query tracker: observation window in seconds before triggering planning #[arg(long, default_value = "100")] tracker_observation_window_secs: u64, + + // ---- SimpleMapStore persistence ---- + // + // When --persistence-enabled is set, the store is constructed via + // SimpleMapStore::with_persistence_per_key with the other + // --persistence-* flags as the config. Forces LockStrategy::PerKey + // regardless of --lock-strategy; the Global variant is + // intentionally left in-memory-only. + + /// Enable the disk-backed persistence layer for SimpleMapStore + #[arg(long)] + persistence_enabled: bool, + + /// Root directory for persistence (manifest + parts/). Required + /// when --persistence-enabled. + #[arg(long)] + persistence_dir: Option, + + /// Primary memory budget for in-memory sealed epochs, in MiB. + /// When exceeded, the background flusher evicts oldest-first. + #[arg(long, default_value = "2048")] + persistence_memory_limit_mb: usize, + + /// Hot-window length in seconds. Any sealed epoch whose end_ts is + /// older than (now - this) is flushed on the next flusher tick, + /// regardless of memory pressure. 0 disables time-based flushing. + #[arg(long, default_value = "3600")] + persistence_hot_window_secs: u64, + + /// Cold-tier TTL in seconds. Parts whose max_ts is older than + /// (now - this) are removed from disk on the next flusher tick. + /// 0 disables disk retention. + #[arg(long, default_value = "604800")] + persistence_delete_older_than_secs: u64, + + /// Cadence of the background flusher loop, in milliseconds. + #[arg(long, default_value = "1000")] + persistence_flush_interval_ms: u64, + + /// Tier-2 part-cache byte budget, in MiB. 0 disables the cache. + /// Defaults to `min(10% * memory_limit_mb, 512)`. + #[arg(long)] + persistence_part_cache_mb: Option, } #[tokio::main] @@ -207,11 +250,67 @@ async fn main() -> Result<()> { // Get cleanup policy from inference config let cleanup_policy = inference_config.cleanup_policy; info!("Using cleanup policy: {:?}", cleanup_policy); - let store = Arc::new(SimpleMapStore::new_with_strategy( - streaming_config.clone(), - cleanup_policy, - args.lock_strategy, - )); + let store = if args.persistence_enabled { + use query_engine_rust::stores::simple_map_store::persistence::SimpleMapStorePersistenceConfig; + let disk_path = args + .persistence_dir + .clone() + .expect("--persistence-enabled requires --persistence-dir"); + let memory_limit_bytes = args.persistence_memory_limit_mb * 1024 * 1024; + let hot_window_ms = if args.persistence_hot_window_secs == 0 { + None + } else { + Some(args.persistence_hot_window_secs * 1000) + }; + let delete_older_than_ms = if args.persistence_delete_older_than_secs == 0 { + None + } else { + Some(args.persistence_delete_older_than_secs * 1000) + }; + let part_cache_bytes = args + .persistence_part_cache_mb + .map(|mb| mb * 1024 * 1024) + .unwrap_or_else(|| { + let ten_pct = (memory_limit_bytes / 10) as u64; + ten_pct.min(512 * 1024 * 1024) + }); + let persistence_cfg = SimpleMapStorePersistenceConfig { + memory_limit_bytes, + memory_low_watermark_bytes: memory_limit_bytes * 8 / 10, + hard_cap_bytes: memory_limit_bytes * 125 / 100, + hot_window_ms, + delete_older_than_ms, + flush_interval: std::time::Duration::from_millis(args.persistence_flush_interval_ms), + disk_path: std::path::PathBuf::from(&disk_path), + part_cache_bytes, + }; + info!( + "Persistence enabled: disk_path={}, memory_limit={} MiB, hot_window={:?} s, delete_older_than={:?} s, flush_interval={} ms, part_cache={} MiB", + disk_path, + args.persistence_memory_limit_mb, + persistence_cfg.hot_window_ms.map(|ms| ms / 1000), + persistence_cfg.delete_older_than_ms.map(|ms| ms / 1000), + args.persistence_flush_interval_ms, + persistence_cfg.part_cache_bytes / (1024 * 1024), + ); + if !matches!(args.lock_strategy, LockStrategy::PerKey) { + info!("--persistence-enabled forces LockStrategy::PerKey (ignoring --lock-strategy)"); + } + Arc::new( + SimpleMapStore::with_persistence_per_key( + streaming_config.clone(), + cleanup_policy, + persistence_cfg, + ) + .expect("SimpleMapStore::with_persistence_per_key failed"), + ) + } else { + Arc::new(SimpleMapStore::new_with_strategy( + streaming_config.clone(), + cleanup_policy, + args.lock_strategy, + )) + }; // // Setup PromSketchStore (shared between engine and remote write server) // let promsketch_store = if args.enable_prometheus_remote_write { diff --git a/asap-query-engine/src/stores/simple_map_store/mod.rs b/asap-query-engine/src/stores/simple_map_store/mod.rs index b3243feb..5caab2cf 100644 --- a/asap-query-engine/src/stores/simple_map_store/mod.rs +++ b/asap-query-engine/src/stores/simple_map_store/mod.rs @@ -65,6 +65,28 @@ impl SimpleMapStore { } } } + + /// Persistence-enabled constructor. Always returns the `PerKey` + /// variant — persistence only targets per-key locking; the + /// `Global` variant is intentionally left in-memory-only. + /// + /// Runs recovery on the disk path, starts the background flusher + /// thread, and returns a store whose sealed epochs are flushed + /// to disk on memory / time pressure. Query paths transparently + /// read back from disk when in-memory state misses. + pub fn with_persistence_per_key( + streaming_config: Arc, + cleanup_policy: CleanupPolicy, + persistence_cfg: persistence::SimpleMapStorePersistenceConfig, + ) -> persistence::PersistResult { + Ok(SimpleMapStore::PerKey( + SimpleMapStorePerKey::with_persistence( + streaming_config, + cleanup_policy, + persistence_cfg, + )?, + )) + } } #[async_trait::async_trait] From 736a7d6fdde770cea16d416a45f6b3e0613d4a6c Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Tue, 14 Apr 2026 13:01:09 -0400 Subject: [PATCH 11/12] test(persistence): add #[ignore] perf harness for storage + flusher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four scenarios covering the hot paths exposed by the persistence layer, runnable via: cargo test --release -p query_engine_rust --lib \ tests::persistence_perf_tests -- --ignored --nocapture 1. insert_throughput_in_memory_vs_persistent — raw insert rate with the in-memory store, with persistence enabled but flusher idle (isolates insert-path overhead), and with persistence + aggressive flushing (measures back-pressure cost). 2. query_latency_memory_only_vs_disk_through — p50/p90/p99 range query latency with all data in memory vs. everything flushed and read back via the moka part cache + disk mmap. 3. flush_throughput_sustained — insert a large batch, then poll diagnostic_info until total_time_map_entries drops to the current epoch's capacity (i.e., everything that can be flushed has been), and report drain rate + MiB/s. The drain-done criterion uses time-map entries because current-epoch bytes are counted in total_sketch_bytes but never get flushed while hot. 4. memory_bound_adherence_under_overload — push 10x the memory limit through the store and record the peak tracked bytes. Verifies the flusher keeps memory roughly bounded without fully draining to zero. SumAccumulator throughout for simplicity; production sketch types (KLL/CountMin/HydraKLL) would have different numbers because their Arroyo serialize path is more expensive. Co-Authored-By: Claude Opus 4.6 (1M context) --- asap-query-engine/src/tests/mod.rs | 1 + .../src/tests/persistence_perf_tests.rs | 511 ++++++++++++++++++ 2 files changed, 512 insertions(+) create mode 100644 asap-query-engine/src/tests/persistence_perf_tests.rs diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index f8ae48dd..e8e2dd5a 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -7,6 +7,7 @@ pub mod prometheus_forwarding_tests; pub mod query_equivalence_tests; pub mod sql_pattern_matching_tests; pub mod persistence_integration_tests; +pub mod persistence_perf_tests; pub mod store_correctness_tests; pub mod trait_design_tests; diff --git a/asap-query-engine/src/tests/persistence_perf_tests.rs b/asap-query-engine/src/tests/persistence_perf_tests.rs new file mode 100644 index 00000000..2267b458 --- /dev/null +++ b/asap-query-engine/src/tests/persistence_perf_tests.rs @@ -0,0 +1,511 @@ +//! Performance harness for the `SimpleMapStore` persistence layer. +//! +//! All tests are `#[ignore]` so they don't slow down the normal +//! `cargo test` run. Exercise them with: +//! +//! ```text +//! cargo test --release -p query_engine_rust --lib \ +//! tests::persistence_perf_tests -- --ignored --nocapture +//! ``` +//! +//! Numbers are illustrative, not load-bearing — ext4 on a laptop SSD +//! is nowhere near the profile of a production deployment, but the +//! *shapes* (memory-bound adherence, relative overheads, disk read +//! costs vs. in-memory) are informative for tuning the flusher. +//! +//! ## Scenarios +//! +//! 1. [`insert_throughput_in_memory_vs_persistent`] — raw insert +//! throughput with and without persistence enabled. Measures the +//! flusher's overhead on the write path. +//! 2. [`query_latency_memory_only_vs_disk_through`] — p50 / p99 +//! range-query latency when results come from memory only vs. +//! when everything has been flushed to disk. +//! 3. [`flush_throughput_sustained`] — how fast the background +//! flusher can turn sealed epochs into parts and evict them, in +//! entries/sec and MiB/sec. +//! 4. [`memory_bound_adherence_under_overload`] — verifies the +//! `mem_bytes_sealed` counter stays bounded when the insert rate +//! vastly exceeds the flusher's steady-state throughput. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use promql_utilities::data_model::KeyByLabelNames; +use tempfile::TempDir; + +use crate::data_model::{ + AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, WindowType, +}; +use crate::precompute_operators::SumAccumulator; +use crate::stores::simple_map_store::per_key::SimpleMapStorePerKey; +use crate::stores::simple_map_store::persistence::SimpleMapStorePersistenceConfig; +use crate::stores::Store; +use crate::{AggregateCore, AggregationConfig}; + +// ========================================================================= +// Fixtures +// ========================================================================= + +fn streaming_config(agg_id: u64, retention: Option) -> Arc { + let cfg = AggregationConfig::new( + agg_id, + AggregationType::Sum, + String::new(), + HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + String::new(), + 60, + 60, + WindowType::Tumbling, + String::new(), + "cpu_usage".to_string(), + retention, + None, + None, + None, + ); + let mut map = HashMap::new(); + map.insert(agg_id, cfg); + Arc::new(StreamingConfig::new(map)) +} + +fn persistence_cfg( + dir: &TempDir, + memory_limit_bytes: usize, + hot_window_ms: Option, + flush_interval: Duration, +) -> SimpleMapStorePersistenceConfig { + SimpleMapStorePersistenceConfig { + memory_limit_bytes, + memory_low_watermark_bytes: memory_limit_bytes * 8 / 10, + hard_cap_bytes: memory_limit_bytes * 125 / 100, + hot_window_ms, + delete_older_than_ms: None, + flush_interval, + disk_path: dir.path().to_path_buf(), + part_cache_bytes: 64 * 1024 * 1024, + } +} + +fn sum_item( + agg_id: u64, + start: u64, + end: u64, + value: f64, +) -> (PrecomputedOutput, Box) { + ( + PrecomputedOutput::new(start, end, None, agg_id), + Box::new(SumAccumulator::with_sum(value)), + ) +} + +/// Generate N distinct-window items. Each has `start_ts = i * 1000` +/// so the rotator fires regularly. +fn gen_items(agg_id: u64, n: usize) -> Vec<(PrecomputedOutput, Box)> { + let mut out = Vec::with_capacity(n); + for i in 0..n { + let start = (i as u64) * 1000; + let end = start + 1000; + out.push(sum_item(agg_id, start, end, i as f64)); + } + out +} + +/// Insert items in batches of `batch_size`, returning wall-clock time. +fn insert_all( + store: &S, + items: Vec<(PrecomputedOutput, Box)>, + batch_size: usize, +) -> Duration { + let start = Instant::now(); + for chunk in items.chunks(batch_size) { + let batch: Vec<_> = chunk.iter().map(|(o, a)| (o.clone(), a.clone())).collect(); + store.insert_precomputed_output_batch(batch).unwrap(); + } + start.elapsed() +} + +fn fmt_rate(count: usize, d: Duration) -> String { + let secs = d.as_secs_f64(); + if secs == 0.0 { + return "∞".to_string(); + } + let rate = count as f64 / secs; + if rate >= 1_000_000.0 { + format!("{:.2} M/s", rate / 1_000_000.0) + } else if rate >= 1_000.0 { + format!("{:.1} K/s", rate / 1_000.0) + } else { + format!("{:.0}/s", rate) + } +} + +// ========================================================================= +// Scenario 1: insert throughput, memory-only vs. with-persistence +// ========================================================================= + +#[test] +#[ignore] +fn insert_throughput_in_memory_vs_persistent() { + const N: usize = 200_000; + const BATCH: usize = 1_000; + + println!("\n== Insert throughput ({} items, batch={}) ==", N, BATCH); + + // -- baseline: in-memory, NoCleanup -- + { + let store = SimpleMapStorePerKey::new( + streaming_config(1, None), + CleanupPolicy::NoCleanup, + ); + let items = gen_items(1, N); + let d = insert_all(&store, items, BATCH); + println!( + " in-memory (NoCleanup): {} total, {:?} wall, {} inserts/sec", + N, + d, + fmt_rate(N, d) + ); + } + + // -- with persistence, flusher idle (no flush, just the overhead + // of mem_bytes_sealed tracking + always-rotate) -- + { + let tmp = TempDir::new().unwrap(); + // Large memory limit, no hot window — flusher has nothing to + // do, so this isolates the insert-path overhead of persistence + // being enabled. + let cfg = persistence_cfg( + &tmp, + 16 * 1024 * 1024 * 1024, // 16 GiB ceiling — unreachable in this test + None, + Duration::from_secs(3600), + ); + let store = SimpleMapStorePerKey::with_persistence( + streaming_config(1, Some(1024)), + CleanupPolicy::NoCleanup, + cfg, + ) + .unwrap(); + let items = gen_items(1, N); + let d = insert_all(&store, items, BATCH); + println!( + " persistent (flusher idle): {} total, {:?} wall, {} inserts/sec", + N, + d, + fmt_rate(N, d) + ); + } + + // -- with persistence + aggressive flushing -- + { + let tmp = TempDir::new().unwrap(); + let cfg = persistence_cfg( + &tmp, + 4 * 1024 * 1024, // 4 MiB limit — forces constant pressure + Some(0), // flush everything ASAP + Duration::from_millis(25), + ); + let store = SimpleMapStorePerKey::with_persistence( + streaming_config(1, Some(512)), + CleanupPolicy::NoCleanup, + cfg, + ) + .unwrap(); + let items = gen_items(1, N); + let d = insert_all(&store, items, BATCH); + println!( + " persistent (flusher active): {} total, {:?} wall, {} inserts/sec", + N, + d, + fmt_rate(N, d) + ); + } +} + +// ========================================================================= +// Scenario 2: query latency memory-only vs. disk-through +// ========================================================================= + +#[test] +#[ignore] +fn query_latency_memory_only_vs_disk_through() { + const POPULATE: usize = 20_000; + const QUERIES: usize = 1_000; + + println!( + "\n== Query latency ({} items populated, {} queries) ==", + POPULATE, QUERIES + ); + + // -- in-memory baseline -- + { + let store = SimpleMapStorePerKey::new( + streaming_config(1, None), + CleanupPolicy::NoCleanup, + ); + let items = gen_items(1, POPULATE); + insert_all(&store, items, 1_000); + + let lats = run_queries(&store, QUERIES, POPULATE); + report_latencies(" in-memory", &lats); + } + + // -- disk read-through (everything flushed) -- + { + let tmp = TempDir::new().unwrap(); + let cfg = persistence_cfg( + &tmp, + 4 * 1024 * 1024, + Some(0), + Duration::from_millis(10), + ); + let store = SimpleMapStorePerKey::with_persistence( + streaming_config(1, Some(256)), + CleanupPolicy::NoCleanup, + cfg, + ) + .unwrap(); + let items = gen_items(1, POPULATE); + insert_all(&store, items, 1_000); + + // Wait until sealed epochs have been flushed and evicted. + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline { + let diag = store.diagnostic_info(); + if diag.total_sketch_bytes == 0 && diag.total_time_map_entries < POPULATE / 20 { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + + let lats = run_queries(&store, QUERIES, POPULATE); + report_latencies(" disk-through", &lats); + } +} + +fn run_queries(store: &S, n_queries: usize, populated: usize) -> Vec { + let mut lats = Vec::with_capacity(n_queries); + // Deterministic range spread: each query covers ~10% of the window. + let window = (populated as u64) * 1000; + let step = window / n_queries as u64; + let span = window / 10; + for i in 0..n_queries { + let start = (i as u64) * step; + let end = start + span; + let t = Instant::now(); + let _ = store + .query_precomputed_output("cpu_usage", 1, start, end) + .unwrap(); + lats.push(t.elapsed()); + } + lats +} + +fn report_latencies(label: &str, lats: &[Duration]) { + let mut sorted: Vec = lats.iter().map(|d| d.as_micros()).collect(); + sorted.sort_unstable(); + let p = |frac: f64| -> u128 { + let idx = ((sorted.len() as f64 - 1.0) * frac) as usize; + sorted[idx] + }; + let avg: u128 = sorted.iter().sum::() / sorted.len() as u128; + println!( + "{}: p50={}µs p90={}µs p99={}µs max={}µs avg={}µs n={}", + label, + p(0.50), + p(0.90), + p(0.99), + sorted.last().copied().unwrap_or(0), + avg, + sorted.len() + ); +} + +// ========================================================================= +// Scenario 3: sustained flush throughput +// ========================================================================= + +#[test] +#[ignore] +fn flush_throughput_sustained() { + const N: usize = 100_000; + + println!("\n== Flush throughput (N={}) ==", N); + + let tmp = TempDir::new().unwrap(); + // Generous memory limit so the flusher is driven by hot_window, + // not by pressure — this measures the flusher's steady-state, + // not its back-pressure behavior. + let cfg = persistence_cfg( + &tmp, + 128 * 1024 * 1024, + Some(0), // flush as fast as sealed epochs arrive + Duration::from_millis(10), + ); + let store = SimpleMapStorePerKey::with_persistence( + streaming_config(1, Some(512)), + CleanupPolicy::NoCleanup, + cfg, + ) + .unwrap(); + + let items = gen_items(1, N); + let insert_start = Instant::now(); + insert_all(&store, items, 1_000); + let insert_d = insert_start.elapsed(); + + // Wait for the flusher to drain sealed epochs. We can't wait for + // `total_sketch_bytes == 0` because the (unsealed) current epoch + // contributes to that counter and it's never flushed while hot. + // Instead, wait until the in-memory count of time-map entries + // drops to ≤ one epoch's capacity — i.e., everything that CAN be + // flushed has been flushed. + // + // "One epoch's worth" = 512 (the retention configured above). + let drain_start = Instant::now(); + let deadline = drain_start + Duration::from_secs(20); + let epoch_capacity_hint = 512usize; + loop { + let diag = store.diagnostic_info(); + if diag.total_time_map_entries <= epoch_capacity_hint { + break; + } + if Instant::now() >= deadline { + println!( + " WARNING: flusher didn't drain in time; {} time-map entries still in memory", + diag.total_time_map_entries + ); + break; + } + std::thread::sleep(Duration::from_millis(5)); + } + let drain_d = drain_start.elapsed(); + let flushed = N.saturating_sub(epoch_capacity_hint); + + // Measure on-disk bytes. + let parts_dir = tmp.path().join("parts"); + let total_bytes = dir_size(&parts_dir); + + println!( + " insert: {:?} ({} items/s)", + insert_d, + fmt_rate(N, insert_d) + ); + println!( + " drain: {:?} ({} items flushed, {} items/s)", + drain_d, + flushed, + fmt_rate(flushed, drain_d) + ); + println!( + " on-disk: {:.2} MiB ({:.2} MiB/s while draining)", + total_bytes as f64 / (1024.0 * 1024.0), + (total_bytes as f64 / (1024.0 * 1024.0)) / drain_d.as_secs_f64() + ); +} + +fn dir_size(path: &std::path::Path) -> u64 { + let mut total = 0u64; + if !path.exists() { + return 0; + } + let walk = |p: &std::path::Path, out: &mut u64| { + if let Ok(entries) = std::fs::read_dir(p) { + for e in entries.flatten() { + let meta = match e.metadata() { + Ok(m) => m, + Err(_) => continue, + }; + if meta.is_file() { + *out += meta.len(); + } else if meta.is_dir() { + if let Ok(inner) = std::fs::read_dir(e.path()) { + for inner_e in inner.flatten() { + if let Ok(m) = inner_e.metadata() { + if m.is_file() { + *out += m.len(); + } + } + } + } + } + } + } + }; + walk(path, &mut total); + total +} + +// ========================================================================= +// Scenario 4: memory-bound adherence under overload +// ========================================================================= + +#[test] +#[ignore] +fn memory_bound_adherence_under_overload() { + // Push 10x more than the high-water mark; verify the tracked + // memory counter stays close to it throughout. + const N: usize = 50_000; + const LIMIT_BYTES: usize = 256 * 1024; + + println!( + "\n== Memory-bound adherence (N={}, limit={} KiB) ==", + N, + LIMIT_BYTES / 1024 + ); + + let tmp = TempDir::new().unwrap(); + let cfg = persistence_cfg( + &tmp, + LIMIT_BYTES, + Some(0), + Duration::from_millis(10), + ); + let store = SimpleMapStorePerKey::with_persistence( + streaming_config(1, Some(128)), + CleanupPolicy::NoCleanup, + cfg, + ) + .unwrap(); + + let items = gen_items(1, N); + + let mut peak_tracked: usize = 0; + let start = Instant::now(); + for chunk in items.chunks(500) { + let batch: Vec<_> = chunk.iter().map(|(o, a)| (o.clone(), a.clone())).collect(); + store.insert_precomputed_output_batch(batch).unwrap(); + let diag = store.diagnostic_info(); + peak_tracked = peak_tracked.max(diag.total_sketch_bytes); + } + let d = start.elapsed(); + + // Let the flusher drain everything. + let deadline = Instant::now() + Duration::from_secs(10); + while Instant::now() < deadline { + if store.diagnostic_info().total_sketch_bytes == 0 { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } + let final_tracked = store.diagnostic_info().total_sketch_bytes; + + println!( + " inserted {} items in {:?} ({})", + N, + d, + fmt_rate(N, d) + ); + println!( + " peak tracked: {} KiB (high-water = {} KiB, ratio = {:.2}x)", + peak_tracked / 1024, + LIMIT_BYTES / 1024, + peak_tracked as f64 / LIMIT_BYTES as f64 + ); + println!(" final tracked after drain: {} bytes", final_tracked); +} From 96b13dcc791bd264c40b2ee33af33867f986a09c Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Tue, 14 Apr 2026 13:14:47 -0400 Subject: [PATCH 12/12] feat(persistence): enforce hard_cap_bytes via insert back-pressure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The v1 design reserved hard_cap_bytes as the insert-path blocking ceiling but never wired it up — insert_for_store_key happily grew mem_bytes_sealed past the limit. Perf test memory_bound_adherence caught this at 1.53x the nominal high-water mark on a 10x overload. This commit wires it up: * FlusherHandle::wait_for_memory_under(mem_counter, cap, max_wait) parks the calling thread on the existing pressure_cv/pressure_mutex until `mem_counter.load() < cap` or the timeout expires. Kicks the flusher once on entry so it runs immediately instead of waiting for its own interval tick. Returns false on timeout so a stuck flusher degrades (logged warning + insert proceeds) rather than hanging the ingest thread. * PerKeyInner gains `hard_cap_bytes: usize`, populated from the persistence config in with_persistence and set to usize::MAX in the in-memory-only constructor path (no blocking). * insert_for_store_key checks mem_bytes_sealed against hard_cap_bytes BEFORE taking the per-agg RwLock::write. Running the check ahead of the lock means back-pressure on one agg-id doesn't stall unrelated queries on the same agg while we wait. * INSERT_BACK_PRESSURE_TIMEOUT = 30s — long enough for a healthy flusher to always complete one tick, short enough that a genuinely stuck flusher surfaces as a warning rather than a deadlock. * PersistenceState._flusher renamed to .flusher since the insert path now reaches it as well as Drop. New integration test hard_cap_back_pressure_blocks_inserts_until_flusher_drains: - Configures memory_limit=512 B, hard_cap=640 B, 200ms flush tick. - Inserts 200 items (~3200 B total) in tight batches. - Asserts that at least one insert took ≥3 ms (unambiguous signal of a condvar wait on a hot path where normal inserts are <1 ms). - Reports median / max / slow-count in the failure message for debugging if the back-pressure regresses. Perf re-run of memory_bound_adherence_under_overload confirms the fix: before: peak 390 KiB (1.53x high-water) after: peak 320 KiB (1.25x high-water = hard_cap_bytes exactly) The 1.25x ratio is the intended behavior — inserts block the moment the counter reaches hard_cap, so it tracks the cap with bounded overshoot from the in-flight batch. Insert throughput under overload drops from 2.80 M/s to 1.06 M/s, which is the cost of back-pressure under sustained overload and exactly what we want: a bounded memory envelope with a deterministic latency impact, not unbounded growth. Full suite: 451 passed, 0 failed, 9 ignored. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/stores/simple_map_store/per_key.rs | 71 ++++++++++++++-- .../simple_map_store/persistence/flusher.rs | 54 +++++++++++- .../tests/persistence_integration_tests.rs | 84 +++++++++++++++++++ 3 files changed, 202 insertions(+), 7 deletions(-) diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index e2f9e3c2..f2eb00bb 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -13,7 +13,7 @@ use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, RwLock}; -use std::time::Instant; +use std::time::{Duration, Instant}; use tracing::{debug, error, info, warn}; use super::persistence::{ @@ -42,6 +42,14 @@ fn epoch_approx_bytes(epoch: &SealedEpoch) -> usize { /// flushable. const PERSISTENCE_DEFAULT_EPOCH_CAPACITY: usize = 1024; +/// Maximum time an insert batch will wait on the flusher's pressure +/// condvar when `mem_bytes_sealed` has exceeded `hard_cap_bytes`. 30 +/// seconds is long enough that a healthy flusher always makes it +/// through a single tick, but short enough that a stuck flusher +/// degrades gracefully (logged warning + continued insert) rather +/// than hanging the ingest thread forever. +const INSERT_BACK_PRESSURE_TIMEOUT: Duration = Duration::from_secs(30); + /// Per-aggregation_id data protected by RwLock struct StoreKeyData { /// Label interning table (Optimization 1) @@ -223,6 +231,13 @@ pub struct PerKeyInner { /// trigger. Approximate — per-type estimates are not guaranteed /// accurate, only proportional. mem_bytes_sealed: AtomicUsize, + + /// Hard memory ceiling. When `mem_bytes_sealed` reaches this, + /// `insert_precomputed_output_batch` blocks on the flusher's + /// `pressure_cv` until the flusher catches up — the v1 + /// mechanism for bounding RAM under sustained overload. Set to + /// `usize::MAX` by the in-memory-only constructor (no blocking). + hard_cap_bytes: usize, } /// Persistence-related state owned by the outer store. Dropping this @@ -230,9 +245,10 @@ pub struct PerKeyInner { struct PersistenceState { manifest: Arc, cache: PartCache, - /// Held so `Drop` stops the thread. Not accessed directly after - /// construction. - _flusher: FlusherHandle, + /// Owned by the store. Dropping it shuts the thread down; the + /// insert path also calls `wait_for_memory_under` on it when the + /// hard cap is hit. + flusher: FlusherHandle, #[allow(dead_code)] parts_root: PathBuf, } @@ -259,6 +275,7 @@ impl SimpleMapStorePerKey { cleanup_policy, persistence_enabled: false, mem_bytes_sealed: AtomicUsize::new(0), + hard_cap_bytes: usize::MAX, }), persistence: None, } @@ -295,6 +312,11 @@ impl SimpleMapStorePerKey { let parts_root = persistence::flusher::parts_root(&persistence_cfg.disk_path); let cache = PartCache::new(parts_root.clone(), persistence_cfg.part_cache_bytes); + // Capture hard_cap before `persistence_cfg` is moved into the + // flusher; PerKeyInner needs it for back-pressure enforcement + // on the insert path. + let hard_cap_bytes = persistence_cfg.hard_cap_bytes; + let inner = Arc::new(PerKeyInner { store: DashMap::new(), earliest_timestamps: DashMap::new(), @@ -304,6 +326,7 @@ impl SimpleMapStorePerKey { cleanup_policy, persistence_enabled: true, mem_bytes_sealed: AtomicUsize::new(0), + hard_cap_bytes, }); let flusher = FlusherHandle::start( @@ -317,7 +340,7 @@ impl SimpleMapStorePerKey { persistence: Some(PersistenceState { manifest, cache, - _flusher: flusher, + flusher, parts_root, }), }) @@ -411,6 +434,42 @@ impl SimpleMapStorePerKey { let inserted_delta = items.len() as u64; let persistence_enabled = self.inner.persistence_enabled; + // ---- Back-pressure (persistence only) ---- + // + // If sealed memory has reached the hard cap, block the insert + // on the flusher's pressure condvar until the flusher makes + // progress. This has to happen BEFORE we take the per-agg + // RwLock::write so unrelated queries on the same agg aren't + // stalled by the wait. The wait is bounded so a stuck flusher + // logs-and-degrades rather than deadlocking the ingest path. + // + // Uses Ordering::Acquire on the load so the check synchronizes + // with the flusher's Relaxed decrements on evict — we might + // see stale values, but the wait loop re-checks after every + // notify_all anyway. + if persistence_enabled { + if let Some(state) = self.persistence.as_ref() { + let cap = self.inner.hard_cap_bytes; + let current = self.inner.mem_bytes_sealed.load(Ordering::Acquire); + if current >= cap { + let ok = state.flusher.wait_for_memory_under( + &self.inner.mem_bytes_sealed, + cap, + INSERT_BACK_PRESSURE_TIMEOUT, + ); + if !ok { + warn!( + "insert back-pressure timed out for agg_id {}: \ + mem={} bytes, hard_cap={}, proceeding anyway", + aggregation_id, + self.inner.mem_bytes_sealed.load(Ordering::Relaxed), + cap + ); + } + } + } + } + // Opt 4: compute batch minimum timestamp before acquiring any lock. let batch_min_ts = items .iter() @@ -666,7 +725,7 @@ impl Drop for SimpleMapStorePerKey { // ref count hits zero, guaranteeing the flusher cannot observe // a half-destroyed store. if let Some(mut state) = self.persistence.take() { - state._flusher.shutdown(); + state.flusher.shutdown(); } } } diff --git a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs index d87f7870..8d108cf9 100644 --- a/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs +++ b/asap-query-engine/src/stores/simple_map_store/persistence/flusher.rs @@ -19,7 +19,7 @@ use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::thread::{self, JoinHandle}; -use std::time::{Instant, SystemTime, UNIX_EPOCH}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tracing::{debug, error, info, warn}; @@ -106,6 +106,58 @@ impl FlusherHandle { self.inner.pressure_cv.notify_all(); } + /// Block the calling thread until `mem_counter.load() < cap` or + /// `max_wait` elapses, whichever comes first. Wakes the flusher + /// once on entry and then sleeps on the pressure condvar until + /// the flusher's per-tick `notify_all` fires — typical wait is + /// one flush tick. + /// + /// Returns `true` if the wait finished with memory under the cap, + /// `false` if the timeout expired first (so the caller can log + /// and proceed rather than hanging forever on a stuck flusher). + /// + /// Safe to call from the insert hot path — the only synchronization + /// held during the wait is the flusher's dedicated pressure mutex, + /// which the flusher thread itself never takes (it only calls + /// `notify_all` on the condvar, which requires no lock). + pub fn wait_for_memory_under( + &self, + mem_counter: &std::sync::atomic::AtomicUsize, + cap: usize, + max_wait: Duration, + ) -> bool { + if mem_counter.load(Ordering::Relaxed) < cap { + return true; + } + // Kick the flusher once so it tries to drain right now instead + // of waiting for its own interval tick. + self.inner.pressure_cv.notify_all(); + + let deadline = Instant::now() + max_wait; + let mut guard = match self.inner.pressure_mutex.lock() { + Ok(g) => g, + Err(_) => return false, // poisoned — give up rather than deadlock + }; + loop { + if mem_counter.load(Ordering::Relaxed) < cap { + return true; + } + if self.inner.shutdown.load(Ordering::Acquire) { + return false; + } + let now = Instant::now(); + if now >= deadline { + return false; + } + let remaining = deadline - now; + let (g, _) = match self.inner.pressure_cv.wait_timeout(guard, remaining) { + Ok(v) => v, + Err(_) => return false, + }; + guard = g; + } + } + /// Access to the manifest for the query read-through path. pub fn manifest(&self) -> Arc { Arc::clone(&self.inner.manifest) diff --git a/asap-query-engine/src/tests/persistence_integration_tests.rs b/asap-query-engine/src/tests/persistence_integration_tests.rs index d4052fa6..1116f31d 100644 --- a/asap-query-engine/src/tests/persistence_integration_tests.rs +++ b/asap-query-engine/src/tests/persistence_integration_tests.rs @@ -201,3 +201,87 @@ fn construct_and_drop_shuts_flusher_cleanly() { // Dropping the store should not deadlock or panic. drop(store); } + +#[test] +fn hard_cap_back_pressure_blocks_inserts_until_flusher_drains() { + // Construct a store with a very small hard cap and a flusher that + // is deliberately slow (long flush interval). The first few + // inserts will push mem_bytes_sealed past the cap; subsequent + // inserts must block in `wait_for_memory_under` until the flusher + // evicts a sealed epoch. + // + // Test strategy: measure wall-clock time of an insert that we + // *know* will hit the cap. If back-pressure is wired up, it + // must be longer than the flusher's tick interval (because it + // waits at least one tick for the condvar notify). If it's not + // wired up, the insert returns in a handful of microseconds. + + let dir = TempDir::new().unwrap(); + let cfg = make_streaming_config(1); + let persistence = SimpleMapStorePersistenceConfig { + // Very small memory limit — a few hundred bytes — so the + // insert path hits the cap after the first handful of items. + memory_limit_bytes: 512, + memory_low_watermark_bytes: 256, + hard_cap_bytes: 640, + hot_window_ms: Some(0), + delete_older_than_ms: None, + // Flusher tick is 200ms — long enough that a blocking insert + // is clearly distinguishable from a non-blocking one. + flush_interval: Duration::from_millis(200), + disk_path: dir.path().to_path_buf(), + part_cache_bytes: 0, + }; + let store = + SimpleMapStorePerKey::with_persistence(cfg, CleanupPolicy::NoCleanup, persistence) + .expect("with_persistence"); + + // Push well past the cap — 200 items × ~16 bytes each, vs. a + // 640-byte cap — so the insert path is forced to block on the + // flusher's condvar at least once. num_aggregates_to_retain=2 + // means every 2 distinct windows triggers a seal, so sealed + // epochs accumulate fast. + // + // A normal insert on the hot path returns in <1 ms; anything + // ≥3 ms unambiguously indicates a condvar wait fired. We also + // capture the median as a baseline to make the assertion + // message informative when it fails. + let mut elapsed_all: Vec = Vec::with_capacity(200); + for i in 0..200u64 { + let start = i * 1_000; + let end = start + 1_000; + let batch = vec![sum_entry(1, start, end, i as f64)]; + let t = Instant::now(); + store + .insert_precomputed_output_batch(batch) + .expect("insert"); + elapsed_all.push(t.elapsed()); + } + let mut sorted = elapsed_all.clone(); + sorted.sort(); + let median = sorted[sorted.len() / 2]; + let max = *sorted.last().unwrap(); + let slow_count = sorted + .iter() + .filter(|d| **d >= Duration::from_millis(3)) + .count(); + assert!( + slow_count > 0, + "expected ≥1 insert to block on back-pressure; all 200 returned fast \ + (median={:?}, max={:?}, slow≥3ms count={})", + median, + max, + slow_count + ); + + // Sanity: after the loop, wait a bit and confirm the flusher + // made progress. + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline { + let diag = store.diagnostic_info(); + if diag.total_time_map_entries <= 4 { + break; + } + std::thread::sleep(Duration::from_millis(20)); + } +}