docs: design for SimpleMapStore persistence (mem limit + disk flush)

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"

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<String>,
+
+    /// 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<u64>,
 }
 
 #[tokio::main]
@@ -89,12 +119,62 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
     );
 
     // Create the store
-    let store: Arc<dyn query_engine_rust::stores::Store> =
+    let store: Arc<dyn query_engine_rust::stores::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, 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 {

asap-query-engine/src/data_model/traits.rs

@@ -44,6 +44,20 @@ pub trait AggregateCore: SerializableToSink + Send + Sync {
         key: &Option<KeyByLabelValues>,
         query_kwargs: &HashMap<String, String>,
     ) -> Result<f64, Box<dyn std::error::Error + Send + Sync>>;
+
+    /// 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

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<String>,
+
+    /// 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<u64>,
 }
 
 #[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 {

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<Vec<crate::KeyByLabelValues>> {
         None
     }

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<Vec<crate::KeyByLabelValues>> {
         None
     }

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>()
+            + (self.added.len() + self.removed.len()) * BYTES_PER_ENTRY
+    }
+
     fn get_keys(&self) -> Option<Vec<KeyByLabelValues>> {
         if !self.removed.is_empty() {
             panic!("DeltaSetAggregatorAccumulator does not support get_keys when removed items are present");

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<Vec<crate::KeyByLabelValues>> {
         None
     }

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::<Self>()
+    }
+
     fn get_keys(&self) -> Option<Vec<crate::KeyByLabelValues>> {
         None
     }