feat(precompute): event-time barrier for Remote Write producers

README.md

@@ -53,6 +53,18 @@ are dropped. Remote Write gets 503 when admitting it would leave more than
 stopped, which also makes `/-/ready` fail. `/api/v1/status/asap` reports the
 engine's sample and pane counters as JSON.
 
+A producer can close panes without waiting for later samples:
+`POST /api/v1/precompute/watermark?time_ms=T` states that every sample at or
+before `T` has been written. After the writes acknowledged before it have
+been ingested, every pane ending at or before `T` (capped at the wall clock)
+is closed, so the last pane of a series that stopped sending is built and
+served. A repeated or lower `T` does nothing and input stays open. A sample
+that arrives later for a pane the barrier closed (one ending at or before
+`T`) counts as late; a later sample in a pane still open is accepted, even if
+it is at or before `T`. The response is `{"barrier_ms": B}`, the barrier in
+force: `T` capped at the wall clock, or an earlier higher barrier. With several producers,
+send the barrier only after all of them have written through `T`.
+
 An instant query whose text is one of the plan's queries is answered by the
 plan (`x-asap-execution: plan`) at its evaluation time `T`. Stored outputs
 built in an `IngestionRun` are read from the SDS as the panes ending at

data_plane/src/engine.rs

@@ -53,6 +53,10 @@ enum Command {
         raw: BTreeMap<u64, Vec<RawSample>>,
         reply: Reply<Vec<ResultSample>>,
     },
+    AdvanceWatermark {
+        barrier_ms: i64,
+        reply: tokio::sync::oneshot::Sender<i64>,
+    },
     Stop,
 }
 
@@ -223,6 +227,21 @@ impl EngineHandle {
         *self.stats.lock().unwrap_or_else(|e| e.into_inner())
     }
 
+    /// Queues an event-time barrier behind every write already acknowledged
+    /// and waits until it has run: every pane ending at or before
+    /// `barrier_ms` (capped at the wall clock) is then closed. Returns the
+    /// barrier in force; see [`Precompute::advance_watermark_at`].
+    pub async fn advance_watermark(&self, barrier_ms: i64) -> Result<i64, IngestError> {
+        if !self.is_running() {
+            return Err(IngestError::Stopped);
+        }
+        let (reply, done) = tokio::sync::oneshot::channel();
+        self.commands
+            .send(Command::AdvanceWatermark { barrier_ms, reply })
+            .map_err(|_| IngestError::Stopped)?;
+        done.await.map_err(|_| IngestError::Stopped)
+    }
+
     /// For tests: stops the engine thread after the commands queued before.
     /// Samples queued between this call and the stop are acknowledged and
     /// then lost; ingestion after the stop fails with
@@ -321,6 +340,11 @@ impl Engine {
                         .unwrap_or_else(|e| e.into_inner()) = self.query.query_run_stats;
                     let _ = reply.send(answer);
                 }
+                Command::AdvanceWatermark { barrier_ms, reply } => {
+                    let barrier_ms = self.precompute.advance_watermark(barrier_ms);
+                    *shared.stats.lock().unwrap_or_else(|e| e.into_inner()) = self.precompute.stats;
+                    let _ = reply.send(barrier_ms);
+                }
                 Command::Stop => break,
             }
         }

data_plane/src/lib.rs

@@ -106,6 +106,7 @@ pub fn router(state: Arc<DataPlane>) -> Router {
         .route("/api/v1/query", get(query).post(query))
         .route("/api/v1/query_range", get(query_range).post(query_range))
         .route("/api/v1/status/asap", get(status))
+        .route("/api/v1/precompute/watermark", post(watermark))
         .route("/-/ready", get(ready))
         .with_state(state)
 }
@@ -159,6 +160,54 @@ async fn write(State(state): State<Arc<DataPlane>>, body: Bytes) -> Response {
     }
 }
 
