fix: add option to disable consumption billing metrics worker to prevent OOM on large KV stores (#32059)

sdk/helper/consts/consts.go

@@ -44,6 +44,13 @@ const (
 
 	VaultDisableUserLockout = "VAULT_DISABLE_USER_LOCKOUT"
 
+	// VaultDisableConsumptionBilling is the environment variable that disables
+	// the periodic consumption billing metrics worker. When set to a non-empty
+	// value, the worker that walks the entire KV store every 10 minutes is
+	// not registered during post-unseal, preventing the memory and performance
+	// impact of the periodic KV store enumeration on large installations.
+	VaultDisableConsumptionBilling = "VAULT_DISABLE_CONSUMPTION_BILLING"
+
 	PerformanceReplicationPathTarget = "performance"
 
 	DRReplicationPathTarget = "dr"

vault/billing/billing_counts.go

@@ -109,6 +109,12 @@ type ConsumptionBilling struct {
 }
 
 type BillingConfig struct {
+	// DisableConsumptionBilling disables the periodic consumption billing metrics
+	// worker that walks the entire KV store every 10 minutes. When set to true,
+	// the billing worker is not registered during post-unseal, preventing the
+	// memory and performance impact of the periodic KV store enumeration.
+	// This can be set via the VAULT_DISABLE_CONSUMPTION_BILLING environment variable.
+	DisableConsumptionBilling bool
 	// For testing purposes. The cadence at which billing metrics are updated
 	MetricsUpdateCadence time.Duration
 	// For testing purposes. The cadence at which plugin counts are sent from perf standby to active
@@ -117,94 +123,3 @@ type BillingConfig struct {
 	// If nil, the default functions from the time package are used
 	TestOverrideClock timeutil.Clock
 }
-
-func GetMonthlyBillingMetricPath(localPrefix string, now time.Time, billingMetric string) string {
-	// Normalize to avoid double slashes since our prefixes include trailing "/".
-	// Example: localPrefix="replicated/", billingMetric="maxKvCounts/" =>
-	// "replicated/2026/01/maxKvCounts/"
-	year := now.Year()
-	month := int(now.Month())
-	return fmt.Sprintf(BillingMonthStorageFormat, localPrefix, year, month, billingMetric)
-}
-
-func GetMonthlyBillingPath(localPrefix string, now time.Time) string {
-	return fmt.Sprintf(BillingMonthStorageFormat, localPrefix, now.Year(), int(now.Month()), "")
-}
-
-func GetAttributionMaxPath(localPathPrefix string, month time.Time, attributionMetricName string) string {
-	return GetMonthlyBillingMetricPath(localPathPrefix, month, AttributionMaxPrefix+attributionMetricName)
-}
-
-type DataProtectionCallCounts struct {
-	Transit   *atomic.Uint64 `json:"transit,omitempty"`
-	Transform *atomic.Uint64 `json:"transform,omitempty"`
-	GcpKms    *atomic.Uint64 `json:"gcpkms,omitempty"`
-}
-
-// IdentityTokenUnits tracks billing metrics for identity and authentication services
-type IdentityTokenUnits struct {
-	// OidcTokenDuration tracks the token duration units (seconds, not duration-adjusted) for billing purposes in memory.
-	// This value is normalized before flushing to storage and is reset to 0 after flush in UpdateOidcDurationAdjustedCount.
-	OidcTokenDuration *uberatomic.Float64 `json:"oidc,omitempty"`
-
-	// SpiffeJwt stores duration-adjusted JWT token units as float64
-	// We need to use the uberAtomic package to store atomic float64 values
-	SpiffeJwt *uberatomic.Float64 `json:"spiffe_jwt,omitempty"`
-}
-
-var _ logical.ConsumptionBillingManager = (*ConsumptionBilling)(nil)
-
-func (s *ConsumptionBilling) WriteBillingData(ctx context.Context, mountType string, data map[string]interface{}) error {
-	if s == nil {
-		return nil
-	}
-
-	switch mountType {
-	case "transit":
-		val, ok := data["count"].(uint64)
-		if !ok {
-			err := fmt.Errorf("invalid value type for transit")
-			return err
-		}
-
-		s.DataProtectionCallCounts.Transit.Add(val)
-	case "transform":
-		val, ok := data["count"].(uint64)
-		if !ok {
-			err := fmt.Errorf("invalid value type for transform")
-			return err
-		}
-
-		s.DataProtectionCallCounts.Transform.Add(val)
-	case "spiffe":
-		// SPIFFE JWT uses float64 for duration-adjusted units
-		val, ok := data["units"].(float64)
-		if !ok {
-			err := fmt.Errorf("invalid value type for spiffe")
-			return err
-		}
-
-		s.IdentityTokenUnits.SpiffeJwt.Add(val)
-	case "gcpkms":
-		val, ok := data["count"].(uint64)
-		if !ok {
-			err := fmt.Errorf("invalid value type for gcp kms")
-			return err
-		}
-
-		s.DataProtectionCallCounts.GcpKms.Add(val)
-	case "external-ca":
-		// External CA uses float64 for duration-adjusted units
-		val, ok := data["units"].(float64)
-		if !ok {
-			err := fmt.Errorf("invalid value type for external-ca")
-			return err
-		}
-
-		s.ExternalCaCertUnits.Add(val)
-	default:
-		err := fmt.Errorf("unknown metric type: %s", mountType)
-		return err
-	}
-	return nil
-}

vault/consumption_billing.go

@@ -48,6 +48,11 @@ func (c *Core) setupConsumptionBilling(ctx context.Context) error {
 	}
 	c.consumptionBillingLock.Unlock()
 
+	if c.billingConfig.DisableConsumptionBilling {
+		c.logger.Warn("consumption billing is disabled by configuration, billing metrics worker will not be started")
+		return nil
+	}
+
 	c.postUnsealFuncs = append(c.postUnsealFuncs, func() {
 		c.consumptionBillingMetricsWorker(ctx)
 		// Start the perf standby plugin counts worker if this is a perf standby