From 9e39a5f2a445c4a7414c19fa2b3cf79d5bd67e0c Mon Sep 17 00:00:00 2001
From: miagilepner <mia.epner@hashicorp.com>
Date: Wed, 1 May 2024 18:16:23 +0200
Subject: [PATCH] Revert "Client count generation simplification (#26692)"
 (#26736)

This reverts commit b84af55a20444fc7a73914db10dc6c3a60c3a331.
---
 .../logical_system_activity_write_testonly.go | 65 +++++++++++--------
 ...cal_system_activity_write_testonly_test.go |  1 -
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/vault/logical_system_activity_write_testonly.go b/vault/logical_system_activity_write_testonly.go
index 6cea57996f17..62f7366887fe 100644
--- a/vault/logical_system_activity_write_testonly.go
+++ b/vault/logical_system_activity_write_testonly.go
@@ -333,27 +333,35 @@ func (m *multipleMonthsActivityClients) addRepeatedClients(monthsAgo int32, c *g
 	return nil
 }
 
-func (m *multipleMonthsActivityClients) timestampForMonth(i int, now time.Time) time.Time {
-	if i > 0 {
-		return timeutil.StartOfMonth(timeutil.MonthsPreviousTo(i, now))
-	}
-	return now
-}
-
 func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[generation.WriteOptions]struct{}, activityLog *ActivityLog) ([]string, error) {
 	now := time.Now().UTC()
 	paths := []string{}
 
 	_, writePQ := opts[generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES]
 	_, writeDistinctClients := opts[generation.WriteOptions_WRITE_DISTINCT_CLIENTS]
+	_, writeEntities := opts[generation.WriteOptions_WRITE_ENTITIES]
 	_, writeIntentLog := opts[generation.WriteOptions_WRITE_INTENT_LOGS]
 
+	pqOpts := pqOptions{}
+	if writePQ || writeDistinctClients {
+		pqOpts.byNamespace = make(map[string]*processByNamespace)
+		pqOpts.byMonth = make(map[int64]*processMonth)
+		pqOpts.activePeriodEnd = m.latestTimestamp(now, true)
+		pqOpts.endTime = timeutil.EndOfMonth(m.latestTimestamp(pqOpts.activePeriodEnd, false))
+		pqOpts.activePeriodStart = m.earliestTimestamp(now)
+	}
+
 	var earliestTimestamp, latestTimestamp time.Time
 	for i, month := range m.months {
 		if month.generationParameters == nil {
 			continue
 		}
-		timestamp := m.timestampForMonth(i, now)
+		var timestamp time.Time
+		if i > 0 {
+			timestamp = timeutil.StartOfMonth(timeutil.MonthsPreviousTo(i, now))
+		} else {
+			timestamp = now
+		}
 		if earliestTimestamp.IsZero() || timestamp.Before(earliestTimestamp) {
 			earliestTimestamp = timestamp
 		}
@@ -365,31 +373,32 @@ func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[gene
 			return nil, err
 		}
 		for segmentIndex, segment := range segments {
-			if segment == nil {
-				// skip the index
-				continue
+			if writeEntities || writeIntentLog {
+				if segment == nil {
+					// skip the index
+					continue
+				}
+				entityPath, err := activityLog.saveSegmentEntitiesInternal(ctx, segmentInfo{
+					startTimestamp:       timestamp.Unix(),
+					currentClients:       &activity.EntityActivityLog{Clients: segment},
+					clientSequenceNumber: uint64(segmentIndex),
+					tokenCount:           &activity.TokenCount{},
+				}, true)
+				if err != nil {
+					return nil, err
+				}
+				paths = append(paths, entityPath)
 			}
-			entityPath, err := activityLog.saveSegmentEntitiesInternal(ctx, segmentInfo{
-				startTimestamp:       timestamp.Unix(),
-				currentClients:       &activity.EntityActivityLog{Clients: segment},
-				clientSequenceNumber: uint64(segmentIndex),
-				tokenCount:           &activity.TokenCount{},
-			}, true)
+		}
+
+		if (writePQ || writeDistinctClients) && i > 0 {
+			reader := newProtoSegmentReader(segments)
+			err = activityLog.segmentToPrecomputedQuery(ctx, timestamp, reader, pqOpts)
 			if err != nil {
 				return nil, err
 			}
-			paths = append(paths, entityPath)
-		}
-	}
-	if writePQ || writeDistinctClients {
-		// start with the oldest month of data, and create precomputed queries
-		// up to that month
-		for i := len(m.months) - 1; i > 0; i-- {
-			activityLog.precomputedQueryWorker(ctx, &ActivityIntentLog{
-				PreviousMonth: m.timestampForMonth(i, now).Unix(),
-				NextMonth:     0,
-			})
 		}
+
 	}
 	if writeIntentLog {
 		err := activityLog.writeIntentLog(ctx, earliestTimestamp.UTC().Unix(), latestTimestamp.UTC())
diff --git a/vault/logical_system_activity_write_testonly_test.go b/vault/logical_system_activity_write_testonly_test.go
index e20bc0abfba7..f3f25773540f 100644
--- a/vault/logical_system_activity_write_testonly_test.go
+++ b/vault/logical_system_activity_write_testonly_test.go
@@ -76,7 +76,6 @@ func TestSystemBackend_handleActivityWriteData(t *testing.T) {
 			name:      "correctly formatted data succeeds",
 			operation: logical.UpdateOperation,
 			input:     map[string]interface{}{"input": `{"write":["WRITE_PRECOMPUTED_QUERIES"],"data":[{"current_month":true,"all":{"clients":[{"count":5}]}}]}`},
-			wantPaths: 1,
 		},
 		{
 			name:      "entities with multiple segments",