From f235bc5320ff384dac35ce7f4dcd65e6ed2ea586 Mon Sep 17 00:00:00 2001 From: Vault Automation Date: Tue, 17 Feb 2026 15:48:10 -0500 Subject: [PATCH] VAULT-41471 Write Data To Census (#12349) (#12391) * Added census metric names and new file build doesn't complete * Added metrics. Working on tests * Tests passing * Added tests * Added a nil check * Try to fix race condition * Fix some nits * Some more nits * Added a go doc * Fix some data races * Remove billing storage lock all together * Revert "Remove billing storage lock all together" This reverts commit cdf4a518c343d8fd0fbb5de13fd860999ca2d131. * Fix failing test * Remove billing storage lock all together * Removed schema json file * Revert "Remove billing storage lock all together" This reverts commit 76970c7d4f7fe98acb993a737aa7410216caf2a8. Co-authored-by: divyaac --- vault/consumption_billing.go | 7 +- vault/consumption_billing_testing_util.go | 50 +++++++-- vault/consumption_billing_util.go | 120 ++++++++++++++++++---- vault/core.go | 2 +- 4 files changed, 150 insertions(+), 29 deletions(-) diff --git a/vault/consumption_billing.go b/vault/consumption_billing.go index 42b39efa8a..c61f8da147 100644 --- a/vault/consumption_billing.go +++ b/vault/consumption_billing.go @@ -13,7 +13,10 @@ import ( "github.com/hashicorp/vault/vault/billing" ) -var ErrCouldNotGetBillingSubView = fmt.Errorf("could not get billing sub view") +var ( + ErrCouldNotGetBillingSubView = fmt.Errorf("could not get billing sub view") + ErrConsumptionBillingNotInitialized = fmt.Errorf("consumption billing is not initialized") +) func (c *Core) setupConsumptionBilling(ctx context.Context) error { // We need replication (post unseal) to start before we run the consumption billing metrics worker @@ -142,6 +145,8 @@ func (c *Core) deletePreviousMonthBillingMetrics(ctx context.Context, currentMon func (c *Core) resetInMemoryBillingMetrics() error { // Reset Transit/Tranform DP counts c.logger.Info("resetting in memory billing metrics") + c.consumptionBillingLock.Lock() + defer c.consumptionBillingLock.Unlock() c.consumptionBilling.DataProtectionCallCounts.Transit.Store(0) c.consumptionBilling.DataProtectionCallCounts.Transform.Store(0) c.consumptionBilling.KmipSeenEnabledThisMonth.Store(false) diff --git a/vault/consumption_billing_testing_util.go b/vault/consumption_billing_testing_util.go index 6cd4e1152b..671c439b14 100644 --- a/vault/consumption_billing_testing_util.go +++ b/vault/consumption_billing_testing_util.go @@ -4,25 +4,63 @@ package vault func (c *Core) ResetInMemoryTransitDataProtectionCallCounts() { - c.consumptionBilling.DataProtectionCallCounts.Transit.Store(0) + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + cb.DataProtectionCallCounts.Transit.Store(0) + } } func (c *Core) GetInMemoryTransitDataProtectionCallCounts() uint64 { - return c.consumptionBilling.DataProtectionCallCounts.Transit.Load() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + return cb.DataProtectionCallCounts.Transit.Load() + } + return 0 } func (c *Core) ResetInMemoryTransformDataProtectionCallCounts() { - c.consumptionBilling.DataProtectionCallCounts.Transform.Store(0) + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + cb.DataProtectionCallCounts.Transform.Store(0) + } } func (c *Core) GetInMemoryTransformDataProtectionCallCounts() uint64 { - return c.consumptionBilling.DataProtectionCallCounts.Transform.Load() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + return cb.DataProtectionCallCounts.Transform.Load() + } + return 0 } func (c *Core) SetInMemoryTransitDataProtectionCallCounts(count uint64) { - c.consumptionBilling.DataProtectionCallCounts.Transit.Store(count) + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + cb.DataProtectionCallCounts.Transit.Store(count) + } } func (c *Core) SetInMemoryTransformDataProtectionCallCounts(count uint64) { - c.consumptionBilling.DataProtectionCallCounts.Transform.Store(count) + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb != nil { + cb.DataProtectionCallCounts.Transform.Store(count) + } } diff --git a/vault/consumption_billing_util.go b/vault/consumption_billing_util.go index b28867822f..aec6595bfb 100644 --- a/vault/consumption_billing_util.go +++ b/vault/consumption_billing_util.go @@ -49,8 +49,16 @@ func (c *Core) getStoredThirdPartyPluginCountsLocked(ctx context.Context, localP // Note that this count is per cluster. It does NOT de-duplicate across clusters. For that reason, // we will always store the count at the "local" prefix. func (c *Core) UpdateMaxThirdPartyPluginCounts(ctx context.Context, currentMonth time.Time) (int, error) { - c.consumptionBilling.BillingStorageLock.Lock() - defer c.consumptionBilling.BillingStorageLock.Unlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.Lock() + defer cb.BillingStorageLock.Unlock() previousThirdPartyPluginCounts, err := c.getStoredThirdPartyPluginCountsLocked(ctx, billing.LocalPrefix, currentMonth) if err != nil { @@ -69,8 +77,16 @@ func (c *Core) UpdateMaxThirdPartyPluginCounts(ctx context.Context, currentMonth } func (c *Core) GetStoredThirdPartyPluginCounts(ctx context.Context, month time.Time) (int, error) { - c.consumptionBilling.BillingStorageLock.RLock() - defer c.consumptionBilling.BillingStorageLock.RUnlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.RLock() + defer cb.BillingStorageLock.RUnlock() return c.getStoredThirdPartyPluginCountsLocked(ctx, billing.LocalPrefix, month) } @@ -144,15 +160,31 @@ func (c *Core) getStoredMaxKvCountsLocked(ctx context.Context, localPathPrefix s } func (c *Core) GetStoredHWMKvCounts(ctx context.Context, localPathPrefix string, month time.Time) (int, error) { - c.consumptionBilling.BillingStorageLock.RLock() - defer c.consumptionBilling.BillingStorageLock.RUnlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.RLock() + defer cb.BillingStorageLock.RUnlock() return c.getStoredMaxKvCountsLocked(ctx, localPathPrefix, month) } // UpdateMaxKvCounts updates the HWM kv counts for the given month, and returns the value that was stored. func (c *Core) UpdateMaxKvCounts(ctx context.Context, localPathPrefix string, currentMonth time.Time) (int, error) { - c.consumptionBilling.BillingStorageLock.Lock() - defer c.consumptionBilling.BillingStorageLock.Unlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.Lock() + defer cb.BillingStorageLock.Unlock() local := localPathPrefix == billing.LocalPrefix @@ -208,8 +240,16 @@ func (c *Core) storeMaxRoleCountsLocked(ctx context.Context, maxRoleCounts *Role } func (c *Core) UpdateMaxRoleCounts(ctx context.Context, localPathPrefix string, currentMonth time.Time) (*RoleCounts, error) { - c.consumptionBilling.BillingStorageLock.Lock() - defer c.consumptionBilling.BillingStorageLock.Unlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return nil, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.Lock() + defer cb.BillingStorageLock.Unlock() local := localPathPrefix == billing.LocalPrefix currentRoleCounts := c.getRoleCountsInternal(local, !local, true) @@ -251,8 +291,16 @@ func (c *Core) UpdateMaxRoleCounts(ctx context.Context, localPathPrefix string, } func (c *Core) GetStoredHWMRoleCounts(ctx context.Context, localPathPrefix string, month time.Time) (*RoleCounts, error) { - c.consumptionBilling.BillingStorageLock.RLock() - defer c.consumptionBilling.BillingStorageLock.RUnlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return nil, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.RLock() + defer cb.BillingStorageLock.RUnlock() return c.getStoredRoleCountsLocked(ctx, localPathPrefix, month) } @@ -333,22 +381,36 @@ func (c *Core) getStoredTransitCallCountsLocked(ctx context.Context, localPathPr } func (c *Core) GetStoredTransitCallCounts(ctx context.Context, month time.Time) (uint64, error) { - c.consumptionBilling.BillingStorageLock.RLock() - defer c.consumptionBilling.BillingStorageLock.RUnlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.RLock() + defer cb.BillingStorageLock.RUnlock() return c.getStoredTransitCallCountsLocked(ctx, billing.LocalPrefix, month) } func (c *Core) UpdateTransitCallCounts(ctx context.Context, currentMonth time.Time) (uint64, error) { - c.consumptionBilling.BillingStorageLock.Lock() - defer c.consumptionBilling.BillingStorageLock.Unlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + if cb == nil { + return 0, ErrConsumptionBillingNotInitialized + } + cb.BillingStorageLock.Lock() + defer cb.BillingStorageLock.Unlock() storedTransitCount, err := c.getStoredTransitCallCountsLocked(ctx, billing.LocalPrefix, currentMonth) if err != nil { return 0, err } // Sum the current count with the stored count - transitCount := c.consumptionBilling.DataProtectionCallCounts.Transit.Swap(0) + storedTransitCount + transitCount := cb.DataProtectionCallCounts.Transit.Swap(0) + storedTransitCount err = c.storeTransitCallCountsLocked(ctx, transitCount, billing.LocalPrefix, currentMonth) if err != nil { @@ -392,8 +454,16 @@ func (c *Core) getStoredKmipEnabledLocked(ctx context.Context, localPathPrefix s } func (c *Core) GetStoredKmipEnabled(ctx context.Context, currentMonth time.Time) (bool, error) { - c.consumptionBilling.BillingStorageLock.RLock() - defer c.consumptionBilling.BillingStorageLock.RUnlock() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return false, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.RLock() + defer cb.BillingStorageLock.RUnlock() return c.getStoredKmipEnabledLocked(ctx, billing.LocalPrefix, currentMonth) } @@ -404,8 +474,16 @@ func (c *Core) GetStoredKmipEnabled(ctx context.Context, currentMonth time.Time) // 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() + c.consumptionBillingLock.RLock() + cb := c.consumptionBilling + c.consumptionBillingLock.RUnlock() + + if cb == nil { + return false, ErrConsumptionBillingNotInitialized + } + + cb.BillingStorageLock.Lock() + defer cb.BillingStorageLock.Unlock() // Check if KMIP is currently enabled, including replicated mounts kmipEnabled, err := c.IsKMIPEnabled(ctx) diff --git a/vault/core.go b/vault/core.go index ced1a765e5..2551b0b1ef 100644 --- a/vault/core.go +++ b/vault/core.go @@ -449,7 +449,7 @@ type Core struct { // consumptionBilling is used to track use case consumption-based billing metrics consumptionBilling *billing.ConsumptionBilling - // consumptionBillingLock protects the consumptionBillingConfig + // consumptionBillingLock protects the consumptionBilling struct consumptionBillingLock sync.RWMutex // metricsCh is used to stop the metrics streaming