From e735f12fd44bfd83e346d6a44630f33c72206f29 Mon Sep 17 00:00:00 2001
From: Hamid Ghaf <83242695+hghaf099@users.noreply.github.com>
Date: Wed, 24 Jan 2024 12:29:33 -0800
Subject: [PATCH] manual licese reporting persist snapshots CE (#25021)

---
 vault/activity_log.go |  2 +-
 vault/census.go       | 17 +++++++++--------
 vault/core.go         | 20 +++++++++++---------
 vault/testing.go      |  1 -
 4 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/vault/activity_log.go b/vault/activity_log.go
index 1914e833b5..5f855eab13 100644
--- a/vault/activity_log.go
+++ b/vault/activity_log.go
@@ -1155,7 +1155,7 @@ func (c *Core) setupActivityLogLocked(ctx context.Context, wg *sync.WaitGroup) e
 			close(manager.retentionDone)
 		}(manager.retentionMonths)
 
-		manager.CensusReportDone = make(chan bool)
+		manager.CensusReportDone = make(chan bool, 1)
 		go c.activityLog.CensusReport(ctx, c.CensusAgent(), c.BillingStart())
 	}
 
diff --git a/vault/census.go b/vault/census.go
index c608ca47eb..03a9790147 100644
--- a/vault/census.go
+++ b/vault/census.go
@@ -10,11 +10,12 @@ import "time"
 // CensusAgent is a stub for OSS
 type CensusReporter interface{}
 
-// setupCensusAgent is a stub for OSS.
-func (c *Core) setupCensusAgent() error                { return nil }
-func (c *Core) BillingStart() time.Time                { return time.Time{} }
-func (c *Core) AutomatedLicenseReportingEnabled() bool { return false }
-func (c *Core) CensusAgent() CensusReporter            { return nil }
-func (c *Core) ReloadCensus() error                    { return nil }
-func (c *Core) teardownCensusAgent() error             { return nil }
-func (c *Core) ManualLicenseReportingEnabled() bool    { return false }
+func (c *Core) setupCensusManager() error                   { return nil }
+func (c *Core) BillingStart() time.Time                     { return time.Time{} }
+func (c *Core) AutomatedLicenseReportingEnabled() bool      { return false }
+func (c *Core) CensusAgent() CensusReporter                 { return nil }
+func (c *Core) ReloadCensus() error                         { return nil }
+func (c *Core) teardownCensusManager() error                { return nil }
+func (c *Core) StartManualCensusSnapshots()                 {}
+func (c *Core) ManualLicenseReportingEnabled() bool         { return false }
+func (c *Core) ManualCensusSnapshotInterval() time.Duration { return time.Duration(0) }
diff --git a/vault/core.go b/vault/core.go
index f6b5cd0add..518b168aa3 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -636,8 +636,6 @@ type Core struct {
 	// it is protected by activityLogLock
 	activityLogConfig ActivityLogCoreConfig
 
-	censusConfig atomic.Value
-
 	// activeTime is set on active nodes indicating the time at which this node
 	// became active.
 	activeTime time.Time
@@ -833,9 +831,6 @@ type CoreConfig struct {
 	LicensePath     string
 	LicensingConfig *LicensingConfig
 
-	// Configured Census Agent
-	CensusAgent CensusReporter
-
 	DisablePerformanceStandby bool
 	DisableIndexing           bool
 	DisableKeyEncodingChecks  bool
@@ -2422,8 +2417,10 @@ func (s standardUnsealStrategy) unseal(ctx context.Context, logger log.Logger, c
 	}
 
 	if !c.IsDRSecondary() {
-		if err := c.setupCensusAgent(); err != nil {
-			logger.Error("skipping reporting for nil agent", "error", err)
+		if !c.perfStandby {
+			if err := c.setupCensusManager(); err != nil {
+				logger.Error("skipping license reporting for nil agent", "error", err)
+			}
 		}
 
 		// not waiting on wg to avoid changing existing behavior
@@ -2431,6 +2428,11 @@ func (s standardUnsealStrategy) unseal(ctx context.Context, logger log.Logger, c
 		if err := c.setupActivityLog(ctx, &wg); err != nil {
 			return err
 		}
+
+		if !c.perfStandby {
+			c.StartManualCensusSnapshots()
+		}
+
 	} else {
 		broker, err := NewAuditBroker(logger)
 		if err != nil {
@@ -2792,8 +2794,8 @@ func (c *Core) preSeal() error {
 		result = multierror.Append(result, fmt.Errorf("error stopping expiration: %w", err))
 	}
 	c.stopActivityLog()
-	// Clean up the censusAgent on seal
-	if err := c.teardownCensusAgent(); err != nil {
+	// Clean up census on seal
+	if err := c.teardownCensusManager(); err != nil {
 		result = multierror.Append(result, fmt.Errorf("error tearing down reporting agent: %w", err))
 	}
 
diff --git a/vault/testing.go b/vault/testing.go
index cc3fb89d69..3e3dbe2088 100644
--- a/vault/testing.go
+++ b/vault/testing.go
@@ -232,7 +232,6 @@ func TestCoreWithSealAndUINoCleanup(t testing.T, opts *CoreConfig) *Core {
 	conf.PluginDirectory = opts.PluginDirectory
 	conf.DetectDeadlocks = opts.DetectDeadlocks
 	conf.Experiments = opts.Experiments
-	conf.CensusAgent = opts.CensusAgent
 	conf.AdministrativeNamespacePath = opts.AdministrativeNamespacePath
 	conf.ImpreciseLeaseRoleTracking = opts.ImpreciseLeaseRoleTracking