[VAULT-15398] Client count tests (#22635)

* fix bugs in client count data generation

* add new tests for client counts

* fix package name

sdk/helper/clientcountutil/clientcountutil.go

@@ -289,7 +289,7 @@ func (d *ActivityLogDataGenerator) Write(ctx context.Context, writeOptions ...ge
 	if err != nil {
 		return nil, err
 	}
-	resp, err := d.client.Logical().WriteBytesWithContext(ctx, "sys/internal/counters/activity/write", data)
+	resp, err := d.client.Logical().WriteWithContext(ctx, "sys/internal/counters/activity/write", map[string]interface{}{"input": string(data)})
 	if err != nil {
 		return nil, err
 	}

vault/external_tests/activity_testonly/activity_testonly_test.go

@@ -0,0 +1,239 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+//go:build testonly
+
+package activity_testonly
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/vault/api"
+	"github.com/hashicorp/vault/helper/testhelpers"
+	"github.com/hashicorp/vault/helper/testhelpers/minimal"
+	"github.com/hashicorp/vault/helper/timeutil"
+	vaulthttp "github.com/hashicorp/vault/http"
+	"github.com/hashicorp/vault/sdk/helper/clientcountutil"
+	"github.com/hashicorp/vault/sdk/helper/clientcountutil/generation"
+	"github.com/hashicorp/vault/vault"
+	"github.com/mitchellh/mapstructure"
+	"github.com/stretchr/testify/require"
+)
+
+// Test_ActivityLog_LoseLeadership writes data for this month, then causes the
+// active node to lose leadership. Once a new node becomes the leader, then the
+// test queries for the current month data and verifies that the data from
+// before the leadership transfer is returned
+func Test_ActivityLog_LoseLeadership(t *testing.T) {
+	t.Parallel()
+	cluster := vault.NewTestCluster(t, nil, &vault.TestClusterOptions{
+		HandlerFunc: vaulthttp.Handler,
+	})
+	cluster.Start()
+	defer cluster.Cleanup()
+
+	testhelpers.WaitForActiveNodeAndStandbys(t, cluster)
+	active := testhelpers.DeriveStableActiveCore(t, cluster)
+	client := active.Client
+	_, err := client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "enable",
+	})
+	require.NoError(t, err)
+
+	_, err = clientcountutil.NewActivityLogData(client).
+		NewCurrentMonthData().
+		NewClientsSeen(10).
+		Write(context.Background(), generation.WriteOptions_WRITE_ENTITIES)
+	require.NoError(t, err)
+	now := time.Now().UTC()
+
+	testhelpers.EnsureCoreSealed(t, active)
+	newActive := testhelpers.WaitForActiveNode(t, cluster)
+	standby := active
+	testhelpers.WaitForStandbyNode(t, standby)
+	testhelpers.EnsureCoreUnsealed(t, cluster, standby)
+
+	resp, err := newActive.Client.Logical().ReadWithData("sys/internal/counters/activity", map[string][]string{
+		"end_time":   {timeutil.EndOfMonth(now).Format(time.RFC3339)},
+		"start_time": {timeutil.StartOfMonth(now).Format(time.RFC3339)},
+	})
+	monthResponse := getMonthsData(t, resp)
+	require.Len(t, monthResponse, 1)
+	require.Equal(t, 10, monthResponse[0].NewClients.Counts.Clients)
+}
+
+// Test_ActivityLog_ClientsOverlapping writes data for the previous month and
+// current month. In the previous month, 7 new clients are seen. In the current
+// month, there are 5 repeated and 2 new clients. The test queries over the
+// previous and current months, and verifies that the repeated clients are not
+// considered new
+func Test_ActivityLog_ClientsOverlapping(t *testing.T) {
+	t.Parallel()
+	cluster := minimal.NewTestSoloCluster(t, nil)
+	client := cluster.Cores[0].Client
+	_, err := client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "enable",
+	})
+	require.NoError(t, err)
+	_, err = clientcountutil.NewActivityLogData(client).
+		NewPreviousMonthData(1).
+		NewClientsSeen(7).
+		NewCurrentMonthData().
+		RepeatedClientsSeen(5).
+		NewClientsSeen(2).
+		Write(context.Background(), generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES, generation.WriteOptions_WRITE_ENTITIES)
+	require.NoError(t, err)
+
+	now := time.Now().UTC()
+
+	// query from the beginning of the previous month to the end of this month
+	resp, err := client.Logical().ReadWithData("sys/internal/counters/activity", map[string][]string{
+		"end_time":   {timeutil.EndOfMonth(now).Format(time.RFC3339)},
+		"start_time": {timeutil.StartOfMonth(timeutil.MonthsPreviousTo(1, now)).Format(time.RFC3339)},
+	})
+	require.NoError(t, err)
+	monthsResponse := getMonthsData(t, resp)
+	require.Len(t, monthsResponse, 2)
+	for _, month := range monthsResponse {
+		ts, err := time.Parse(time.RFC3339, month.Timestamp)
+		require.NoError(t, err)
+		// This month should have a total of 7 clients
+		// 2 of those will be considered new
+		if ts.UTC().Equal(timeutil.StartOfMonth(now)) {
+			require.Equal(t, month.Counts.Clients, 7)
+			require.Equal(t, month.NewClients.Counts.Clients, 2)
+		} else {
+			// All clients will be considered new for the previous month
+			require.Equal(t, month.Counts.Clients, 7)
+			require.Equal(t, month.NewClients.Counts.Clients, 7)
+
+		}
+	}
+}
+
+// Test_ActivityLog_ClientsNewCurrentMonth writes data for the past month and
+// current month with 5 repeated clients and 2 new clients in the current month.
+// The test then queries the activity log for only the current month, and
+// verifies that all 7 clients seen this month are considered new.
+func Test_ActivityLog_ClientsNewCurrentMonth(t *testing.T) {
+	t.Parallel()
+	cluster := minimal.NewTestSoloCluster(t, nil)
+	client := cluster.Cores[0].Client
+	_, err := client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "enable",
+	})
+	require.NoError(t, err)
+	_, err = clientcountutil.NewActivityLogData(client).
+		NewPreviousMonthData(1).
+		NewClientsSeen(5).
+		NewCurrentMonthData().
+		RepeatedClientsSeen(5).
+		NewClientsSeen(2).
+		Write(context.Background(), generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES, generation.WriteOptions_WRITE_ENTITIES)
+	require.NoError(t, err)
+
+	now := time.Now().UTC()
+
+	// query from the beginning of this month to the end of this month
+	resp, err := client.Logical().ReadWithData("sys/internal/counters/activity", map[string][]string{
+		"end_time":   {timeutil.EndOfMonth(now).Format(time.RFC3339)},
+		"start_time": {timeutil.StartOfMonth(now).Format(time.RFC3339)},
+	})
+	require.NoError(t, err)
+	monthsResponse := getMonthsData(t, resp)
+	require.Len(t, monthsResponse, 1)
+	require.Equal(t, 7, monthsResponse[0].NewClients.Counts.Clients)
+}
+
+// Test_ActivityLog_Disable writes data for a past month and a current month and
+// then disables the activity log. The test then queries for a timeframe that
+// includes both the disabled and enabled dates. The test verifies that the past
+// month's data is returned, but there is no current month data.
+func Test_ActivityLog_Disable(t *testing.T) {
+	t.Parallel()
+	cluster := minimal.NewTestSoloCluster(t, nil)
+	client := cluster.Cores[0].Client
+	_, err := client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "enable",
+	})
+	require.NoError(t, err)
+	_, err = clientcountutil.NewActivityLogData(client).
+		NewPreviousMonthData(1).
+		NewClientsSeen(5).
+		NewCurrentMonthData().
+		NewClientsSeen(5).
+		Write(context.Background(), generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES, generation.WriteOptions_WRITE_ENTITIES)
+	require.NoError(t, err)
+
+	_, err = client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "disable",
+	})
+	now := time.Now().UTC()
+	// query from the beginning of the previous month to the end of this month
+	resp, err := client.Logical().ReadWithData("sys/internal/counters/activity", map[string][]string{
+		"end_time":   {timeutil.EndOfMonth(now).Format(time.RFC3339)},
+		"start_time": {timeutil.StartOfMonth(timeutil.MonthsPreviousTo(1, now)).Format(time.RFC3339)},
+	})
+	require.NoError(t, err)
+	monthsResponse := getMonthsData(t, resp)
+
+	// we only expect data for the previous month
+	require.Len(t, monthsResponse, 1)
+	lastMonthResp := monthsResponse[0]
+	ts, err := time.Parse(time.RFC3339, lastMonthResp.Timestamp)
+	require.NoError(t, err)
+	require.Equal(t, ts.UTC(), timeutil.StartOfPreviousMonth(now.UTC()))
+}
+
+// Test_ActivityLog_EmptyDataMonths writes data for only the current month,
+// then queries a timeframe of several months in the past to now. The test
+// verifies that empty months of data are returned for the past, and the current
+// month data is correct.
+func Test_ActivityLog_EmptyDataMonths(t *testing.T) {
+	t.Parallel()
+	cluster := minimal.NewTestSoloCluster(t, nil)
+	client := cluster.Cores[0].Client
+	_, err := client.Logical().Write("sys/internal/counters/config", map[string]interface{}{
+		"enabled": "enable",
+	})
+	require.NoError(t, err)
+	_, err = clientcountutil.NewActivityLogData(client).
+		NewCurrentMonthData().
+		NewClientsSeen(10).
+		Write(context.Background(), generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES, generation.WriteOptions_WRITE_ENTITIES)
+	require.NoError(t, err)
+
+	now := time.Now().UTC()
+	// query from the beginning of 3 months ago to the end of this month
+	resp, err := client.Logical().ReadWithData("sys/internal/counters/activity", map[string][]string{
+		"end_time":   {timeutil.EndOfMonth(now).Format(time.RFC3339)},
+		"start_time": {timeutil.StartOfMonth(timeutil.MonthsPreviousTo(3, now)).Format(time.RFC3339)},
+	})
+	require.NoError(t, err)
+	monthsResponse := getMonthsData(t, resp)
+
+	require.Len(t, monthsResponse, 4)
+	for _, month := range monthsResponse {
+		ts, err := time.Parse(time.RFC3339, month.Timestamp)
+		require.NoError(t, err)
+		// current month should have data
+		if ts.UTC().Equal(timeutil.StartOfMonth(now)) {
+			require.Equal(t, month.Counts.Clients, 10)
+		} else {
+			// other months should be empty
+			require.Nil(t, month.Counts)
+		}
+	}
+}
+
+func getMonthsData(t *testing.T, resp *api.Secret) []vault.ResponseMonth {
+	t.Helper()
+	monthsRaw, ok := resp.Data["months"]
+	require.True(t, ok)
+	monthsResponse := make([]vault.ResponseMonth, 0)
+	err := mapstructure.Decode(monthsRaw, &monthsResponse)
+	require.NoError(t, err)
+	return monthsResponse
+}

