Stats: Stop counting the same user multiple times (#26777)

* Stats: Cache based stats implementation

* Stats: Correct logic and add larger scale test

* Stats: linter

* Stats: SQL implementation

* Stats: cleanup SQL

* Stats: Tab -> Spaces

* Update pkg/services/sqlstore/stats.go

Co-authored-by: Sofia Papagiannaki <papagian@users.noreply.github.com>

* Stats: Quote 'user' table with dialect.Quote

* Stats: Ensure test is run as integration test

* Stats: Use boolean value

...because if (v) { true } else { false } is unnecessary at best.

Co-authored-by: Sofia Papagiannaki <papagian@users.noreply.github.com>

pkg/models/stats.go

@@ -92,13 +92,15 @@ type GetSystemUserCountStatsQuery struct {
 	Result *SystemUserCountStats
 }
 
-type ActiveUserStats struct {
-	ActiveUsers   int64
-	ActiveAdmins  int64
-	ActiveEditors int64
-	ActiveViewers int64
+type UserStats struct {
+	Users   int64
+	Admins  int64
+	Editors int64
+	Viewers int64
 }
 
-type GetActiveUserStatsQuery struct {
-	Result *ActiveUserStats
+type GetUserStatsQuery struct {
+	MustUpdate bool
+	Active     bool
+	Result     UserStats
 }

pkg/services/sqlstore/stats.go

@@ -2,6 +2,7 @@ package sqlstore
 
 import (
 	"context"
+	"strconv"
 	"time"
 
 	"github.com/grafana/grafana/pkg/bus"
@@ -13,12 +14,12 @@ func init() {
 	bus.AddHandler("sql", GetDataSourceStats)
 	bus.AddHandler("sql", GetDataSourceAccessStats)
 	bus.AddHandler("sql", GetAdminStats)
-	bus.AddHandler("sql", GetActiveUserStats)
+	bus.AddHandler("sql", GetUserStats)
 	bus.AddHandlerCtx("sql", GetAlertNotifiersUsageStats)
 	bus.AddHandlerCtx("sql", GetSystemUserCountStats)
 }
 
-var activeUserTimeLimit = time.Hour * 24 * 30
+const activeUserTimeLimit = time.Hour * 24 * 30
 
 func GetAlertNotifiersUsageStats(ctx context.Context, query *models.GetAlertNotifierUsageStatsQuery) error {
 	var rawSql = `SELECT COUNT(*) AS count, type FROM ` + dialect.Quote("alert_notification") + ` GROUP BY type`
@@ -80,9 +81,7 @@ func GetSystemStats(query *models.GetSystemStatsQuery) error {
 	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("team") + `) AS teams,`)
 	sb.Write(`(SELECT COUNT(id) FROM ` + dialect.Quote("user_auth_token") + `) AS auth_tokens,`)
 
-	sb.Write(roleCounterSQL("Viewer", "viewers", false)+`,`, activeUserDeadlineDate)
-	sb.Write(roleCounterSQL("Editor", "editors", false)+`,`, activeUserDeadlineDate)
-	sb.Write(roleCounterSQL("Admin", "admins", false)+``, activeUserDeadlineDate)
+	sb.Write(roleCounterSQL())
 
 	var stats models.SystemStats
 	_, err := x.SQL(sb.GetSqlString(), sb.params...).Get(&stats)
@@ -95,22 +94,15 @@ func GetSystemStats(query *models.GetSystemStatsQuery) error {
 	return nil
 }
 
-func roleCounterSQL(role string, alias string, onlyActive bool) string {
-	var sqlQuery = `
-		(
-			SELECT COUNT(DISTINCT u.id)
-			FROM ` + dialect.Quote("user") + ` AS u, org_user
-			WHERE u.last_seen_at > ? AND ( org_user.user_id=u.id AND org_user.role='` + role + `' )
-		) AS active_` + alias
-
-	if !onlyActive {
-		sqlQuery += `,
-		(
-			SELECT COUNT(DISTINCT u.id)
-			FROM ` + dialect.Quote("user") + ` AS u, org_user
-			WHERE ( org_user.user_id=u.id AND org_user.role='` + role + `' )
-		) AS ` + alias
-	}
+func roleCounterSQL() string {
+	_ = updateUserRoleCountsIfNecessary(false)
+	sqlQuery :=
+		strconv.FormatInt(userStatsCache.total.Admins, 10) + ` AS admins, ` +
+			strconv.FormatInt(userStatsCache.total.Editors, 10) + ` AS editors, ` +
+			strconv.FormatInt(userStatsCache.total.Viewers, 10) + ` AS viewers, ` +
+			strconv.FormatInt(userStatsCache.active.Admins, 10) + ` AS active_admins, ` +
+			strconv.FormatInt(userStatsCache.active.Editors, 10) + ` AS active_editors, ` +
+			strconv.FormatInt(userStatsCache.active.Viewers, 10) + ` AS active_viewers`
 
 	return sqlQuery
 }
@@ -159,9 +151,7 @@ func GetAdminStats(query *models.GetAdminStatsQuery) error {
 			SELECT COUNT(*)
 			FROM ` + dialect.Quote("user") + ` WHERE last_seen_at > ?
 		) AS active_users,
-		` + roleCounterSQL("Admin", "admins", false) + `,
-		` + roleCounterSQL("Editor", "editors", false) + `,
-		` + roleCounterSQL("Viewer", "viewers", false) + `,
+		` + roleCounterSQL() + `,
 		(
 			SELECT COUNT(*)
 			FROM ` + dialect.Quote("user_auth_token") + ` WHERE rotated_at > ?
@@ -192,23 +182,105 @@ func GetSystemUserCountStats(ctx context.Context, query *models.GetSystemUserCou
 	})
 }
 
-func GetActiveUserStats(query *models.GetActiveUserStatsQuery) error {
-	activeUserDeadlineDate := time.Now().Add(-activeUserTimeLimit)
-	sb := &SqlBuilder{}
+func GetUserStats(query *models.GetUserStatsQuery) error {
+	err := updateUserRoleCountsIfNecessary(query.MustUpdate)
+	if err != nil {
+		return err
+	}
 
-	sb.Write(`SELECT `)
-	sb.Write(`(SELECT COUNT(*) FROM `+dialect.Quote("user")+` WHERE last_seen_at > ?) AS active_users,`, activeUserDeadlineDate)
-	sb.Write(roleCounterSQL("Viewer", "viewers", true)+`,`, activeUserDeadlineDate)
-	sb.Write(roleCounterSQL("Editor", "editors", true)+`,`, activeUserDeadlineDate)
-	sb.Write(roleCounterSQL("Admin", "admins", true)+``, activeUserDeadlineDate)
+	if query.Active {
+		query.Result = userStatsCache.active
+	} else {
+		query.Result = userStatsCache.total
+	}
 
-	var stats models.ActiveUserStats
-	_, err := x.SQL(sb.GetSqlString(), sb.params...).Get(&stats)
+	return nil
+}
+
+func updateUserRoleCountsIfNecessary(forced bool) error {
+	memoizationPeriod := time.Now().Add(-userStatsCacheLimetime)
+	if forced || userStatsCache.memoized.Before(memoizationPeriod) {
+		err := updateUserRoleCounts()
 		if err != nil {
 			return err
 		}
+	}
+	return nil
+}
 
-	query.Result = &stats
+type memoUserStats struct {
+	active models.UserStats
+	total  models.UserStats
 
+	memoized time.Time
+}
+
+var (
+	userStatsCache         = memoUserStats{}
+	userStatsCacheLimetime = 5 * time.Minute
+)
+
+func updateUserRoleCounts() error {
+	query := `
+SELECT role AS bitrole, active, COUNT(role) AS count FROM
+    (SELECT active, SUM(role) AS role
+     FROM (SELECT
+               u.id,
+               CASE org_user.role
+                   WHEN 'Admin' THEN 4
+                   WHEN 'Editor' THEN 2
+                   ELSE 1
+                   END AS role,
+               u.last_seen_at>? AS active
+           FROM ` + dialect.Quote("user") + ` AS u LEFT JOIN org_user ON org_user.user_id = u.id
+           GROUP BY u.id, u.last_seen_at, org_user.role) AS t2
+     GROUP BY active, id) AS t1
+GROUP BY active, role;`
+
+	activeUserDeadline := time.Now().Add(-activeUserTimeLimit)
+
+	type rolebitmap struct {
+		Active  bool
+		Bitrole int64
+		Count   int64
+	}
+
+	bitmap := []rolebitmap{}
+	err := x.SQL(query, activeUserDeadline).Find(&bitmap)
+	if err != nil {
+		return err
+	}
+
+	memo := memoUserStats{memoized: time.Now()}
+	for _, role := range bitmap {
+		roletype := models.ROLE_VIEWER
+		if role.Bitrole&0b100 != 0 {
+			roletype = models.ROLE_ADMIN
+		} else if role.Bitrole&0b10 != 0 {
+			roletype = models.ROLE_EDITOR
+		}
+
+		memo.total = addToStats(memo.total, roletype, role.Count)
+		if role.Active {
+			memo.active = addToStats(memo.active, roletype, role.Count)
+		}
+	}
+
+	userStatsCache = memo
 	return nil
 }
+
+func addToStats(base models.UserStats, role models.RoleType, count int64) models.UserStats {
+	base.Users += count
+
+	switch role {
+	case models.ROLE_ADMIN:
+		base.Admins += count
+	case models.ROLE_EDITOR:
+		base.Editors += count
+	default:
+		base.Viewers += count
+	}
+
+	return base
+}

pkg/services/sqlstore/stats_integration_test.go

+// +build integration
+
+package sqlstore
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIntegration_GetUserStats(t *testing.T) {
+	InitTestDB(t)
+
+	cmd := &models.CreateUserCommand{
+		Email:   "admin@test.com",
+		Name:    "Admin",
+		Login:   "admin",
+		OrgName: mainOrgName,
+		IsAdmin: true,
+	}
+	err := CreateUser(context.Background(), cmd)
+	require.NoError(t, err)
+	firstUser := cmd.Result
+
+	{
+		defaultAutoAssign := setting.AutoAssignOrg
+		defaultOrgID := setting.AutoAssignOrgId
+		defaultRole := setting.AutoAssignOrgRole
+
+		setting.AutoAssignOrg = true
+		setting.AutoAssignOrgId = int(firstUser.OrgId)
+		setting.AutoAssignOrgRole = "Editor"
+
+		defer func() {
+			setting.AutoAssignOrg = defaultAutoAssign
+			setting.AutoAssignOrgId = defaultOrgID
+			setting.AutoAssignOrgRole = defaultRole
+		}()
+	}
+
+	users := make([]models.User, 5)
+
+	for i := range users {
+		cmd := &models.CreateUserCommand{
+			Email: fmt.Sprintf("usertest%v@test.com", i),
+			Name:  fmt.Sprintf("user name %v", i),
+			Login: fmt.Sprintf("user_test_%v_login", i),
+			OrgId: firstUser.OrgId,
+		}
+		err := CreateUser(context.Background(), cmd)
+		require.NoError(t, err)
+		users[i] = cmd.Result
+	}
+
+	query := models.GetUserStatsQuery{
+		MustUpdate: true,
+	}
+	err = GetUserStats(&query)
+	require.NoError(t, err)
+	assert.EqualValues(t, models.UserStats{
+		Users:   6,
+		Admins:  1,
+		Editors: 5,
+		Viewers: 0,
+	}, query.Result)
+}

pkg/services/sqlstore/stats_test.go

@@ -11,18 +11,16 @@ import (
 )
 
 func TestStatsDataAccess(t *testing.T) {
-	t.Run("Testing Stats Data Access", func(t *testing.T) {
 	InitTestDB(t)
-
-		t.Run("Get system stats should not results in error", func(t *testing.T) {
 	populateDB(t)
 
+	t.Run("Get system stats should not results in error", func(t *testing.T) {
 		query := models.GetSystemStatsQuery{}
 		err := GetSystemStats(&query)
 		require.NoError(t, err)
 		assert.Equal(t, int64(3), query.Result.Users)
-			assert.Equal(t, 1, query.Result.Editors)
-			assert.Equal(t, 1, query.Result.Viewers)
+		assert.Equal(t, 0, query.Result.Editors)
+		assert.Equal(t, 0, query.Result.Viewers)
 		assert.Equal(t, 3, query.Result.Admins)
 	})
 
@@ -57,14 +55,16 @@ func TestStatsDataAccess(t *testing.T) {
 	})
 
 	t.Run("Get active user count stats should not result in error", func(t *testing.T) {
-			query := models.GetActiveUserStatsQuery{}
-			err := GetActiveUserStats(&query)
+		query := models.GetUserStatsQuery{
+			MustUpdate: true,
+			Active:     true,
+		}
+		err := GetUserStats(&query)
 		require.NoError(t, err)
-			assert.Equal(t, int64(1), query.Result.ActiveUsers)
-			assert.Equal(t, int64(1), query.Result.ActiveAdmins)
-			assert.Equal(t, int64(0), query.Result.ActiveEditors)
-			assert.Equal(t, int64(0), query.Result.ActiveViewers)
-		})
+		assert.Equal(t, int64(1), query.Result.Users)
+		assert.Equal(t, int64(1), query.Result.Admins)
+		assert.Equal(t, int64(0), query.Result.Editors)
+		assert.Equal(t, int64(0), query.Result.Viewers)
 	})
 }
 
@@ -127,4 +127,12 @@ func populateDB(t *testing.T) {
 	}
 	err = UpdateUserLastSeenAt(updateUserLastSeenAtCmd)
 	require.NoError(t, err)
+
+	// force renewal of user stats
+	query := models.GetUserStatsQuery{
+		MustUpdate: true,
+		Active:     true,
+	}
+	err = GetUserStats(&query)
+	require.NoError(t, err)
 }