From a1e68134c37f0b6e65655366cd0a6a6b3050075b 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 ceadca10791f..7d42bdfe6dd7 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 c608ca47eb75..03a979014784 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 86b4e9059e76..68c24c3ed335 100644
--- a/vault/core.go
+++ b/vault/core.go
@@ -649,8 +649,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
@@ -831,9 +829,6 @@ type CoreConfig struct {
 	LicensePath     string
 	LicensingConfig *LicensingConfig
 
-	// Configured Census Agent
-	CensusAgent CensusReporter
-
 	DisablePerformanceStandby bool
 	DisableIndexing           bool
 	DisableKeyEncodingChecks  bool
@@ -2466,8 +2461,10 @@ func (s standardUnsealStrategy) unseal(ctx context.Context, logger log.Logger, c
 			return err
 		}
 
-		if err := c.setupCensusAgent(); err != nil {
-			c.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
@@ -2475,6 +2472,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 {
 		var err error
 		disableEventLogger, err := parseutil.ParseBool(os.Getenv(featureFlagDisableEventLogger))
@@ -2725,8 +2727,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 177b84274aee..3073e569642b 100644
--- a/vault/testing.go
+++ b/vault/testing.go
@@ -235,7 +235,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