Revert "temp"

This reverts commit 4932979.

Author: divyaac

vault/activity_log.go

@@ -1830,29 +1830,54 @@ 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, _ = a.populateNamespaceAndMonthlyBreakdowns()
+		partialByMonth, partialByNamespace = a.populateNamespaceAndMonthlyBreakdowns()
 		a.fragmentLock.RUnlock()
 
-		// 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)
+		// 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)
 		if err != nil {
 			return nil, err
 		}
 
-		// Combine the existing months precomputed query with the current month data
-		pq.CombineWithCurrentMonth(currentMonth)
-	}
+		// 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
+		}
 
-	// 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
+		// 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)
+			}
+		}
 	}
 
 	// Sort clients within each namespace
@@ -1862,6 +1887,34 @@ 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)
@@ -1878,6 +1931,8 @@ 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,7 +9,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"math"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -26,7 +25,6 @@ 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"
@@ -5086,193 +5084,3 @@ 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)
-			}
-		})
-	}
-}