Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions src/invocation-plane-services/grpc-proxy/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
11 changes: 11 additions & 0 deletions src/invocation-plane-services/grpc-proxy/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand Down Expand Up @@ -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=
Expand All @@ -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=
Expand Down Expand Up @@ -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=
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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",
],
)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

// 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 {
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
}
Loading
Loading