From 3d5df344c10afd2827c35a9624e551dbc2ec2ad0 Mon Sep 17 00:00:00 2001 From: MrAlders0n <55921894+MrAlders0n@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:28:23 -0400 Subject: [PATCH] Finish presence flushes before the pools close on shutdown (#285) * Finish presence flushes before the pools close on shutdown * Ticker flushes outlive a shutdown that starts mid-flush; test and comment fixes --- cmd/beacon/main.go | 8 +- internal/presence/coalescer.go | 26 +++++- internal/presence/coalescer_test.go | 121 ++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 7 deletions(-) diff --git a/cmd/beacon/main.go b/cmd/beacon/main.go index 967197f..54887e6 100644 --- a/cmd/beacon/main.go +++ b/cmd/beacon/main.go @@ -196,8 +196,9 @@ func main() { // ── Presence write coalescing ──────────────────────────────────────────── // Ingest writes go through the coalescer; reads use the API store. coalescer := presence.New(ingestStore, resolved.PresenceFlushInterval, resolved.PresencePacketTTL) + finishPresence := func(context.Context) {} if mode.Coalescer { - go coalescer.Run(ctx) + finishPresence = coalescer.Start(ctx) } // ── Redis cache layer ──────────────────────────────────────────────────── @@ -609,9 +610,8 @@ func main() { } ingestWorkers.Wait() backgroundJobs.Wait() - if mode.Coalescer { - coalescer.Flush(shutdownCtx) - } + // Run's final flush, then one for the bumps written after it, both before the deferred pools.Close. + finishPresence(shutdownCtx) } // proxyLimitWarning explains why per-IP limits go site-wide behind an unlisted reverse proxy. diff --git a/internal/presence/coalescer.go b/internal/presence/coalescer.go index 221aac9..3cf6b2a 100644 --- a/internal/presence/coalescer.go +++ b/internal/presence/coalescer.go @@ -234,15 +234,35 @@ func (c *Coalescer) DeleteOldObservers(ctx context.Context, cutoff time.Time) ([ return c.Store.DeleteOldObservers(ctx, cutoff) } -// Run flushes on the configured interval until ctx is cancelled, then does a -// final flush so a clean shutdown loses nothing. +// Start runs Run in the background until ctx ends. The returned finish waits for Run's final flush to end, then +// flushes the bumps written since. Call it once, after every writer has stopped and before the store's connections +// close: when it returns, no flush is left running. +func (c *Coalescer) Start(ctx context.Context) (finish func(context.Context)) { + done := make(chan struct{}) + go func() { + defer close(done) + c.Run(ctx) + }() + return func(fctx context.Context) { + <-done // Run's final flush is bounded by its own timeout + c.Flush(fctx) + } +} + +// Run flushes every flush interval until ctx ends, then flushes once more and returns. A flush runs on its own +// context, so a shutdown that starts mid-flush lets it finish instead of dropping its rows. func (c *Coalescer) Run(ctx context.Context) { ticker := time.NewTicker(c.flushInterval) defer ticker.Stop() for { select { case <-ticker.C: - c.Flush(ctx) + if ctx.Err() != nil { + continue // both ready: take the final flush + } + flushCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), max(c.flushInterval, 10*time.Second)) + c.Flush(flushCtx) + cancel() case <-ctx.Done(): flushCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) c.Flush(flushCtx) diff --git a/internal/presence/coalescer_test.go b/internal/presence/coalescer_test.go index 6dd2267..0e9185b 100644 --- a/internal/presence/coalescer_test.go +++ b/internal/presence/coalescer_test.go @@ -139,3 +139,124 @@ func TestObserverLastIATAFlush(t *testing.T) { t.Fatalf("flushed %v at %v, want YYZ at %s", gotIATAs, gotAts, heard) } } + +// shutdownStore blocks the first observer flush until release closes, and reports a flush that starts once the pool is +// closed. +type shutdownStore struct { + retentionStore + release chan struct{} + closed atomic.Bool + flushes atomic.Int32 + counts atomic.Int32 + t *testing.T +} + +func (s *shutdownStore) TouchObservers(ctx context.Context, ids []uuid.UUID, seen []time.Time, counts []int32, iatas []string, iataAts []time.Time) error { + if s.closed.Load() { + s.t.Error("observer flush after the pool closed") + } + if s.flushes.Add(1) == 1 { + <-s.release + } + for _, n := range counts { + s.counts.Add(n) + } + if s.closed.Load() { + s.t.Error("observer flush still running when the pool closed") + } + return nil +} + +// On shutdown, finish returns only after Run's final flush has ended, and it flushes the bumps written after that +// flush took its snapshot, so closing the pool right after finish drops nothing. +func TestCoalescerShutdownFlushesBeforePoolClose(t *testing.T) { + store := &shutdownStore{retentionStore: retentionStore{id: uuid.New()}, release: make(chan struct{}), t: t} + c := New(store, time.Hour, time.Hour) + bump := func() { + t.Helper() + if _, _, err := c.UpsertObserver(context.Background(), []byte("fixture"), "YVR"); err != nil { + t.Fatal(err) + } + } + bump() // identity + bump() // one bump for Run's final flush + ctx, cancel := context.WithCancel(context.Background()) + finish := c.Start(ctx) + cancel() + // Run's final flush has taken its snapshot and blocks in TouchObservers. + deadline := time.Now().Add(5 * time.Second) + for store.flushes.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if store.flushes.Load() == 0 { + t.Fatal("Run's final flush never started") + } + bump() // a writer that stopped late + + finished := make(chan struct{}) + go func() { + finish(context.Background()) + close(finished) + }() + select { + case <-finished: + t.Fatal("finish returned while Run's final flush was still running") + case <-time.After(100 * time.Millisecond): + } + close(store.release) + select { + case <-finished: + case <-time.After(5 * time.Second): + t.Fatal("finish did not return") + } + store.closed.Store(true) // main's deferred pools.Close + time.Sleep(50 * time.Millisecond) + if n := store.counts.Load(); n != 2 { + t.Errorf("flushed %d observations, want 2 (Run's final flush and the late one)", n) + } +} + +// A ticker flush in flight when shutdown starts finishes on its own context: its rows are written, not dropped. +func TestCoalescerTickerFlushSurvivesShutdown(t *testing.T) { + store := &shutdownStore{retentionStore: retentionStore{id: uuid.New()}, release: make(chan struct{}), t: t} + var ctxErr atomic.Value + c := New(&ctxStore{shutdownStore: store, ctxErr: &ctxErr}, 20*time.Millisecond, time.Hour) + for range 2 { + if _, _, err := c.UpsertObserver(context.Background(), []byte("fixture"), "YVR"); err != nil { + t.Fatal(err) + } + } + ctx, cancel := context.WithCancel(context.Background()) + finish := c.Start(ctx) + deadline := time.Now().Add(5 * time.Second) + for store.flushes.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if store.flushes.Load() == 0 { + t.Fatal("no ticker flush started") + } + cancel() // shutdown while the ticker flush is blocked + time.Sleep(20 * time.Millisecond) + close(store.release) + finish(context.Background()) + if err, _ := ctxErr.Load().(error); err != nil { + t.Errorf("ticker flush saw a cancelled context: %v", err) + } + if n := store.counts.Load(); n != 1 { + t.Errorf("flushed %d observations, want 1", n) + } +} + +// ctxStore records whether the first observer flush's context was cancelled by the time it wrote. +type ctxStore struct { + *shutdownStore + ctxErr *atomic.Value +} + +func (s *ctxStore) TouchObservers(ctx context.Context, ids []uuid.UUID, seen []time.Time, counts []int32, iatas []string, iataAts []time.Time) error { + err := s.shutdownStore.TouchObservers(ctx, ids, seen, counts, iatas, iataAts) + if ctx.Err() != nil { + s.ctxErr.CompareAndSwap(nil, ctx.Err()) + } + return err +}