+/// `POST /api/v1/precompute/watermark?time_ms=T`: the producer states that
+/// every sample at or before `T` has been written. Once the engine has run
+/// every write acknowledged before this request, every pane with
+/// `end <= T` is closed, so the last pane of a series that stopped sending
+/// is built without a later sample. `T` is capped at the wall clock; a
+/// repeated or lower `T` does nothing. A sample that arrives later for a
+/// pane the barrier closed counts as late. Input stays open. With several
+/// producers, call it only after all of them have written through `T`.
+/// Answers `{"barrier_ms": B}`, the barrier in force (`T` capped at the wall
+/// clock, or a higher earlier barrier); 400 `bad_data` with the reason when
+/// `time_ms` is missing or not an integer.
+async fn watermark(State(state): State<Arc<DataPlane>>, RawQuery(query): RawQuery) -> Response {
+    let time_ms = form_urlencoded::parse(query.unwrap_or_default().as_bytes())
+        .find(|(name, _)| name == "time_ms")
+        .map(|(_, value)| value.into_owned());
+    let time_ms = match time_ms.as_deref().map(str::parse::<i64>) {
+        Some(Ok(time_ms)) => time_ms,
+        Some(Err(_)) => {
+            let reason = format!(
+                "invalid parameter \"time_ms\": {:?} is not an integer number of milliseconds",
+                time_ms.unwrap_or_default()
+            );
+            return (
+                StatusCode::BAD_REQUEST,
+                axum::Json(prometheus::bad_data(&reason)),
+            )
+                .into_response();
+        }
+        None => {
+            let reason = "missing parameter \"time_ms\" (integer milliseconds)";
+            return (
+                StatusCode::BAD_REQUEST,
+                axum::Json(prometheus::bad_data(reason)),
+            )
+                .into_response();
+        }
+    };
+    match state.engine.advance_watermark(time_ms).await {
+        Ok(barrier_ms) => {
+            axum::Json(serde_json::json!({ "barrier_ms": barrier_ms })).into_response()
+        }
+        Err(error) => {
+            tracing::error!(%error, "event-time barrier refused");
+            (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response()
+        }
+    }
+}
+
 /// Ready while the engine thread runs; a stopped engine never ingests
 /// again, so the data plane must be restarted.
 async fn ready(State(state): State<Arc<DataPlane>>) -> StatusCode {

data_plane/src/precompute.rs

@@ -208,6 +208,9 @@ pub struct Precompute {
     /// passes it, so a series stamped ahead of the others cannot close
     /// panes early.
     newest_clock_ms: Option<i64>,
+    /// The newest event-time barrier a producer declared: every pane ending
+    /// at or before it is closed.
+    barrier_ms: Option<i64>,
     pub store: SummaryStore<Batch>,
     pub stats: PrecomputeStats,
 }
@@ -260,6 +263,7 @@ impl Precompute {
             max_future_ms: millis(options.max_future_ms),
             newest_sample_ms: None,
             newest_clock_ms: None,
+            barrier_ms: None,
             store: SummaryStore::new(),
             stats: PrecomputeStats::default(),
         })
@@ -281,21 +285,42 @@ impl Precompute {
 
     /// The watermark: no sample older than this is expected any more. It is
     /// the newest routed sample, capped at the wall clock, minus the allowed
-    /// lateness.
+    /// lateness, or the newest barrier if that is later.
     pub fn watermark_ms(&self) -> Option<i64> {
         let newest = match self.newest_clock_ms {
-            Some(clock) => self.newest_sample_ms?.min(clock),
-            None => self.newest_sample_ms?,
+            Some(clock) => self.newest_sample_ms.map(|newest| newest.min(clock)),
+            None => self.newest_sample_ms,
         };
-        Some(newest.saturating_sub(self.allowed_lateness_ms))
+        let from_samples = newest.map(|newest| newest.saturating_sub(self.allowed_lateness_ms));
+        from_samples.max(self.barrier_ms)
+    }
+
+    /// [`Self::advance_watermark_at`] the current wall-clock time.
+    pub fn advance_watermark(&mut self, barrier_ms: i64) -> i64 {
+        self.advance_watermark_at(barrier_ms, wall_clock_ms())
+    }
+
+    /// An event-time barrier: the producer states that every sample at or
+    /// before `barrier_ms` has been written. Closes every pane with
+    /// `end <= barrier_ms` (panes are `(end - width, end]`), so a pane whose
+    /// series stopped sending closes without a later sample. The barrier is
+    /// capped at `now_ms`; a repeated or lower barrier does nothing. A sample
+    /// that arrives later for a pane the barrier closed counts as late.
+    /// Returns the barrier in force.
+    pub fn advance_watermark_at(&mut self, barrier_ms: i64, now_ms: i64) -> i64 {
+        self.newest_clock_ms = self.newest_clock_ms.max(Some(now_ms));
+        let barrier_ms = barrier_ms.min(now_ms);
+        if self.barrier_ms.is_some_and(|current| current >= barrier_ms) {
+            return self.barrier_ms.unwrap_or(barrier_ms);
+        }
+        self.barrier_ms = Some(barrier_ms);
+        self.close_panes();
+        barrier_ms
     }
 
     /// [`Self::ingest_at`] the current wall-clock time.
     pub fn ingest(&mut self, samples: &[RawSample]) {
-        let now_ms = std::time::SystemTime::now()
-            .duration_since(std::time::UNIX_EPOCH)
-            .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX));
-        self.ingest_at(samples, now_ms);
+        self.ingest_at(samples, wall_clock_ms());
     }
 
     /// Buffers `samples` into their panes, then builds every pane the
