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..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, @@ -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 send stateful session reconnect 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..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{ @@ -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" // the reconnect request 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..1f258fe16 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() @@ -206,6 +264,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] + 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 + // 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, conns, &lastConnErr) + } + clientConn, err := getClientConnFromProxy(ctx, work, p.h3) if err != nil { return traceError(span, err) @@ -221,13 +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)) }() - // 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) + conns.add() go func() { - defer wg.Done() + defer conns.done() err := p.listener.ServeConn(handlerConn) if err != nil { _ = handlerConn.Close() @@ -238,67 +312,81 @@ 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) + 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, 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 { + 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)) + // 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 + } + // 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)) + } + 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 conns.done() + clientConn, err := getClientConnFromProxy(ctx, &reconnectWork, 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..f77b3bbe6 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,198 @@ 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". + // 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 { + _, 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") + 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. + proxyCancel() + closeStalledConns() + select { + case <-proxyDone: + case <-time.After(15 * time.Second): + 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") + } +}