Add HTTP ingress drain control

crates/core/protobuf/node_ctl_svc.proto

@@ -33,6 +33,16 @@ service NodeCtlSvc {
   // Trigger manual compaction of RocksDB databases on this node.
   rpc TriggerCompaction(TriggerCompactionRequest)
       returns (TriggerCompactionResponse);
+
+  // Stops accepting new HTTP ingress traffic and reports HTTP drain progress.
+  // The first call returns after listeners are closed and graceful shutdown has
+  // been initiated on every accepted connection. Repeat calls poll for the
+  // DRAINED state. DRAINED reports only HTTP listener and connection
+  // quiescence; it is not a process-shutdown or worker-drain barrier. The
+  // transition is scoped to the current node generation and is irreversible
+  // until the process restarts.
+  rpc DrainHttpIngress(DrainHttpIngressRequest)
+      returns (DrainHttpIngressResponse);
 }
 
 enum ClusterFeature {
@@ -135,6 +145,29 @@ message TriggerCompactionResponse {
   repeated DatabaseCompactionResult results = 1;
 }
 
+message DrainHttpIngressRequest {
+  // The node process that the caller intends to drain. The request is rejected
+  // if this does not match the currently running generation.
+  restate.common.GenerationalNodeId expected_node_id = 1;
+}
+
+enum HttpIngressDrainStatus {
+  HTTP_INGRESS_DRAIN_STATUS_UNSPECIFIED = 0;
+  HTTP_INGRESS_DRAIN_STATUS_NOT_PRESENT = 1;
+  HTTP_INGRESS_DRAIN_STATUS_ACTIVE = 2;
+  HTTP_INGRESS_DRAIN_STATUS_DRAINING = 3;
+  // No accepted HTTP ingress connections remain; worker and process activity
+  // are not covered.
+  HTTP_INGRESS_DRAIN_STATUS_DRAINED = 4;
+}
+
+message DrainHttpIngressResponse {
+  restate.common.GenerationalNodeId node_id = 1;
+  HttpIngressDrainStatus status = 2;
+  uint64 in_flight_requests = 3;
+  uint64 in_flight_connections = 4;
+}
+
 message DatabaseCompactionResult {
   string db_name = 1;
   bool success = 2;

crates/ingress-http/src/drain.rs

@@ -0,0 +1,245 @@
+// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use tokio::sync::watch;
+use tokio_util::sync::CancellationToken;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum IngressDrainStatus {
+    Active,
+    Draining,
+    Drained,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct IngressDrainProgress {
+    pub status: IngressDrainStatus,
+    pub in_flight_requests: u64,
+    pub in_flight_connections: u64,
+}
+
+#[derive(Debug, Clone)]
+pub struct IngressDrainHandle {
+    inner: Arc<Inner>,
+}
+
+#[derive(Debug)]
+struct Inner {
+    drain: CancellationToken,
+    state: watch::Sender<State>,
+    in_flight_requests: AtomicU64,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct State {
+    status: IngressDrainStatus,
+    admission_closed: bool,
+    in_flight_connections: u64,
+    connections_awaiting_goaway: u64,
+}
+
+impl Default for IngressDrainHandle {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl IngressDrainHandle {
+    pub fn new() -> Self {
+        Self {
+            inner: Arc::new(Inner {
+                drain: CancellationToken::new(),
+                state: watch::Sender::new(State {
+                    status: IngressDrainStatus::Active,
+                    admission_closed: false,
+                    in_flight_connections: 0,
+                    connections_awaiting_goaway: 0,
+                }),
+                in_flight_requests: AtomicU64::new(0),
+            }),
+        }
+    }
+
+    /// Starts the irreversible drain and waits until the admission barrier has been crossed.
+    ///
+    /// Once this returns, listeners have been closed and every accepted connection has initiated
+    /// graceful shutdown. Existing requests may still be running.
+    pub async fn drain(&self) -> IngressDrainProgress {
+        self.start_draining();
+
+        let mut state = self.inner.state.subscribe();
+        state
+            .wait_for(|state| state.admission_closed && state.connections_awaiting_goaway == 0)
+            .await
+            .expect("the ingress drain state sender is retained by the handle");
+        self.to_progress(*state.borrow())
+    }
+
+    pub fn progress(&self) -> IngressDrainProgress {
+        self.to_progress(*self.inner.state.borrow())
+    }
+
+    pub(crate) fn cancellation_token(&self) -> CancellationToken {
+        self.inner.drain.clone()
+    }
+
+    pub(crate) fn start_draining(&self) {
+        self.inner.state.send_if_modified(|state| {
+            if state.status != IngressDrainStatus::Active {
+                return false;
+            }
+            state.status = IngressDrainStatus::Draining;
+            state.connections_awaiting_goaway = state.in_flight_connections;
+            true
+        });
+        self.inner.drain.cancel();
+    }
+
+    pub(crate) fn admission_closed(&self) {
+        self.inner.state.send_if_modified(|state| {
+            if state.admission_closed {
+                return false;
+            }
+            state.admission_closed = true;
+            true
+        });
+    }
+
+    pub(crate) fn drained(&self) {
+        self.inner.state.send_if_modified(|state| {
+            if state.status == IngressDrainStatus::Drained {
+                return false;
+            }
+            debug_assert!(state.admission_closed);
+            debug_assert_eq!(state.in_flight_connections, 0);
+            state.status = IngressDrainStatus::Drained;
+            true
+        });
+    }
+
+    pub(crate) fn connection_started(&self) -> ConnectionGuard {
+        self.inner.state.send_modify(|state| {
+            state.in_flight_connections += 1;
+            if state.status != IngressDrainStatus::Active {
+                state.connections_awaiting_goaway += 1;
+            }
+        });
+        ConnectionGuard {
+            drain: self.clone(),
+            goaway_initiated: false,
+        }
+    }
+
+    pub(crate) fn request_started(&self) -> RequestGuard {
+        self.inner
+            .in_flight_requests
+            .fetch_add(1, Ordering::Relaxed);
+        RequestGuard(self.clone())
+    }
+
+    fn to_progress(&self, state: State) -> IngressDrainProgress {
+        IngressDrainProgress {
+            status: state.status,
+            in_flight_requests: self.inner.in_flight_requests.load(Ordering::Relaxed),
+            in_flight_connections: state.in_flight_connections,
+        }
+    }
+}
+
+pub(crate) struct ConnectionGuard {
+    drain: IngressDrainHandle,
+    goaway_initiated: bool,
+}
+
+impl ConnectionGuard {
+    pub(crate) fn goaway_initiated(&mut self) {
+        if self.goaway_initiated {
+            return;
+        }
+        self.goaway_initiated = true;
+        self.drain.inner.state.send_modify(|state| {
+            debug_assert!(state.connections_awaiting_goaway > 0);
+            state.connections_awaiting_goaway -= 1;
+        });
+    }
+}
+
+impl Drop for ConnectionGuard {
+    fn drop(&mut self) {
+        self.drain.inner.state.send_modify(|state| {
+            debug_assert!(state.in_flight_connections > 0);
+            state.in_flight_connections -= 1;
+            if state.status != IngressDrainStatus::Active && !self.goaway_initiated {
+                debug_assert!(state.connections_awaiting_goaway > 0);
+                state.connections_awaiting_goaway -= 1;
+            }
+        });
+    }
+}
+
+pub(crate) struct RequestGuard(IngressDrainHandle);
+
+impl Drop for RequestGuard {
+    fn drop(&mut self) {
+        let previous = self
+            .0
+            .inner
+            .in_flight_requests
+            .fetch_sub(1, Ordering::Relaxed);
+        debug_assert!(previous > 0);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn drain_waits_for_admission_barrier_and_is_idempotent() {
+        let drain = IngressDrainHandle::new();
+        let mut connection = drain.connection_started();
+        let request = drain.request_started();
+
+        let pending_drain = tokio::spawn({
+            let drain = drain.clone();
+            async move { drain.drain().await }
+        });
+        tokio::task::yield_now().await;
+
+        drain.admission_closed();
+        assert!(
+            tokio::time::timeout(Duration::from_millis(10), pending_drain)
+                .await
+                .is_err(),
+            "drain must wait until existing connections have initiated GOAWAY"
+        );
+
+        connection.goaway_initiated();
+        let progress = tokio::time::timeout(Duration::from_secs(1), drain.drain())
+            .await
+            .unwrap();
+        assert_eq!(progress.status, IngressDrainStatus::Draining);
+        assert_eq!(progress.in_flight_connections, 1);
+        assert_eq!(progress.in_flight_requests, 1);
+
+        drop(request);
+        drop(connection);
+        drain.drained();
+        let progress = drain.drain().await;
+        assert_eq!(progress.status, IngressDrainStatus::Drained);
+        assert_eq!(progress.in_flight_connections, 0);
+        assert_eq!(progress.in_flight_requests, 0);
+    }
+}

crates/ingress-http/src/layers/in_flight_requests.rs

@@ -0,0 +1,83 @@
+// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use pin_project_lite::pin_project;
+use tower::{Layer, Service};
+
+use crate::drain::{IngressDrainHandle, RequestGuard};
+
+#[derive(Clone)]
+pub(crate) struct InFlightRequestsLayer(IngressDrainHandle);
+
+impl InFlightRequestsLayer {
+    pub(crate) fn new(drain: IngressDrainHandle) -> Self {
+        Self(drain)
+    }
+}
+
+impl<S> Layer<S> for InFlightRequestsLayer {
+    type Service = InFlightRequests<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        InFlightRequests {
+            inner,
+            drain: self.0.clone(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct InFlightRequests<S> {
+    inner: S,
+    drain: IngressDrainHandle,
+}
+
+impl<S, Request> Service<Request> for InFlightRequests<S>
+where
+    S: Service<Request>,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = InFlightRequestFuture<S::Future>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, request: Request) -> Self::Future {
+        InFlightRequestFuture {
+            inner: self.inner.call(request),
+            _guard: self.drain.request_started(),
+        }
+    }
+}
+
+pin_project! {
+    pub(crate) struct InFlightRequestFuture<F> {
+        #[pin]
+        inner: F,
+        _guard: RequestGuard,
+    }
+}
+
+impl<F> Future for InFlightRequestFuture<F>
+where
+    F: Future,
+{
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        self.project().inner.poll(cx)
+    }
+}

crates/ingress-http/src/layers/mod.rs

@@ -8,5 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+pub(crate) mod in_flight_requests;
 pub mod load_shed;
 pub mod tracing_context_extractor;

crates/ingress-http/src/lib.rs

@@ -8,12 +8,14 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+mod drain;
 mod handler;
 mod layers;
 mod metric_definitions;
 mod rpc_request_dispatcher;
 mod server;
 
+pub use drain::{IngressDrainHandle, IngressDrainProgress, IngressDrainStatus};
 pub use rpc_request_dispatcher::InvocationClientRequestDispatcher;
 pub use server::{HyperServerIngress, IngressServerError};