add backlog for re-reporting in case of failure

mafredri · mafredri · commit 7dc4ff226baf · 2023-08-14T11:43:02.000Z
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceappstats.sql b/coderd/database/queries/workspaceappstats.sql
@@ -30,4 +30,8 @@ DO
 	WHERE
 		workspace_app_stats.user_id = EXCLUDED.user_id
 		AND workspace_app_stats.agent_id = EXCLUDED.agent_id
-		AND workspace_app_stats.session_id = EXCLUDED.session_id;
+		AND workspace_app_stats.session_id = EXCLUDED.session_id
+		-- Since stats are updated in place as time progresses, we only
+		-- want to update this row if it's fresh.
+		AND workspace_app_stats.session_ended_at <= EXCLUDED.session_ended_at
+		AND workspace_app_stats.requests <= EXCLUDED.requests;
diff --git a/coderd/workspaceapps/stats.go b/coderd/workspaceapps/stats.go
@@ -142,6 +142,7 @@ type StatsCollector struct {
 	mu               sync.Mutex                       // Protects following.
 	statsBySessionID map[uuid.UUID]*StatsReport       // Track unique sessions.
 	groupedStats     map[statsGroupKey][]*StatsReport // Rolled up stats for sessions in close proximity.
+	backlog          []StatsReport                    // Stats that have not been reported yet (due to error).
 }
 
 type StatsCollectorOptions struct {
@@ -242,7 +243,7 @@ func (sc *StatsCollector) rollup() []StatsReport {
 				rolledUp:         true,
 			}
 		}
