Keep the analytics flusher alive through drain

cmd/meshtender/main.go

@@ -101,9 +101,16 @@ func run(logger *slog.Logger) error {
 	}
 
 	// First-party traffic analytics: wrap the whole dispatcher so every host is
-	// captured, with a background goroutine doing the writes + rollups.
+	// captured, with a background goroutine doing the writes + rollups. It runs on
+	// its own context — NOT the signal context — so it keeps consuming events while
+	// in-flight requests drain during shutdown; we stop and drain it afterwards.
 	rec := analytics.New(st, cfg)
-	go rec.Run(ctx)
+	analyticsCtx, stopAnalytics := context.WithCancel(context.Background())
+	analyticsDone := make(chan struct{})
+	go func() {
+		defer close(analyticsDone)
+		rec.Run(analyticsCtx)
+	}()
 
 	httpSrv := &http.Server{
 		Addr:              cfg.Addr,
@@ -126,14 +133,25 @@ func run(logger *slog.Logger) error {
 		}
 	}()
 
+	var srvErr error
 	select {
 	case <-ctx.Done():
 		logger.Info("shutting down")
-	case err := <-errCh:
-		return err
+	case srvErr = <-errCh:
+		// Server failed to listen/serve; fall through to orderly teardown.
 	}
 
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-	return httpSrv.Shutdown(shutdownCtx)
+	// Drain in-flight HTTP requests first — they may still record analytics events,
+	// so the flusher must stay alive through the drain.
+	if srvErr == nil {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		srvErr = httpSrv.Shutdown(shutdownCtx)
+		cancel()
+	}
+
+	// Now stop the flusher and wait for its final flush to complete before the
+	// deferred st.Close() closes the pool underneath it.
+	stopAnalytics()
+	<-analyticsDone
+	return srvErr
 }

internal/analytics/analytics.go

@@ -69,7 +69,17 @@ func (rec *Recorder) Run(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
-			// Final flush on a fresh context, since ctx is already cancelled.
+			// Drain anything still queued (e.g. events recorded while in-flight
+			// requests were draining) into the batch, then final-flush on a fresh
+			// context since ctx is already cancelled.
+			for drained := true; drained; {
+				select {
+				case e := <-rec.ch:
+					batch = append(batch, e)
+				default:
+					drained = false
+				}
+			}
 			fc, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			write(fc)
 			cancel()

internal/analytics/main_test.go

@@ -0,0 +1,14 @@
+package analytics
+
+import (
+	"os"
+	"testing"
+
+	"github.com/jleight/meshtender/internal/testdb"
+)
+
+// TestMain wires testdb's container teardown. The container only starts if a test
+// actually calls testdb.Fresh, so the nil-store handler tests stay DB-free.
+func TestMain(m *testing.M) {
+	os.Exit(testdb.RunMain(m))
+}

internal/analytics/shutdown_test.go

@@ -0,0 +1,64 @@
+package analytics
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/jleight/meshtender/internal/config"
+	"github.com/jleight/meshtender/internal/store"
+	"github.com/jleight/meshtender/internal/testdb"
+)
+
+func analyticsMigrate(dsn string) error {
+	ctx := context.Background()
+	s, err := store.New(ctx, dsn)
+	if err != nil {
+		return err
+	}
+	defer s.Close()
+	return s.Migrate(ctx)
+}
+
+// TestRunFlushesQueuedEventsOnShutdown: when Run's context is cancelled, it drains
+// events still queued in the channel and persists them in the final flush, rather
+// than dropping them. This is what lets main.go stop the flusher after the HTTP
+// drain without losing the events recorded during that window.
+func TestRunFlushesQueuedEventsOnShutdown(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	st, err := store.New(ctx, testdb.Fresh(t, analyticsMigrate))
+	if err != nil {
+		t.Fatalf("store: %v", err)
+	}
+	defer st.Close()
+
+	rec := New(st, &config.Config{PrimaryHost: "app.x", RootHost: "x", AuthHost: "auth.x", WWWHost: "www.x"})
+
+	runCtx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		rec.Run(runCtx)
+	}()
+
+	// Queue events (below flushBatch, so nothing is written until shutdown), then
+	// cancel. The final drain+flush must persist all of them.
+	const n = 5
+	for i := 0; i < n; i++ {
+		rec.ch <- store.AnalyticsEvent{
+			Ts: time.Now(), Surface: "app", Host: "app.x", Path: "/dashboard",
+			Method: "GET", Status: 200, Visitor: "v",
+		}
+	}
+	cancel()
+	<-done
+
+	var got int
+	if err := st.Pool().QueryRow(ctx, `SELECT count(*) FROM analytics_events`).Scan(&got); err != nil {
+		t.Fatalf("count events: %v", err)
+	}
+	if got != n {
+		t.Fatalf("persisted %d events on shutdown, want %d (queued events were dropped)", got, n)
+	}
+}