From 704ad38200a99d5418bbefb9d707be9fbc700f22 Mon Sep 17 00:00:00 2001 From: Max Xing Date: Tue, 18 Aug 2026 16:14:52 -0700 Subject: [PATCH] fix(ratelimiter): keep rate-limit counters across pod replacement Counters live only in Olric's memory, and Olric seeds a member on write: its balancer moves a partition when the holder stops owning it, and nothing back-fills a backup owner that joined while a counter sat idle. Replacing every member in turn therefore dropped the counters, and callers regained a full budget with nothing logged. A rolling upgrade did this to every counter at once, so enforcement was effectively off for the duration. Four parts, none sufficient alone: - Keep a backup of each counter. The Olric replica count now defaults to 2 in the service rather than being derived by the chart, so every deployment gets it without environment-specific config. Read and write quorums stay at one so a degraded cluster keeps serving checks. - Hand the counters over on SIGTERM. The pod reports unready, stops serving gRPC, then re-writes each entry so the put path replicates it to the current owners. Incr with a zero delta is used rather than Get plus Put because it is atomic per key and preserves the TTL. The hand-off is budgeted to fit the default 30s termination grace period. - Report unready until this member has joined the cluster. A member on its own counts against an empty view, and cannot receive counters from a member that is leaving. This also paces rolling updates: with maxUnavailable at zero, the next pod is not replaced until the new member is in the cluster. - Pin maxSurge to one and maxUnavailable to zero, so a live member always exists to receive the hand-off. The percentage defaults allow an unavailable pod at some replica counts. Verified on a local self-hosted cluster with a 10-per-hour limit, spending 4 then probing 10 more. Graceful kill of one or all pods, rollout restart, scale down and up, and an abrupt kill of one pod all hold at 10 admitted. Killing every pod at the same instant still admits 14: nothing is persisted, so there is no copy left to recover from. Closes #975 Signed-off-by: Max Xing --- .../nvcf-ratelimiter/templates/_helpers.tpl | 2 +- .../templates/deployment.yaml | 16 ++- .../ratelimiter/nvcf-ratelimiter/values.yaml | 8 ++ .../ratelimiter/cmd/info_test.go | 5 +- .../ratelimiter/cmd/main.go | 75 +++++++++++- .../ratelimiter/olric_store.go | 45 +++++++ .../ratelimiter/olric_store_test.go | 115 ++++++++++++++++++ .../ratelimiter/rate_limiter.go | 17 ++- 8 files changed, 276 insertions(+), 7 deletions(-) create mode 100644 src/invocation-plane-services/ratelimiter/olric_store_test.go 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