-
+		rollupChanged := false
 		newGroup := []*StatsReport{rolledUp} // Must be first in slice for future iterations (see group[0] above).
 		for _, stat := range group {
 			if !stat.SessionEndedAt.IsZero() && stat.SessionEndedAt.Sub(stat.SessionStartedAt) <= sc.opts.RollupWindow {
@@ -251,6 +252,7 @@ func (sc *StatsCollector) rollup() []StatsReport {
 					rolledUp.SessionID = stat.SessionID // Borrow the first session ID, useful in tests.
 				}
 				rolledUp.Requests += stat.Requests
+				rollupChanged = true
 				continue
 			}
 			if stat.SessionEndedAt.IsZero() && sc.opts.Now().Sub(stat.SessionStartedAt) <= sc.opts.RollupWindow {
@@ -273,7 +275,7 @@ func (sc *StatsCollector) rollup() []StatsReport {
 				newGroup = append(newGroup, stat) // Keep it for future updates.
 			}
 		}
-		if rolledUp.Requests > 0 {
+		if rollupChanged {
 			report = append(report, *rolledUp)
 		}
 
@@ -300,12 +302,33 @@ func (sc *StatsCollector) flush(ctx context.Context) (err error) {
 		}
 	}()
 
+	// We keep the backlog as a simple slice so that we don't need to
+	// attempt to merge it with the stats we're about to report. This
+	// is because the rollup is a one-way operation and the backlog may
+	// contain stats that are still in the statsBySessionID map and will
+	// be reported again in the future. It is possible to merge the
+	// backlog and the stats we're about to report, but it's not worth
+	// the complexity.
+	if len(sc.backlog) > 0 {
+		err = sc.opts.Reporter.Report(ctx, sc.backlog)
+		if err != nil {
+			return xerrors.Errorf("report workspace app stats from backlog failed: %w", err)
+		}
+		sc.backlog = nil
+	}
+
 	stats := sc.rollup()
 	if len(stats) == 0 {
 		return nil
 	}
 
-	return sc.opts.Reporter.Report(ctx, stats)
+	err = sc.opts.Reporter.Report(ctx, stats)
+	if err != nil {
+		sc.backlog = stats
+		return xerrors.Errorf("report workspace app stats failed: %w", err)
+	}
+
+	return nil
 }
 
 func (sc *StatsCollector) Close() error {
diff --git a/coderd/workspaceapps/stats_test.go b/coderd/workspaceapps/stats_test.go
@@ -11,15 +11,18 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/exp/slices"
+	"golang.org/x/xerrors"
 
 	"github.com/coder/coder/coderd/database"
 	"github.com/coder/coder/coderd/workspaceapps"
 	"github.com/coder/coder/testutil"
 )
 
 type fakeReporter struct {
-	mu sync.Mutex
-	s  []workspaceapps.StatsReport
+	mu   sync.Mutex
+	s    []workspaceapps.StatsReport
+	err  error
+	errN int
 }
 
 func (r *fakeReporter) stats() []workspaceapps.StatsReport {
@@ -28,8 +31,25 @@ func (r *fakeReporter) stats() []workspaceapps.StatsReport {
 	return r.s
 }
 
+func (r *fakeReporter) errors() int {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.errN
+}
+
+func (r *fakeReporter) setError(err error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.err = err
+}
+
 func (r *fakeReporter) Report(_ context.Context, stats []workspaceapps.StatsReport) error {
 	r.mu.Lock()
+	if r.err != nil {
+		r.errN++
+		r.mu.Unlock()
+		return r.err
+	}
 	r.s = append(r.s, stats...)
 	r.mu.Unlock()
 	return nil
@@ -296,7 +316,6 @@ func TestStatsCollector(t *testing.T) {
 			var gotStats []workspaceapps.StatsReport
 			require.Eventually(t, func() bool {
 				gotStats = reporter.stats()
-				t.Logf("len(reporter.stats()) = %d, len(tt.want) = %d", len(gotStats), len(tt.want))
 				return len(gotStats) == len(tt.want)
 			}, testutil.WaitMedium, testutil.IntervalFast)
 
@@ -329,6 +348,63 @@ func TestStatsCollector(t *testing.T) {
 	}
 }
 
+func TestStatsCollector_backlog(t *testing.T) {
+	t.Parallel()
+
+	rollupWindow := time.Minute
+	flush := make(chan chan<- struct{}, 1)
+
+	start := database.Now().Truncate(time.Minute).UTC()
+	var now atomic.Pointer[time.Time]
+	now.Store(&start)
+
+	reporter := &fakeReporter{}
+	collector := workspaceapps.NewStatsCollector(workspaceapps.StatsCollectorOptions{
+		Reporter:       reporter,
+		ReportInterval: time.Hour,
+		RollupWindow:   rollupWindow,
+
+		Flush: flush,
+		Now:   func() time.Time { return *now.Load() },
+	})
+
+	reporter.setError(xerrors.New("some error"))
+
+	// The first collected stat is "rolled up" and moved into the
+	// backlog during the first flush. On the second flush nothing is
+	// rolled up due to being unable to report the backlog.
+	for i := 0; i < 2; i++ {
+		collector.Collect(workspaceapps.StatsReport{
+			SessionID:        uuid.New(),
+			SessionStartedAt: start,
+			SessionEndedAt:   start.Add(10 * time.Second),
+			Requests:         1,
+		})
+		start = start.Add(time.Minute)
+		now.Store(&start)
+
+		flushDone := make(chan struct{}, 1)
+		flush <- flushDone
+		<-flushDone
+	}
+
+	// Flush was performed 2 times, 2 reports should have failed.
+	wantErrors := 2
+	assert.Equal(t, wantErrors, reporter.errors())
+	assert.Empty(t, reporter.stats())
+
+	reporter.setError(nil)
+
+	// Flush again, this time the backlog should be reported in addition
+	// to the second collected stat being rolled up and reported.
+	flushDone := make(chan struct{}, 1)
+	flush <- flushDone
+	<-flushDone
+
+	assert.Equal(t, wantErrors, reporter.errors())
+	assert.Len(t, reporter.stats(), 2)
+}
+
 func TestStatsCollector_Close(t *testing.T) {
 	t.Parallel()
 
