feat(atenet): graceful termination of atenet-router

cmd/atenet/README.md

@@ -23,6 +23,14 @@ likely be split in the future for better scalability.)
 * Upstream: Envoy's `ORIGINAL_DST` actor cluster dials the actor's in-worker
   `atunnel` ingress server on the worker pod's port 443 over mTLS, using the
   address `atenet router`'s ext_proc resolved into `x-ate-original-dst`.
+* Termination: the router drains gracefully on SIGTERM (readiness flip →
+  endpoint propagation → Envoy admin-API drain → ext_proc drain), and the
+  Envoy container's `preStop` hook waits for the router's drain-complete
+  marker on a pod-shared emptyDir — so established connections and parked
+  requests finish instead of resetting. The whole sequence must fit within
+  `terminationGracePeriodSeconds` (see the manifest comments). Upgrades are
+  whole-system swaps (#473) rather than per-Deployment rolling updates; the
+  drain is what makes the old system's termination lossless.
 
 RBAC permissions:
 * read, list on ActorTemplate

cmd/atenet/internal/router/README.md

@@ -19,6 +19,13 @@ Router has several responsibilities:
   worker-pool saturation, retrying the resume until the actor is routable or a
   bounded wait elapses, instead of failing fast. See
   [docs/request-parking.md](../../../../../docs/request-parking.md).
+* Drains gracefully on SIGTERM: flips `/readyz` so the Service stops sending
+  new connections, waits out endpoint propagation (`--drain-delay`), drives
+  Envoy's admin API to drain established connections, gracefully stops the
+  ext_proc server so parked requests finish normally (`--drain-timeout`,
+  derived from the parking budget), then writes a drain-complete marker that
+  releases the Envoy container's `preStop` hook. See `drain.go` and
+  `envoydrain.go`.
 
 ## status page
 

cmd/atenet/internal/router/cmd.go

@@ -76,6 +76,13 @@ func NewRouterCmd() *cobra.Command {
 	cmd.Flags().Float64Var(&cfg.ParkedRequest.RetryFactor, "parked-request-retry-factor", defaultParkedRequestRetryFactor, "Multiplier applied to the retry delay after each attempt; must be >= 1")
 	cmd.Flags().Float64Var(&cfg.ParkedRequest.RetryJitter, "parked-request-retry-jitter", defaultParkedRequestRetryJitter, "Random fraction in [0, 1) added to each retry delay to de-synchronize parked requests")
 	cmd.Flags().IntVar(&cfg.ExtProcMaxRequests, "extproc-max-requests", 0, "Circuit-breaker max_requests for Envoy's ext_proc cluster; 0 (the default) derives it as twice --parked-request-max (minimum 1024). Explicit values must be >= --parked-request-max: every parked request holds one slot for its full wait, and the excess is fast-path headroom")
+	// Graceful shutdown knobs. The router sits behind a Service, so
+	// route-drain window is needed: after SIGTERM the readiness flip
+	// must propagate to the Service endpoints before the drain starts.
+	cmd.Flags().DurationVar(&cfg.DrainDelay, "drain-delay", 13*time.Second, "How long to keep serving after SIGTERM before starting the drain, covering readiness-probe detection and Service endpoint propagation")
+	cmd.Flags().DurationVar(&cfg.DrainTimeout, "drain-timeout", 0, "Deadline for the ext_proc drain on shutdown; streams still open past it (parked requests included) are forcefully cancelled. 0 (the default) derives --parked-request-budget + the actor route timeout + margin so parked requests always finish normally. Explicit values must be >= --parked-request-budget")
+	cmd.Flags().StringVar(&cfg.EnvoyAdminAddr, "envoy-admin-address", "127.0.0.1:9901", "Envoy admin interface the shutdown sequence drives to drain the sidecar (healthcheck/fail, drain_listeners, stats polling). Ignored with --atenet-router=agentgateway")
+	cmd.Flags().StringVar(&cfg.DrainCompleteFile, "drain-complete-file", defaultDrainCompleteFile, "Marker file created (on a pod-shared emptyDir) once the shutdown drain completes; the Envoy container's preStop hook polls for it so Envoy exits as soon as — and no sooner than — the drain is done. Removed at startup to defuse stale markers. Empty disables the handshake")
 
 	return cmd
 }

cmd/atenet/internal/router/config.go

@@ -99,6 +99,21 @@ type routerConfig struct {
 	// excess is fast-path headroom for requests to already-running actors.
 	// 0 derives it from the parking lot — see extProcMaxRequests.
 	ExtProcMaxRequests int
+
+	// DrainDelay is how long the router serves after SIGTERM before draining,
+	// allowing readiness flip propagation to Service endpoints. DrainTimeout
+	// bounds the ext_proc drain (0 derives it automatically — see drainTimeout).
+	DrainDelay   time.Duration
+	DrainTimeout time.Duration
+
+	// EnvoyAdminAddr is the Envoy admin interface the drain sequence drives
+	// (healthcheck/fail, drain_listeners, stats polling). Same-pod loopback.
+	EnvoyAdminAddr string
+
+	// DrainCompleteFile is the marker file written once shutdown completes,
+	// releasing Envoy's preStop hook on the shared emptyDir. Removed at startup to
+	// defuse stale markers. Empty disables the handshake.
+	DrainCompleteFile string
 }
 
 func (c routerConfig) atenetRouter() atenetRouter {
@@ -128,6 +143,24 @@ func (c routerConfig) extProcMaxRequests() int {
 	return derived
 }
 
+// drainTimeoutMargin is the slack added on top of the bounded in-flight work
+// when deriving the drain timeout, mirroring the +5s Envoy ext_proc
+// MessageTimeout margin so the router always sheds before a hard cut.
+const drainTimeoutMargin = 5 * time.Second
+
+// drainTimeout resolves the effective ext_proc drain deadline: an explicit
+// flag wins; 0 derives park budget + the DEFAULT route timeout + margin. The
+// derivation deliberately ignores a configured --route-timeout so a raised
+// route ceiling cannot silently stretch shutdown past the pod's grace period
+// (see defaultRouteTimeout); operators pair a long route timeout with an
+// explicit --drain-timeout instead.
+func (c routerConfig) drainTimeout(parkCfg ParkedRequestConfig) time.Duration {
+	if c.DrainTimeout > 0 {
+		return c.DrainTimeout
+	}
+	return parkCfg.Budget + defaultRouteTimeout + drainTimeoutMargin
+}
+
 // validate rejects flag combinations that would make the router misbehave
 // rather than merely differ.
 func (c routerConfig) validate() error {
@@ -147,5 +180,15 @@ func (c routerConfig) validate() error {
 		return fmt.Errorf("--extproc-max-requests (%d) must be >= --parked-request-max (%d): a circuit breaker below the parking lot silently truncates it with Envoy-generated 503s",
 			c.ExtProcMaxRequests, c.ParkedRequest.Max)
 	}
+	if c.DrainDelay < 0 {
+		return fmt.Errorf("--drain-delay must not be negative, got %s", c.DrainDelay)
+	}
+	if c.DrainTimeout < 0 {
+		return fmt.Errorf("--drain-timeout must not be negative, got %s (0 derives it from --parked-request-budget)", c.DrainTimeout)
+	}
+	if c.DrainTimeout > 0 && c.ParkedRequest.enabled() && c.DrainTimeout < c.ParkedRequest.normalized().Budget {
+		return fmt.Errorf("--drain-timeout (%s) must be >= --parked-request-budget (%s): a drain shorter than the parking budget resets parked requests on shutdown instead of letting them finish",
+			c.DrainTimeout, c.ParkedRequest.normalized().Budget)
+	}
 	return nil
 }

cmd/atenet/internal/router/config_test.go

@@ -17,6 +17,7 @@ package router
 import (
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestRouterConfigValidate(t *testing.T) {
@@ -60,6 +61,33 @@ func TestRouterConfigValidate(t *testing.T) {
 			name: "parking disabled ignores the relation",
 			cfg:  routerConfig{ExtProcMaxRequests: 8, ParkedRequest: ParkedRequestConfig{Max: 0}},
 		},
+		{
+			name:    "drain-timeout below the parking budget rejected",
+			cfg:     routerConfig{ParkedRequest: ParkedRequestConfig{Budget: 5 * time.Second, Max: 1024}, DrainTimeout: 2 * time.Second},
+			wantErr: "must be >= --parked-request-budget",
+		},
+		{
+			name: "drain-timeout equal to the parking budget accepted",
+			cfg:  routerConfig{ParkedRequest: ParkedRequestConfig{Budget: 5 * time.Second, Max: 1024}, DrainTimeout: 5 * time.Second},
+		},
+		{
+			name: "drain-timeout above the parking budget accepted",
+			cfg:  routerConfig{ParkedRequest: ParkedRequestConfig{Budget: 5 * time.Second, Max: 1024}, DrainTimeout: 30 * time.Second},
+		},
+		{
+			name: "short drain-timeout with parking disabled accepted",
+			cfg:  routerConfig{ParkedRequest: ParkedRequestConfig{Max: 0}, DrainTimeout: time.Second},
+		},
+		{
+			name:    "negative drain-timeout rejected",
+			cfg:     routerConfig{ParkedRequest: ParkedRequestConfig{Max: defaultParkedRequestMax}, DrainTimeout: -time.Second},
+			wantErr: "--drain-timeout must not be negative",
+		},
+		{
+			name:    "negative drain-delay rejected",
+			cfg:     routerConfig{ParkedRequest: ParkedRequestConfig{Max: defaultParkedRequestMax}, DrainDelay: -time.Second},
+			wantErr: "--drain-delay must not be negative",
+		},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
@@ -117,6 +145,50 @@ func TestRouterConfigExtProcMaxRequests(t *testing.T) {
 	}
 }
 
+func TestRouterConfigDrainTimeout(t *testing.T) {
+	tests := []struct {
+		name    string
+		cfg     routerConfig
+		parkCfg ParkedRequestConfig
+		want    time.Duration
+	}{
+		{
+			name:    "auto derives budget + route timeout + margin",
+			cfg:     routerConfig{DrainTimeout: 0},
+			parkCfg: ParkedRequestConfig{Budget: 5 * time.Second, Max: 1024}.normalized(),
+			want:    5*time.Second + defaultRouteTimeout + drainTimeoutMargin,
+		},
+		{
+			name:    "auto scales with a larger budget",
+			cfg:     routerConfig{DrainTimeout: 0},
+			parkCfg: ParkedRequestConfig{Budget: 30 * time.Second, Max: 1024}.normalized(),
+			want:    30*time.Second + defaultRouteTimeout + drainTimeoutMargin,
+		},
+		{
+			name: "parking disabled still derives from the normalized default budget",
+			cfg:  routerConfig{DrainTimeout: 0},
+			// normalized() fills Budget even when Max disables parking, so the
+			// derived drain still covers a later re-enable without a restart
+			// surprise.
+			parkCfg: ParkedRequestConfig{Max: 0}.normalized(),
+			want:    defaultParkedRequestBudget + defaultRouteTimeout + drainTimeoutMargin,
+		},
+		{
+			name:    "explicit value wins over derivation",
+			cfg:     routerConfig{DrainTimeout: 42 * time.Second},
+			parkCfg: ParkedRequestConfig{Budget: 5 * time.Second, Max: 1024}.normalized(),
+			want:    42 * time.Second,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := tc.cfg.drainTimeout(tc.parkCfg); got != tc.want {
+				t.Errorf("drainTimeout() = %s, want %s", got, tc.want)
+			}
+		})
+	}
+}
+
 func TestSetOtlpCollector(t *testing.T) {
 	// No collector address may keep the router from starting. The address
 	// defaults to OTEL_EXPORTER_OTLP_ENDPOINT, which also feeds the router's

cmd/atenet/internal/router/drain.go

@@ -0,0 +1,162 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package router
+
+import (
+	"context"
+	"errors"
+	"io/fs"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"time"
+
+	"github.com/agent-substrate/substrate/internal/serverboot"
+)
+
+// defaultDrainCompleteFile is where the drain sequence leaves its completion
+// marker. The manifest mounts an emptyDir here in both containers; the
+// dataplane container's preStop hook polls the same path.
+const defaultDrainCompleteFile = "/var/run/atenet/drain-complete"
+
+// removeStaleDrainMarker deletes a leftover drain-complete marker at startup.
+// The emptyDir the marker lives on survives container restarts within the
+// pod, and a stale marker would let the dataplane container's preStop hook
+// exit the instant a later drain begins — before any connection has drained.
+func removeStaleDrainMarker(ctx context.Context, path string) {
+	if path == "" {
+		return
+	}
+	if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
+		slog.WarnContext(ctx, "Failed to remove stale drain-complete marker", slog.String("path", path), slog.Any("err", err))
+	}
+}
+
+// writeDrainMarker creates the drain-complete marker, releasing the dataplane
+// container's preStop hook. Failure is logged, never fatal: the kubelet still bounds the
+// hook at terminationGracePeriodSeconds, so a missing marker degrades to a
+// slower exit, not a wedge.
+func writeDrainMarker(ctx context.Context, path string) {
+	if path == "" {
+		return
+	}
+	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
+		slog.WarnContext(ctx, "Failed to create drain-complete marker directory", slog.String("path", path), slog.Any("err", err))
+		return
+	}
+	if err := os.WriteFile(path, []byte(time.Now().UTC().Format(time.RFC3339)+"\n"), 0o644); err != nil {
+		slog.WarnContext(ctx, "Failed to write drain-complete marker", slog.String("path", path), slog.Any("err", err))
+		return
+	}
+	slog.InfoContext(ctx, "Drain-complete marker written", slog.String("path", path))
+}
+
+// grpcStopper is the subset of *grpc.Server the drain sequence drives.
+type grpcStopper interface {
+	// GracefulStop stops accepting new connections and RPCs and blocks until
+	// all in-flight RPCs (parked ext_proc streams, most of all) finish.
+	GracefulStop()
+	// Stop cancels in-flight RPCs and closes all connections immediately.
+	Stop()
+}
+
+// dataplaneDrainer actively drains the dataplane proxy sidecar: established
+// connections finish and their in-flight requests complete before the
+// implementation returns. envoyDrainer is the Envoy implementation (driving
+// the admin API); the orchestrator below stays agnostic of which proxy is
+// deployed.
+type dataplaneDrainer interface {
+	// Drain blocks until the dataplane has quiesced or ctx expires; an error
+	// reports an incomplete drain and the shutdown sequence continues.
+	Drain(ctx context.Context) error
+}
+
+// drainParams wires the shutdown sequence. The order is forced by the
+// ext_proc filter being failClosed in the dataplane: the ext_proc server must
+// outlive the dataplane's drain, because any request the dataplane still
+// accepts during its drain window needs ext_proc answering.
+type drainParams struct {
+	readiness *serverboot.Readiness
+	// delay is the route-drain window: after the readiness flip, how long to
+	// keep serving while the Service endpoints drop this pod.
+	delay time.Duration
+	// dataplane, when non-nil, is drained after the delay and before the
+	// ext_proc server stops, bounded by dataplaneWindow. nil means the
+	// deployed dataplane offers the router no drain hook and manages its own
+	// termination; the sequence then proceeds directly to the ext_proc drain.
+	dataplane       dataplaneDrainer
+	dataplaneWindow time.Duration
+	// extproc is the ext_proc gRPC server; timeout bounds its graceful drain
+	// (sized >= the parking budget so parked requests finish normally).
+	extproc grpcStopper
+	timeout time.Duration
+	// stopRest cancels the work context, stopping the remaining subsystems
+	// (xDS, controller, health checker, statusz) once no traffic depends on
+	// them.
+	stopRest func()
+}
+
+// drainOnShutdown drives graceful shutdown when ctx is cancelled (SIGTERM or
+// interrupt): flip readiness (Service stops sending new connections), wait
+// out the propagation delay, drain the dataplane (established connections
+// finish), then drain ext_proc so parked requests complete — force-stopping
+// past the timeout — and finally stop everything else. The returned channel
+// closes once the sequence completes, so Run can block on it before letting
+// the deferred tracer/meter flushes run.
+func drainOnShutdown(ctx context.Context, p drainParams) <-chan struct{} {
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		<-ctx.Done()
+		slog.InfoContext(ctx, "Shutdown signal received; draining")
+		p.readiness.MarkNotReady()
+		time.Sleep(p.delay)
+
+		if p.dataplane != nil {
+			slog.InfoContext(ctx, "Draining dataplane", slog.Duration("window", p.dataplaneWindow))
+			dpCtx, cancel := context.WithTimeout(context.Background(), p.dataplaneWindow)
+			if err := p.dataplane.Drain(dpCtx); err != nil {
+				// TODO: Add a shutdown-outcome metric (clean vs
+				// dataplane-drain-incomplete vs ext_proc force-stopped) so
+				// unclean shutdowns are visible in dashboards, not only logs.
+				slog.WarnContext(ctx, "Dataplane drain incomplete; continuing shutdown", slog.Any("err", err))
+			} else {
+				slog.InfoContext(ctx, "Dataplane drained")
+			}
+			cancel()
+		}
+
+		slog.InfoContext(ctx, "Starting ext_proc drain")
+		drainComplete := make(chan struct{})
+		go func() {
+			p.extproc.GracefulStop()
+			close(drainComplete)
+		}()
+		select {
+		case <-drainComplete:
+			slog.InfoContext(ctx, "ext_proc drain completed within deadline")
+		case <-time.After(p.timeout):
+			// TODO: Count this in the shutdown-outcome metric above: a
+			// force-stop here means in-flight ext_proc streams (parked
+			// requests included) were cancelled — the unclean-shutdown signal
+			// operators most need to see.
+			slog.WarnContext(ctx, "ext_proc drain deadline exceeded; forcing stop")
+			p.extproc.Stop()
+		}
+
+		p.stopRest()
+	}()
+	return done
+}