fix: implement fake DeleteOldWorkspaceAgentStats (#11076)

mtojek · web-flow · commit 941e3873a8d1 · 2023-12-07T14:08:16.000+01:00
diff --git a/coderd/database/dbmem/dbmem.go b/coderd/database/dbmem/dbmem.go
@@ -1180,8 +1180,22 @@ func (q *FakeQuerier) DeleteOldWorkspaceAgentLogs(_ context.Context) error {
 	return nil
 }
 
-func (*FakeQuerier) DeleteOldWorkspaceAgentStats(_ context.Context) error {
-	// no-op
+func (q *FakeQuerier) DeleteOldWorkspaceAgentStats(_ context.Context) error {
+	q.mutex.Lock()
+	defer q.mutex.Unlock()
+
+	now := dbtime.Now()
+	sixMonthInterval := 6 * 30 * 24 * time.Hour
+	sixMonthsAgo := now.Add(-sixMonthInterval)
+
+	var validStats []database.WorkspaceAgentStat
+	for _, stat := range q.workspaceAgentStats {
+		if stat.CreatedAt.Before(sixMonthsAgo) {
+			continue
+		}
+		validStats = append(validStats, stat)
+	}
+	q.workspaceAgentStats = validStats
 	return nil
 }
 
diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go
@@ -34,6 +34,57 @@ func TestPurge(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestDeleteOldWorkspaceAgentStats(t *testing.T) {
+	t.Parallel()
+
+	db, _ := dbtestutil.NewDB(t)
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	now := dbtime.Now()
+
+	// given
+	// Let's use RxBytes to identify stat entries.
+	// Stat inserted 6 months + 1 hour ago, should be deleted.
+	first := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
+		CreatedAt:                 now.Add(-6*30*24*time.Hour - time.Hour),
+		ConnectionMedianLatencyMS: 1,
+		RxBytes:                   1111,
+	})
+
+	// Stat inserted 6 months - 1 hour ago, should not be deleted.
+	second := dbgen.WorkspaceAgentStat(t, db, database.WorkspaceAgentStat{
+		CreatedAt:                 now.Add(-5*30*24*time.Hour + time.Hour),
+		ConnectionMedianLatencyMS: 1,
+		RxBytes:                   2222,
+	})
+
+	// when
+	closer := dbpurge.New(ctx, logger, db)
+	defer closer.Close()
+
+	// then
+	var stats []database.GetWorkspaceAgentStatsRow
+	var err error
+	require.Eventually(t, func() bool {
+		// Query all stats created not earlier than 7 months ago
+		stats, err = db.GetWorkspaceAgentStats(ctx, now.Add(-7*30*24*time.Hour))
+		if err != nil {
+			return false
+		}
+		return !containsWorkspaceAgentStat(stats, first) &&
+			containsWorkspaceAgentStat(stats, second)
+	}, testutil.WaitShort, testutil.IntervalFast, stats)
+}
+
+func containsWorkspaceAgentStat(stats []database.GetWorkspaceAgentStatsRow, needle database.WorkspaceAgentStat) bool {
+	return slices.ContainsFunc(stats, func(s database.GetWorkspaceAgentStatsRow) bool {
+		return s.WorkspaceRxBytes == needle.RxBytes
+	})
+}
+
 func TestDeleteOldWorkspaceAgentLogs(t *testing.T) {
 	t.Parallel()
 
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceagentstats.sql b/coderd/database/queries/workspaceagentstats.sql
@@ -90,7 +90,7 @@ ORDER BY
 	date ASC;
 
 -- name: DeleteOldWorkspaceAgentStats :exec
-DELETE FROM workspace_agent_stats WHERE created_at < NOW() - INTERVAL '6 months';
+DELETE FROM workspace_agent_stats WHERE created_at < NOW() - INTERVAL '180 days';
 
 -- name: GetDeploymentWorkspaceAgentStats :one
 WITH agent_stats AS (
