From aebc8b92966afbfc3ff5b175ff3e13037bb21ab2 Mon Sep 17 00:00:00 2001 From: balaji Date: Wed, 19 Aug 2026 19:14:18 -0700 Subject: [PATCH 1/3] fix(grpc-proxy): detect stateful sessions whose worker is gone Rejoining an existing stateful session published the reconnect message with a core NATS publish, which succeeds whether or not a worker is subscribed. When the session's worker was gone the message was silently discarded, the proxy reported a healthy rejoin, and the client was left holding a session cookie it presented again on every retry. No error was raised on any path, so the session never recovered without operator action. The recovery path already existed: the director clears the client's request id cookie on ErrSessionNotFound, and a client that drops the cookie opens a fresh session on its next request. It was unreachable because the rejoin could not tell a dead session from a live one. Send the reconnect as a request instead. A no-responders answer means nothing can serve the session, so report ErrSessionNotFound and let the existing cookie-clearing path recover it. Interest is confirmed twice before acting, because a single answer can reflect a momentary gap in interest propagation and the cost of believing it is severing a live session. A subscribed worker that does not answer is treated as live, which is how the previous publish behaved. On the worker side, acknowledge reconnects the way the polling listener already does, and subscribe to the reconnect subject before the first CONNECT rather than after it, so a rejoin arriving during session establishment is not mistaken for a dead session. The two sides are independent. The proxy change alone is sufficient and costs one short probe deadline per rejoin against workers that predate the acknowledgement; the worker change removes that cost. Rollout order does not matter. This mirrors polling_request in the invocation service, which already maps a no-responders answer onto "no worker picked this up". Adds nvcf_grpc_proxy_service_stateful_rejoin_total{result}, pre- initialised across all four outcomes. Promotes github.com/nats-io/nats-server/v2 from an indirect dependency to a direct test dependency at v2.11.6, matching the version already used elsewhere in the repo, so the tests exercise real no-responders semantics rather than a mock. Apache-2.0, already present in MODULE.bazel. Co-Authored-By: Balaji Ganesan --- .../grpc-proxy/go.mod | 5 + .../grpc-proxy/go.sum | 11 ++ .../grpc-proxy/proxy/invocation/BUILD.bazel | 10 +- .../proxy/invocation/function_invoker.go | 82 +++++++- .../invocation/join_existing_session_test.go | 177 ++++++++++++++++++ .../grpc-proxy/proxy/metrics/metrics.go | 31 +++ src/libraries/go/worker/proxy/proxy.go | 132 +++++++------ .../go/worker/proxy/proxy_e2e_test.go | 113 +++++++++++ 8 files changed, 499 insertions(+), 62 deletions(-) create mode 100644 src/invocation-plane-services/grpc-proxy/proxy/invocation/join_existing_session_test.go diff --git a/src/invocation-plane-services/grpc-proxy/go.mod b/src/invocation-plane-services/grpc-proxy/go.mod index 2163b7c0f..74ad21111 100644 --- a/src/invocation-plane-services/grpc-proxy/go.mod +++ b/src/invocation-plane-services/grpc-proxy/go.mod @@ -20,6 +20,7 @@ require ( github.com/hashicorp/go-metrics v0.5.4 github.com/hellofresh/health-go/v5 v5.5.5 github.com/jellydator/ttlcache/v3 v3.4.0 + github.com/nats-io/nats-server/v2 v2.11.6 github.com/nats-io/nats.go v1.43.0 github.com/nats-io/nkeys v0.4.11 github.com/prometheus/client_golang v1.23.2 @@ -95,6 +96,7 @@ require ( github.com/go-ole/go-ole v1.3.0 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect github.com/gogo/protobuf v1.3.2 // indirect + github.com/google/go-tpm v0.9.5 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 // indirect github.com/hashicorp/go-cleanhttp v0.5.2 // indirect github.com/hashicorp/go-immutable-radix v1.3.1 // indirect @@ -106,6 +108,7 @@ require ( github.com/magiconair/properties v1.8.10 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect + github.com/minio/highwayhash v1.0.3 // indirect github.com/mitchellh/go-homedir v1.1.0 // indirect github.com/moby/docker-image-spec v1.3.1 // indirect github.com/moby/patternmatcher v0.6.0 // indirect @@ -115,6 +118,7 @@ require ( github.com/moby/term v0.5.0 // indirect github.com/morikuni/aec v1.0.0 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/nats-io/jwt/v2 v2.7.4 // indirect github.com/nats-io/nuid v1.0.1 // indirect github.com/oklog/oklog v0.3.2 // indirect github.com/oklog/run v1.2.0 // indirect @@ -160,6 +164,7 @@ require ( golang.org/x/oauth2 v0.36.0 // indirect golang.org/x/sys v0.45.0 // indirect golang.org/x/text v0.37.0 // indirect + golang.org/x/time v0.12.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260226221140-a57be14db171 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260226221140-a57be14db171 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect diff --git a/src/invocation-plane-services/grpc-proxy/go.sum b/src/invocation-plane-services/grpc-proxy/go.sum index 849f315d0..3c07d3618 100644 --- a/src/invocation-plane-services/grpc-proxy/go.sum +++ b/src/invocation-plane-services/grpc-proxy/go.sum @@ -27,6 +27,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuy github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0= github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0= github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho= +github.com/antithesishq/antithesis-sdk-go v0.4.3-default-no-op h1:+OSa/t11TFhqfrX0EOSqQBDJ0YlpmK0rDSiB19dg9M0= +github.com/antithesishq/antithesis-sdk-go v0.4.3-default-no-op/go.mod h1:IUpT2DPAKh6i/YhSbt6Gl3v2yvUZjmKncl7U91fup7E= github.com/armon/go-metrics v0.4.1/go.mod h1:E6amYzXo6aW1tqzoZGT755KkbgrJsSdpwZ+3JqfkOG4= github.com/aws/aws-sdk-go-v2 v1.41.2 h1:LuT2rzqNQsauaGkPK/7813XxcZ3o3yePY0Iy891T2ls= github.com/aws/aws-sdk-go-v2 v1.41.2/go.mod h1:IvvlAZQXvTXznUPfRVfryiG1fbzE2NGK6m9u39YQ+S4= @@ -191,6 +193,8 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/go-tpm v0.9.5 h1:ocUmnDebX54dnW+MQWGQRbdaAcJELsa6PqZhJ48KwVU= +github.com/google/go-tpm v0.9.5/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= @@ -264,6 +268,8 @@ github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stg github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= +github.com/minio/highwayhash v1.0.3 h1:kbnuUMoHYyVl7szWjSxJnxw11k2U709jqFPPmIUyD6Q= +github.com/minio/highwayhash v1.0.3/go.mod h1:GGYsuwP/fPD6Y9hMiXuapVvlIUEhFhMTh0rxU3ik1LQ= github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0= @@ -288,6 +294,10 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U= +github.com/nats-io/jwt/v2 v2.7.4 h1:jXFuDDxs/GQjGDZGhNgH4tXzSUK6WQi2rsj4xmsNOtI= +github.com/nats-io/jwt/v2 v2.7.4/go.mod h1:me11pOkwObtcBNR8AiMrUbtVOUGkqYjMQZ6jnSdVUIA= +github.com/nats-io/nats-server/v2 v2.11.6 h1:4VXRjbTUFKEB+7UoaKL3F5Y83xC7MxPoIONOnGgpkHw= +github.com/nats-io/nats-server/v2 v2.11.6/go.mod h1:2xoztlcb4lDL5Blh1/BiukkKELXvKQ5Vy29FPVRBUYs= github.com/nats-io/nats.go v1.43.0 h1:uRFZ2FEoRvP64+UUhaTokyS18XBCR/xM2vQZKO4i8ug= github.com/nats-io/nats.go v1.43.0/go.mod h1:iRWIPokVIFbVijxuMQq4y9ttaBTMe0SFdlZfMDd+33g= github.com/nats-io/nkeys v0.4.11 h1:q44qGV008kYd9W1b1nEBkNzvnWxtRSQ7A8BoqRrcfa0= @@ -531,6 +541,7 @@ golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBc golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.43.0 h1:S4RLU2sB31O/NCl+zFN9Aru9A/Cq2aqKpTZJ6B+DwT4= 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..25a0d5d42 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel @@ -57,11 +57,19 @@ alias( go_test( name = "invocation_test", - srcs = ["nats_test.go"], + srcs = [ + "join_existing_session_test.go", + "nats_test.go", + ], embed = [":invocation"], deps = [ + "//src/invocation-plane-services/grpc-proxy/nvcf/pb", + "@com_github_google_uuid//:uuid", + "@com_github_nats_io_nats_go//:nats_go", + "@com_github_nats_io_nats_server_v2//server", "@com_github_nats_io_nkeys//:nkeys", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", + "@org_golang_google_protobuf//proto", ], ) diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go index 6e12f52da..c71c29c18 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go @@ -304,21 +304,93 @@ func (f *FunctionInvoker) tryRegionForLLS(ctx context.Context, region string, re return err } +const ( + // reconnectAckTimeout bounds one attempt to hand a session back to its + // worker. Only the no-responders answer is acted on and the NATS server + // produces that from interest state without waiting for anyone, so this + // deadline is reached only when a worker is subscribed but predates the + // acknowledgement in its reconnect listener. Kept short because until + // those workers roll over it is added to every rejoin. + reconnectAckTimeout = 500 * time.Millisecond + // reconnectNoRespondersRetryDelay separates the two liveness probes. A + // single no-responders answer can reflect a momentary gap in interest + // propagation rather than a dead worker, and the cost of believing it is + // severing a live session, so it has to be seen twice. + reconnectNoRespondersRetryDelay = 250 * time.Millisecond +) + +// joinExistingSession asks the worker holding requestId to open another tunnel. +// +// This is a request rather than a publish so that a session whose worker is +// gone is detectable. A publish always succeeds, so the proxy reported a +// healthy rejoin, waited for a worker that could not arrive, and left the +// client holding a session cookie it would present again on every retry, with +// no error surfaced anywhere. That is the state that could only be cleared by +// restarting the function. The stateless path already works this way: see +// polling_request in the invocation service, which maps NoResponders onto +// "no worker picked this up". func (f *FunctionInvoker) joinExistingSession(ctx context.Context, requestId uuid.UUID, proxyAuthResponse *pb.ProxyAuthResponse, workerAuthToken string) error { marshalledInvokeFunctionRequest, err := f.marshalStatefulSessionRequest(requestId, proxyAuthResponse, workerAuthToken) if err != nil { return err } subject := reconnectSubject(requestId) - err = f.nc.PublishMsg(&nats.Msg{ + span := trace.SpanFromContext(ctx) + + result, err := f.probeSessionWorker(ctx, subject, marshalledInvokeFunctionRequest) + if result == metrics.RejoinNoResponders { + // Confirm before acting: see reconnectNoRespondersRetryDelay. + SleepWithContext(ctx, reconnectNoRespondersRetryDelay) + if ctx.Err() != nil { + return ctx.Err() + } + result, err = f.probeSessionWorker(ctx, subject, marshalledInvokeFunctionRequest) + } + + metrics.StatefulRejoinTotal.WithLabelValues(result).Inc() + span.SetAttributes(attribute.String("rejoin_result", result)) + + switch result { + case metrics.RejoinAcked, metrics.RejoinAssumedLive: + return nil + case metrics.RejoinNoResponders: + // Nothing is subscribed to this session's reconnect subject, so no + // worker can ever serve it. Reporting it as a missing session makes + // the director clear the client's request id cookie, and the client's + // next request opens a fresh session without any operator action. + zap.L().Info("no worker subscribed for existing stateful session, asking client to start a new one", + zap.Stringer("request_id", requestId), + zap.String("function_id", proxyAuthResponse.FunctionId)) + return fmt.Errorf("%w for request id %s", ErrSessionNotFound, requestId) + default: + return fmt.Errorf("failed to publish function invocation request to nats: %w", err) + } +} + +// probeSessionWorker sends the reconnect message once and classifies the answer. +func (f *FunctionInvoker) probeSessionWorker(ctx context.Context, subject string, payload []byte) (string, error) { + requestCtx, cancel := context.WithTimeout(ctx, reconnectAckTimeout) + defer cancel() + + _, err := f.nc.RequestMsgWithContext(requestCtx, &nats.Msg{ Subject: subject, Header: otelHeaders(ctx), - Data: marshalledInvokeFunctionRequest, + Data: payload, }) - if err != nil { - return fmt.Errorf("failed to publish function invocation request to nats: %w", err) + switch { + case err == nil: + return metrics.RejoinAcked, nil + case errors.Is(err, nats.ErrNoResponders): + return metrics.RejoinNoResponders, err + case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil: + // Our own deadline, not the caller's. A worker is subscribed but did + // not acknowledge, which is how a worker built before the reconnect + // acknowledgement behaves. Interest is sufficient proof the session is + // live, so treat it exactly as the previous publish did. + return metrics.RejoinAssumedLive, nil + default: + return metrics.RejoinFailed, err } - return nil } func otelHeaders(ctx context.Context) nats.Header { diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/join_existing_session_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/join_existing_session_test.go new file mode 100644 index 000000000..89164553d --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/join_existing_session_test.go @@ -0,0 +1,177 @@ +/* +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" + "errors" + "testing" + "time" + + "github.com/google/uuid" + natsserver "github.com/nats-io/nats-server/v2/server" + "github.com/nats-io/nats.go" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "nvcf-grpc-proxy/nvcf/pb" +) + +// startEmbeddedNats starts an in-process NATS server on an OS-assigned +// ephemeral port so these tests can run alongside the rest of the suite +// without a fixed-port collision. Torn down via t.Cleanup. +func startEmbeddedNats(t *testing.T) string { + t.Helper() + + s, err := natsserver.NewServer(&natsserver.Options{ + Host: "127.0.0.1", + Port: -1, + NoSigs: true, + }) + require.NoError(t, err) + + s.Start() + t.Cleanup(func() { + s.Shutdown() + s.WaitForShutdown() + }) + + require.True(t, s.ReadyForConnections(10*time.Second), "embedded nats server did not become ready") + return s.ClientURL() +} + +func newTestInvoker(t *testing.T) (*FunctionInvoker, *nats.Conn) { + t.Helper() + + nc, err := nats.Connect(startEmbeddedNats(t)) + require.NoError(t, err) + t.Cleanup(nc.Close) + + // no_responders is what lets a rejoin tell a dead session from a live one. + // If the server or client ever stopped supporting it the rest of these + // assertions would still pass for the wrong reason, so check it directly. + require.True(t, nc.HeadersSupported(), "no-responders detection needs header support") + + return &FunctionInvoker{ + nc: nc, + region: "region-1", + connectPaths: ConnectPaths{HTTP1: "http://10.0.0.1:10086/v1/proxy"}, + }, nc +} + +// A session whose worker is gone has nothing subscribed to its reconnect +// subject. That has to surface as ErrSessionNotFound, because that is the only +// error the director turns into a cookie-clearing response, which is what lets +// the client open a fresh session without the function being restarted. +func TestJoinExistingSessionNoWorkerSubscribed(t *testing.T) { + invoker, _ := newTestInvoker(t) + requestId := uuid.New() + + start := time.Now() + err := invoker.joinExistingSession(context.Background(), requestId, &pb.ProxyAuthResponse{FunctionId: "fn-1"}, "worker-token") + + require.Error(t, err) + require.ErrorIs(t, err, ErrSessionNotFound) + assert.Contains(t, err.Error(), requestId.String()) + // No responders is answered from interest state, so detection must not + // cost the full probe deadline even with the confirmation probe. + assert.Less(t, time.Since(start), 2*reconnectAckTimeout+reconnectNoRespondersRetryDelay, + "dead session should be detected without waiting out both probe deadlines") +} + +// A worker that acknowledges the reconnect is live, so the rejoin succeeds and +// the worker receives the connection config it needs to CONNECT back. +func TestJoinExistingSessionWorkerAcks(t *testing.T) { + invoker, nc := newTestInvoker(t) + requestId := uuid.New() + + received := make(chan *pb.WorkerInvokeFunctionRequest, 1) + sub, err := nc.Subscribe(reconnectSubject(requestId), func(msg *nats.Msg) { + var work pb.WorkerInvokeFunctionRequest + if err := proto.Unmarshal(msg.Data, &work); err != nil { + return + } + // mirrors the worker's reconnect listener + _ = msg.Respond(nil) + received <- &work + }) + require.NoError(t, err) + require.NoError(t, nc.Flush()) + t.Cleanup(func() { _ = sub.Unsubscribe() }) + + start := time.Now() + err = invoker.joinExistingSession(context.Background(), requestId, &pb.ProxyAuthResponse{FunctionId: "fn-1"}, "worker-token") + require.NoError(t, err) + assert.Less(t, time.Since(start), reconnectAckTimeout, "an acknowledged rejoin should not wait on the deadline") + + select { + case work := <-received: + assert.Equal(t, requestId.String(), work.RequestId) + require.Len(t, work.StatefulConfig.ConnectionConfigs, 1) + assert.Equal(t, "worker-token", + work.StatefulConfig.ConnectionConfigs[0].GetHttp1Config().ProxyAuthorizationToken) + case <-time.After(5 * time.Second): + t.Fatal("worker never received the reconnect message") + } +} + +// A worker built before the acknowledgement is still subscribed, and interest +// alone proves the session is live. The rejoin must succeed rather than be +// mistaken for a dead session, otherwise deploying the proxy ahead of the +// worker would sever every live session. +func TestJoinExistingSessionSubscribedWorkerWithoutAck(t *testing.T) { + invoker, nc := newTestInvoker(t) + requestId := uuid.New() + + received := make(chan struct{}, 1) + sub, err := nc.Subscribe(reconnectSubject(requestId), func(msg *nats.Msg) { + received <- struct{}{} // deliberately no Respond + }) + require.NoError(t, err) + require.NoError(t, nc.Flush()) + t.Cleanup(func() { _ = sub.Unsubscribe() }) + + err = invoker.joinExistingSession(context.Background(), requestId, &pb.ProxyAuthResponse{FunctionId: "fn-1"}, "worker-token") + require.NoError(t, err, "a subscribed worker that does not ack must be treated as live") + + select { + case <-received: + case <-time.After(5 * time.Second): + t.Fatal("worker never received the reconnect message") + } +} + +// A caller that goes away mid-rejoin must not be reported as a dead session: +// that would clear a client cookie on the strength of our own cancellation. +func TestJoinExistingSessionCancelledCallerIsNotASessionLoss(t *testing.T) { + invoker, nc := newTestInvoker(t) + requestId := uuid.New() + + sub, err := nc.Subscribe(reconnectSubject(requestId), func(msg *nats.Msg) {}) + require.NoError(t, err) + require.NoError(t, nc.Flush()) + t.Cleanup(func() { _ = sub.Unsubscribe() }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err = invoker.joinExistingSession(ctx, requestId, &pb.ProxyAuthResponse{FunctionId: "fn-1"}, "worker-token") + require.Error(t, err) + assert.NotErrorIs(t, err, ErrSessionNotFound) + assert.True(t, errors.Is(err, context.Canceled), "expected the caller's cancellation, got %v", err) +} 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..a4b2d3054 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go @@ -324,6 +324,23 @@ var ConnectResults = []string{ ConnectHijackFailed, } +// Outcomes of asking the worker that owns an existing session to open a new +// tunnel. RejoinNoResponders is the one that matters: it is the only signal +// the proxy has that a session can never be served again. +const ( + RejoinAcked = "acked" // worker acknowledged the reconnect + RejoinAssumedLive = "assumed_live" // subscribed but did not ack, worker predates the ack + RejoinNoResponders = "no_responders" // nothing subscribed, session is unrecoverable + RejoinFailed = "failed" // publish itself failed +) + +var RejoinResults = []string{ + RejoinAcked, + RejoinAssumedLive, + RejoinNoResponders, + RejoinFailed, +} + var ( // WorkerConnectTotal counts every CONNECT attempt by outcome. Splitting // the 403s into expired / unknown / mismatch is the point: they are @@ -385,6 +402,17 @@ var ( Help: "worker tunnels still attached to a client connection when it closed", Buckets: []float64{0, 1, 2, 3, 5, 10, 25, 50, 100}, }) + + // StatefulRejoinTotal counts attempts to hand an existing session back to + // its worker. A sustained no_responders rate means clients are holding + // cookies for sessions whose workers are gone; before this was measured + // those rejoins were published into a dead subject and silently dropped. + StatefulRejoinTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Namespace: RootNamespace, + Name: "stateful_rejoin_total", + Help: "attempts to rejoin an existing stateful session, by outcome", + }, []string{"result"}) ) func init() { @@ -406,6 +434,9 @@ func init() { for _, code := range CloseCodes { WorkerConnectionCloseCodeTotal.WithLabelValues(code) } + for _, result := range RejoinResults { + StatefulRejoinTotal.WithLabelValues(result) + } } var nc atomic.Pointer[nats.Conn] diff --git a/src/libraries/go/worker/proxy/proxy.go b/src/libraries/go/worker/proxy/proxy.go index 699fe7558..19eb5c55e 100644 --- a/src/libraries/go/worker/proxy/proxy.go +++ b/src/libraries/go/worker/proxy/proxy.go @@ -206,6 +206,26 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ // doesn't make it in time. go p.keepaliveReconnectRegistration(ctx, work, trackingRegion) + // exit once all connections have completed successfully, + // or 30 seconds after the last connection fails in order to wait for reconnects + var lastConnErr atomic.Pointer[error] + wg := sync.WaitGroup{} + + // Subscribe before the first CONNECT, not after it. The proxy decides + // whether a session is still alive by whether anything is subscribed to + // this subject, so any window where the session exists but the + // subscription does not is a window where a rejoin can be misread as a + // dead session. Subscribing here closes that window for the whole of + // session establishment, including the CONNECT retries below. + subscription, err := p.nc.SubscribeSync("stateful_session.reconnect." + work.RequestId) + if err != nil { + // Not fatal. Without the subscription the session cannot be rejoined, + // but it can still serve the connection it is about to open. + zap.L().Error("failed to listen for stateful session reconnects", zap.String("req id", work.RequestId), zap.Error(err)) + } else { + go p.serveStatefulReconnects(ctx, span, work, subscription, &wg, &lastConnErr) + } + clientConn, err := getClientConnFromProxy(ctx, work, p.h3) if err != nil { return traceError(span, err) @@ -221,10 +241,6 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ zap.L().Info("stateful work request shutting down", zap.String("req id", work.RequestId)) }() - // exit once all connections have completed successfully, - // or 30 seconds after the last connection fails in order to wait for reconnects - var lastConnErr atomic.Pointer[error] - wg := sync.WaitGroup{} wg.Add(1) go func() { defer wg.Done() @@ -238,67 +254,71 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ p.disconnectCallback(ctx, work.RequestId, err) }() - go func() { - zap.L().Info("listening for stateful reconnects", zap.String("req id", work.RequestId)) - subscription, err := p.nc.SubscribeSync("stateful_session.reconnect." + work.RequestId) + wg.Wait() + if err := lastConnErr.Load(); err != nil { + return *err + } + return nil +} + +func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span, work *pb.WorkerInvokeFunctionRequest, subscription *nats.Subscription, wg *sync.WaitGroup, lastConnErr *atomic.Pointer[error]) { + zap.L().Info("listening for stateful reconnects", zap.String("req id", work.RequestId)) + defer func() { _ = subscription.Unsubscribe() }() + for ctx.Err() == nil { + var msg *nats.Msg + err := backoff.Retry(func() error { + nextMsg, err := subscription.NextMsgWithContext(ctx) + if err != nil { + if !errors.Is(err, context.Canceled) { + zap.L().Warn("failed to get next stateful session reconnect message", zap.String("req id", work.RequestId), zap.Error(err)) + } + return err + } + msg = nextMsg + return nil + }, backoff.WithContext(backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(0)), ctx)) if err != nil { - zap.L().Error("failed to listen for stateful session reconnects", zap.String("req id", work.RequestId), zap.Error(err)) return } - defer func() { _ = subscription.Unsubscribe() }() - for ctx.Err() == nil { - var msg *nats.Msg - err = backoff.Retry(func() error { - nextMsg, err := subscription.NextMsgWithContext(ctx) - if err != nil { - if !errors.Is(err, context.Canceled) { - zap.L().Warn("failed to get next stateful session reconnect message", zap.String("req id", work.RequestId), zap.Error(err)) - } - return err - } - msg = nextMsg - return nil - }, backoff.WithContext(backoff.NewExponentialBackOff(backoff.WithMaxElapsedTime(0)), ctx)) + var work pb.WorkerInvokeFunctionRequest + err = proto.Unmarshal(msg.Data, &work) + if err != nil { + zap.L().Warn("malformed stateful session reconnect message", zap.String("req id", work.RequestId), zap.Error(err)) + continue + } + // Acknowledge receipt, as the polling listener does. The proxy uses + // this to tell a live session from one whose worker is gone; without + // it the proxy has to fall back on subscription interest alone and + // waits out its probe deadline on every rejoin. A failed ack is not + // a reason to drop the reconnect: the proxy treats an unanswered + // probe as live, so carry on and serve it. + if err := msg.Respond(nil); err != nil { + zap.L().Warn("failed to ack stateful session reconnect", zap.String("req id", work.RequestId), zap.Error(err)) + } + wg.Add(1) + go func() { + defer wg.Done() + clientConn, err := getClientConnFromProxy(ctx, &work, p.h3) if err != nil { + _ = traceError(span, err) return } - var work pb.WorkerInvokeFunctionRequest - err = proto.Unmarshal(msg.Data, &work) + handlerConn, err := handlerisolationconn.NewHandlerConn(clientConn, p.handlerPool) if err != nil { - zap.L().Warn("malformed stateful session reconnect message", zap.String("req id", work.RequestId), zap.Error(err)) - continue + _ = traceError(span, err) + return } - wg.Add(1) - go func() { - defer wg.Done() - clientConn, err := getClientConnFromProxy(ctx, &work, p.h3) - if err != nil { - _ = traceError(span, err) - return - } - handlerConn, err := handlerisolationconn.NewHandlerConn(clientConn, p.handlerPool) - if err != nil { - _ = traceError(span, err) - return - } - zap.L().Info("connected to nvcf stateful proxy (reconnect)", zap.String("req id", work.RequestId)) - err = p.listener.ServeConn(handlerConn) - if err != nil { - _ = handlerConn.Close() - } - err = traceError(span, err) - lastConnErr.Store(&err) - zap.L().Info("connection closed. triggering callback.", zap.String("req id", work.RequestId)) - p.disconnectCallback(ctx, work.RequestId, err) - }() - } - }() - - wg.Wait() - if err := lastConnErr.Load(); err != nil { - return *err + zap.L().Info("connected to nvcf stateful proxy (reconnect)", zap.String("req id", work.RequestId)) + err = p.listener.ServeConn(handlerConn) + if err != nil { + _ = handlerConn.Close() + } + err = traceError(span, err) + lastConnErr.Store(&err) + zap.L().Info("connection closed. triggering callback.", zap.String("req id", work.RequestId)) + p.disconnectCallback(ctx, work.RequestId, err) + }() } - return nil } func (p *HttpProxy) keepaliveReconnectRegistration(ctx context.Context, work *pb.WorkerInvokeFunctionRequest, trackingRegion string) { diff --git a/src/libraries/go/worker/proxy/proxy_e2e_test.go b/src/libraries/go/worker/proxy/proxy_e2e_test.go index e298749ed..baaa9bcf9 100644 --- a/src/libraries/go/worker/proxy/proxy_e2e_test.go +++ b/src/libraries/go/worker/proxy/proxy_e2e_test.go @@ -25,6 +25,7 @@ import ( "net/http/httptest" "net/http/httputil" "strings" + "sync" "testing" "time" @@ -185,3 +186,115 @@ func TestProxy_EndToEnd(t *testing.T) { t.Fatal("Proxy did not return after teardown") } } + +// The proxy decides whether a stateful session still exists by asking whether +// anything answers on its reconnect subject: no responders means no worker can +// ever serve it, so the client is told to start a new session. Two properties +// have to hold for that to be safe, and this test pins both while the initial +// CONNECT is deliberately left hanging. +// +// The session must be answerable before its first tunnel is up, otherwise a +// rejoin arriving during session establishment reads as a dead session and +// severs a session that was only starting. And the listener must reply, so the +// proxy gets its answer immediately rather than waiting out a probe deadline. +func TestProxy_ReconnectSubjectIsAnsweredBeforeFirstTunnel(t *testing.T) { + setupLogger() + allowInsecure(t) + + cluster, err := testutils.NewNatsSuperCluster(t) + require.NoError(t, err) + defer cluster.Shutdown() + + nc, err := nats.Connect(cluster.Clusters[0].Servers[0].ClientURL()) + require.NoError(t, err) + defer nc.Close() + js, err := jetstream.New(nc) + require.NoError(t, err) + + region := cluster.Clusters[0].Region + + // A listener that accepts and then never answers, so the worker's first + // CONNECT stays outstanding for the whole test and the window under test + // stays open rather than being raced. + stalledProxy, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer func() { _ = stalledProxy.Close() }() + + // Accepted conns are tracked so they can be closed during teardown. + // tcpConnect's response read is not context aware, so cancelling the proxy + // context alone will not unblock a CONNECT that is waiting on a reply. + var stalledMu sync.Mutex + var stalledConns []net.Conn + closeStalledConns := func() { + stalledMu.Lock() + defer stalledMu.Unlock() + for _, conn := range stalledConns { + _ = conn.Close() + } + stalledConns = nil + } + defer closeStalledConns() + go func() { + for { + conn, acceptErr := stalledProxy.Accept() + if acceptErr != nil { + return + } + stalledMu.Lock() + stalledConns = append(stalledConns, conn) + stalledMu.Unlock() + } + }() + + httpProxy, err := NewHttpProxy(nc, js, uuid.New().String(), uuid.New().String(), + func(request *httputil.ProxyRequest) {}, nil, nil, nil) + require.NoError(t, err) + defer httpProxy.Close() + + requestId := uuid.New().String() + work := &pb.WorkerInvokeFunctionRequest{ + RequestId: requestId, + NcaId: "nca-1", + StatefulConfig: &pb.WorkerInvokeFunctionRequest_StatefulConfig{ + ConnectionConfigs: []*pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig{ + { + Config: &pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_Http1Config{ + Http1Config: &pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_HTTP1ConnectionConfig{ + ProxyURI: "http://" + stalledProxy.Addr().String() + "/v1/proxy", + ProxyAuthorizationToken: "dummy-token", + }, + }, + }, + }, + }, + } + + proxyCtx, proxyCancel := context.WithCancel(t.Context()) + proxyDone := make(chan error, 1) + go func() { + proxyDone <- httpProxy.Proxy(proxyCtx, work, region) + }() + + reconnectBody, err := proto.Marshal(work) + require.NoError(t, err) + + // Retry only to absorb the scheduling gap before Proxy runs at all. What is + // being asserted is that the answer, once the session is running, is an + // acknowledgement and never "no responders". + var lastErr error + require.Eventually(t, func() bool { + _, lastErr = nc.Request("stateful_session.reconnect."+requestId, reconnectBody, 2*time.Second) + return lastErr == nil + }, 15*time.Second, 50*time.Millisecond, + "reconnect subject should be acknowledged while the first CONNECT is still outstanding, last error: %v", lastErr) + + // Cancel first so the CONNECT retry loop stops, then drop the stalled conns + // to release the read that cancellation cannot reach. + proxyCancel() + closeStalledConns() + select { + case <-proxyDone: + case <-time.After(15 * time.Second): + t.Fatal("Proxy did not return after teardown") + } +} From 5f17d211668f8dd127f03be5baba2db19a3244b8 Mon Sep 17 00:00:00 2001 From: balaji Date: Wed, 19 Aug 2026 20:25:17 -0700 Subject: [PATCH 2/3] fix(grpc-proxy): remove the WaitGroup race in stateful session teardown Review flagged that the reconnect listener calls wg.Add from its own goroutine while Proxy is in wg.Wait. A reconnect arriving at the instant the last connection finishes takes the counter from zero back to one while Wait is returning, which is the misuse the race detector reports. Holding a WaitGroup reference for the listener itself, the obvious fix, deadlocks: the listener only exits once the session context is cancelled, and that cancellation is deferred until after Proxy returns, so the wait and the listener would wait on each other. Verified by applying it and watching a session with one closed tunnel never return. Track connections with a small lock instead, so registering a connection and declaring the session over happen under the same lock and the listener is never counted. A reconnect that loses the race is refused rather than served against a finished session. Adds a test for the natural teardown path, where the connections close and nothing cancels the context. Every existing test cancels at teardown, which is why the deadlock above went unnoticed. Also corrects two stale references to the old publish in the error text and the metric comment. Co-Authored-By: Balaji Ganesan --- .../proxy/invocation/function_invoker.go | 4 +- .../grpc-proxy/proxy/metrics/metrics.go | 22 ++--- src/libraries/go/worker/proxy/proxy.go | 80 +++++++++++++++-- .../go/worker/proxy/proxy_e2e_test.go | 89 ++++++++++++++++++- 4 files changed, 171 insertions(+), 24 deletions(-) diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go index c71c29c18..cba3f65d5 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/function_invoker.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, @@ -363,7 +363,7 @@ func (f *FunctionInvoker) joinExistingSession(ctx context.Context, requestId uui zap.String("function_id", proxyAuthResponse.FunctionId)) return fmt.Errorf("%w for request id %s", ErrSessionNotFound, requestId) default: - return fmt.Errorf("failed to publish function invocation request to nats: %w", err) + return fmt.Errorf("failed to send stateful session reconnect request to nats: %w", err) } } 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 a4b2d3054..30d2ce969 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, @@ -301,15 +301,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{ @@ -331,7 +331,7 @@ const ( RejoinAcked = "acked" // worker acknowledged the reconnect RejoinAssumedLive = "assumed_live" // subscribed but did not ack, worker predates the ack RejoinNoResponders = "no_responders" // nothing subscribed, session is unrecoverable - RejoinFailed = "failed" // publish itself failed + RejoinFailed = "failed" // the reconnect request itself failed ) var RejoinResults = []string{ diff --git a/src/libraries/go/worker/proxy/proxy.go b/src/libraries/go/worker/proxy/proxy.go index 19eb5c55e..6b7202ffa 100644 --- a/src/libraries/go/worker/proxy/proxy.go +++ b/src/libraries/go/worker/proxy/proxy.go @@ -194,6 +194,64 @@ func (p *HttpProxy) Close() error { return p.handler.Close() } +// sessionConns tracks the connections serving one stateful session. +// +// The reconnect listener can register a connection at any moment, including +// the instant the last existing one finishes. A sync.WaitGroup cannot express +// that safely: the counter would go from zero back to one while Wait was +// already returning, which is the misuse the race detector reports. Holding a +// WaitGroup reference for the listener itself does not work either, because +// the listener only exits once the session context is cancelled and that +// cancellation happens after the wait returns, so the two deadlock. +// +// Registering a connection and declaring the session over therefore happen +// under one lock, and the listener is never counted. +type sessionConns struct { + mu sync.Mutex + cond *sync.Cond + active int + sealed bool +} + +func newSessionConns() *sessionConns { + s := &sessionConns{} + s.cond = sync.NewCond(&s.mu) + return s +} + +// add registers a connection, reporting false once the session has been sealed +// and no further connections will be served. +func (s *sessionConns) add() bool { + s.mu.Lock() + defer s.mu.Unlock() + if s.sealed { + return false + } + s.active++ + return true +} + +func (s *sessionConns) done() { + s.mu.Lock() + defer s.mu.Unlock() + s.active-- + if s.active == 0 { + s.cond.Broadcast() + } +} + +// waitAndSeal blocks until every registered connection has finished, then seals +// under the same lock so a reconnect cannot register after the session has been +// declared over. +func (s *sessionConns) waitAndSeal() { + s.mu.Lock() + defer s.mu.Unlock() + for s.active > 0 { + s.cond.Wait() + } + s.sealed = true +} + func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequest, trackingRegion string) error { ctx, cancel := context.WithCancel(ctx) defer cancel() @@ -209,7 +267,7 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ // exit once all connections have completed successfully, // or 30 seconds after the last connection fails in order to wait for reconnects var lastConnErr atomic.Pointer[error] - wg := sync.WaitGroup{} + conns := newSessionConns() // Subscribe before the first CONNECT, not after it. The proxy decides // whether a session is still alive by whether anything is subscribed to @@ -223,7 +281,7 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ // but it can still serve the connection it is about to open. zap.L().Error("failed to listen for stateful session reconnects", zap.String("req id", work.RequestId), zap.Error(err)) } else { - go p.serveStatefulReconnects(ctx, span, work, subscription, &wg, &lastConnErr) + go p.serveStatefulReconnects(ctx, span, work, subscription, conns, &lastConnErr) } clientConn, err := getClientConnFromProxy(ctx, work, p.h3) @@ -241,9 +299,9 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ zap.L().Info("stateful work request shutting down", zap.String("req id", work.RequestId)) }() - wg.Add(1) + conns.add() go func() { - defer wg.Done() + defer conns.done() err := p.listener.ServeConn(handlerConn) if err != nil { _ = handlerConn.Close() @@ -254,14 +312,14 @@ func (p *HttpProxy) Proxy(ctx context.Context, work *pb.WorkerInvokeFunctionRequ p.disconnectCallback(ctx, work.RequestId, err) }() - wg.Wait() + conns.waitAndSeal() if err := lastConnErr.Load(); err != nil { return *err } return nil } -func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span, work *pb.WorkerInvokeFunctionRequest, subscription *nats.Subscription, wg *sync.WaitGroup, lastConnErr *atomic.Pointer[error]) { +func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span, work *pb.WorkerInvokeFunctionRequest, subscription *nats.Subscription, conns *sessionConns, lastConnErr *atomic.Pointer[error]) { zap.L().Info("listening for stateful reconnects", zap.String("req id", work.RequestId)) defer func() { _ = subscription.Unsubscribe() }() for ctx.Err() == nil { @@ -295,9 +353,15 @@ func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span if err := msg.Respond(nil); err != nil { zap.L().Warn("failed to ack stateful session reconnect", zap.String("req id", work.RequestId), zap.Error(err)) } - wg.Add(1) + if !conns.add() { + // The session finished while this reconnect was in flight, so + // there is nothing left to attach it to. Stop listening rather + // than serve a connection nobody is waiting on. + zap.L().Info("stateful session already finished, ignoring reconnect", zap.String("req id", work.RequestId)) + return + } go func() { - defer wg.Done() + defer conns.done() clientConn, err := getClientConnFromProxy(ctx, &work, p.h3) if err != nil { _ = traceError(span, err) diff --git a/src/libraries/go/worker/proxy/proxy_e2e_test.go b/src/libraries/go/worker/proxy/proxy_e2e_test.go index baaa9bcf9..f77b3bbe6 100644 --- a/src/libraries/go/worker/proxy/proxy_e2e_test.go +++ b/src/libraries/go/worker/proxy/proxy_e2e_test.go @@ -281,12 +281,23 @@ func TestProxy_ReconnectSubjectIsAnsweredBeforeFirstTunnel(t *testing.T) { // Retry only to absorb the scheduling gap before Proxy runs at all. What is // being asserted is that the answer, once the session is running, is an // acknowledgement and never "no responders". + // testify runs the condition on its own goroutine, so the last error is + // captured under a lock and reported after Eventually returns rather than + // being passed as a message argument, which would be evaluated up front and + // always read nil. + var lastErrMu sync.Mutex var lastErr error require.Eventually(t, func() bool { - _, lastErr = nc.Request("stateful_session.reconnect."+requestId, reconnectBody, 2*time.Second) - return lastErr == nil + _, reqErr := nc.Request("stateful_session.reconnect."+requestId, reconnectBody, 2*time.Second) + lastErrMu.Lock() + lastErr = reqErr + lastErrMu.Unlock() + return reqErr == nil }, 15*time.Second, 50*time.Millisecond, - "reconnect subject should be acknowledged while the first CONNECT is still outstanding, last error: %v", lastErr) + "reconnect subject should be acknowledged while the first CONNECT is still outstanding") + lastErrMu.Lock() + require.NoError(t, lastErr) + lastErrMu.Unlock() // Cancel first so the CONNECT retry loop stops, then drop the stalled conns // to release the read that cancellation cannot reach. @@ -298,3 +309,75 @@ func TestProxy_ReconnectSubjectIsAnsweredBeforeFirstTunnel(t *testing.T) { t.Fatal("Proxy did not return after teardown") } } + +// A session ends when its connections close, with nothing cancelling the work +// request's context. Proxy has to return on its own at that point, so the +// reconnect listener must never be something the session waits on: the listener +// only exits once that context is cancelled, and the cancellation happens after +// Proxy returns. Counting it would deadlock every stateful session, and no +// other test covers this because they all cancel at teardown. +func TestProxy_ReturnsWhenConnectionsCloseWithoutExternalCancel(t *testing.T) { + setupLogger() + allowInsecure(t) + + cluster, err := testutils.NewNatsSuperCluster(t) + require.NoError(t, err) + defer cluster.Shutdown() + + nc, err := nats.Connect(cluster.Clusters[0].Servers[0].ClientURL()) + require.NoError(t, err) + defer nc.Close() + js, err := jetstream.New(nc) + require.NoError(t, err) + region := cluster.Clusters[0].Region + + serverConns, _, _, server := mockGrpcProxy() + defer server.Close() + + inferenceServer := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = io.Copy(w, r.Body) + })) + protocols := &http.Protocols{} + protocols.SetUnencryptedHTTP2(true) + protocols.SetHTTP1(true) + inferenceServer.Config.Protocols = protocols + inferenceServer.Start() + defer inferenceServer.Close() + + httpProxy, err := NewHttpProxy(nc, js, uuid.New().String(), uuid.New().String(), + func(request *httputil.ProxyRequest) { + request.Out.URL.Scheme = "http" + request.Out.URL.Host = inferenceServer.Listener.Addr().String() + }, nil, nil, + // No-op disconnect callback so the test does not sit through the + // reconnect wait; the teardown path under test is the same. + func(ctx context.Context, reqId string, err error) {}) + require.NoError(t, err) + defer httpProxy.Close() + + work := &pb.WorkerInvokeFunctionRequest{ + RequestId: uuid.New().String(), + StatefulConfig: &pb.WorkerInvokeFunctionRequest_StatefulConfig{ + ConnectionConfigs: []*pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig{ + {Config: &pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_Http3Config{ + Http3Config: &pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_HTTP3ConnectionConfig{ + ProxyURI: "https://localhost:10084/v1/proxy", + ProxyAuthorizationToken: "dummy-token", + }}}, + }, + }, + } + + // Deliberately not cancelled anywhere in this test. + proxyDone := make(chan error, 1) + go func() { proxyDone <- httpProxy.Proxy(context.Background(), work, region) }() + + serverConn := <-serverConns + _ = serverConn.Close() + + select { + case <-proxyDone: + case <-time.After(20 * time.Second): + t.Fatal("Proxy did not return after its only tunnel closed") + } +} From 311fb44f7907fb00ad31d16f812ea675a9bddc5b Mon Sep 17 00:00:00 2001 From: balaji Date: Thu, 20 Aug 2026 09:23:31 -0700 Subject: [PATCH 3/3] fix(worker): keep the session id in the malformed-reconnect log The unmarshalled reconnect payload shadowed the work parameter, so the malformed-payload warning logged the request id of the payload that had just failed to parse. That field was therefore always empty for exactly the message an operator needs to trace. Name the payload apart from the session so the log keeps the session id and the two stay distinguishable at every use. Co-Authored-By: Balaji Ganesan --- src/libraries/go/worker/proxy/proxy.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/libraries/go/worker/proxy/proxy.go b/src/libraries/go/worker/proxy/proxy.go index 6b7202ffa..1f258fe16 100644 --- a/src/libraries/go/worker/proxy/proxy.go +++ b/src/libraries/go/worker/proxy/proxy.go @@ -338,8 +338,12 @@ func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span if err != nil { return } - var work pb.WorkerInvokeFunctionRequest - err = proto.Unmarshal(msg.Data, &work) + // Named apart from the work parameter on purpose. Shadowing it meant the + // malformed-payload log below reported the request id of the payload + // that had just failed to parse, so the field was always empty for + // exactly the message an operator needs to trace. + var reconnectWork pb.WorkerInvokeFunctionRequest + err = proto.Unmarshal(msg.Data, &reconnectWork) if err != nil { zap.L().Warn("malformed stateful session reconnect message", zap.String("req id", work.RequestId), zap.Error(err)) continue @@ -362,7 +366,7 @@ func (p *HttpProxy) serveStatefulReconnects(ctx context.Context, span trace.Span } go func() { defer conns.done() - clientConn, err := getClientConnFromProxy(ctx, &work, p.h3) + clientConn, err := getClientConnFromProxy(ctx, &reconnectWork, p.h3) if err != nil { _ = traceError(span, err) return