diff --git a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel index 82b929724..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,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", diff --git a/src/invocation-plane-services/grpc-proxy/proxy/director.go b/src/invocation-plane-services/grpc-proxy/proxy/director.go index ffcbad43a..ffa51b0a0 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/director.go +++ b/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. 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/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..1ee2ecc5e --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation_gate_test.go @@ -0,0 +1,109 @@ +/* +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" + "sync/atomic" + "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()) + + 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 + // returns can only mean the drain did not wait. + finished.Store(true) + gate.end() + }() + + gate.closeAndDrain(5 * time.Second) + + 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. +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()) +} 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()) +}