diff --git a/src/invocation-plane-services/grpc-proxy/go.mod b/src/invocation-plane-services/grpc-proxy/go.mod index 2163b7c0f..8090eaca0 100644 --- a/src/invocation-plane-services/grpc-proxy/go.mod +++ b/src/invocation-plane-services/grpc-proxy/go.mod @@ -23,7 +23,7 @@ require ( 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 - github.com/quic-go/quic-go v0.59.1 + github.com/quic-go/quic-go v0.61.0 github.com/samber/lo v1.51.0 github.com/spf13/cobra v1.10.2 github.com/spf13/pflag v1.0.10 @@ -39,8 +39,8 @@ require ( go.opentelemetry.io/otel/trace v1.44.0 go.uber.org/atomic v1.11.0 go.uber.org/zap v1.27.0 - golang.org/x/net v0.55.0 - golang.org/x/sync v0.20.0 + golang.org/x/net v0.56.0 + golang.org/x/sync v0.22.0 google.golang.org/grpc v1.81.1 google.golang.org/protobuf v1.36.11 ) @@ -155,11 +155,11 @@ require ( go.uber.org/multierr v1.11.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/crypto v0.51.0 // indirect - golang.org/x/mod v0.35.0 // indirect + golang.org/x/crypto v0.54.0 // indirect + golang.org/x/mod v0.37.0 // indirect 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/sys v0.47.0 // indirect + golang.org/x/text v0.40.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..8013136b5 100644 --- a/src/invocation-plane-services/grpc-proxy/go.sum +++ b/src/invocation-plane-services/grpc-proxy/go.sum @@ -351,10 +351,12 @@ github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA= github.com/prometheus/procfs v0.20.1 h1:XwbrGOIplXW/AU3YhIhLODXMJYyC1isLFfYCsTEycfc= github.com/prometheus/procfs v0.20.1/go.mod h1:o9EMBZGRyvDrSPH1RqdxhojkuXstoe4UlK79eF5TGGo= +github.com/quic-go/go-ossfuzz-seeds v0.1.0 h1:APacT+iIaNF6fd8AGEiN3bT/Jtkd2jz4v4TzM7MFjy0= +github.com/quic-go/go-ossfuzz-seeds v0.1.0/go.mod h1:3IOHRbJIc+L6YKMwfDtJAM9Vj9k0YY4muhuyUYk5tbk= github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= -github.com/quic-go/quic-go v0.59.1 h1:0Gmua0HW1Tv7ANR7hUYwRyD0MG5OJfgvYSZasGZzBic= -github.com/quic-go/quic-go v0.59.1/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= +github.com/quic-go/quic-go v0.61.0 h1:ui88A53s8MSVYLC56en0KQ17HARk+9986Dn0SBfKNvA= +github.com/quic-go/quic-go v0.61.0/go.mod h1:9So2anK4Tp22URSQq00k+Vo2PNkle96ycDPDHL4s9vs= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= @@ -473,8 +475,8 @@ golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnf golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= -golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI= -golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8= +golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= +golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= @@ -482,8 +484,8 @@ golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHl golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= -golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -496,8 +498,8 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200625001655-4c5254603344/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= -golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= +golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= @@ -509,8 +511,8 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -531,15 +533,15 @@ 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.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= -golang.org/x/term v0.43.0/go.mod h1:lrhlHNdQJHO+1qVYiHfFKVuVioJIheAc3fBSMFYEIsk= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0= +golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= diff --git a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel index 82b929724..6af7ac884 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel +++ b/src/invocation-plane-services/grpc-proxy/proxy/BUILD.bazel @@ -88,6 +88,7 @@ go_test( "director_test.go", "hijack_test.go", "info_test.go", + "pending_work_test.go", "proxy_test.go", ], embed = [":proxy"], @@ -103,6 +104,8 @@ go_test( "@com_github_gorilla_websocket//:websocket", "@com_github_hellofresh_health_go_v5//:health-go", "@com_github_jellydator_ttlcache_v3//:ttlcache", + "@com_github_stretchr_testify//assert", + "@com_github_stretchr_testify//require", "@org_golang_google_grpc//:grpc", "@org_golang_google_grpc//codes", "@org_golang_google_grpc//credentials/insecure", diff --git a/src/invocation-plane-services/grpc-proxy/proxy/director.go b/src/invocation-plane-services/grpc-proxy/proxy/director.go index ffcbad43a..483be0fb1 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/director.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/director.go @@ -69,10 +69,26 @@ type issuedTokenInfo struct { mintedAt time.Time } +// pendingWorkInfo identifies a stateful work request this pod has issued a +// worker token for but has not yet seen a CONNECT for. The token lives only in +// this pod's memory, so if the pod goes away the queued request can never +// authenticate; the entry is what lets shutdown find it and drop it. +type pendingWorkInfo struct { + functionVersionId string +} + +// pendingWorkPurger removes a queued stateful work request. Implemented by the +// function invoker and asserted optionally, so an invoker that cannot reach the +// work queue (tests, alternative implementations) simply skips the purge. +type pendingWorkPurger interface { + PurgePendingWork(ctx context.Context, requestId uuid.UUID, functionVersionId string) error +} + type StreamDirector struct { shuttingDown *atomic.Bool workerAuth *ttlcache.Cache[string, workerAuthInfo] // auth -> request + function info issuedTokens *ttlcache.Cache[string, issuedTokenInfo] // diagnostic only, see issuedTokenInfo + pendingWork *ttlcache.Cache[uuid.UUID, pendingWorkInfo] workers *ttlcache.Cache[workerConnectionKey, *worker.WorkerConnection] functionInvoker FunctionInvoker cors *cors.Cors @@ -108,6 +124,20 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector { ) go issuedTokenCache.Start() + // Sessions waiting on a worker CONNECT. Retention is deliberately much + // longer than the token TTL: the point is to still know about a request + // whose token has already aged out, because that request is still sitting + // in the work queue. Bounded so it cannot grow without limit. + pendingWorkCache := ttlcache.New( + // Deliberately reuses the issued-token retention and capacity rather + // than introducing its own. This tracks the same population from the + // same call site, and the codebase does not need another timeout. + ttlcache.WithTTL[uuid.UUID, pendingWorkInfo](issuedTokenRetention), + ttlcache.WithCapacity[uuid.UUID, pendingWorkInfo](issuedTokenCacheCapacity), + ttlcache.WithDisableTouchOnHit[uuid.UUID, pendingWorkInfo](), + ) + go pendingWorkCache.Start() + // Set immediately before DeleteAll in Close so the eviction handler can // report shutdown rather than attributing a drain to a client or worker. shuttingDown := &atomic.Bool{} @@ -241,6 +271,7 @@ func NewStreamDirector(functionInvoker FunctionInvoker) *StreamDirector { workers: cache, shuttingDown: shuttingDown, issuedTokens: issuedTokenCache, + pendingWork: pendingWorkCache, workerAuth: workerAuthCache, functionInvoker: functionInvoker, cors: cors.New(middleware.DefaultCorsOptions), @@ -323,14 +354,77 @@ const ( issuedTokenCacheCapacity = 50000 ) +// purgePendingWork drops the work requests for sessions that never got a +// worker CONNECT. +// +// Best effort by design. A session records its pending work just before the +// invocation publishes the work request, so a shutdown landing precisely +// between those two steps will miss that one request. Closing that window +// needs an admission gate and a drain timeout, which is more machinery and +// another tunable than the gap justifies: a missed request is simply left as +// it is today, and today every one of them is left. +// +// Their tokens exist only in this pod's memory, so once it is +// gone every one of them is guaranteed to be rejected; leaving them queued +// means each is still pulled, still takes a concurrency slot, and still fails. +// +// Sessions with a worker already attached are deliberately not touched. Those +// can reattach through another pod, and their work request has already left the +// queue anyway. +func (s *StreamDirector) purgePendingWork() { + purger, ok := s.functionInvoker.(pendingWorkPurger) + if !ok { + return + } + pending := s.pendingWork.Items() + if len(pending) == 0 { + return + } + + ctx, cancel := context.WithTimeout(context.Background(), consts.Timeout) + defer cancel() + + var purged, failed int + for requestId, item := range pending { + if ctx.Err() != nil { + // Out of budget. Report what is left rather than trailing off + // silently, so a shutdown that could not finish is visible. + failed += len(pending) - purged - failed + break + } + if err := purger.PurgePendingWork(ctx, requestId, item.Value().functionVersionId); err != nil { + failed++ + // Expected when this service has no rights on the work queue, so + // this stays a warning: the purge is an optimisation and shutdown + // is still correct without it. + zap.L().Warn("failed to purge pending stateful work request on shutdown", + zap.Stringer("request_id", requestId), + zap.String("function_version_id", item.Value().functionVersionId), + zap.Error(err)) + continue + } + purged++ + } + + metrics.PendingWorkPurgedTotal.WithLabelValues(metrics.PurgeSucceeded).Add(float64(purged)) + metrics.PendingWorkPurgedTotal.WithLabelValues(metrics.PurgeFailed).Add(float64(failed)) + zap.L().Info("purged pending stateful work requests on shutdown", + zap.Int("purged", purged), + zap.Int("failed", failed)) +} + func (s *StreamDirector) Close() error { // Mark first: DeleteAll evicts every entry, and without this those // evictions would be misreported as client or worker initiated. s.shuttingDown.Store(true) + // Before the caches go away, drop the queued work this pod can no longer + // authenticate. Best effort: a failure here must not hold up shutdown. + s.purgePendingWork() s.workers.DeleteAll() s.workers.Stop() s.workerAuth.Stop() s.issuedTokens.Stop() + s.pendingWork.Stop() if s.functionInvoker != nil { if closer, ok := s.functionInvoker.(io.Closer); ok { _ = closer.Close() @@ -532,6 +626,11 @@ func (s *StreamDirector) getAndInitWorkerConnection(ctx context.Context, conn *w functionVersionId: apiFuncVersion, mintedAt: now, }, ttlcache.DefaultTTL) + // Remembered until the worker CONNECTs back, so that a shutdown can + // find the requests whose tokens are about to be lost with this pod + // and drop them from the work queue instead of leaving them to be + // pulled and rejected. + s.pendingWork.Set(requestId, pendingWorkInfo{functionVersionId: apiFuncVersion}, ttlcache.DefaultTTL) // Diagnostic shadow record, longer lived than the auth entry, so a // later rejection can say "expired N seconds ago" instead of just // "not found". Never consulted when granting access. diff --git a/src/invocation-plane-services/grpc-proxy/proxy/hijack.go b/src/invocation-plane-services/grpc-proxy/proxy/hijack.go index 747a3fd18..129250ef1 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/hijack.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/hijack.go @@ -174,6 +174,9 @@ func (s *StreamDirector) HijackHandler(w http.ResponseWriter, r *http.Request) { // already been sent away, so no point in keeping the auth around. // also need to make sure we don't allow reconnects to guard against replay attacks with 0-rtt. s.workerAuth.Delete(auth) + // The worker is attached, so this request is no longer queued and must not + // be purged if this pod shuts down. The session can reattach elsewhere. + s.pendingWork.Delete(parsedRequestId) } // networkPeerAddress returns the transport-level remote host of the tunnel connection and diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel b/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel index 8d5ef783b..fb8eb0445 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/BUILD.bazel @@ -21,6 +21,7 @@ go_library( "function_invoker.go", "nats.go", "nats_tracing.go", + "pending_work.go", "sleep.go", "stateful_session_request_registration.go", ], @@ -57,9 +58,13 @@ alias( go_test( name = "invocation_test", - srcs = ["nats_test.go"], + srcs = [ + "nats_test.go", + "pending_work_test.go", + ], embed = [":invocation"], deps = [ + "@com_github_google_uuid//:uuid", "@com_github_nats_io_nkeys//:nkeys", "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go new file mode 100644 index 000000000..d0837b775 --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work.go @@ -0,0 +1,65 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ +package invocation + +import ( + "context" + "fmt" + + "github.com/google/uuid" + "github.com/nats-io/nats.go/jetstream" +) + +// The work queue a stateful invocation lands in. These are owned by the +// invocation service, which publishes to requestWorkSubject and removes a +// cancelled request with the same subject-filtered purge used here. The +// formats are duplicated rather than shared because the owning service is +// written in another language; they must not drift. +func requestWorkStream(region, functionVersionId string) string { + return fmt.Sprintf("rq_%s_%s", region, functionVersionId) +} + +func requestWorkSubject(region, functionVersionId string, requestId uuid.UUID) string { + return fmt.Sprintf("rq.%s.%s.%s", region, functionVersionId, requestId) +} + +// PurgePendingWork drops a stateful work request that is still queued. +// +// It is called for sessions this pod issued a worker token for that never came +// back to CONNECT, at the point the pod is shutting down. Those tokens only +// exist in this pod's memory, so every one of those queued requests is already +// guaranteed to fail authentication whenever a worker eventually pulls it. Left +// in place they are pulled anyway, occupy a worker concurrency slot, and fail, +// which is what keeps a saturated function at zero goodput long after the +// restart that caused it. +// +// Purging by subject only removes messages still held by the stream. A session +// that already has a worker attached had its message delivered, so an +// established session is unaffected and remains free to reattach to another +// pod. +func (f *FunctionInvoker) PurgePendingWork(ctx context.Context, requestId uuid.UUID, functionVersionId string) error { + streamName := requestWorkStream(f.region, functionVersionId) + stream, err := f.js.Stream(ctx, streamName) + if err != nil { + return fmt.Errorf("failed to look up work stream %s: %w", streamName, err) + } + subject := requestWorkSubject(f.region, functionVersionId, requestId) + if err := stream.Purge(ctx, jetstream.WithPurgeSubject(subject)); err != nil { + return fmt.Errorf("failed to purge work subject %s: %w", subject, err) + } + return nil +} diff --git a/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go new file mode 100644 index 000000000..ff878c626 --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/invocation/pending_work_test.go @@ -0,0 +1,50 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ +package invocation + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" +) + +// The work queue is owned by the invocation service, which is written in +// another language, so these formats are duplicated rather than shared. If they +// drift the purge silently targets a subject nothing was ever published to and +// removes nothing, with no error to show for it. These cases pin the exact +// strings against the owning service's request_stream_name and request_subject. +func TestRequestWorkStreamAndSubjectMatchTheInvocationService(t *testing.T) { + requestId := uuid.MustParse("11111111-2222-3333-4444-555555555555") + region := "us-west-2" + versionId := "66666666-7777-8888-9999-000000000000" + + assert.Equal(t, "rq_us-west-2_66666666-7777-8888-9999-000000000000", + requestWorkStream(region, versionId)) + assert.Equal(t, "rq.us-west-2.66666666-7777-8888-9999-000000000000.11111111-2222-3333-4444-555555555555", + requestWorkSubject(region, versionId, requestId)) +} + +// The subject has to fall inside the stream's own subject space, otherwise a +// filtered purge matches nothing. +func TestRequestWorkSubjectIsCoveredByTheStreamSubjectSpace(t *testing.T) { + region := "eu-west-1" + versionId := uuid.New().String() + + subject := requestWorkSubject(region, versionId, uuid.New()) + assert.Regexp(t, `^rq\.`+region+`\.`+versionId+`\.`, subject) +} diff --git a/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go b/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go index 187b19c87..eb2942d12 100644 --- a/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go +++ b/src/invocation-plane-services/grpc-proxy/proxy/metrics/metrics.go @@ -6,7 +6,7 @@ Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at - http://www.apache.org/licenses/LICENSE-2.0 + http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, @@ -189,6 +189,14 @@ var ( }) ) +// Outcomes of the shutdown purge of queued stateful work requests. +const ( + PurgeSucceeded = "succeeded" + PurgeFailed = "failed" +) + +var PurgeResults = []string{PurgeSucceeded, PurgeFailed} + // Close reasons reported when a worker tunnel goes away. The three "deleted" // variants matter: a bare `deleted` cannot distinguish the client hanging up // from the worker hanging up from a proxy shutdown, and which side went first @@ -301,15 +309,15 @@ var WorkerConnectionCloseReasons = []string{ // Outcomes of a worker CONNECT to /v1/proxy. Every terminal path in // HijackHandler maps to exactly one of these. const ( - ConnectAccepted = "accepted" - ConnectNotHijackable = "rejected_not_hijackable" // 500 - ConnectMissingAuth = "rejected_missing_auth" // 401 - ConnectMissingRequestID = "rejected_missing_requestid" // 400 - ConnectInvalidRequestID = "rejected_invalid_requestid" // 400 - ConnectTokenExpired = "rejected_token_expired" // 403, token was issued but has aged out - ConnectTokenUnknown = "rejected_token_unknown" // 403, token was never issued by this pod - ConnectRequestIDMismatch = "rejected_requestid_mismatch"// 403, token valid but bound to another request - ConnectHijackFailed = "rejected_hijack_failed" // 500 + ConnectAccepted = "accepted" + ConnectNotHijackable = "rejected_not_hijackable" // 500 + ConnectMissingAuth = "rejected_missing_auth" // 401 + ConnectMissingRequestID = "rejected_missing_requestid" // 400 + ConnectInvalidRequestID = "rejected_invalid_requestid" // 400 + ConnectTokenExpired = "rejected_token_expired" // 403, token was issued but has aged out + ConnectTokenUnknown = "rejected_token_unknown" // 403, token was never issued by this pod + ConnectRequestIDMismatch = "rejected_requestid_mismatch" // 403, token valid but bound to another request + ConnectHijackFailed = "rejected_hijack_failed" // 500 ) var ConnectResults = []string{ @@ -375,6 +383,17 @@ var ( Buckets: []float64{1, 5, 8, 10, 15, 30, 45, 60, 120, 300, 600, 1800, 3600}, }) + // PendingWorkPurgedTotal counts stateful work requests dropped from the + // work queue during shutdown because this pod held the only copy of their + // worker token. A persistent failed count usually means this service lacks + // purge rights on the work queue rather than a transient NATS error. + PendingWorkPurgedTotal = promauto.NewCounterVec( + prometheus.CounterOpts{ + Namespace: RootNamespace, + Name: "pending_work_purged_total", + Help: "queued stateful work requests dropped at shutdown, by outcome", + }, []string{"result"}) + // ClientConnectionWorkerTunnelsAtClose records how many worker tunnels a // client connection was still holding when it closed. Anything above zero // means that close tore down live tunnels. @@ -406,6 +425,9 @@ func init() { for _, code := range CloseCodes { WorkerConnectionCloseCodeTotal.WithLabelValues(code) } + for _, result := range PurgeResults { + PendingWorkPurgedTotal.WithLabelValues(result) + } } var nc atomic.Pointer[nats.Conn] diff --git a/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go b/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go new file mode 100644 index 000000000..ec5261784 --- /dev/null +++ b/src/invocation-plane-services/grpc-proxy/proxy/pending_work_test.go @@ -0,0 +1,124 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ +package proxy + +import ( + "context" + "errors" + "net" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/jellydator/ttlcache/v3" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "nvcf-grpc-proxy/proxy/invocation" +) + +// recordingPurger stands in for the function invoker, capturing which requests +// shutdown asked to drop. +type recordingPurger struct { + mu sync.Mutex + purged map[uuid.UUID]string + err error +} + +func newRecordingPurger() *recordingPurger { + return &recordingPurger{purged: map[uuid.UUID]string{}} +} + +func (p *recordingPurger) InvokeStatefulFunction(_ context.Context, _ net.Conn, _, _, _ string, _ *uuid.UUID, _ func(string, uuid.UUID, string, string)) (invocation.Result, context.CancelFunc, error) { + return invocation.Result{}, nil, errors.New("not used") +} + +func (p *recordingPurger) PurgePendingWork(_ context.Context, requestId uuid.UUID, functionVersionId string) error { + p.mu.Lock() + defer p.mu.Unlock() + if p.err != nil { + return p.err + } + p.purged[requestId] = functionVersionId + return nil +} + +func (p *recordingPurger) purgedRequests() map[uuid.UUID]string { + p.mu.Lock() + defer p.mu.Unlock() + out := map[uuid.UUID]string{} + for k, v := range p.purged { + out[k] = v + } + return out +} + +// A request whose worker never came back to CONNECT can never authenticate once +// this pod is gone, because the token only ever existed here. Shutdown has to +// take it out of the work queue, otherwise it is still pulled, still occupies a +// worker slot, and still fails. +func TestClosePurgesWorkForSessionsAwaitingConnect(t *testing.T) { + purger := newRecordingPurger() + director := NewStreamDirector(purger) + + awaiting := uuid.New() + director.pendingWork.Set(awaiting, pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) + + purged := purger.purgedRequests() + require.Len(t, purged, 1) + assert.Equal(t, "version-1", purged[awaiting]) +} + +// A session with a worker already attached is not queued any more and can +// reattach through another pod. Purging it would sever a session that was going +// to survive the restart, so shutdown must leave it alone. +func TestClosePurgesNothingOnceTheWorkerHasConnected(t *testing.T) { + purger := newRecordingPurger() + director := NewStreamDirector(purger) + + connected := uuid.New() + director.pendingWork.Set(connected, pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + // what HijackHandler does once the worker's CONNECT is accepted + director.pendingWork.Delete(connected) + + require.NoError(t, director.Close()) + + assert.Empty(t, purger.purgedRequests()) +} + +// The purge is an optimisation. If this service has no rights on the work queue +// the calls fail, and shutdown still has to complete. +func TestClosePurgeFailureDoesNotBlockShutdown(t *testing.T) { + purger := newRecordingPurger() + purger.err = errors.New("nats: permissions violation for stream purge") + director := NewStreamDirector(purger) + + director.pendingWork.Set(uuid.New(), pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) +} + +// An invoker with no route to the work queue simply skips the purge rather than +// failing shutdown. +func TestClosePurgeSkippedWhenInvokerCannotPurge(t *testing.T) { + director := NewStreamDirector((*mockInvoker)(&invocation.Result{})) + director.pendingWork.Set(uuid.New(), pendingWorkInfo{functionVersionId: "version-1"}, ttlcache.DefaultTTL) + + require.NoError(t, director.Close()) +}