Vault 2823 cc namespace

changelog/12393.txt

@@ -0,0 +1,3 @@
+```release-note: improvement
+core: observe the client counts broken down by namespace for partial month client count
+```

vault/activity_log.go

@@ -67,6 +67,11 @@ type segmentInfo struct {
 	entitySequenceNumber uint64
 }
 
+type clients struct {
+	distinctEntities uint64
+	nonEntityTokens  uint64
+}
+
 // ActivityLog tracks unique entity counts and non-entity token counts.
 // It handles assembling log fragments (and sending them to the active
 // node), writing log segments, and precomputing queries.
@@ -143,6 +148,10 @@ type ActivityLog struct {
 
 	// for testing: is config currently being invalidated. protected by l
 	configInvalidationInProgress bool
+
+	// All known active entity count by namespace ID this month; use fragmentLock read-locked
+	// to check whether it already exists.
+	entityCountByNamespaceID map[string]uint64
 }
 
 // These non-persistent configuration options allow us to disable
@@ -164,17 +173,18 @@ func NewActivityLog(core *Core, logger log.Logger, view *BarrierView, metrics me
 	}
 
 	a := &ActivityLog{
-		core:            core,
-		configOverrides: &core.activityLogConfig,
-		logger:          logger,
-		view:            view,
-		metrics:         metrics,
-		nodeID:          hostname,
-		newFragmentCh:   make(chan struct{}, 1),
-		sendCh:          make(chan struct{}, 1), // buffered so it can be triggered by fragment size
-		writeCh:         make(chan struct{}, 1), // same for full segment
-		doneCh:          make(chan struct{}, 1),
-		activeEntities:  make(map[string]struct{}),
+		core:                     core,
+		configOverrides:          &core.activityLogConfig,
+		logger:                   logger,
+		view:                     view,
+		metrics:                  metrics,
+		nodeID:                   hostname,
+		newFragmentCh:            make(chan struct{}, 1),
+		sendCh:                   make(chan struct{}, 1), // buffered so it can be triggered by fragment size
+		writeCh:                  make(chan struct{}, 1), // same for full segment
+		doneCh:                   make(chan struct{}, 1),
+		activeEntities:           make(map[string]struct{}),
+		entityCountByNamespaceID: make(map[string]uint64),
 		currentSegment: segmentInfo{
 			startTimestamp: 0,
 			currentEntities: &activity.EntityActivityLog{
@@ -531,7 +541,7 @@ func (a *ActivityLog) loadPriorEntitySegment(ctx context.Context, startTime time
 	// Or the feature has been disabled.
 	if a.enabled && startTime.Unix() == a.currentSegment.startTimestamp {
 		for _, ent := range out.Entities {
-			a.activeEntities[ent.EntityID] = struct{}{}
+			a.AddEntity(ent)
 		}
 	}
 	a.fragmentLock.Unlock()
@@ -571,7 +581,7 @@ func (a *ActivityLog) loadCurrentEntitySegment(ctx context.Context, startTime ti
 	}
 
 	for _, ent := range out.Entities {
-		a.activeEntities[ent.EntityID] = struct{}{}
+		a.AddEntity(ent)
 	}
 
 	return nil
@@ -684,6 +694,7 @@ func (a *ActivityLog) resetCurrentLog() {
 
 	a.fragment = nil
 	a.activeEntities = make(map[string]struct{})
+	a.entityCountByNamespaceID = make(map[string]uint64)
 	a.standbyFragmentsReceived = make([]*activity.LogFragment, 0)
 }
 
@@ -1095,6 +1106,7 @@ func (a *ActivityLog) perfStandbyFragmentWorker() {
 			// clear active entity set
 			a.fragmentLock.Lock()
 			a.activeEntities = make(map[string]struct{})
+			a.entityCountByNamespaceID = make(map[string]uint64)
 			a.fragmentLock.Unlock()
 
 			// Set timer for next month.
@@ -1309,6 +1321,7 @@ func (a *ActivityLog) AddEntityToFragment(entityID string, namespaceID string, t
 			NamespaceID: namespaceID,
 			Timestamp:   timestamp,
 		})
+	a.entityCountByNamespaceID[namespaceID] += 1
 	a.activeEntities[entityID] = struct{}{}
 }
 
@@ -1353,7 +1366,7 @@ func (a *ActivityLog) receivedFragment(fragment *activity.LogFragment) {
 	}
 
 	for _, e := range fragment.Entities {
-		a.activeEntities[e.EntityID] = struct{}{}
+		a.AddEntity(e)
 	}
 
 	a.standbyFragmentsReceived = append(a.standbyFragmentsReceived, fragment)
@@ -1792,26 +1805,91 @@ func (c *Core) activeEntityGaugeCollector(ctx context.Context) ([]metricsutil.Ga
 
 // partialMonthClientCount returns the number of clients used so far this month.
 // If activity log is not enabled, the response will be nil
-func (a *ActivityLog) partialMonthClientCount(ctx context.Context) map[string]interface{} {
+func (a *ActivityLog) partialMonthClientCount(ctx context.Context) (map[string]interface{}, error) {
 	a.fragmentLock.RLock()
 	defer a.fragmentLock.RUnlock()
 
 	if !a.enabled {
 		// nothing to count
-		return nil
+		return nil, nil
 	}
+	byNamespace := make([]*ClientCountInNamespace, 0)
+	responseData := make(map[string]interface{})
+	totalEntities := 0
+	totalTokens := 0
+
+	clientCountTable := createClientCountTable(a.entityCountByNamespaceID, a.currentSegment.tokenCount.CountByNamespaceID)
 
-	entityCount := len(a.activeEntities)
-	var tokenCount int
-	for _, countByNS := range a.currentSegment.tokenCount.CountByNamespaceID {
-		tokenCount += int(countByNS)
+	queryNS, err := namespace.FromContext(ctx)
+	if err != nil {
+		return responseData, err
 	}
-	clientCount := entityCount + tokenCount
 
-	responseData := make(map[string]interface{})
-	responseData["distinct_entities"] = entityCount
-	responseData["non_entity_tokens"] = tokenCount
-	responseData["clients"] = clientCount
+	for nsID, clients := range clientCountTable {
+
+		ns, err := NamespaceByID(ctx, nsID, a.core)
+		if err != nil {
+			return responseData, err
+		}
+
+		if a.includeInResponse(queryNS, ns) {
+			var displayPath string
+			if ns == nil {
+				displayPath = fmt.Sprintf("deleted namespace %q", nsID)
+			} else {
+				displayPath = ns.Path
+			}
+
+			byNamespace = append(byNamespace, &ClientCountInNamespace{
+				NamespaceID:   nsID,
+				NamespacePath: displayPath,
+				Counts: ClientCountResponse{
+					DistinctEntities: int(clients.distinctEntities),
+					NonEntityTokens:  int(clients.nonEntityTokens),
+					Clients:          int(clients.distinctEntities + clients.nonEntityTokens),
+				},
+			})
+			totalEntities += int(clients.distinctEntities)
+			totalTokens += int(clients.nonEntityTokens)
+
+		}
+	}
+
+	responseData["by_namespace"] = byNamespace
+	responseData["distinct_entities"] = totalEntities
+	responseData["non_entity_tokens"] = totalTokens
+	responseData["clients"] = totalEntities + totalTokens
+
+	return responseData, nil
+}
+
+//createClientCountTable maps the entitycount and token count to the namespace id
+func createClientCountTable(entityMap map[string]uint64, tokenMap map[string]uint64) map[string]*clients {
+	//add distinct entity count
+	clientCountTable := make(map[string]*clients)
+	for nsID, count := range entityMap {
+		if _, ok := clientCountTable[nsID]; !ok {
+			clientCountTable[nsID] = &clients{distinctEntities: 0, nonEntityTokens: 0}
+		}
+		clientCountTable[nsID].distinctEntities += count
+
+	}
+	//add non-entity token count
+	for nsID, count := range tokenMap {
+		if _, ok := clientCountTable[nsID]; !ok {
+			clientCountTable[nsID] = &clients{distinctEntities: 0, nonEntityTokens: 0}
+		}
+		clientCountTable[nsID].nonEntityTokens += count
+
+	}
+	return clientCountTable
+
+}
+
+func (a *ActivityLog) AddEntity(e *activity.EntityRecord) {
+	if _, ok := a.activeEntities[e.EntityID]; !ok {
+		a.activeEntities[e.EntityID] = struct{}{}
+		a.entityCountByNamespaceID[e.NamespaceID] += 1
+	}
 
-	return responseData
 }

vault/activity_log_test.go

@@ -1390,6 +1390,80 @@ func setupActivityRecordsInStorage(t *testing.T, base time.Time, includeEntities
 	return a, entityRecords, tokenRecords
 }
 
+//setupActivityRecordsInStorageRootNS creates entity and non-entity tokens only with root namespace
+func setupActivityRecordsInStorageRootNS(t *testing.T, base time.Time, includeEntities, includeTokens bool) (*ActivityLog, []*activity.EntityRecord, map[string]uint64) {
+	t.Helper()
+
+	core, _, _ := TestCoreUnsealed(t)
+	a := core.activityLog
+	monthsAgo := base.AddDate(0, -3, 0)
+
+	var entityRecords []*activity.EntityRecord
+	if includeEntities {
+		entityRecords = []*activity.EntityRecord{
+			{
+				EntityID:    "11111111-1111-1111-1111-111111111111",
+				NamespaceID: "root",
+				Timestamp:   time.Now().Unix(),
+			},
+			{
+				EntityID:    "22222222-2222-2222-2222-222222222222",
+				NamespaceID: "root",
+				Timestamp:   time.Now().Unix(),
+			},
+			{
+				EntityID:    "33333333-2222-2222-2222-222222222222",
+				NamespaceID: "root",
+				Timestamp:   time.Now().Unix(),
+			},
+		}
+
+		testEntities1 := &activity.EntityActivityLog{
+			Entities: entityRecords[:1],
+		}
+		entityData1, err := proto.Marshal(testEntities1)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		testEntities2 := &activity.EntityActivityLog{
+			Entities: entityRecords[1:2],
+		}
+		entityData2, err := proto.Marshal(testEntities2)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+		testEntities3 := &activity.EntityActivityLog{
+			Entities: entityRecords[2:],
+		}
+		entityData3, err := proto.Marshal(testEntities3)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+
+		WriteToStorage(t, core, ActivityLogPrefix+"entity/"+fmt.Sprint(monthsAgo.Unix())+"/0", entityData1)
+		WriteToStorage(t, core, ActivityLogPrefix+"entity/"+fmt.Sprint(base.Unix())+"/0", entityData2)
+		WriteToStorage(t, core, ActivityLogPrefix+"entity/"+fmt.Sprint(base.Unix())+"/1", entityData3)
+	}
+
+	var tokenRecords map[string]uint64
+	if includeTokens {
+		tokenRecords = make(map[string]uint64)
+		tokenRecords["root"] = 4
+		tokenCount := &activity.TokenCount{
+			CountByNamespaceID: tokenRecords,
+		}
+
+		tokenData, err := proto.Marshal(tokenCount)
+		if err != nil {
+			t.Fatalf(err.Error())
+		}
+
+		WriteToStorage(t, core, ActivityLogPrefix+"directtokens/"+fmt.Sprint(base.Unix())+"/0", tokenData)
+	}
+
+	return a, entityRecords, tokenRecords
+}
+
 func TestActivityLog_refreshFromStoredLog(t *testing.T) {
 	a, expectedEntityRecords, expectedTokenCounts := setupActivityRecordsInStorage(t, time.Now().UTC(), true, true)
 	a.SetEnable(true)
@@ -2499,9 +2573,9 @@ func TestActivityLog_Deletion(t *testing.T) {
 func TestActivityLog_partialMonthClientCount(t *testing.T) {
 	timeutil.SkipAtEndOfMonth(t)
 
-	ctx := context.Background()
+	ctx := namespace.RootContext(nil)
 	now := time.Now().UTC()
-	a, entities, tokenCounts := setupActivityRecordsInStorage(t, timeutil.StartOfMonth(now), true, true)
+	a, entities, tokenCounts := setupActivityRecordsInStorageRootNS(t, timeutil.StartOfMonth(now), true, true)
 
 	a.SetEnable(true)
 	var wg sync.WaitGroup
@@ -2520,7 +2594,10 @@ func TestActivityLog_partialMonthClientCount(t *testing.T) {
 
 	expectedClientCount := partialMonthEntityCount + partialMonthTokenCount
 
-	results := a.partialMonthClientCount(ctx)
+	results, err := a.partialMonthClientCount(ctx)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
 	if results == nil {
 		t.Fatal("no results to test")
 	}

vault/activity_log_testing_util.go

@@ -36,6 +36,29 @@ func (c *Core) InjectActivityLogDataThisMonth(t *testing.T) (map[string]struct{}
 	return activeEntities, tokens
 }
 
+// InjectActivityLogDataThisMonth populates the in-memory client store
+// with some entities and tokens for root namespace , overriding what was already there
+// It is currently used for API integration tests
+func (c *Core) InjectActivityLogDataThisMonthRootNS(t *testing.T) (map[string]uint64, map[string]uint64) {
+	t.Helper()
+
+	entitiesByNS := map[string]uint64{
+		"root": 1,
+	}
+	tokens := map[string]uint64{
+		"root": 5,
+	}
+
+	c.activityLog.l.Lock()
+	defer c.activityLog.l.Unlock()
+	c.activityLog.fragmentLock.Lock()
+	defer c.activityLog.fragmentLock.Unlock()
+
+	c.activityLog.currentSegment.tokenCount.CountByNamespaceID = tokens
+	c.activityLog.entityCountByNamespaceID = entitiesByNS
+	return entitiesByNS, tokens
+}
+
 // Return the in-memory activeEntities from an activity log
 func (c *Core) GetActiveEntities() map[string]struct{} {
 	out := make(map[string]struct{})

vault/external_tests/activity/activity_test.go

@@ -87,7 +87,7 @@ func TestActivityLog_MonthlyActivityApi(t *testing.T) {
 	validateClientCounts(t, resp, 0, 0)
 
 	// inject some data and query the API
-	entities, tokens := core.InjectActivityLogDataThisMonth(t)
+	entities, tokens := core.InjectActivityLogDataThisMonthRootNS(t)
 	expectedEntities := len(entities)
 	var expectedTokens int
 	for _, tokenCount := range tokens {

vault/logical_system_activity.go

@@ -142,7 +142,12 @@ func (b *SystemBackend) handleMonthlyActivityCount(ctx context.Context, req *log
 		return logical.ErrorResponse("no activity log present"), nil
 	}
 
-	results := a.partialMonthClientCount(ctx)
+	results, err := a.partialMonthClientCount(ctx)
+	if err != nil {
+		return &logical.Response{
+			Data: results,
+		}, err
+	}
 	if results == nil {
 		return logical.RespondWithStatusCode(nil, req, http.StatusNoContent)
 	}