fix(coderd/batchstats): fix init race and close flush (#9248)

mafredri · web-flow · commit ed2b1236c0ae · 2023-08-23T11:58:25.000+03:00
diff --git a/coderd/batchstats/batcher.go b/coderd/batchstats/batcher.go
@@ -105,6 +105,8 @@ func New(ctx context.Context, opts ...Option) (*Batcher, func(), error) {
 		b.tickCh = b.ticker.C
 	}
 
+	b.initBuf(b.batchSize)
+
 	cancelCtx, cancelFunc := context.WithCancel(ctx)
 	done := make(chan struct{})
 	go func() {
@@ -172,7 +174,6 @@ func (b *Batcher) Add(
 
 // Run runs the batcher.
 func (b *Batcher) run(ctx context.Context) {
-	b.initBuf(b.batchSize)
 	// nolint:gocritic // This is only ever used for one thing - inserting agent stats.
 	authCtx := dbauthz.AsSystemRestricted(ctx)
 	for {
@@ -184,7 +185,13 @@ func (b *Batcher) run(ctx context.Context) {
 			b.flush(authCtx, true, "reaching capacity")
 		case <-ctx.Done():
 			b.log.Debug(ctx, "context done, flushing before exit")
-			b.flush(authCtx, true, "exit")
+
+			// We must create a new context here as the parent context is done.
+			ctxTimeout, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+			defer cancel() //nolint:revive // We're returning, defer is fine.
+
+			// nolint:gocritic // This is only ever used for one thing - inserting agent stats.
+			b.flush(dbauthz.AsSystemRestricted(ctxTimeout), true, "exit")
 			return
 		}
 	}
diff --git a/coderd/batchstats/batcher_internal_test.go b/coderd/batchstats/batcher_internal_test.go
@@ -31,7 +31,7 @@ func TestBatchStats(t *testing.T) {
 	deps1 := setupDeps(t, store)
 	deps2 := setupDeps(t, store)
 	tick := make(chan time.Time)
-	flushed := make(chan int)
+	flushed := make(chan int, 1)
 
 	b, closer, err := New(ctx,
 		WithStore(store),
