temp

Author: divyaac

vault/activity_log.go

@@ -1830,54 +1830,29 @@ func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.T
 		pq = storedQuery
 	}
 
-	// Calculate the namespace response breakdowns and totals for entities and tokens from the initial
-	// namespace data.
-	totalCounts, byNamespaceResponse, err := a.calculateByNamespaceResponseForQuery(ctx, pq.Namespaces)
-	if err != nil {
-		return nil, err
-	}
-
-	// If we need to add the current month's client counts into the total, compute the namespace
-	// breakdown for the current month as well.
 	var partialByMonth map[int64]*processMonth
-	var partialByNamespace map[string]*processByNamespace
-	var byNamespaceResponseCurrent []*ResponseNamespace
-	var totalCurrentCounts *ResponseCounts
 	if computePartial {
 		// Traverse through current month's activitylog data and group clients
 		// into months and namespaces
 		a.fragmentLock.RLock()
-		partialByMonth, partialByNamespace = a.populateNamespaceAndMonthlyBreakdowns()
+		partialByMonth, _ = a.populateNamespaceAndMonthlyBreakdowns()
 		a.fragmentLock.RUnlock()
 
-		// Convert the byNamespace breakdowns into structs that are
-		// consumable by the /activity endpoint, so as to reuse code between these two
-		// endpoints.
-		byNamespaceComputation := a.transformALNamespaceBreakdowns(partialByNamespace)
-
-		// Calculate the namespace response breakdowns and totals for entities
-		// and tokens from current month namespace data.
-		totalCurrentCounts, byNamespaceResponseCurrent, err = a.calculateByNamespaceResponseForQuery(ctx, byNamespaceComputation)
+		// Estimate the current month totals. These record contains is complete with all the
+		// current month data, grouped by namespace and mounts
+		currentMonth, err := a.computeCurrentMonthForBillingPeriod(ctx, partialByMonth, startTime, endTime)
 		if err != nil {
 			return nil, err
 		}
 
-		// Create a mapping of namespace id to slice index, so that we can efficiently update our results without
-		// having to traverse the entire namespace response slice every time.
-		nsrMap := make(map[string]int)
-		for i, nr := range byNamespaceResponse {
-			nsrMap[nr.NamespaceID] = i
-		}
+		// Combine the existing months precomputed query with the current month data
+		pq.CombineWithCurrentMonth(currentMonth)
+	}
 
-		// Rather than blindly appending, which will create duplicates, check our existing counts against the current
-		// month counts, and append or update as necessary. We also want to account for mounts and their counts.
-		for _, nrc := range byNamespaceResponseCurrent {
-			if ndx, ok := nsrMap[nrc.NamespaceID]; ok {
-				byNamespaceResponse[ndx].Add(nrc)
-			} else {
-				byNamespaceResponse = append(byNamespaceResponse, nrc)
-			}
-		}
+	// Convert the namespace data into a protobuf format that can be returned in the response
+	totalCounts, byNamespaceResponse, err := a.calculateByNamespaceResponseForQuery(ctx, pq.Namespaces)
+	if err != nil {
+		return nil, err
 	}
 
 	// Sort clients within each namespace
@@ -1887,34 +1862,6 @@ func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.T
 		totalCounts, byNamespaceResponse = a.limitNamespacesInALResponse(byNamespaceResponse, limitNamespaces)
 	}
 
-	distinctEntitiesResponse := totalCounts.EntityClients
-	if computePartial {
-		currentMonth, err := a.computeCurrentMonthForBillingPeriod(ctx, partialByMonth, startTime, endTime)
-		if err != nil {
-			return nil, err
-		}
-
-		// Add the namespace attribution for the current month to the newly computed current month value. Note
-		// that transformMonthBreakdowns calculates a superstruct of the required namespace struct due to its
-		// primary use-case being for precomputedQueryWorker, but we will reuse this code for brevity and extract
-		// the namespaces from it.
-		currentMonthNamespaceAttribution := a.transformMonthBreakdowns(partialByMonth)
-
-		// Ensure that there is only one element in this list -- if not, warn.
-		if len(currentMonthNamespaceAttribution) > 1 {
-			a.logger.Warn("more than one month worth of namespace and mount attribution calculated for "+
-				"current month values", "number of months", len(currentMonthNamespaceAttribution))
-		}
-		if len(currentMonthNamespaceAttribution) == 0 {
-			a.logger.Warn("no month data found, returning query with no namespace attribution for current month")
-		} else {
-			currentMonth.Namespaces = currentMonthNamespaceAttribution[0].Namespaces
-			currentMonth.NewClients.Namespaces = currentMonthNamespaceAttribution[0].NewClients.Namespaces
-		}
-		pq.Months = append(pq.Months, currentMonth)
-		distinctEntitiesResponse += pq.Months[len(pq.Months)-1].NewClients.Counts.EntityClients
-	}
-
 	// Now populate the response based on breakdowns.
 	responseData := make(map[string]interface{})
 	responseData["start_time"] = pq.StartTime.Format(time.RFC3339)
