diff --git a/coderd/batchstats/batcher.go b/coderd/batchstats/batcher.go index e53a24cb4c1d1..b3b881f2133e9 100644 --- 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 index b5d18e81327fd..8c1c367f7db5b 100644 --- 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),