fix batcher maybe

coadler · coadler · commit 1fee5b9ca0d1 · 2023-08-08T00:16:58.000Z
diff --git a/coderd/batchstats/batcher.go b/coderd/batchstats/batcher.go
@@ -199,15 +199,6 @@ func (b *Batcher) flush(ctx context.Context, forced bool, reason string) {
 	defer func() {
 		b.flushForced.Store(false)
 		b.mu.Unlock()
-		// Notify that a flush has completed. This only happens in tests.
-		if b.flushed != nil {
-			select {
-			case <-ctx.Done():
-				close(b.flushed)
-			default:
-				b.flushed <- count
-			}
-		}
 		if count > 0 {
 			elapsed := time.Since(start)
 			b.log.Debug(ctx, "flush complete",
@@ -217,6 +208,15 @@ func (b *Batcher) flush(ctx context.Context, forced bool, reason string) {
 				slog.F("reason", reason),
 			)
 		}
+		// Notify that a flush has completed. This only happens in tests.
+		if b.flushed != nil {
+			select {
+			case <-ctx.Done():
+				close(b.flushed)
+			default:
+				b.flushed <- count
+			}
+		}
 	}()
 
 	if len(b.buf.ID) == 0 {
diff --git a/coderd/batchstats/batcher_internal_test.go b/coderd/batchstats/batcher_internal_test.go
@@ -46,7 +46,7 @@ func TestBatchStats(t *testing.T) {
 
 	// Given: no data points are added for workspace
 	// When: it becomes time to report stats
-	t1 := time.Now()
+	t1 := database.Now()
 	// Signal a tick and wait for a flush to complete.
 	tick <- t1
 	f := <-flushed
@@ -59,7 +59,7 @@ func TestBatchStats(t *testing.T) {
 	require.Empty(t, stats, "should have no stats for workspace")
 
 	// Given: a single data point is added for workspace
-	t2 := time.Now()
+	t2 := database.Now()
 	t.Logf("inserting 1 stat")
 	require.NoError(t, b.Add(deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randAgentSDKStats(t)))
 
@@ -77,7 +77,7 @@ func TestBatchStats(t *testing.T) {
 
 	// Given: a lot of data points are added for both workspaces
 	// (equal to batch size)
-	t3 := time.Now()
+	t3 := database.Now()
 	done := make(chan struct{})
 
 	go func() {
@@ -105,15 +105,15 @@ func TestBatchStats(t *testing.T) {
 	require.Len(t, stats, 2, "should have stats for both workspaces")
 
 	// Ensures that a subsequent flush pushes all the remaining data
-	t4 := time.Now()
+	t4 := database.Now()
 	tick <- t4
 	f2 := <-flushed
 	t.Logf("flush 4 completed")
 	expectedCount := defaultBufferSize - f
 	require.Equal(t, expectedCount, f2, "did not flush expected remaining rows")
 
 	// Ensure that a subsequent flush does not push stale data.
-	t5 := time.Now()
+	t5 := database.Now()
 	tick <- t5
 	f = <-flushed
 	require.Zero(t, f, "expected zero stats to have been flushed")
