diff --git a/cmd/meshtender/main.go b/cmd/meshtender/main.go index 7c0ab04..b01fc13 100644 --- a/cmd/meshtender/main.go +++ b/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 } diff --git a/internal/analytics/analytics.go b/internal/analytics/analytics.go index 1f0b869..d40f6cb 100644 --- a/internal/analytics/analytics.go +++ b/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() diff --git a/internal/analytics/main_test.go b/internal/analytics/main_test.go new file mode 100644 index 0000000..c0e3de5 --- /dev/null +++ b/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)) +} diff --git a/internal/analytics/shutdown_test.go b/internal/analytics/shutdown_test.go new file mode 100644 index 0000000..b5db919 --- /dev/null +++ b/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) + } +}