@@ -1931,8 +1878,6 @@ func (a *ActivityLog) handleQuery(ctx context.Context, startTime, endTime time.T
 	}
 
 	responseData["by_namespace"] = byNamespaceResponse
-	totalCounts.Add(totalCurrentCounts)
-	totalCounts.DistinctEntities = distinctEntitiesResponse
 	responseData["total"] = totalCounts
 
 	// Create and populate the month response structs based on the monthly breakdown.

vault/activity_log_test.go

@@ -9,6 +9,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -25,6 +26,7 @@ import (
 	"github.com/go-test/deep"
 	"github.com/golang/protobuf/proto"
 	"github.com/hashicorp/go-uuid"
+	"github.com/hashicorp/vault/builtin/credential/userpass"
 	"github.com/hashicorp/vault/helper/constants"
 	"github.com/hashicorp/vault/helper/namespace"
 	"github.com/hashicorp/vault/helper/timeutil"
@@ -5084,3 +5086,193 @@ func TestActivityLog_reportPrecomputedQueryMetrics(t *testing.T) {
 		hasMetric(t, data, "identity.pki_acme.active.reporting_period", 3, nil)
 	})
 }
+
+// TestHandleQuery_MultipleMounts creates a cluster with
+// two userpass mounts. It then tests verifies that
+// the total new counts are calculated within a reasonably level of accuracy for
+// various numbers of clients in each mount.
+func TestHandleQuery_MultipleMounts(t *testing.T) {
+	tests := map[string]struct {
+		twoMonthsAgo          [][]int
+		oneMonthAgo           [][]int
+		currentMonth          [][]int
+		expectedNewClients    int
+		repeatPreviousMonth   bool
+		expectedTotalAccuracy float64
+	}{
+		"low volume, all mounts": {
+			twoMonthsAgo: [][]int{
+				{20, 20},
+			},
+			oneMonthAgo: [][]int{
+				{30, 30},
+			},
+			currentMonth: [][]int{
+				{40, 40},
+			},
+			repeatPreviousMonth:   true,
+			expectedNewClients:    80,
+			expectedTotalAccuracy: 1,
+		},
+		"medium volume, all mounts": {
+			twoMonthsAgo: [][]int{
+				{200, 200},
+			},
+			oneMonthAgo: [][]int{
+				{300, 300},
+			},
+			currentMonth: [][]int{
+				{400, 400},
+			},
+			repeatPreviousMonth:   true,
+			expectedNewClients:    800,
+			expectedTotalAccuracy: 0.98,
+		},
+		"higher volume, all mounts": {
+			twoMonthsAgo: [][]int{
+				{200, 200},
+			},
+			oneMonthAgo: [][]int{
+				{300, 300},
+			},
+			currentMonth: [][]int{
+				{2000, 5000},
+			},
+			repeatPreviousMonth:   true,
+			expectedNewClients:    7000,
+			expectedTotalAccuracy: 0.95,
+		},
+		"higher volume, no repeats": {
+			twoMonthsAgo: [][]int{
+				{200, 200},
+			},
+			oneMonthAgo: [][]int{
+				{300, 300},
+			},
+			currentMonth: [][]int{
+				{4000, 6000},
+			},
+			repeatPreviousMonth:   false,
+			expectedNewClients:    10000,
+			expectedTotalAccuracy: 0.98,
+		},
+	}
+
+	for i, tt := range tests {
+		testname := fmt.Sprintf("%s", i)
+		t.Run(testname, func(t *testing.T) {
+			// Normalize to start of month to prevent end of month weirdness
+			startOfMonth := timeutil.StartOfMonth(time.Now().UTC())
+
+			storage := &logical.InmemStorage{}
+			coreConfig := &CoreConfig{
+				CredentialBackends: map[string]logical.Factory{
+					"userpass": userpass.Factory,
+				},
+				Physical: storage.Underlying(),
+			}
+
+			cluster := NewTestCluster(t, coreConfig, nil)
+			cluster.Start()
+			defer cluster.Cleanup()
+			core := cluster.Cores[0].Core
+			TestWaitActive(t, core)
+
+			a := core.activityLog
+			ctx := namespace.RootContext(nil)
+			var err error
+
+			namespaces := make([]*namespace.Namespace, 0, 6)
+			mounts := make(map[string][]*MountEntry)
+			ns := namespace.RootNamespace
+			namespaces = append(namespaces, ns)
+
+			// Add two userpass mounts to the root namespace
+			for i := 0; i < 1; i++ {
+				me1 := &MountEntry{
+					Table: credentialTableType,
+					Path:  "up1/",
+					Type:  "userpass",
+				}
+				err = core.enableCredential(namespace.ContextWithNamespace(ctx, namespaces[i]), me1)
+				require.NoError(t, err)
+
+				me2 := &MountEntry{
+					Table: credentialTableType,
+					Path:  "up2/",
+					Type:  "userpass",
+				}
+				err = core.enableCredential(namespace.ContextWithNamespace(ctx, namespaces[i]), me2)
+				require.NoError(t, err)
+				mounts[namespaces[i].ID] = []*MountEntry{me1, me2}
+			}
+
+			// Generate data for two months ago
+			clientPrefix := 1
+			var entityRecordsMonth2 []*activity.EntityRecord
+			for namespaceId, mountSlice := range mounts {
+				for mountIndex, mount := range mountSlice {
+					entityRecordsMonth2 = append(entityRecordsMonth2, generateClientData(0, tt.twoMonthsAgo[0][mountIndex], mount.Accessor, namespaceId, fmt.Sprintf("%d", clientPrefix))...)
+				}
+			}
+
+			// Generate data for a month ago
+			clientPrefix += 1
+			var entityRecordsMonth1 []*activity.EntityRecord
+			for namespaceId, mountSlice := range mounts {
+				for mountIndex, mount := range mountSlice {
+					entityRecordsMonth2 = append(entityRecordsMonth1, generateClientData(0, tt.oneMonthAgo[0][mountIndex], mount.Accessor, namespaceId, fmt.Sprintf("%d", clientPrefix))...)
+				}
+			}
+
+			startOfTwoMonthsAgo := timeutil.StartOfMonth(startOfMonth.AddDate(0, -2, 0)).UTC()
+			startOfOneMonthAgo := timeutil.StartOfMonth(startOfMonth.AddDate(0, -1, 0)).UTC()
+			generatePreviousMonthClientData(t, core, startOfTwoMonthsAgo, entityRecordsMonth2)
+			generatePreviousMonthClientData(t, core, startOfOneMonthAgo, entityRecordsMonth1)
+
+			// Generate the current months data
+			clientPrefix += 1
+			var currentMonthData []*activity.EntityRecord
+			for namespaceId, mountSlice := range mounts {
+				for mountIndex, mount := range mountSlice {
+					currentMonthData = append(currentMonthData, generateClientData(0, tt.currentMonth[0][mountIndex], mount.Accessor, namespaceId, fmt.Sprintf("%d", clientPrefix))...)
+					clientPrefix += 1
+					mountIndex++
+				}
+			}
+			generateCurrentMonthClientData(t, core, entityRecordsMonth2, currentMonthData)
+
+			endOfCurrentMonth := timeutil.EndOfMonth(time.Now().UTC())
+			actual, err := a.handleQuery(ctx, startOfTwoMonthsAgo, endOfCurrentMonth, 0)
+
+			// Ensure that the month response is the same as the totals, because all clients
+			// are new clients and there will be no approximation in the single month partial
+			// case
+			monthsRaw, ok := actual["months"]
+			if !ok {
+				t.Fatalf("malformed results. got %v", actual)
+			}
+			monthsResponse := make([]ResponseMonth, 0)
+			err = mapstructure.Decode(monthsRaw, &monthsResponse)
+
+			currentMonthClients := monthsResponse[len(monthsResponse)-1]
+
+			// New verify that the new client totals for ALL namespaces are approximately accurate
+			newClientsError := math.Abs((float64)(currentMonthClients.NewClients.Counts.Clients - tt.expectedNewClients))
+			newClientsErrorMargin := newClientsError / (float64)(tt.expectedNewClients)
+			expectedAccuracyCalc := (1 - tt.expectedTotalAccuracy) * 100 / 100
+			if newClientsErrorMargin > expectedAccuracyCalc {
+				t.Fatalf("bad accuracy: expected %+v, found %+v", expectedAccuracyCalc, newClientsErrorMargin)
+			}
+
+			// Verify that the totals for the clients are visibly sensible (that is the total of all the individual new clients per namespace)
+			total := 0
+			for _, newClientCounts := range currentMonthClients.NewClients.Namespaces {
+				total += newClientCounts.Counts.Clients
+			}
+			if diff := math.Abs(float64(currentMonthClients.NewClients.Counts.Clients - total)); diff >= 1 {
+				t.Fatalf("total expected was %d but got %d", currentMonthClients.NewClients.Counts.Clients, total)
+			}
+		})
+	}
+}