coder · f0ssel · Dec 7, 2023 · Dec 5, 2023 · Dec 6, 2023 · Dec 6, 2023
diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go
@@ -2049,6 +2049,13 @@ func (q *querier) GetWorkspaceResourcesCreatedAfter(ctx context.Context, created
 	return q.db.GetWorkspaceResourcesCreatedAfter(ctx, createdAt)
 }
 
+func (q *querier) GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx context.Context, templateIds []uuid.UUID) ([]database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow, error) {
+	if err := q.authorizeContext(ctx, rbac.ActionRead, rbac.ResourceSystem); err != nil {
+		return nil, err
+	}
+	return q.db.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, templateIds)
+}
+
 func (q *querier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) {
 	prep, err := prepareSQLFilter(ctx, q.auth, rbac.ActionRead, rbac.ResourceWorkspace.Type)
 	if err != nil {

diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -4462,6 +4462,36 @@ func (q *FakeQuerier) GetWorkspaceResourcesCreatedAfter(_ context.Context, after
 	return resources, nil
 }
 
+func (q *FakeQuerier) GetWorkspaceUniqueOwnerCountByTemplateIDs(_ context.Context, templateIds []uuid.UUID) ([]database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	workspaceOwners := make(map[uuid.UUID]map[uuid.UUID]struct{})
+	for _, workspace := range q.workspaces {
+		if workspace.Deleted {
+			continue
+		}
+		if !slices.Contains(templateIds, workspace.TemplateID) {
+			continue
+		}
+		_, ok := workspaceOwners[workspace.TemplateID]
+		if !ok {
+			workspaceOwners[workspace.TemplateID] = make(map[uuid.UUID]struct{})
+		}
+		workspaceOwners[workspace.TemplateID][workspace.OwnerID] = struct{}{}
+	}
+	resp := make([]database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow, 0)
+	for _, templateID := range templateIds {
+		count := len(workspaceOwners[templateID])
+		resp = append(resp, database.GetWorkspaceUniqueOwnerCountByTemplateIDsRow{
+			TemplateID:      templateID,
+			UniqueOwnersSum: int64(count),
+		})
+	}
+
+	return resp, nil
+}
+
 func (q *FakeQuerier) GetWorkspaces(ctx context.Context, arg database.GetWorkspacesParams) ([]database.GetWorkspacesRow, error) {
 	if err := validateDatabaseType(arg); err != nil {
 		return nil, err

diff --git a/coderd/database/dbmetrics/dbmetrics.go b/coderd/database/dbmetrics/dbmetrics.go
diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go
diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaces.sql b/coderd/database/queries/workspaces.sql
@@ -287,6 +287,15 @@ WHERE
 	AND LOWER("name") = LOWER(@name)
 ORDER BY created_at DESC;
 
+-- name: GetWorkspaceUniqueOwnerCountByTemplateIDs :many
+SELECT
+	template_id, COUNT(DISTINCT owner_id) AS unique_owners_sum
+FROM
+	workspaces
+WHERE
+	template_id = ANY(@template_ids :: uuid[]) AND deleted = false
+GROUP BY template_id;
+
 -- name: InsertWorkspace :one
 INSERT INTO
 	workspaces (

diff --git a/coderd/metricscache/metricscache.go b/coderd/metricscache/metricscache.go
@@ -52,6 +52,7 @@ type Cache struct {
 	deploymentDAUResponses   atomic.Pointer[map[int]codersdk.DAUsResponse]
 	templateDAUResponses     atomic.Pointer[map[int]map[uuid.UUID]codersdk.DAUsResponse]
 	templateUniqueUsers      atomic.Pointer[map[uuid.UUID]int]
+	templateWorkspaceOwners  atomic.Pointer[map[uuid.UUID]int]
 	templateAverageBuildTime atomic.Pointer[map[uuid.UUID]database.GetTemplateAverageBuildTimeRow]
 	deploymentStatsResponse  atomic.Pointer[codersdk.DeploymentStats]
 
@@ -206,6 +207,7 @@ func (c *Cache) refreshTemplateDAUs(ctx context.Context) error {
 	var (
 		templateDAUs              = make(map[int]map[uuid.UUID]codersdk.DAUsResponse, len(templates))
 		templateUniqueUsers       = make(map[uuid.UUID]int)
+		templateWorkspaceOwners   = make(map[uuid.UUID]int)
 		templateAverageBuildTimes = make(map[uuid.UUID]database.GetTemplateAverageBuildTimeRow)
 	)
 
@@ -214,7 +216,9 @@ func (c *Cache) refreshTemplateDAUs(ctx context.Context) error {
 		return xerrors.Errorf("deployment daus: %w", err)
 	}
 
+	ids := make([]uuid.UUID, 0, len(templates))
 	for _, template := range templates {
+		ids = append(ids, template.ID)
 		for _, tzOffset := range templateTimezoneOffsets {
 			rows, err := c.database.GetTemplateDAUs(ctx, database.GetTemplateDAUsParams{
 				TemplateID: template.ID,
@@ -249,6 +253,17 @@ func (c *Cache) refreshTemplateDAUs(ctx context.Context) error {
 		}
 		templateAverageBuildTimes[template.ID] = templateAvgBuildTime
 	}
+
+	owners, err := c.database.GetWorkspaceUniqueOwnerCountByTemplateIDs(ctx, ids)
+	if err != nil {
+		return xerrors.Errorf("get workspace unique owner count by template ids: %w", err)
+	}
+
+	for _, owner := range owners {
+		templateWorkspaceOwners[owner.TemplateID] = int(owner.UniqueOwnersSum)
+	}
+
+	c.templateWorkspaceOwners.Store(&templateWorkspaceOwners)
 	c.templateDAUResponses.Store(&templateDAUs)
 	c.templateUniqueUsers.Store(&templateUniqueUsers)
 	c.templateAverageBuildTime.Store(&templateAverageBuildTimes)
@@ -469,6 +484,21 @@ func (c *Cache) TemplateBuildTimeStats(id uuid.UUID) codersdk.TemplateBuildTimeS
 	}
 }
 
+func (c *Cache) TemplateWorkspaceOwners(id uuid.UUID) (int, bool) {
+	m := c.templateWorkspaceOwners.Load()
+	if m == nil {
+		// Data loading.
+		return -1, false
+	}
+
+	resp, ok := (*m)[id]
+	if !ok {
+		// Probably no data.
+		return -1, false
+	}
+	return resp, true
+}
+
 func (c *Cache) DeploymentStats() (codersdk.DeploymentStats, bool) {
 	deploymentStats := c.deploymentStatsResponse.Load()
 	if deploymentStats == nil {

diff --git a/coderd/metricscache/metricscache_test.go b/coderd/metricscache/metricscache_test.go
@@ -254,6 +254,74 @@ func TestCache_TemplateUsers(t *testing.T) {
 	}
 }
 
+func TestCache_TemplateWorkspaceOwners(t *testing.T) {
+	t.Parallel()
+	var ()
+
+	var (
+		db    = dbmem.New()
+		cache = metricscache.New(db, slogtest.Make(t, nil), metricscache.Intervals{
+			TemplateDAUs: testutil.IntervalFast,
+		})
+	)
+
+	defer cache.Close()
+
+	user1 := dbgen.User(t, db, database.User{})
+	user2 := dbgen.User(t, db, database.User{})
+	template := dbgen.Template(t, db, database.Template{
+		Provisioner: database.ProvisionerTypeEcho,
+	})
+	require.Eventuallyf(t, func() bool {
+		count, ok := cache.TemplateWorkspaceOwners(template.ID)
+		return ok && count == 0
+	}, testutil.WaitShort, testutil.IntervalMedium,
+		"TemplateWorkspaceOwners never populated 0 owners",
+	)
+
+	dbgen.Workspace(t, db, database.Workspace{
+		TemplateID: template.ID,
+		OwnerID:    user1.ID,
+	})
+
+	require.Eventuallyf(t, func() bool {
+		count, _ := cache.TemplateWorkspaceOwners(template.ID)
+		return count == 1
+	}, testutil.WaitShort, testutil.IntervalMedium,
+		"TemplateWorkspaceOwners never populated 1 owner",
+	)
+
+	workspace2 := dbgen.Workspace(t, db, database.Workspace{
+		TemplateID: template.ID,
+		OwnerID:    user2.ID,
+	})
+
+	require.Eventuallyf(t, func() bool {
+		count, _ := cache.TemplateWorkspaceOwners(template.ID)
+		return count == 2
+	}, testutil.WaitShort, testutil.IntervalMedium,
+		"TemplateWorkspaceOwners never populated 2 owners",
+	)
+
+	// 3rd workspace should not be counted since we have the same owner as workspace2.
+	dbgen.Workspace(t, db, database.Workspace{
+		TemplateID: template.ID,
+		OwnerID:    user1.ID,
+	})
+
+	db.UpdateWorkspaceDeletedByID(context.Background(), database.UpdateWorkspaceDeletedByIDParams{
+		ID:      workspace2.ID,
+		Deleted: true,
+	})
+
+	require.Eventuallyf(t, func() bool {
+		count, _ := cache.TemplateWorkspaceOwners(template.ID)
+		return count == 1
+	}, testutil.WaitShort, testutil.IntervalMedium,
+		"TemplateWorkspaceOwners never populated 1 owner after delete",
+	)
+}
+
 func clockTime(t time.Time, hour, minute, sec int) time.Time {
 	return time.Date(t.Year(), t.Month(), t.Day(), hour, minute, sec, t.Nanosecond(), t.Location())
 }

diff --git a/coderd/templates.go b/coderd/templates.go
@@ -831,7 +831,12 @@ func (api *API) convertTemplate(
 	template database.Template,
 ) codersdk.Template {
 	templateAccessControl := (*(api.Options.AccessControlStore.Load())).GetTemplateAccessControl(template)
-	activeCount, _ := api.metricsCache.TemplateUniqueUsers(template.ID)
+
+	owners := 0
+	o, ok := api.metricsCache.TemplateWorkspaceOwners(template.ID)
+	if ok {
+		owners = o
+	}
 
 	buildTimeStats := api.metricsCache.TemplateBuildTimeStats(template.ID)
 
@@ -849,7 +854,7 @@ func (api *API) convertTemplate(
 		DisplayName:                    template.DisplayName,
 		Provisioner:                    codersdk.ProvisionerType(template.Provisioner),
 		ActiveVersionID:                template.ActiveVersionID,
-		ActiveUserCount:                activeCount,
+		ActiveUserCount:                owners,
 		BuildTimeStats:                 buildTimeStats,
 		Description:                    template.Description,
 		Icon:                           template.Icon,