diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index ee2e772..7fbc9ce 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -74,6 +74,10 @@ pub struct PrecomputeEngine { receivers: Option>>, /// Shared ingest agg_configs, swappable at runtime. ingest_agg_configs: Arc>>>, + /// Test-support wall-clock override, applied to every spawned worker. + /// See `with_now_ms_fn`. `None` in production — each worker keeps its + /// default `SystemTime::now`-backed clock. + now_ms_fn: Option i64 + Send + Sync>>, } impl PrecomputeEngine { @@ -119,6 +123,7 @@ impl PrecomputeEngine { senders, receivers: Some(receivers), ingest_agg_configs, + now_ms_fn: None, } } @@ -127,6 +132,16 @@ impl PrecomputeEngine { self.diagnostics.clone() } + /// Test-support builder: override the wall-clock source used by every + /// worker's wall-clock fallback (`flush_all`). Lets integration tests + /// drive the fallback deterministically through the real ingest pipeline + /// instead of racing a real clock with `tokio::time::sleep`. Must be + /// called before `run()`. Production code never calls this. + pub fn with_now_ms_fn(mut self, f: impl Fn() -> i64 + Send + Sync + 'static) -> Self { + self.now_ms_fn = Some(Arc::new(f)); + self + } + /// Return a handle for applying runtime config updates to this engine. /// Must be called before `run()`. pub fn handle(&self) -> PrecomputeEngineHandle { @@ -160,7 +175,7 @@ impl PrecomputeEngine { // Spawn workers let mut worker_handles = Vec::with_capacity(num_workers); for (id, rx) in receivers.into_iter().enumerate() { - let worker = Worker::new( + let mut worker = Worker::new( id, rx, self.output_sink.clone(), @@ -177,6 +192,10 @@ impl PrecomputeEngine { self.diagnostics.worker_watermarks[id].clone(), self.diagnostics.worker_watermarks.to_vec(), ); + if let Some(now_ms_fn) = &self.now_ms_fn { + let now_ms_fn = now_ms_fn.clone(); + worker.set_now_ms_fn(Box::new(move || now_ms_fn())); + } let handle = tokio::spawn(async move { worker.run().await; }); diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index 35edc88..da3c7b9 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -103,7 +103,8 @@ pub struct Worker { wall_clock_grace_period_ms: i64, /// Injectable clock returning current wall-clock time in milliseconds /// since the unix epoch. Production uses `SystemTime::now`; tests - /// override with a deterministic fake via `set_now_ms_fn`. + /// override with a deterministic fake via `set_now_ms_fn` (directly, or + /// via `PrecomputeEngine::with_now_ms_fn` for integration tests). now_ms_fn: Box i64 + Send + Sync>, } @@ -145,12 +146,14 @@ impl Worker { } } - /// Test/diagnostic-only setter for the wall-clock source. Replaces the - /// default `SystemTime::now`-backed clock with a deterministic fake so - /// unit tests can drive the wall-clock fallback in `flush_all` without - /// `std::thread::sleep`. Production code never calls this. - #[cfg(test)] - pub fn set_now_ms_fn(&mut self, f: Box i64 + Send + Sync>) { + /// Test-support setter for the wall-clock source. Replaces the default + /// `SystemTime::now`-backed clock with a deterministic fake so tests can + /// drive the wall-clock fallback in `flush_all` without real sleeping. + /// Crate-visible (not `#[cfg(test)]`-gated) so `PrecomputeEngine:: + /// with_now_ms_fn` can reach it from integration tests in `tests/`, + /// which link the library normally and don't see `#[cfg(test)]` items. + /// Production code never calls this. + pub(crate) fn set_now_ms_fn(&mut self, f: Box i64 + Send + Sync>) { self.now_ms_fn = f; }