Entity and alias counts (#9262)

* Added gauge collectors for entity counts.
* Entity and alias gauges.
* Locking around accessor to core.identityStore and core.mount.

vault/core_metrics.go

@@ -323,3 +323,73 @@ func (c *Core) kvSecretGaugeCollector(ctx context.Context) ([]metricsutil.GaugeL
 
 	return results, nil
 }
+
+func (c *Core) entityGaugeCollector(ctx context.Context) ([]metricsutil.GaugeLabelValues, error) {
+	// Protect against concurrent changes during seal
+	c.stateLock.RLock()
+	identityStore := c.identityStore
+	c.stateLock.RUnlock()
+	if identityStore == nil {
+		return []metricsutil.GaugeLabelValues{}, errors.New("nil identity store")
+	}
+
+	byNamespace, err := identityStore.countEntitiesByNamespace(ctx)
+	if err != nil {
+		return []metricsutil.GaugeLabelValues{}, err
+	}
+
+	// No check for expiration here; the bulk of the work should be in
+	// counting the entities.
+	allNamespaces := c.collectNamespaces()
+	values := make([]metricsutil.GaugeLabelValues, len(allNamespaces))
+	for i := range values {
+		values[i].Labels = []metrics.Label{
+			metricsutil.NamespaceLabel(allNamespaces[i]),
+		}
+		values[i].Value = float32(byNamespace[allNamespaces[i].ID])
+	}
+
+	return values, nil
+}
+
+func (c *Core) entityGaugeCollectorByMount(ctx context.Context) ([]metricsutil.GaugeLabelValues, error) {
+	c.stateLock.RLock()
+	identityStore := c.identityStore
+	c.stateLock.RUnlock()
+	if identityStore == nil {
+		return []metricsutil.GaugeLabelValues{}, errors.New("nil identity store")
+	}
+
+	byAccessor, err := identityStore.countEntitiesByMountAccessor(ctx)
+	if err != nil {
+		return []metricsutil.GaugeLabelValues{}, err
+	}
+
+	values := make([]metricsutil.GaugeLabelValues, 0)
+	for accessor, count := range byAccessor {
+		// Terminate if taking too long to do the translation
+		select {
+		case <-ctx.Done():
+			return values, errors.New("context cancelled")
+		default:
+			break
+		}
+
+		c.stateLock.RLock()
+		mountEntry := c.router.MatchingMountByAccessor(accessor)
+		c.stateLock.RUnlock()
+		if mountEntry == nil {
+			continue
+		}
+		values = append(values, metricsutil.GaugeLabelValues{
+			Labels: []metrics.Label{
+				metricsutil.NamespaceLabel(mountEntry.namespace),
+				{"auth_method", mountEntry.Type},
+				{"mount_point", "auth/" + mountEntry.Path},
+			},
+			Value: float32(count),
+		})
+	}
+
+	return values, nil
+}

vault/core_metrics_test.go

@@ -175,5 +175,93 @@ func TestCoreMetrics_KvSecretGaugeError(t *testing.T) {
 	if counter.Count != 1 {
 		t.Errorf("Counter number of samples %v is not 1.", counter.Count)
 	}
-
+}
+
+func metricLabelsMatch(t *testing.T, actual []metrics.Label, expected map[string]string) {
+	t.Helper()
+
+	if len(actual) != len(expected) {
+		t.Errorf("Expected %v labels, got %v: %v", len(expected), len(actual), actual)
+	}
+
+	for _, l := range actual {
+		if v, ok := expected[l.Name]; ok {
+			if v != l.Value {
+				t.Errorf("Mismatched value %v=%v, expected %v", l.Name, l.Value, v)
+			}
+		} else {
+			t.Errorf("Unexpected label %v", l.Name)
+		}
+	}
+}
+
+func TestCoreMetrics_EntityGauges(t *testing.T) {
+	ctx := namespace.RootContext(nil)
+	is, ghAccessor, core := testIdentityStoreWithGithubAuth(ctx, t)
+
+	// Create an entity
+	alias1 := &logical.Alias{
+		MountType:     "github",
+		MountAccessor: ghAccessor,
+		Name:          "githubuser",
+	}
+
+	entity, err := is.CreateOrFetchEntity(ctx, alias1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create a second alias for the same entity
+	registerReq := &logical.Request{
+		Operation: logical.UpdateOperation,
+		Path:      "entity-alias",
+		Data: map[string]interface{}{
+			"name":           "githubuser2",
+			"canonical_id":   entity.ID,
+			"mount_accessor": ghAccessor,
+		},
+	}
+
+	resp, err := is.HandleRequest(ctx, registerReq)
+	if err != nil || (resp != nil && resp.IsError()) {
+		t.Fatalf("err:%v resp:%#v", err, resp)
+	}
+
+	glv, err := core.entityGaugeCollector(ctx)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if len(glv) != 1 {
+		t.Fatalf("Wrong number of gauges %v, expected %v", len(glv), 1)
+	}
+
+	if glv[0].Value != 1.0 {
+		t.Errorf("Entity count %v, expected %v", glv[0].Value, 1.0)
+	}
+
+	metricLabelsMatch(t, glv[0].Labels,
+		map[string]string{
+			"namespace": "root",
+		})
+
+	glv, err = core.entityGaugeCollectorByMount(ctx)
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if len(glv) != 1 {
+		t.Fatalf("Wrong number of gauges %v, expected %v", len(glv), 1)
+	}
+
+	if glv[0].Value != 2.0 {
+		t.Errorf("Alias count %v, expected %v", glv[0].Value, 2.0)
+	}
+
+	metricLabelsMatch(t, glv[0].Labels,
+		map[string]string{
+			"namespace":   "root",
+			"auth_method": "github",
+			"mount_point": "auth/github/",
+		})
 }

vault/identity_store_util.go

@@ -2131,3 +2131,62 @@ func (i *IdentityStore) countEntities() (int, error) {
 
 	return count, nil
 }
+
+// Sum up the number of entities belonging to each namespace (keyed by ID)
+func (i *IdentityStore) countEntitiesByNamespace(ctx context.Context) (map[string]int, error) {
+	txn := i.db.Txn(false)
+	iter, err := txn.Get(entitiesTable, "id")
+	if err != nil {
+		return nil, err
+	}
+
+	byNamespace := make(map[string]int)
+	val := iter.Next()
+	for val != nil {
+		// Check if runtime exceeded.
+		select {
+		case <-ctx.Done():
+			return byNamespace, errors.New("context cancelled")
+		default:
+			break
+		}
+
+		// Count in the namespace attached to the entity.
+		entity := val.(*identity.Entity)
+		byNamespace[entity.NamespaceID] = byNamespace[entity.NamespaceID] + 1
+		val = iter.Next()
+	}
+
+	return byNamespace, nil
+}
+
+// Sum up the number of entities belonging to each mount point (keyed by accessor)
+func (i *IdentityStore) countEntitiesByMountAccessor(ctx context.Context) (map[string]int, error) {
+	txn := i.db.Txn(false)
+	iter, err := txn.Get(entitiesTable, "id")
+	if err != nil {
+		return nil, err
+	}
+
+	byMountAccessor := make(map[string]int)
+	val := iter.Next()
+	for val != nil {
+		// Check if runtime exceeded.
+		select {
+		case <-ctx.Done():
+			return byMountAccessor, errors.New("context cancelled")
+		default:
+			break
+		}
+
+		// Count each alias separately; will translate to mount point and type
+		// in the caller.
+		entity := val.(*identity.Entity)
+		for _, alias := range entity.Aliases {
+			byMountAccessor[alias.MountAccessor] = byMountAccessor[alias.MountAccessor] + 1
+		}
+		val = iter.Next()
+	}
+
+	return byMountAccessor, nil
+}