vault/logical_system_activity_write_testonly.go

@@ -325,7 +325,7 @@ func (m *multipleMonthsActivityClients) addRepeatedClients(monthsAgo int32, c *g
 }
 
 func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[generation.WriteOptions]struct{}, activityLog *ActivityLog) ([]string, error) {
-	now := timeutil.StartOfMonth(time.Now().UTC())
+	now := time.Now().UTC()
 	paths := []string{}
 
 	_, writePQ := opts[generation.WriteOptions_WRITE_PRECOMPUTED_QUERIES]
@@ -337,8 +337,8 @@ func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[gene
 	if writePQ || writeDistinctClients {
 		pqOpts.byNamespace = make(map[string]*processByNamespace)
 		pqOpts.byMonth = make(map[int64]*processMonth)
-		pqOpts.activePeriodEnd = m.latestTimestamp(now)
-		pqOpts.endTime = timeutil.EndOfMonth(pqOpts.activePeriodEnd)
+		pqOpts.activePeriodEnd = m.latestTimestamp(now, true)
+		pqOpts.endTime = timeutil.EndOfMonth(m.latestTimestamp(pqOpts.activePeriodEnd, false))
 		pqOpts.activePeriodStart = m.earliestTimestamp(now)
 	}
 
@@ -382,7 +382,7 @@ func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[gene
 			}
 		}
 
-		if writePQ || writeDistinctClients {
+		if (writePQ || writeDistinctClients) && i > 0 {
 			reader := newProtoSegmentReader(segments)
 			err = activityLog.segmentToPrecomputedQuery(ctx, timestamp, reader, pqOpts)
 			if err != nil {
@@ -402,12 +402,13 @@ func (m *multipleMonthsActivityClients) write(ctx context.Context, opts map[gene
 	if err != nil {
 		return nil, err
 	}
+	wg.Wait()
 	return paths, nil
 }
 
-func (m *multipleMonthsActivityClients) latestTimestamp(now time.Time) time.Time {
+func (m *multipleMonthsActivityClients) latestTimestamp(now time.Time, includeCurrentMonth bool) time.Time {
 	for i, month := range m.months {
-		if month.generationParameters != nil {
+		if month.generationParameters != nil && (i != 0 || includeCurrentMonth) {
 			return timeutil.StartOfMonth(timeutil.MonthsPreviousTo(i, now))
 		}
 	}

vault/logical_system_activity_write_testonly_test.go

@@ -502,7 +502,7 @@ func Test_handleActivityWriteData(t *testing.T) {
 		},
 	}
 
-	t.Run("write entitites", func(t *testing.T) {
+	t.Run("write entities", func(t *testing.T) {
 		core, _, _ := TestCoreUnsealed(t)
 		marshaled, err := protojson.Marshal(&generation.ActivityLogMockInput{
 			Data:  data,
@@ -601,15 +601,15 @@ func Test_handleActivityWriteData(t *testing.T) {
 
 		now := time.Now().UTC()
 		start := timeutil.StartOfMonth(timeutil.MonthsPreviousTo(3, now))
-		end := timeutil.EndOfMonth(now)
+		end := timeutil.EndOfMonth(timeutil.MonthsPreviousTo(1, now))
 		pq, err := core.activityLog.queryStore.Get(context.Background(), start, end)
 		require.NoError(t, err)
 		require.NotNil(t, pq)
 		require.Equal(t, end, pq.EndTime)
 		require.Equal(t, start, pq.StartTime)
 		require.Len(t, pq.Namespaces, 1)
-		require.Equal(t, uint64(12), pq.Namespaces[0].Entities)
-		require.Len(t, pq.Months, 4)
+		require.Equal(t, uint64(10), pq.Namespaces[0].Entities)
+		require.Len(t, pq.Months, 3)
 	})
 	t.Run("write intent logs", func(t *testing.T) {
 		core, _, _ := TestCoreUnsealed(t)