VAULT-41207: KMIP Metrics (#12116) (#12208)

* add a new method to identify whether KMIP is enabled

* add a new prefix for the new metric

* add new methods to store and update the metric

* update the kmip usage in billing

* move the method to ent file since kmip is ent only feature

* add unit tests at the core metrics level

* add new unit tests to test the billing methods for the new metric

* add persistence to test cases

* add external tests for primary and secondary

* account for DR secondaries, add clarifying comments, fix tests

* fmt

* move call of update into update local hwm metric method

* feedback: simplify update method by removing operation to get stored value

* feedback: optimize kmip usage detection by adding atomic tracker to detect usage once kmip mount is enabled

* fmt

* feedback: remove check on DR secondary inside update method but leave it at Get method for now

* feedback: change kmip prefix to a more flexible structure with sub item

* feedback: rename atomic tracker for kmip usage

* feedback: simplify the kmip identifier method

* revert back on kmip path prefix changes

* feedback: move the atomic bool into consumption billing struct

* feedback: remove DR check in Get method since dr needs to have billing data replicated

* add another external test to test local mount detection in perf secondary

* add a no-op oss stub for kmip enabled method

Co-authored-by: Amir Aslamov <amir.aslamov@hashicorp.com>

vault/billing/billing_counts.go

@@ -22,6 +22,7 @@ const (
 	TransitDataProtectionCallCountsPrefix = "transitDataProtectionCallCounts/"
 	LocalPrefix                           = "local/"
 	ThirdPartyPluginsPrefix               = "thirdPartyPluginCounts/"
+	KmipEnabledPrefix                     = "kmipEnabled/"
 
 	BillingWriteInterval = 10 * time.Minute
 )
@@ -35,6 +36,10 @@ type ConsumptionBilling struct {
 	BillingConfig            BillingConfig
 	DataProtectionCallCounts DataProtectionCallCounts
 	Logger                   log.Logger
+
+	// KmipSeenEnabledThisMonth tracks whether KMIP has been enabled during the current billing month.
+	// This is used to avoid scanning all mounts every 10 minutes for KMIP billing detection.
+	KmipSeenEnabledThisMonth atomic.Bool
 }
 
 type BillingConfig struct {

vault/consumption_billing.go

@@ -57,6 +57,8 @@ func (c *Core) consumptionBillingMetricsWorker(ctx context.Context) {
 			case <-endOfMonth.C:
 				// Reset the timer for the next month
 				endOfMonth.Reset(time.Until(timeutil.StartOfNextMonth(time.Now())))
+				// Reset KMIP enabled flag for the new billing month
+				c.consumptionBilling.KmipSeenEnabledThisMonth.Store(false)
 				// On month boundary, we need to flush the current in-memory counts to storage
 				if err := c.updateBillingMetrics(ctx); err != nil {
 					c.logger.Error("error updating billing metrics at month boundary", "error", err)
@@ -126,11 +128,18 @@ func (c *Core) UpdateLocalHWMMetrics(ctx context.Context, currentMonth time.Time
 	} else {
 		c.logger.Info("updated local max external plugin counts", "prefix", billing.LocalPrefix, "currentMonth", currentMonth)
 	}
+	if _, err := c.UpdateKmipEnabled(ctx, currentMonth); err != nil {
+		c.logger.Error("error updating local kmip enabled", "error", err)
+	} else {
+		c.logger.Info("updated local kmip enabled", "prefix", billing.LocalPrefix, "currentMonth", currentMonth)
+	}
 
 	return nil
 }
 
+// UpdateLocalAggregatedMetrics updates local metrics that are aggregated across all replicated clusters
 func (c *Core) UpdateLocalAggregatedMetrics(ctx context.Context, currentMonth time.Time) error {
+	// Update aggregrated count of data protection calls
 	if _, err := c.UpdateDataProtectionCallCounts(ctx, currentMonth); err != nil {
 		return fmt.Errorf("could not store transit data protection call counts: %w", err)
 	}

vault/consumption_billing_util.go

@@ -329,3 +329,59 @@ func (c *Core) UpdateDataProtectionCallCounts(ctx context.Context, currentMonth
 
 	return storedDataProtectionCallCounts, nil
 }
+
+func (c *Core) storeKmipEnabledLocked(ctx context.Context, localPathPrefix string, currentMonth time.Time, kmipEnabled bool) error {
+	billingPath := billing.GetMonthlyBillingPath(localPathPrefix, currentMonth, billing.KmipEnabledPrefix)
+	entry, err := logical.StorageEntryJSON(billingPath, kmipEnabled)
+	if err != nil {
+		return err
+	}
+	return c.GetBillingSubView().Put(ctx, entry)
+}
+
+func (c *Core) getStoredKmipEnabledLocked(ctx context.Context, localPathPrefix string, currentMonth time.Time) (bool, error) {
+	billingPath := billing.GetMonthlyBillingPath(localPathPrefix, currentMonth, billing.KmipEnabledPrefix)
+	entry, err := c.GetBillingSubView().Get(ctx, billingPath)
+	if err != nil {
+		return false, err
+	}
+	if entry == nil {
+		return false, nil
+	}
+	var kmipEnabled bool
+	if err := entry.DecodeJSON(&kmipEnabled); err != nil {
+		return false, err
+	}
+	return kmipEnabled, nil
+}
+
+func (c *Core) GetStoredKmipEnabled(ctx context.Context, currentMonth time.Time) (bool, error) {
+	c.consumptionBilling.BillingStorageLock.RLock()
+	defer c.consumptionBilling.BillingStorageLock.RUnlock()
+	return c.getStoredKmipEnabledLocked(ctx, billing.LocalPrefix, currentMonth)
+}
+
+// UpdateKmipEnabled updates the KMIP enabled status for the current month.
+// Note that each cluster is billed independently, so we only store the status at the local prefix.
+// Additionally, KMIP usage detection covers both local and replicated mounts, meaning if primary has KMIP,
+// secondary also detects it and gets charged. This is intentional, as the KMIP usage is per cluster.
+// We only store true when KMIP is enabled; we never store false. This means storing true multiple times
+// is idempotent and safe.
+func (c *Core) UpdateKmipEnabled(ctx context.Context, currentMonth time.Time) (bool, error) {
+	c.consumptionBilling.BillingStorageLock.Lock()
+	defer c.consumptionBilling.BillingStorageLock.Unlock()
+
+	// Check if KMIP is currently enabled, including replicated mounts
+	kmipEnabled, err := c.IsKMIPEnabled(ctx)
+	if err != nil {
+		return false, err
+	}
+
+	if kmipEnabled {
+		if err := c.storeKmipEnabledLocked(ctx, billing.LocalPrefix, currentMonth, true); err != nil {
+			return false, err
+		}
+	}
+
+	return kmipEnabled, nil
+}

vault/core_metrics_oss.go

@@ -0,0 +1,15 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+//go:build !enterprise
+
+package vault
+
+import (
+	"context"
+)
+
+// IsKMIPEnabled is a stub for OSS. KMIP is an enterprise feature.
+func (c *Core) IsKMIPEnabled(ctx context.Context) (bool, error) {
+	return false, nil
+}