From 3b5a38d17a90f58ad87d57aac2bb75959ac97c00 Mon Sep 17 00:00:00 2001 From: balaji Date: Wed, 19 Aug 2026 21:40:51 -0700 Subject: [PATCH 1/4] fix(grpc-proxy): purge unclaimed stateful work on shutdown A worker CONNECT token lives only in the memory of the pod that minted it, but the work request it belongs to is durable and waits in the JetStream work queue until a worker has a slot to pull it. When the pod goes away, every request it issued that has not been pulled yet is already doomed: a worker pulls it, takes a concurrency slot, is rejected with 403, and hands the slot back having achieved nothing. Nothing removed those requests, so on a saturated function this repeats for as long as the backlog takes to drain while clients retry and refill it, which is the extended near-zero-goodput window seen after a restart. Track sessions from the point a token is issued until the worker CONNECTs back, and on shutdown purge the work requests still waiting. This is the same subject-filtered purge the invocation service uses in cancel_request. Only sessions still waiting for a worker are purged. A session with a worker attached is not tied to the pod that started it: on reconnect the config is rebuilt from the answering pod's address with a fresh token, so the worker reattaches elsewhere and the session survives a rolling update. Purging those would sever sessions that were going to live. Purging by subject only removes what the stream still holds, so an established session is untouched for that reason too. The purge is best effort and bounded. Whether this service may purge the work queue is granted outside this repository, so a rejected purge is logged and shutdown continues rather than failing. Adds nvcf_grpc_proxy_service_pending_work_purged_total{result}. A persistent failed count is the signal that the permission is missing. Co-Authored-By: Balaji Ganesan --- .../grpc-proxy/proxy/BUILD.bazel | 3 + .../grpc-proxy/proxy/director.go | 97 ++++++++++++++ .../grpc-proxy/proxy/hijack.go | 3 + .../grpc-proxy/proxy/invocation/BUILD.bazel | 7 +- .../proxy/invocation/pending_work.go | 65 +++++++++ .../proxy/invocation/pending_work_test.go | 50 +++++++ .../grpc-proxy/proxy/metrics/metrics.go | 42 ++++-- .../grpc-proxy/proxy/pending_work_test.go | 124 ++++++++++++++++++ 8 files changed, 380 insertions(+), 11 deletions(-) create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go diff --git a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel index 82b929724..6af7ac884 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel +++ b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel @@ -88,6 +88,7 @@ go_test( "director_test.go", "hijack_test.go", "info_test.go", + "pending_work_test.go", "proxy_test.go", ], embed = [":proxy"], @@ -103,6 +104,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", diff --git a/src/invocation-plane-services/grpc-proxy/proxy/director.go b/src/invocation-plane-services/grpc-proxy/proxy/director.go index ffcbad43a..5c04e6838 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/director.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/director.go @@ -69,10 +69,26 @@ 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] workers *ttlcache.Cache[workerConnectionKey, *worker.WorkerConnection] functionInvoker FunctionInvoker cors *cors.Cors @@ -108,6 +124,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 +268,7 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector { workers: cache, shuttingDown: shuttingDown, issuedTokens: issuedTokenCache, + pendingWork: pendingWorkCache, workerAuth: workerAuthCache, functionInvoker: functionInvoker, cors: cors.New(middleware.DefaultCorsOptions), @@ -321,16 +349,80 @@ 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 + // 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) + // 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() @@ -532,6 +624,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. diff --git a/src/invocation-plane-services/grpc-proxy/proxy/hijack.go b/src/invocation-plane-services/grpc-proxy/proxy/hijack.go index 747a3fd18..129250ef1 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/hijack.go +++ b/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 diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel index 8d5ef783b..fb8eb0445 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel +++ b/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", diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go new file mode 100644 index 000000000..d0837b775 --- /dev/null +++ b/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 +} diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go new file mode 100644 index 000000000..ff878c626 --- /dev/null +++ b/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) +} diff --git a/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go b/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go index 187b19c87..eb2942d12 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go @@ -6,7 +6,7 @@ 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 + 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, @@ -189,6 +189,14 @@ var ( }) ) +// Outcomes of the shutdown purge of queued stateful work requests. +const ( + PurgeSucceeded = "succeeded" + PurgeFailed = "failed" +) + +var PurgeResults = []string{PurgeSucceeded, PurgeFailed} + // Close reasons reported when a worker tunnel goes away. The three "deleted" // variants matter: a bare `deleted` cannot distinguish the client hanging up // from the worker hanging up from a proxy shutdown, and which side went first @@ -301,15 +309,15 @@ var WorkerConnectionCloseReasons = []string{ // Outcomes of a worker CONNECT to /v1/proxy. Every terminal path in // HijackHandler maps to exactly one of these. const ( - ConnectAccepted = "accepted" - ConnectNotHijackable = "rejected_not_hijackable" // 500 - ConnectMissingAuth = "rejected_missing_auth" // 401 - ConnectMissingRequestID = "rejected_missing_requestid" // 400 - ConnectInvalidRequestID = "rejected_invalid_requestid" // 400 - ConnectTokenExpired = "rejected_token_expired" // 403, token was issued but has aged out - ConnectTokenUnknown = "rejected_token_unknown" // 403, token was never issued by this pod - ConnectRequestIDMismatch = "rejected_requestid_mismatch"// 403, token valid but bound to another request - ConnectHijackFailed = "rejected_hijack_failed" // 500 + ConnectAccepted = "accepted" + ConnectNotHijackable = "rejected_not_hijackable" // 500 + ConnectMissingAuth = "rejected_missing_auth" // 401 + ConnectMissingRequestID = "rejected_missing_requestid" // 400 + ConnectInvalidRequestID = "rejected_invalid_requestid" // 400 + ConnectTokenExpired = "rejected_token_expired" // 403, token was issued but has aged out + ConnectTokenUnknown = "rejected_token_unknown" // 403, token was never issued by this pod + ConnectRequestIDMismatch = "rejected_requestid_mismatch" // 403, token valid but bound to another request + ConnectHijackFailed = "rejected_hijack_failed" // 500 ) var ConnectResults = []string{ @@ -375,6 +383,17 @@ var ( Buckets: []float64{1, 5, 8, 10, 15, 30, 45, 60, 120, 300, 600, 1800, 3600}, }) + // PendingWorkPurgedTotal counts stateful work requests dropped from the + // work queue during shutdown because this pod held the only copy of their + // worker token. A persistent failed count usually means this service lacks + // purge rights on the work queue rather than a transient NATS error. + PendingWorkPurgedTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Namespace: RootNamespace, + Name: "pending_work_purged_total", + Help: "queued stateful work requests dropped at shutdown, by outcome", + }, []string{"result"}) + // ClientConnectionWorkerTunnelsAtClose records how many worker tunnels a // client connection was still holding when it closed. Anything above zero // means that close tore down live tunnels. @@ -406,6 +425,9 @@ func init() { for _, code := range CloseCodes { WorkerConnectionCloseCodeTotal.WithLabelValues(code) } + for _, result := range PurgeResults { + PendingWorkPurgedTotal.WithLabelValues(result) + } } var nc atomic.Pointer[nats.Conn] diff --git a/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go b/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go new file mode 100644 index 000000000..ec5261784 --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go @@ -0,0 +1,124 @@ +/* +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 proxy + +import ( + "context" + "errors" + "net" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/jellydator/ttlcache/v3" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "nvcf-grpc-proxy/proxy/invocation" +) + +// recordingPurger stands in for the function invoker, capturing which requests +// shutdown asked to drop. +type recordingPurger struct { + mu sync.Mutex + purged map[uuid.UUID]string + err error +} + +func newRecordingPurger() *recordingPurger { + return &recordingPurger{purged: map[uuid.UUID]string{}} +} + +func (p *recordingPurger) InvokeStatefulFunction(_ context.Context, _ net.Conn, _, _, _ string, _ *uuid.UUID, _ func(string, uuid.UUID, string, string)) (invocation.Result, context.CancelFunc, error) { + return invocation.Result{}, nil, errors.New("not used") +} + +func (p *recordingPurger) PurgePendingWork(_ context.Context, requestId uuid.UUID, functionVersionId string) error { + p.mu.Lock() + defer p.mu.Unlock() + if p.err != nil { + return p.err + } + p.purged[requestId] = functionVersionId + return nil +} + +func (p *recordingPurger) purgedRequests() map[uuid.UUID]string { + p.mu.Lock() + defer p.mu.Unlock() + out := map[uuid.UUID]string{} + for k, v := range p.purged { + out[k] = v + } + return out +} + +// A request whose worker never came back to CONNECT can never authenticate once +// this pod is gone, because the token only ever existed here. Shutdown has to +// take it out of the work queue, otherwise it is still pulled, still occupies a +// worker slot, and still fails. +func TestClosePurgesWorkForSessionsAwaitingConnect(t *testing.T) { + purger := newRecordingPurger() + director := NewStreamDirector(purger) + + awaiting := uuid.New() + director.pendingWork.Set(awaiting, pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) + + purged := purger.purgedRequests() + require.Len(t, purged, 1) + assert.Equal(t, "version-1", purged[awaiting]) +} + +// A session with a worker already attached is not queued any more and can +// reattach through another pod. Purging it would sever a session that was going +// to survive the restart, so shutdown must leave it alone. +func TestClosePurgesNothingOnceTheWorkerHasConnected(t *testing.T) { + purger := newRecordingPurger() + director := NewStreamDirector(purger) + + connected := uuid.New() + director.pendingWork.Set(connected, pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + // what HijackHandler does once the worker's CONNECT is accepted + director.pendingWork.Delete(connected) + + require.NoError(t, director.Close()) + + assert.Empty(t, purger.purgedRequests()) +} + +// The purge is an optimisation. If this service has no rights on the work queue +// the calls fail, and shutdown still has to complete. +func TestClosePurgeFailureDoesNotBlockShutdown(t *testing.T) { + purger := newRecordingPurger() + purger.err = errors.New("nats: permissions violation for stream purge") + director := NewStreamDirector(purger) + + director.pendingWork.Set(uuid.New(), pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) +} + +// An invoker with no route to the work queue simply skips the purge rather than +// failing shutdown. +func TestClosePurgeSkippedWhenInvokerCannotPurge(t *testing.T) { + director := NewStreamDirector((*mockInvoker)(&invocation.Result{})) + director.pendingWork.Set(uuid.New(), pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) +} From 069be4636c3bb9fbb153d867cc003f950d551e4a Mon Sep 17 00:00:00 2001 From: balaji Date: Wed, 19 Aug 2026 21:56:43 -0700 Subject: [PATCH 2/4] fix(grpc-proxy): close the shutdown window that orphaned queued work An invocation records its pending work before it publishes the work request. A shutdown landing between those two steps purged nothing and then let the publish leave a request in the queue with no surviving token to authenticate it, which is exactly the state the purge exists to prevent. Marking the director as shutting down did not help: that flag only affects eviction reporting and does not stop an invocation already in progress. Gate admission instead. Shutdown closes the gate and waits, bounded, for invocations already past it to finish publishing, so the purge snapshot sees every request this pod created and nothing can publish after it. An invocation refused at the gate returns Unavailable and the client retries against a live pod, which is correct once the servers have drained. The gate covers only the invocation, never the session that follows, so shutdown is never held for the length of a session, and the drain is bounded so a stuck invocation cannot block exit. Co-Authored-By: Balaji Ganesan --- .../grpc-proxy/proxy/BUILD.bazel | 2 + .../grpc-proxy/proxy/director.go | 19 ++++ .../grpc-proxy/proxy/invocation_gate.go | 100 +++++++++++++++++ .../grpc-proxy/proxy/invocation_gate_test.go | 104 ++++++++++++++++++ 4 files changed, 225 insertions(+) create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/invocation_gate.go create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go diff --git a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel index 6af7ac884..088118faa 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel +++ b/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,7 @@ go_test( "director_test.go", "hijack_test.go", "info_test.go", + "invocation_gate_test.go", "pending_work_test.go", "proxy_test.go", ], diff --git a/src/invocation-plane-services/grpc-proxy/proxy/director.go b/src/invocation-plane-services/grpc-proxy/proxy/director.go index 5c04e6838..ffa51b0a0 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/director.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/director.go @@ -89,6 +89,7 @@ type StreamDirector struct { 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 @@ -269,6 +270,7 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector { shuttingDown: shuttingDown, issuedTokens: issuedTokenCache, pendingWork: pendingWorkCache, + invocations: newInvocationGate(), workerAuth: workerAuthCache, functionInvoker: functionInvoker, cors: cors.New(middleware.DefaultCorsOptions), @@ -356,6 +358,9 @@ const ( 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 @@ -415,6 +420,9 @@ 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() @@ -615,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() diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate.go new file mode 100644 index 000000000..b4107efdf --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate.go @@ -0,0 +1,100 @@ +/* +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 proxy + +import ( + "sync" + "time" +) + +// invocationGate counts stateful invocations that are between recording their +// pending work and publishing the work request, and lets shutdown close +// admission and wait for the ones already inside. +// +// Without it the shutdown purge is not sound. An invocation records its pending +// work first and publishes the work request afterwards, so a purge landing +// between those two steps finds nothing to remove and then the publish leaves a +// request in the queue with no surviving token to authenticate it, which is the +// exact situation the purge exists to prevent. +// +// It covers only the invocation itself, never the session that follows, so +// shutdown is never held up for the length of a session. +type invocationGate struct { + mu sync.Mutex + cond *sync.Cond + active int + closed bool + expired bool +} + +func newInvocationGate() *invocationGate { + g := &invocationGate{} + g.cond = sync.NewCond(&g.mu) + return g +} + +// begin admits an invocation, reporting false once admission has closed. A +// refused caller must not go on to publish a work request. +func (g *invocationGate) begin() bool { + g.mu.Lock() + defer g.mu.Unlock() + if g.closed { + return false + } + g.active++ + return true +} + +func (g *invocationGate) end() { + g.mu.Lock() + defer g.mu.Unlock() + g.active-- + if g.active == 0 { + g.cond.Broadcast() + } +} + +// closeAndDrain stops admitting invocations and waits for those already +// admitted, giving up after timeout so a stuck invocation cannot hold up +// shutdown. Closing happens under the same lock that begin checks, so once this +// returns no further invocation can publish a work request. +func (g *invocationGate) closeAndDrain(timeout time.Duration) { + g.mu.Lock() + g.closed = true + if g.active == 0 { + g.mu.Unlock() + return + } + g.mu.Unlock() + + // sync.Cond has no deadline of its own, so a timer sets the expiry flag and + // wakes the wait. Waiting in place rather than in a helper goroutine means + // an invocation that never finishes cannot leave one behind. + timer := time.AfterFunc(timeout, func() { + g.mu.Lock() + defer g.mu.Unlock() + g.expired = true + g.cond.Broadcast() + }) + defer timer.Stop() + + g.mu.Lock() + defer g.mu.Unlock() + for g.active > 0 && !g.expired { + g.cond.Wait() + } +} diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go new file mode 100644 index 000000000..a0276ce12 --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go @@ -0,0 +1,104 @@ +/* +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 proxy + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestInvocationGateAdmitsUntilClosed(t *testing.T) { + gate := newInvocationGate() + + require.True(t, gate.begin()) + gate.end() + + gate.closeAndDrain(time.Second) + + assert.False(t, gate.begin(), "admission must be closed after drain") +} + +// The point of the gate: an invocation admitted before shutdown finishes +// publishing before the purge takes its snapshot. Draining must therefore block +// until that invocation calls end. +func TestInvocationGateDrainWaitsForAdmittedInvocation(t *testing.T) { + gate := newInvocationGate() + require.True(t, gate.begin()) + + released := make(chan struct{}) + go func() { + time.Sleep(100 * time.Millisecond) + close(released) + gate.end() + }() + + gate.closeAndDrain(5 * time.Second) + + select { + case <-released: + default: + t.Fatal("drain returned before the admitted invocation finished") + } +} + +// A stuck invocation must not hold shutdown open indefinitely. +func TestInvocationGateDrainGivesUpAfterTimeout(t *testing.T) { + gate := newInvocationGate() + require.True(t, gate.begin()) + t.Cleanup(gate.end) + + start := time.Now() + gate.closeAndDrain(200 * time.Millisecond) + elapsed := time.Since(start) + + assert.GreaterOrEqual(t, elapsed, 200*time.Millisecond) + assert.Less(t, elapsed, 5*time.Second, "drain should give up rather than block shutdown") +} + +// Draining with nothing in flight is the common case and must not wait. +func TestInvocationGateDrainReturnsImmediatelyWhenIdle(t *testing.T) { + gate := newInvocationGate() + + start := time.Now() + gate.closeAndDrain(5 * time.Second) + + assert.Less(t, time.Since(start), time.Second) +} + +func TestInvocationGateConcurrentUse(t *testing.T) { + gate := newInvocationGate() + + var wg sync.WaitGroup + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if gate.begin() { + gate.end() + } + }() + } + + gate.closeAndDrain(5 * time.Second) + wg.Wait() + + assert.False(t, gate.begin()) +} From 152d181b90a916ce9d424bb65e1a6e4f7d6f4c1b Mon Sep 17 00:00:00 2001 From: balaji Date: Wed, 19 Aug 2026 22:04:30 -0700 Subject: [PATCH 3/4] test(grpc-proxy): make the invocation drain assertion unambiguous The drain test proved the right thing but read as if it could race: the goroutine closed the channel before calling end, so the assertion looked order-dependent even though the drain cannot return until end runs. Assert on a flag set before end plus the elapsed time instead, so a drain that failed to wait is caught directly rather than inferred. Inverting the original order, as suggested in review, would have introduced a real flake: the drain can return between end and the channel close. Co-Authored-By: Balaji Ganesan --- .../grpc-proxy/proxy/invocation_gate_test.go | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go index a0276ce12..9cffa9a48 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go @@ -18,6 +18,7 @@ package proxy import ( "sync" + "sync/atomic" "testing" "time" @@ -43,20 +44,21 @@ func TestInvocationGateDrainWaitsForAdmittedInvocation(t *testing.T) { gate := newInvocationGate() require.True(t, gate.begin()) - released := make(chan struct{}) + const held = 100 * time.Millisecond + var finished atomic.Bool go func() { - time.Sleep(100 * time.Millisecond) - close(released) + time.Sleep(held) + // Set before end so that observing "not finished" after the drain + // returns can only mean the drain did not wait. + finished.Store(true) gate.end() }() + start := time.Now() gate.closeAndDrain(5 * time.Second) - select { - case <-released: - default: - t.Fatal("drain returned before the admitted invocation finished") - } + assert.True(t, finished.Load(), "drain returned before the admitted invocation finished") + assert.GreaterOrEqual(t, time.Since(start), held, "drain returned without waiting out the invocation") } // A stuck invocation must not hold shutdown open indefinitely. From 2c0f94ff879ffd1f9b61d82a91539495c11a5521 Mon Sep 17 00:00:00 2001 From: balaji Date: Thu, 20 Aug 2026 09:23:33 -0700 Subject: [PATCH 4/4] test(grpc-proxy): remove a timing flake in the drain assertion The start time was taken after launching the goroutine, so the sleep could begin first and the measured elapsed time come out just under the held duration, failing the assertion for no real reason. Co-Authored-By: Balaji Ganesan --- .../grpc-proxy/proxy/invocation_gate_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go index 9cffa9a48..1ee2ecc5e 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go @@ -46,6 +46,10 @@ func TestInvocationGateDrainWaitsForAdmittedInvocation(t *testing.T) { const held = 100 * time.Millisecond var finished atomic.Bool + // Recorded before the goroutine starts. Taking it afterwards allows the + // sleep to begin first, which makes the measured elapsed time slightly + // shorter than held and fails the assertion for no real reason. + start := time.Now() go func() { time.Sleep(held) // Set before end so that observing "not finished" after the drain @@ -54,7 +58,6 @@ func TestInvocationGateDrainWaitsForAdmittedInvocation(t *testing.T) { gate.end() }() - start := time.Now() gate.closeAndDrain(5 * time.Second) assert.True(t, finished.Load(), "drain returned before the admitted invocation finished")