fix(ratelimiter): keep rate-limit counters across pod replacement

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 }} 
+{{- end }}

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 }}

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.

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{

src/invocation-plane-services/ratelimiter/cmd/main.go

@@ -27,14 +27,17 @@ 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"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"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))

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
+// "<prefix>:user:<clientAuthSubject>:<ncaId>:..." 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