fix(grpc-proxy): purge unclaimed stateful work on shutdown

src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "h3.go",
         "health.go",
         "hijack.go",
+        "invocation_gate.go",
         "proxy.go",
         "tracing.go",
     ],
@@ -88,6 +89,8 @@ go_test(
         "director_test.go",
         "hijack_test.go",
         "info_test.go",
+        "invocation_gate_test.go",
+        "pending_work_test.go",
         "proxy_test.go",
     ],
     embed = [":proxy"],
@@ -103,6 +106,8 @@ go_test(
         "@com_github_gorilla_websocket//:websocket",
         "@com_github_hellofresh_health_go_v5//:health-go",
         "@com_github_jellydator_ttlcache_v3//:ttlcache",
+        "@com_github_stretchr_testify//assert",
+        "@com_github_stretchr_testify//require",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//credentials/insecure",

src/invocation-plane-services/grpc-proxy/proxy/director.go

@@ -69,10 +69,27 @@ type issuedTokenInfo struct {
 	mintedAt time.Time
 }
 
+// pendingWorkInfo identifies a stateful work request this pod has issued a
+// worker token for but has not yet seen a CONNECT for. The token lives only in
+// this pod's memory, so if the pod goes away the queued request can never
+// authenticate; the entry is what lets shutdown find it and drop it.
+type pendingWorkInfo struct {
+	functionVersionId string
+}
+
+// pendingWorkPurger removes a queued stateful work request. Implemented by the
+// function invoker and asserted optionally, so an invoker that cannot reach the
+// work queue (tests, alternative implementations) simply skips the purge.
+type pendingWorkPurger interface {
+	PurgePendingWork(ctx context.Context, requestId uuid.UUID, functionVersionId string) error
+}
+
 type StreamDirector struct {
 	shuttingDown    *atomic.Bool
 	workerAuth      *ttlcache.Cache[string, workerAuthInfo]  // auth -> request + function info
 	issuedTokens    *ttlcache.Cache[string, issuedTokenInfo] // diagnostic only, see issuedTokenInfo
+	pendingWork     *ttlcache.Cache[uuid.UUID, pendingWorkInfo]
+	invocations     *invocationGate
 	workers         *ttlcache.Cache[workerConnectionKey, *worker.WorkerConnection]
 	functionInvoker FunctionInvoker
 	cors            *cors.Cors
@@ -108,6 +125,17 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector {
 	)
 	go issuedTokenCache.Start()
 
+	// Sessions waiting on a worker CONNECT. Retention is deliberately much
+	// longer than the token TTL: the point is to still know about a request
+	// whose token has already aged out, because that request is still sitting
+	// in the work queue. Bounded so it cannot grow without limit.
+	pendingWorkCache := ttlcache.New(
+		ttlcache.WithTTL[uuid.UUID, pendingWorkInfo](pendingWorkRetention),
+		ttlcache.WithCapacity[uuid.UUID, pendingWorkInfo](pendingWorkCacheCapacity),
+		ttlcache.WithDisableTouchOnHit[uuid.UUID, pendingWorkInfo](),
+	)
+	go pendingWorkCache.Start()
+
 	// Set immediately before DeleteAll in Close so the eviction handler can
 	// report shutdown rather than attributing a drain to a client or worker.
 	shuttingDown := &atomic.Bool{}
@@ -241,6 +269,8 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector {
 		workers:         cache,
 		shuttingDown:    shuttingDown,
 		issuedTokens:    issuedTokenCache,
+		pendingWork:     pendingWorkCache,
+		invocations:     newInvocationGate(),
 		workerAuth:      workerAuthCache,
 		functionInvoker: functionInvoker,
 		cors:            cors.New(middleware.DefaultCorsOptions),
@@ -321,16 +351,86 @@ const (
 	issuedTokenRetention = 15 * time.Minute
 	// issuedTokenCacheCapacity bounds the diagnostic cache.
 	issuedTokenCacheCapacity = 50000
+	// pendingWorkRetention is how long a session waiting on a worker CONNECT is
+	// remembered. It has to outlast the token by a wide margin, because the
+	// queued request this refers to survives its token and is exactly what
+	// shutdown needs to clean up.
+	pendingWorkRetention = 30 * time.Minute
+	// pendingWorkCacheCapacity bounds the pending work cache.
+	pendingWorkCacheCapacity = 50000
+	// invocationDrainTimeout bounds how long shutdown waits for invocations
+	// that are already past admission to finish publishing their work request.
+	invocationDrainTimeout = 5 * time.Second
+	// pendingWorkPurgeTimeout bounds the whole shutdown purge. Shutdown must not
+	// block on the work queue, so this is a best-effort budget.
+	pendingWorkPurgeTimeout = 10 * time.Second
 )
 
+// purgePendingWork drops the work requests for sessions that never got a
+// worker CONNECT. Their tokens exist only in this pod's memory, so once it is
+// gone every one of them is guaranteed to be rejected; leaving them queued
+// means each is still pulled, still takes a concurrency slot, and still fails.
+//
+// Sessions with a worker already attached are deliberately not touched. Those
+// can reattach through another pod, and their work request has already left the
+// queue anyway.
+func (s *StreamDirector) purgePendingWork() {
+	purger, ok := s.functionInvoker.(pendingWorkPurger)
+	if !ok {
+		return
+	}
+	pending := s.pendingWork.Items()
+	if len(pending) == 0 {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), pendingWorkPurgeTimeout)
+	defer cancel()
+
+	var purged, failed int
+	for requestId, item := range pending {
+		if ctx.Err() != nil {
+			// Out of budget. Report what is left rather than trailing off
+			// silently, so a shutdown that could not finish is visible.
+			failed += len(pending) - purged - failed
+			break
+		}
+		if err := purger.PurgePendingWork(ctx, requestId, item.Value().functionVersionId); err != nil {
+			failed++
+			// Expected when this service has no rights on the work queue, so
+			// this stays a warning: the purge is an optimisation and shutdown
+			// is still correct without it.
+			zap.L().Warn("failed to purge pending stateful work request on shutdown",
+				zap.Stringer("request_id", requestId),
+				zap.String("function_version_id", item.Value().functionVersionId),
+				zap.Error(err))
+			continue
+		}
+		purged++
+	}
+
+	metrics.PendingWorkPurgedTotal.WithLabelValues(metrics.PurgeSucceeded).Add(float64(purged))
+	metrics.PendingWorkPurgedTotal.WithLabelValues(metrics.PurgeFailed).Add(float64(failed))
+	zap.L().Info("purged pending stateful work requests on shutdown",
+		zap.Int("purged", purged),
+		zap.Int("failed", failed))
+}
+
 func (s *StreamDirector) Close() error {
 	// Mark first: DeleteAll evicts every entry, and without this those
 	// evictions would be misreported as client or worker initiated.
 	s.shuttingDown.Store(true)
+	// Stop admitting invocations and let the ones already running finish
+	// publishing, so the purge below sees every request this pod created.
+	s.invocations.closeAndDrain(invocationDrainTimeout)
+	// Before the caches go away, drop the queued work this pod can no longer
+	// authenticate. Best effort: a failure here must not hold up shutdown.
+	s.purgePendingWork()
 	s.workers.DeleteAll()
 	s.workers.Stop()
 	s.workerAuth.Stop()
 	s.issuedTokens.Stop()
+	s.pendingWork.Stop()
 	if s.functionInvoker != nil {
 		if closer, ok := s.functionInvoker.(io.Closer); ok {
 			_ = closer.Close()
@@ -523,6 +623,17 @@ func (s *StreamDirector) getAndInitWorkerConnection(ctx context.Context, conn *w
 		// Capture API-provided function info in closure
 		var apiFunctionId, apiFunctionVersionId string
 
+		// An invocation records its pending work before it publishes the work
+		// request, so a shutdown landing between those two steps would purge
+		// nothing and then leave the published request behind with no token to
+		// authenticate it. Admission is closed before shutdown purges, and
+		// shutdown waits for whatever is already past this point, so no
+		// invocation can publish after the purge has taken its snapshot.
+		if !s.invocations.begin() {
+			return nil, nverrors.NewNVError(errors.New("proxy is shutting down")).WithCode(grpcCodes.Unavailable)
+		}
+		defer s.invocations.end()
+
 		invokeResponse, cancelInvokingWorker, err := s.functionInvoker.InvokeStatefulFunction(ctx, conn, auth, functionId, functionVersionId, requestId, func(workerAuthToken string, requestId uuid.UUID, apiFunc string, apiFuncVersion string) {
 			// Populate workerAuth cache BEFORE worker is notified (atomicity guarantee)
 			now := time.Now()
@@ -532,6 +643,11 @@ func (s *StreamDirector) getAndInitWorkerConnection(ctx context.Context, conn *w
 				functionVersionId: apiFuncVersion,
 				mintedAt:          now,
 			}, ttlcache.DefaultTTL)
+			// Remembered until the worker CONNECTs back, so that a shutdown can
+			// find the requests whose tokens are about to be lost with this pod
+			// and drop them from the work queue instead of leaving them to be
+			// pulled and rejected.
+			s.pendingWork.Set(requestId, pendingWorkInfo{functionVersionId: apiFuncVersion}, ttlcache.DefaultTTL)
 			// Diagnostic shadow record, longer lived than the auth entry, so a
 			// later rejection can say "expired N seconds ago" instead of just
 			// "not found". Never consulted when granting access.

src/invocation-plane-services/grpc-proxy/proxy/hijack.go

@@ -174,6 +174,9 @@ func (s *StreamDirector) HijackHandler(w http.ResponseWriter, r *http.Request) {
 	// already been sent away, so no point in keeping the auth around.
 	// also need to make sure we don't allow reconnects to guard against replay attacks with 0-rtt.
 	s.workerAuth.Delete(auth)
+	// The worker is attached, so this request is no longer queued and must not
+	// be purged if this pod shuts down. The session can reattach elsewhere.
+	s.pendingWork.Delete(parsedRequestId)
 }
 
 // networkPeerAddress returns the transport-level remote host of the tunnel connection and

src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel

@@ -21,6 +21,7 @@ go_library(
         "function_invoker.go",
         "nats.go",
         "nats_tracing.go",
+        "pending_work.go",
         "sleep.go",
         "stateful_session_request_registration.go",
     ],
@@ -57,9 +58,13 @@ alias(
 
 go_test(
     name = "invocation_test",
-    srcs = ["nats_test.go"],
+    srcs = [
+        "nats_test.go",
+        "pending_work_test.go",
+    ],
     embed = [":invocation"],
     deps = [
+        "@com_github_google_uuid//:uuid",
         "@com_github_nats_io_nkeys//:nkeys",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",

src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go

@@ -0,0 +1,65 @@
+/*
+SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+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 invocation
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/uuid"
+	"github.com/nats-io/nats.go/jetstream"
+)
+
+// The work queue a stateful invocation lands in. These are owned by the
+// invocation service, which publishes to requestWorkSubject and removes a
+// cancelled request with the same subject-filtered purge used here. The
+// formats are duplicated rather than shared because the owning service is
+// written in another language; they must not drift.
+func requestWorkStream(region, functionVersionId string) string {
+	return fmt.Sprintf("rq_%s_%s", region, functionVersionId)
+}
+
+func requestWorkSubject(region, functionVersionId string, requestId uuid.UUID) string {
+	return fmt.Sprintf("rq.%s.%s.%s", region, functionVersionId, requestId)
+}
+
+// PurgePendingWork drops a stateful work request that is still queued.
+//
+// It is called for sessions this pod issued a worker token for that never came
+// back to CONNECT, at the point the pod is shutting down. Those tokens only
+// exist in this pod's memory, so every one of those queued requests is already
+// guaranteed to fail authentication whenever a worker eventually pulls it. Left
+// in place they are pulled anyway, occupy a worker concurrency slot, and fail,
+// which is what keeps a saturated function at zero goodput long after the
+// restart that caused it.
+//
+// Purging by subject only removes messages still held by the stream. A session
+// that already has a worker attached had its message delivered, so an
+// established session is unaffected and remains free to reattach to another
+// pod.
+func (f *FunctionInvoker) PurgePendingWork(ctx context.Context, requestId uuid.UUID, functionVersionId string) error {
+	streamName := requestWorkStream(f.region, functionVersionId)
+	stream, err := f.js.Stream(ctx, streamName)
+	if err != nil {
+		return fmt.Errorf("failed to look up work stream %s: %w", streamName, err)
+	}
+	subject := requestWorkSubject(f.region, functionVersionId, requestId)
+	if err := stream.Purge(ctx, jetstream.WithPurgeSubject(subject)); err != nil {
+		return fmt.Errorf("failed to purge work subject %s: %w", subject, err)
+	}
+	return nil
+}

src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go

@@ -0,0 +1,50 @@
+/*
+SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+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 invocation
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+)
+
+// The work queue is owned by the invocation service, which is written in
+// another language, so these formats are duplicated rather than shared. If they
+// drift the purge silently targets a subject nothing was ever published to and
+// removes nothing, with no error to show for it. These cases pin the exact
+// strings against the owning service's request_stream_name and request_subject.
+func TestRequestWorkStreamAndSubjectMatchTheInvocationService(t *testing.T) {
+	requestId := uuid.MustParse("11111111-2222-3333-4444-555555555555")
+	region := "us-west-2"
+	versionId := "66666666-7777-8888-9999-000000000000"
+
+	assert.Equal(t, "rq_us-west-2_66666666-7777-8888-9999-000000000000",
+		requestWorkStream(region, versionId))
+	assert.Equal(t, "rq.us-west-2.66666666-7777-8888-9999-000000000000.11111111-2222-3333-4444-555555555555",
+		requestWorkSubject(region, versionId, requestId))
+}
+
+// The subject has to fall inside the stream's own subject space, otherwise a
+// filtered purge matches nothing.
+func TestRequestWorkSubjectIsCoveredByTheStreamSubjectSpace(t *testing.T) {
+	region := "eu-west-1"
+	versionId := uuid.New().String()
+
+	subject := requestWorkSubject(region, versionId, uuid.New())
+	assert.Regexp(t, `^rq\.`+region+`\.`+versionId+`\.`, subject)
+}