diff --git a/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/_helpers.tpl b/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/_helpers.tpl index d371fcaa4..e5549da4f 100644 --- a/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/_helpers.tpl +++ b/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/_helpers.tpl @@ -149,4 +149,4 @@ Create the name of the service account to use {{- else }} {{- default "default" .Values.rateLimiter.serviceAccount.name }} {{- end }} -{{- end }} \ No newline at end of file +{{- end }} diff --git a/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/deployment.yaml b/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/deployment.yaml index ae6f1a728..c43604964 100644 --- a/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/deployment.yaml +++ b/deploy/helm/ratelimiter/nvcf-ratelimiter/templates/deployment.yaml @@ -22,6 +22,14 @@ metadata: {{- include "nvcf-ratelimiter.labels" . | nindent 4 }} spec: replicas: {{ .Values.rateLimiter.replicaCount }} + # Surge first so a live member exists to receive the outgoing pod's counters. + # Pinned because the percentage defaults allow an unavailable pod at some + # replica counts. + strategy: + type: RollingUpdate + rollingUpdate: + maxSurge: 1 + maxUnavailable: 0 selector: matchLabels: {{- include "nvcf-ratelimiter.selectorLabels" . | nindent 6 }} @@ -32,10 +40,12 @@ spec: {{- with .Values.rateLimiter.podAnnotations }} {{- toYaml . | nindent 8 }} {{- end }} - checksum/config-env: {{ printf "%s|%s|%s" (toYaml .Values.rateLimiter.env) (include "nvcf-ratelimiter.oauth2Issuer" .) (include "nvcf-ratelimiter.audience" .) | sha256sum }} + checksum/config-env: {{ printf "%s|%s|%s|%s" (toYaml .Values.rateLimiter.env) (include "nvcf-ratelimiter.oauth2Issuer" .) (include "nvcf-ratelimiter.audience" .) (toString .Values.rateLimiter.olricReplicaCount) | sha256sum }} labels: {{- include "nvcf-ratelimiter.selectorLabels" . | nindent 8 }} spec: + # Covers the readiness delay plus the counter hand-off. + terminationGracePeriodSeconds: {{ .Values.rateLimiter.terminationGracePeriodSeconds }} {{- with .Values.rateLimiter.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -83,6 +93,10 @@ spec: value: {{ include "nvcf-ratelimiter.oauth2Issuer" . | quote }} - name: AUDIENCE value: {{ include "nvcf-ratelimiter.audience" . | quote }} + {{- with .Values.rateLimiter.olricReplicaCount }} + - name: OLRIC_REPLICA_COUNT + value: {{ . | quote }} + {{- end }} {{- range $key, $value := .Values.rateLimiter.env }} - name: {{ $key }} value: {{ $value | quote }} diff --git a/deploy/helm/ratelimiter/nvcf-ratelimiter/values.yaml b/deploy/helm/ratelimiter/nvcf-ratelimiter/values.yaml index 0baf68b5d..6cd1d07bc 100644 --- a/deploy/helm/ratelimiter/nvcf-ratelimiter/values.yaml +++ b/deploy/helm/ratelimiter/nvcf-ratelimiter/values.yaml @@ -32,6 +32,14 @@ rateLimiter: replicaCount: 1 + # Copies of each counter in the Olric cluster. Empty uses the service default, + # which keeps one backup so losing a member does not reset its counters. + olricReplicaCount: "" + + # Headroom for the counter hand-off on SIGTERM. The hand-off fits the 30s + # Kubernetes default; this leaves room for a slower cluster. + terminationGracePeriodSeconds: 60 + podDisruptionBudget: enabled: false # minAvailable and maxUnavailable are mutually exclusive; set exactly one. diff --git a/src/invocation-plane-services/ratelimiter/cmd/info_test.go b/src/invocation-plane-services/ratelimiter/cmd/info_test.go index 6c41cc0ea..d73b1fcf9 100644 --- a/src/invocation-plane-services/ratelimiter/cmd/info_test.go +++ b/src/invocation-plane-services/ratelimiter/cmd/info_test.go @@ -21,6 +21,7 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "sync/atomic" "testing" golibversion "github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/version" @@ -41,7 +42,7 @@ func TestNewHealthServeMux_Info(t *testing.T) { golibversion.GitHash = "" }) - mux := newHealthServeMux(nil) + mux := newHealthServeMux(nil, &atomic.Bool{}) require.NotNil(t, mux) w := httptest.NewRecorder() @@ -59,7 +60,7 @@ func TestNewHealthServeMux_Info(t *testing.T) { } func TestNewHealthServeMux_Info_RejectsNonGET(t *testing.T) { - mux := newHealthServeMux(nil) + mux := newHealthServeMux(nil, &atomic.Bool{}) require.NotNil(t, mux) for _, method := range []string{ diff --git a/src/invocation-plane-services/ratelimiter/cmd/main.go b/src/invocation-plane-services/ratelimiter/cmd/main.go index 7c30b8465..adf2e4c1f 100644 --- a/src/invocation-plane-services/ratelimiter/cmd/main.go +++ b/src/invocation-plane-services/ratelimiter/cmd/main.go @@ -27,7 +27,9 @@ import ( "os" "os/signal" "reflect" + "sync/atomic" "syscall" + "time" "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging" olricConfig "github.com/olric-data/olric/config" @@ -35,6 +37,7 @@ import ( "github.com/spf13/cobra" "github.com/spf13/pflag" "go.uber.org/zap" + "google.golang.org/grpc" "github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/nvkit/config" "github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/nvkit/logs" @@ -85,6 +88,17 @@ func setupPprof() { }() } +const ( + // Long enough for the endpoints controller to see the failing probe. + drainReadinessDelay = 5 * time.Second + // GracefulStop waits on in-flight RPCs and can outlast the grace period on + // its own, which would skip the drain entirely. + gracefulStopTimeout = 5 * time.Second + // All three must fit the default 30s termination grace period, since not + // every deployment can raise it. + drainTimeout = 15 * time.Second +) + func setupOlricStats(rateLimiter *ratelimiter.RateLimiter) { // Get the DMap from the store for on-demand stats store, ok := rateLimiter.GetStore().(*ratelimiter.Store) @@ -105,6 +119,49 @@ func setupOlricStats(rateLimiter *ratelimiter.RateLimiter) { }() } +// drainAndStop hands this member's counters to the current partition owners +// before Olric closes. Graceful termination only; an abrupt kill relies on the +// Olric replica count instead. +func drainAndStop(draining *atomic.Bool, server *grpc.Server, rateLimiter *ratelimiter.RateLimiter) { + draining.Store(true) + time.Sleep(drainReadinessDelay) + + stopped := make(chan struct{}) + go func() { + server.GracefulStop() + close(stopped) + }() + select { + case <-stopped: + case <-time.After(gracefulStopTimeout): + zap.L().Warn("Graceful stop timed out, forcing shutdown") + server.Stop() + } + + store, ok := rateLimiter.GetStore().(*ratelimiter.Store) + if !ok { + zap.L().Warn("Failed to cast store to *ratelimiter.Store for drain") + return + } + + // Not the command context: the shutdown signal may have cancelled it. + ctx, cancel := context.WithTimeout(context.Background(), drainTimeout) + defer cancel() + + start := time.Now() + drained, failed, err := store.Drain(ctx) + if err != nil { + zap.L().Error("Failed to drain counters", zap.Int("drained", drained), zap.Int("failed", failed), zap.Error(err)) + return + } + if failed > 0 { + zap.L().Error("Drained counters with failures", zap.Int("drained", drained), zap.Int("failed", failed), + zap.Duration("took", time.Since(start))) + return + } + zap.L().Info("Drained counters", zap.Int("drained", drained), zap.Duration("took", time.Since(start))) +} + func InterceptorLogger(l *zap.Logger) logging.Logger { return logging.LoggerFunc(func(ctx context.Context, lvl logging.Level, msg string, fields ...any) { f := make([]zap.Field, 0, len(fields)/2) @@ -142,9 +199,14 @@ func InterceptorLogger(l *zap.Logger) logging.Logger { // newHealthServeMux builds the management HTTP mux serving /health and the // GET /info build-version endpoint. -func newHealthServeMux(rateLimiter *ratelimiter.RateLimiter) *http.ServeMux { +func newHealthServeMux(rateLimiter *ratelimiter.RateLimiter, draining *atomic.Bool) *http.ServeMux { mux := http.NewServeMux() mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) { + // Unready first so the Service drops this pod before counters move. + if draining.Load() { + w.WriteHeader(http.StatusServiceUnavailable) + return + } if err := rateLimiter.Health(); err != nil { zap.L().Error("rate limiter error", zap.Error(err)) w.WriteHeader(http.StatusInternalServerError) @@ -180,6 +242,9 @@ func NewRootCommand() *cobra.Command { if c.SecretsPath == "" { c.SecretsPath = "vault/secrets.json" } + if c.OlricReplicaCount <= 0 { + c.OlricReplicaCount = ratelimiter.DefaultOlricReplicaCount + } if c.OAuth2Issuer == "" { return fmt.Errorf("missing required OAUTH2_ISSUER (expected JWT iss claim on inbound gRPC calls)") } @@ -222,6 +287,10 @@ func NewRootCommand() *cobra.Command { } cfg.DMaps.EvictionPolicy = olricConfig.LRUEviction cfg.DMaps.MaxInuse = 100_000_000 // 100 MB + cfg.ReplicaCount = rateLimiterConfig.OlricReplicaCount + // Quorum of one so a degraded cluster keeps serving checks. + cfg.ReadQuorum = 1 + cfg.WriteQuorum = 1 rateLimiter, err := ratelimiter.NewRateLimiter(*rateLimiterConfig, cfg) if err != nil { return err @@ -231,7 +300,8 @@ func NewRootCommand() *cobra.Command { // Setup Olric stats endpoint for debugging setupOlricStats(rateLimiter) // make a http health endpoint since astro doesn't support gRPC health endpoint - healthServer := newHealthServeMux(rateLimiter) + var draining atomic.Bool + healthServer := newHealthServeMux(rateLimiter, &draining) healthErrChan := make(chan error, 1) go func() { if err := http.ListenAndServe(":8080", healthServer); err != nil { @@ -259,6 +329,7 @@ func NewRootCommand() *cobra.Command { select { case sig := <-sigChan: zap.L().Info("Received signal, shutting down...", zap.String("signal", sig.String())) + drainAndStop(&draining, baseServer, rateLimiter) return nil case err := <-grpcErrChan: zap.L().Error("grpc server error, shutting down...", zap.Error(err)) diff --git a/src/invocation-plane-services/ratelimiter/olric_store.go b/src/invocation-plane-services/ratelimiter/olric_store.go index 446ec215f..8fe5f6234 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store.go +++ b/src/invocation-plane-services/ratelimiter/olric_store.go @@ -20,6 +20,8 @@ package ratelimiter import ( "context" "errors" + "fmt" + "strings" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" @@ -153,6 +155,35 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim return common.GetContextFromState(now, rate, expiration, int64(value)), nil } +// Drain re-writes every entry so it replicates to the current partition owners. +// Olric only seeds a member on write, so without this a pod that joined while a +// counter sat idle holds nothing. +func (store *Store) Drain(ctx context.Context) (drained int, failed int, err error) { + iter, err := store.dmap.Scan(ctx) + if err != nil { + return 0, 0, fmt.Errorf("failed to scan DMap for drain: %w", err) + } + defer iter.Close() + + for iter.Next() { + if ctx.Err() != nil { + return drained, failed, ctx.Err() + } + key := iter.Key() + // Zero delta: atomic per key and keeps the TTL. + if _, err := store.dmap.Incr(ctx, key, 0); err != nil { + if errors.Is(err, olric.ErrKeyNotFound) { + continue + } + failed++ + zap.L().Warn("Failed to drain key", zap.String("key", redactKey(key)), zap.Error(err)) + continue + } + drained++ + } + return drained, failed, nil +} + // Peek returns the limit for the given identifier, without modification on current values. // NOT USED func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) { @@ -173,6 +204,20 @@ func (store *Store) Reset(ctx context.Context, key string, rate limiter.Rate) (l return common.GetContextFromState(now, rate, expiration, 0), nil } +// redactKey hashes the subject in a per-user counter key. Keys are built as +// ":user:::..." for the per-user tier, so the +// segment after the "user" sentinel is the caller identity. +func redactKey(key string) string { + parts := strings.Split(key, ":") + for i, part := range parts { + if part == "user" && i+1 < len(parts) { + parts[i+1] = redactSubject(parts[i+1]) + return strings.Join(parts, ":") + } + } + return key +} + // getCacheKey returns the full path for an identifier. func (store *Store) getCacheKey(key string) string { return store.Prefix + ":" + key diff --git a/src/invocation-plane-services/ratelimiter/olric_store_test.go b/src/invocation-plane-services/ratelimiter/olric_store_test.go new file mode 100644 index 000000000..11d69af85 --- /dev/null +++ b/src/invocation-plane-services/ratelimiter/olric_store_test.go @@ -0,0 +1,115 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 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 ratelimiter + +import ( + "context" + "testing" + "time" + + "github.com/olric-data/olric" + olricConfig "github.com/olric-data/olric/config" + "github.com/stretchr/testify/require" +) + +func newTestStore(t *testing.T) *Store { + t.Helper() + + cfg := olricConfig.New("local") + cfg.DMaps.EvictionPolicy = olricConfig.LRUEviction + cfg.DMaps.MaxInuse = 100_000_000 + + started := make(chan struct{}) + cfg.Started = func() { close(started) } + + db, err := olric.New(cfg) + require.NoError(t, err) + + go func() { + if err := db.Start(); err != nil { + t.Errorf("olric failed to start: %v", err) + } + }() + + select { + case <-started: + case <-time.After(30 * time.Second): + t.Fatal("olric did not start in time") + } + + t.Cleanup(func() { + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _ = db.Shutdown(ctx) + }) + + dmap, err := db.NewEmbeddedClient().NewDMap("test-limiter") + require.NoError(t, err) + + return &Store{Prefix: "test-limiter", dmap: dmap} +} + +func TestDrainReplicatesEveryKeyWithoutChangingCounts(t *testing.T) { + ctx := context.Background() + store := newTestStore(t) + + require.NoError(t, store.dmap.Put(ctx, "a", 4, olric.EX(time.Hour))) + require.NoError(t, store.dmap.Put(ctx, "b", 9, olric.EX(time.Hour))) + + ttlBefore := map[string]int64{} + for _, key := range []string{"a", "b"} { + entry, err := store.dmap.Get(ctx, key) + require.NoError(t, err) + ttlBefore[key] = entry.TTL() + } + + drained, failed, err := store.Drain(ctx) + require.NoError(t, err) + require.Equal(t, 2, drained) + require.Zero(t, failed) + + for key, want := range map[string]int{"a": 4, "b": 9} { + got, err := store.dmap.Get(ctx, key) + require.NoError(t, err) + value, err := got.Int() + require.NoError(t, err) + require.Equal(t, want, value, "drain must not change the counter for %q", key) + // The window must end when it would have without the drain, so compare + // against the deadline recorded before draining rather than just + // asserting some TTL survived. + require.InDelta(t, ttlBefore[key], got.TTL(), float64(2*time.Second/time.Millisecond), + "drain must keep the window deadline for %q", key) + } +} + +func TestDrainOnEmptyStore(t *testing.T) { + store := newTestStore(t) + + drained, failed, err := store.Drain(context.Background()) + require.NoError(t, err) + require.Zero(t, drained) + require.Zero(t, failed) +} + +func TestRedactKeyHidesSubject(t *testing.T) { + // Non per-user keys carry no caller identity and pass through untouched. + require.Equal(t, "limiter:nca-1:version-1:10-H", redactKey("limiter:nca-1:version-1:10-H")) + + // Real per-user keys are prefixed by the store, so the sentinel is not first. + got := redactKey("limiter:user:caller@example.com:nca-1:version-1:10-H") + require.NotContains(t, got, "caller@example.com") + require.Equal(t, "limiter:user:"+redactSubject("caller@example.com")+":nca-1:version-1:10-H", got) +} diff --git a/src/invocation-plane-services/ratelimiter/rate_limiter.go b/src/invocation-plane-services/ratelimiter/rate_limiter.go index 304788bdb..7df68e785 100644 --- a/src/invocation-plane-services/ratelimiter/rate_limiter.go +++ b/src/invocation-plane-services/ratelimiter/rate_limiter.go @@ -70,6 +70,9 @@ import ( const ( defaultLimiterCacheCapacity = 300 defaultIndexedPolicyCacheCapacity = 100 + // Keep a backup of every counter so losing a member does not reset it. + // Olric warns and places fewer backups when the cluster is smaller. + DefaultOlricReplicaCount = 2 ) type FunctionRateLimitConfig struct { @@ -99,6 +102,9 @@ type Config struct { NvcfApiUrl string `mapstructure:"NVCF_API_URL"` CacheTTL int `mapstructure:"CACHE_TTL"` CollectMetrics bool `mapstructure:"COLLECT_METRICS"` + // Copies of each counter kept in the cluster. At 1 the owning member is a + // single point of failure for the counters it holds. + OlricReplicaCount int `mapstructure:"OLRIC_REPLICA_COUNT"` } type CustomClaims struct { @@ -419,8 +425,17 @@ func (r *RateLimiter) Health() error { } return errors.New("olric db shutdown") default: - return nil } + // Not ready until this member is in the cluster: a member on its own counts + // against an empty view, and cannot receive counters from a member leaving. + members, err := r.db.NewEmbeddedClient().Members(context.Background()) + if err != nil { + return fmt.Errorf("failed to list olric members: %w", err) + } + if len(members) == 0 { + return errors.New("olric member list is empty") + } + return nil } // GetStore returns the limiter store for metrics/debugging purposes