@@ -422,6 +447,13 @@ impl Precompute {
     }
 }
 
+/// Milliseconds since the Unix epoch.
+fn wall_clock_ms() -> i64 {
+    std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
+}
+
 /// End of the tumbling pane of width `pane_ms` that contains `timestamp_ms`,
 /// or `None` when it does not fit in an `i64`. Panes are
 /// `(end - pane_ms, end]`, like PromQL ranges: a sample at a pane boundary

data_plane/tests/answers.rs

@@ -391,3 +391,105 @@ async fn rejects_evaluation_times_out_of_range() {
     }
     assert!(state.engine.is_running());
 }
+
+/// After a producer's event-time barrier, the last pane of a series that
+/// stopped sending is closed, and the planned query is served from the
+/// plan instead of forwarded. Input stays open.
+#[tokio::test]
+async fn barrier_serves_the_window_of_a_stopped_series() {
+    let state = Arc::new(
+        DataPlane::new(
+            common::plan_json("example4b", false),
+            fake_prometheus(vec![]).await,
+            15_000,
+        )
+        .unwrap(),
+    );
+    let app = router(state.clone());
+    let write = |from: i64, to: i64| {
+        let request = WriteRequest {
+            timeseries: vec![TimeSeries {
+                labels: vec![
+                    Label {
+                        name: "__name__".into(),
+                        value: "latency_ms".into(),
+                    },
+                    Label {
+                        name: "job".into(),
+                        value: "a".into(),
+                    },
+                ],
+                samples: (from..to)
+                    .step_by(15_000)
+                    .map(|t| Sample {
+                        timestamp: t,
+                        value: (t / 15_000 % 240) as f64,
+                    })
+                    .collect(),
+            }],
+        };
+        Request::post("/api/v1/write")
+            .body(Body::from(encode(&request)))
+            .unwrap()
+    };
+    let written = app
+        .clone()
+        .oneshot(write(MINUTE, 70 * MINUTE))
+        .await
+        .unwrap();
+    assert_eq!(written.status(), StatusCode::NO_CONTENT);
+    let query = "quantile_over_time(0.99, latency_ms[60m])";
+    // The series stopped at 69:45; with 15 s lateness (60, 70 min] is open.
+    let response = instant(&app, query, 70 * MINUTE).await;
+    assert_eq!(response.execution.as_deref(), Some("forwarded"));
+
+    let barrier = app
+        .clone()
+        .oneshot(
+            Request::post(format!(
+                "/api/v1/precompute/watermark?time_ms={}",
+                70 * MINUTE
+            ))
+            .body(Body::empty())
+            .unwrap(),
+        )
+        .await
+        .unwrap();
+    assert_eq!(barrier.status(), StatusCode::OK);
+    let body = barrier.into_body().collect().await.unwrap().to_bytes();
+    let body: Value = serde_json::from_slice(&body).unwrap();
+    assert_eq!(body["barrier_ms"], 70 * MINUTE, "{body}");
+    let response = instant(&app, query, 70 * MINUTE).await;
+    assert_eq!(
+        response.execution.as_deref(),
+        Some("plan"),
+        "{:?}",
+        response.body
+    );
+
+    let written = app
+        .clone()
+        .oneshot(write(70 * MINUTE + 15_000, 72 * MINUTE))
+        .await
+        .unwrap();
+    assert_eq!(written.status(), StatusCode::NO_CONTENT, "input stays open");
+    for uri in [
+        "/api/v1/precompute/watermark",
+        "/api/v1/precompute/watermark?time_ms=soon",
+        "/api/v1/precompute/watermark?time_ms=1.5",
+    ] {
+        let bad = app
+            .clone()
+            .oneshot(Request::post(uri).body(Body::empty()).unwrap())
+            .await
+            .unwrap();
+        assert_eq!(bad.status(), StatusCode::BAD_REQUEST, "{uri}");
+        let body = bad.into_body().collect().await.unwrap().to_bytes();
+        let body: Value = serde_json::from_slice(&body).unwrap();
+        assert_eq!(body["errorType"], "bad_data", "{uri}");
+        assert!(
+            body["error"].as_str().unwrap().contains("time_ms"),
+            "{uri}: {body}"
+        );
+    }
+}