router: grant an in-flight resume a bounded grace past the park budget when a worker is already assigned

cmd/atenet/internal/router/parking.go

@@ -43,6 +43,22 @@ const (
 	defaultParkedRequestRetryJitter   = 0.1
 )
 
+// extProcMessageTimeoutMargin is how much longer Envoy waits on the ext_proc
+// stream than the router can possibly hold a parked request. Without a margin
+// the two deadlines race, and an Envoy win replaces the router's meaningful 503
+// ("no free workers available") with a generic gateway error.
+const extProcMessageTimeoutMargin = 2 * time.Second
+
+// extProcMessageTimeoutFor returns the ext_proc message timeout that covers the
+// router's worst-case hold on a parked request: the parking budget, plus the
+// grace an attempt already in flight when the budget elapses may get (see
+// resumeAttemptGrace), plus the margin above. The resumer's hold and Envoy's
+// patience must move together — raising one without the other reintroduces the
+// race the margin exists to prevent.
+func extProcMessageTimeoutFor(cfg ParkedRequestConfig) time.Duration {
+	return cfg.normalized().Budget + resumeAttemptGrace + extProcMessageTimeoutMargin
+}
+
 // parkOutcome is the terminal disposition of a parked request. It is recorded
 // as the `outcome` label on the parking.wait.duration histogram.
 type parkOutcome string

cmd/atenet/internal/router/resumer.go

@@ -36,6 +36,26 @@ import (
 // conflicts are retried; capacity errors fail immediately.
 const failFastResumeBudget = 15 * time.Second
 
+// resumeAttemptGrace is the extra time a resume attempt still in flight at the
+// budget gets to land, but only if a worker is already assigned to the actor
+// (see workerAssigned). Retries always stop at the budget; the grace covers the
+// one attempt underway, never another round.
+//
+// Canceling an assigned attempt is the case worth avoiding: it does not
+// release the worker, so the pool the request was waiting on stays starved,
+// and a restore that was about to succeed is discarded. Under node contention
+// those overshoot the budget by a few hundred milliseconds routinely. With no
+// worker assigned there is nothing to lose, so the request fails on its budget.
+//
+// budget+grace is the router's worst-case hold on a parked request, which the
+// ext_proc message timeout must cover; see extProcMessageTimeoutFor.
+const resumeAttemptGrace = 3 * time.Second
+
+// resumeProbeTimeout bounds the single GetActor that decides whether a worker
+// is assigned. It is spent inside the grace, not on top of it, so a slow
+// control plane cannot stretch the hold past budget+grace.
+const resumeProbeTimeout = 1 * time.Second
+
 // resumeBackoff builds the backoff between resume attempts while a request is
 // parked, from the configured retry parameters.
 //
@@ -98,6 +118,11 @@ type ActorResumer struct {
 	// budget bounds the total time a single resume operation retries before the
 	// underlying error is returned.
 	budget time.Duration
+	// grace is the extra time a resume attempt already in flight when the
+	// budget elapses is allowed, if a worker has been assigned to the actor by
+	// then. Zero — the fail-fast behavior applied when parking is off — cuts it
+	// at the budget and skips the probe entirely.
+	grace time.Duration
 	// backoff paces the retries within the budget.
 	backoff wait.Backoff
 	// nextID is a counter assigned to each incoming ResumeActor call.
@@ -119,6 +144,7 @@ func withParking(cfg ParkedRequestConfig) resumerOption {
 		r.parkEnabled = cfg.enabled()
 		if r.parkEnabled {
 			r.budget = cfg.Budget
+			r.grace = resumeAttemptGrace
 		}
 		r.backoff = resumeBackoff(cfg.RetryInterval, cfg.RetryFactor, cfg.RetryJitter)
 	}
@@ -157,6 +183,28 @@ func (r *ActorResumer) retryable(err error) bool {
 	}
 }
 
+// workerAssigned reports whether the control plane has already bound a worker
+// to the actor, i.e. whether a resume in flight is past the point where
+// abandoning it would strand that worker. RESUMING is written durably by the
+// assignment step, before the snapshot restore begins; RUNNING is the same
+// answer after it. A probe that fails or times out counts as unassigned, which
+// cancels the attempt exactly as it would without the grace.
+func (r *ActorResumer) workerAssigned(actorRef resources.ActorRef) bool {
+	ctx, cancel := context.WithTimeout(context.Background(), resumeProbeTimeout)
+	defer cancel()
+
+	actor, err := r.apiClient.GetActor(ctx, &ateapipb.GetActorRequest{Actor: actorRef.ToObjectRef()})
+	if err != nil {
+		return false
+	}
+	switch actor.GetStatus() {
+	case ateapipb.Actor_STATUS_RESUMING, ateapipb.Actor_STATUS_RUNNING:
+		return true
+	default:
+		return false
+	}
+}
+
 // ResumeActor ensures the requested actor is running. It deduplicates concurrent
 // requests within the process and, when parking is enabled, holds the request
 // while retrying transient failures until the budget elapses.
@@ -185,9 +233,35 @@ func (r *ActorResumer) ResumeActor(ctx context.Context, actorRef resources.Actor
 		var resumeResp *ateapipb.ResumeActorResponse
 		var lastRetryErr error
 
-		err := wait.ExponentialBackoffWithContext(bgCtx, backoff, func(ctx context.Context) (bool, error) {
+		// bgCtx bounds the retry loop; attemptCtx bounds the RPCs. They part
+		// company only when the budget elapses with an attempt still running:
+		// if a worker is by then assigned to the actor, that attempt gets the
+		// grace to land instead of being cancelled (see resumeAttemptGrace).
+		// With the grace disabled the two coincide, and the budget cuts the RPC
+		// exactly as it did before.
+		attemptCtx := bgCtx
+		if r.grace > 0 {
+			var cancelAttempt context.CancelFunc
+			attemptCtx, cancelAttempt = context.WithCancel(context.Background())
+			defer cancelAttempt()
+
+			stopWatchdog := context.AfterFunc(bgCtx, func() {
+				graceExpiry := time.After(r.grace) // started before the probe, so the probe is spent inside the grace
+				if !r.workerAssigned(actorRef) {
+					cancelAttempt()
+					return
+				}
+				<-graceExpiry
+				cancelAttempt()
+			})
+			// Deferred last, so it runs first on the way out: the flight's own
+			// teardown of bgCtx must not be mistaken for a budget expiry.
+			defer stopWatchdog()
+		}
+
+		err := wait.ExponentialBackoffWithContext(bgCtx, backoff, func(context.Context) (bool, error) {
 			var err error
-			resumeResp, err = r.apiClient.ResumeActor(ctx, &ateapipb.ResumeActorRequest{
+			resumeResp, err = r.apiClient.ResumeActor(attemptCtx, &ateapipb.ResumeActorRequest{
 				Actor: actorRef.ToObjectRef(),
 			})
 			if err == nil {
@@ -212,8 +286,10 @@ func (r *ActorResumer) ResumeActor(ctx context.Context, actorRef resources.Actor
 			// when the deadline lands during an in-flight ResumeActor RPC, gRPC
 			// surfaces a *status* error with code DeadlineExceeded that does not
 			// match the context sentinel, which would misreport budget exhaustion
-			// as a 504. bgCtx is this loop's only deadline source, so checking it
-			// covers both landing spots (mid-RPC and between retries).
+			// as a 504. Every way this loop can run out of time — between
+			// retries, mid-RPC, or when the grace cuts an attempt that never
+			// landed — happens at or after bgCtx's own deadline, so checking
+			// bgCtx covers all of them.
 			if lastRetryErr != nil && (bgCtx.Err() != nil || wait.Interrupted(err)) {
 				return &resumeCallResult{leaderID: reqID, err: &budgetExhaustedError{lastErr: lastRetryErr}}, nil
 			}

cmd/atenet/internal/router/resumer_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"errors"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -31,6 +32,10 @@ import (
 type resumerMockClient struct {
 	ateapipb.ControlClient
 	resumeFn func(ctx context.Context, in *ateapipb.ResumeActorRequest, opts ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error)
+	// status is what GetActor reports; the resumer probes it to decide whether
+	// a worker is already assigned. The zero value (STATUS_UNSPECIFIED) means
+	// unassigned.
+	status ateapipb.Actor_Status
 }
 
 func (m *resumerMockClient) ResumeActor(ctx context.Context, in *ateapipb.ResumeActorRequest, opts ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error) {
@@ -40,6 +45,13 @@ func (m *resumerMockClient) ResumeActor(ctx context.Context, in *ateapipb.Resume
 	return nil, status.Error(codes.Unimplemented, "unimplemented")
 }
 
+func (m *resumerMockClient) GetActor(ctx context.Context, in *ateapipb.GetActorRequest, opts ...grpc.CallOption) (*ateapipb.Actor, error) {
+	return &ateapipb.Actor{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: in.GetActor().GetAtespace(), Name: in.GetActor().GetName()},
+		Status:   m.status,
+	}, nil
+}
+
 func TestActorResumer_ResumeActor(t *testing.T) {
 	const testActorName = "actor-a"
 	const testAtespace = "team-a"
@@ -367,13 +379,15 @@ func TestActorResumer_Parking(t *testing.T) {
 					// last retryable error.
 					return nil, status.Error(codes.FailedPrecondition, "no free workers available")
 				}
-				// Later attempt: block until the park budget cancels the RPC,
-				// then return what a real gRPC client returns — a *status*
-				// error with code DeadlineExceeded that does NOT satisfy
+				// Later attempt: block until the budget cuts the RPC, then
+				// return what a real gRPC client returns — a *status* error
+				// that does NOT satisfy errors.Is(err, context.Canceled) or
 				// errors.Is(err, context.DeadlineExceeded).
 				<-ctx.Done()
 				return nil, status.FromContextError(ctx.Err()).Err()
 			},
+			// No worker assigned, so no grace: the attempt is cut at the budget.
+			status: ateapipb.Actor_STATUS_SUSPENDED,
 		}
 
 		resumer := NewActorResumer(mock, withParking(ParkedRequestConfig{Max: 1, Budget: 300 * time.Millisecond}))
@@ -415,6 +429,137 @@ func TestActorResumer_Parking(t *testing.T) {
 	})
 }
 
+// TestActorResumer_AttemptGrace covers the grace granted to a resume attempt
+// still in flight when the parking budget elapses — and the probe that decides
+// whether it is warranted. Canceling an attempt after a worker has been
+// assigned strands that worker in RESUMING and starves the pool the caller is
+// waiting on, while under node contention the snapshot restore overshoots the
+// budget by a few hundred milliseconds routinely; with no worker assigned there
+// is nothing to protect, and the request fails on its budget as before.
+func TestActorResumer_AttemptGrace(t *testing.T) {
+	const testActorName = "actor-grace"
+	const expectedIP = "10.0.0.99"
+	testActorRef := resources.ActorRef{Atespace: "team-a", Name: testActorName}
+
+	const budget = 300 * time.Millisecond
+	const grace = 2 * time.Second
+
+	// overshoot returns a ResumeActor stub whose first attempt reports a
+	// saturated pool and whose second runs past the budget, landing only if it
+	// is not cancelled first.
+	overshoot := func(lands bool) func(context.Context, *ateapipb.ResumeActorRequest, ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error) {
+		var calls atomic.Int32
+		return func(ctx context.Context, in *ateapipb.ResumeActorRequest, opts ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error) {
+			if calls.Add(1) == 1 {
+				return nil, status.Error(codes.FailedPrecondition, "no free workers available")
+			}
+			wait := budget
+			if !lands {
+				wait = time.Hour // never completes on its own
+			}
+			select {
+			case <-time.After(wait):
+			case <-ctx.Done():
+				return nil, status.FromContextError(ctx.Err()).Err()
+			}
+			return &ateapipb.ResumeActorResponse{
+				Actor:   &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Name: testActorName}, Status: ateapipb.Actor_STATUS_RUNNING, AteomPodIp: expectedIP},
+				Resumed: true,
+			}, nil
+		}
+	}
+
+	newResumer := func(mock *resumerMockClient) *ActorResumer {
+		// The production grace (resumeAttemptGrace) is seconds long by design.
+		r := NewActorResumer(mock, withParking(ParkedRequestConfig{Max: 1, Budget: budget}))
+		r.grace = grace
+		return r
+	}
+
+	t.Run("AssignedWorkerLandsWithinGrace", func(t *testing.T) {
+		// A worker was assigned, so the overshooting restore is worth waiting
+		// for rather than throwing away.
+		mock := &resumerMockClient{resumeFn: overshoot(true), status: ateapipb.Actor_STATUS_RESUMING}
+
+		start := time.Now()
+		actor, _, err := newResumer(mock).ResumeActor(context.Background(), testActorRef)
+		elapsed := time.Since(start)
+		if err != nil {
+			t.Fatalf("expected the overshooting resume to be served, got %v", err)
+		}
+		if actor.GetAteomPodIp() != expectedIP {
+			t.Errorf("expected IP %q, got %q", expectedIP, actor.GetAteomPodIp())
+		}
+		if elapsed < budget {
+			t.Errorf("expected the resume to land after the budget, took only %v", elapsed)
+		}
+		if elapsed >= budget+grace {
+			t.Errorf("expected the resume to land inside the grace, took %v", elapsed)
+		}
+	})
+
+	t.Run("UnassignedActorFailsAtBudget", func(t *testing.T) {
+		// Same overshooting attempt, but the actor never left SUSPENDED: no
+		// worker to strand, nothing to wait for.
+		mock := &resumerMockClient{resumeFn: overshoot(true), status: ateapipb.Actor_STATUS_SUSPENDED}
+
+		start := time.Now()
+		_, _, err := newResumer(mock).ResumeActor(context.Background(), testActorRef)
+		elapsed := time.Since(start)
+		if got := status.Code(err); got != codes.FailedPrecondition {
+			t.Errorf("expected the capacity error, got %v (err=%v)", got, err)
+		}
+		var exhausted *budgetExhaustedError
+		if !errors.As(err, &exhausted) {
+			t.Errorf("expected budget exhaustion, got %T (%v)", err, err)
+		}
+		if elapsed >= budget+grace {
+			t.Errorf("expected the attempt to be cut at the budget, took %v", elapsed)
+		}
+	})
+
+	t.Run("GraceIsBounded", func(t *testing.T) {
+		// A worker is assigned but the resume never lands: the grace, not the
+		// control plane, has to end the wait.
+		mock := &resumerMockClient{resumeFn: overshoot(false), status: ateapipb.Actor_STATUS_RESUMING}
+
+		start := time.Now()
+		_, _, err := newResumer(mock).ResumeActor(context.Background(), testActorRef)
+		elapsed := time.Since(start)
+		if got := status.Code(err); got != codes.FailedPrecondition {
+			t.Errorf("expected the capacity error after the grace, got %v (err=%v)", got, err)
+		}
+		if elapsed < budget+grace {
+			t.Errorf("expected the full grace to be granted, took only %v", elapsed)
+		}
+		if elapsed >= budget+2*grace {
+			t.Errorf("expected the grace to bound the wait, took %v", elapsed)
+		}
+	})
+
+	t.Run("GraceDoesNotExtendRetries", func(t *testing.T) {
+		// A stale RESUMING (left by some earlier abandoned attempt) must not
+		// buy extra retries: with the pool refusing instantly there is no
+		// attempt in flight at the budget, so exhaustion lands on time.
+		mock := &resumerMockClient{
+			resumeFn: func(ctx context.Context, in *ateapipb.ResumeActorRequest, opts ...grpc.CallOption) (*ateapipb.ResumeActorResponse, error) {
+				return nil, status.Error(codes.FailedPrecondition, "no free workers available")
+			},
+			status: ateapipb.Actor_STATUS_RESUMING,
+		}
+
+		start := time.Now()
+		_, _, err := newResumer(mock).ResumeActor(context.Background(), testActorRef)
+		elapsed := time.Since(start)
+		if got := status.Code(err); got != codes.FailedPrecondition {
+			t.Errorf("expected the capacity error, got %v (err=%v)", got, err)
+		}
+		if elapsed >= budget+grace {
+			t.Errorf("expected exhaustion at the budget, took %v", elapsed)
+		}
+	})
+}
+
 // TestActorResumer_CallerCancelDoesNotAbortFlight pins the detached-context
 // contract from both sides: a caller that disconnects while parked gets
 // context.Canceled (classified as the `canceled` outcome) WITHOUT aborting the

cmd/atenet/internal/router/router.go

@@ -25,7 +25,6 @@ import (
 	"os/signal"
 	"strings"
 	"syscall"
-	"time"
 
 	"github.com/spf13/cobra"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -200,9 +199,7 @@ func (s *RouterServer) Run(ctx context.Context) error {
 
 	xdsSrv.SetExtProcMaxRequests(s.cfg.extProcMaxRequests())
 	if parkCfg.enabled() {
-		// Envoy must keep a parked request open at least as long as the router
-		// will hold it; add a margin so the router surfaces its own 503 first.
-		xdsSrv.SetExtProcMessageTimeout(parkCfg.Budget + 5*time.Second)
+		xdsSrv.SetExtProcMessageTimeout(extProcMessageTimeoutFor(parkCfg))
 	}
 
 	xdsSrv.SetTlsConfig(s.cfg.HttpsPort, s.cfg.EnvoyCertPath)

docs/request-parking.md

@@ -43,6 +43,29 @@ exponential backoff until either
   underlying capacity error is returned, surfacing as `503 "actor <id>
   unavailable: no free workers available"`.
 
+**An attempt with a worker already assigned gets a grace.** When the budget
+elapses with a resume still in flight, the router issues one `GetActor`. If the
+actor is `RESUMING` or `RUNNING` — a worker is assigned — that attempt gets a
+bounded **grace** (3s) to land instead of being cancelled. Anything else fails
+on the budget exactly as before. Either way the retry loop stops at the budget:
+the grace covers the one attempt underway, never another round.
+
+The distinction is what makes the wait worth it. The control plane durably marks
+an actor `RESUMING` once it has bound a worker to it, *before* the expensive
+part — the snapshot restore — begins. A cancellation after that point releases
+nothing: the actor is left `RESUMING` with the worker still assigned, so the
+undersized pool that made the request park stays starved for every request
+behind it, and a resume that was about to succeed is discarded. Under node
+contention the restore overshoots the budget by a few hundred milliseconds
+routinely, and the grace is what turns those into a `200` rather than a `503`
+plus a stranded worker. With no worker assigned there is nothing to protect and
+nothing to wait for.
+
+The worst-case hold on a parked request is therefore `budget + 3s` (the probe is
+spent inside the grace, not on top of it), which Envoy's ext_proc message
+timeout is sized to cover: `budget + 3s + 2s` margin, so the router's own 503
+surfaces first.
+
 To bound resource use and provide backpressure, the router admits requests to a
 **parking lot** of fixed capacity (`--parked-request-max`, default `1024`). Each
 in-flight resume occupies one slot. When the lot is full, further requests are