From 1f9024e2ca77de407b357cb8add21de661ae4ca6 Mon Sep 17 00:00:00 2001 From: Tiago Farto Date: Wed, 29 Jul 2026 11:08:56 +0000 Subject: [PATCH 1/3] fix(enrichment): reset leader session state, clear gauge on shutdown (ENG-1670) Review follow-ups on the leader election and backlog gauge: - The leader's `SET idle_session_timeout` leaked onto the pooled connection. SET is session state and the pool hands that exact backend to unrelated callers afterwards, so Postgres could eventually terminate some other component's idle connection. release() now RESETs it before returning the connection. Pinned by a MaxConns=1 test that fails with "30min" if the reset is removed. - The gauge was not withdrawn on graceful shutdown -- the most frequent handover of all. Close released the advisory lock while the values stayed, so the meter provider's final export published a backlog this process no longer owned. Added the matching defer. - Corrected the alerting guidance, which the withdraw-on-demotion change had inverted: a failed poll now removes the series rather than freezing it, so the symptom is a MISSING gauge and a value-based staleness rule would not catch it. Updated the interface doc, the metric description and the escalated log. - Fixed a comment that contradicted migration 020 about the 016 index predicate, and one left over from the transaction-scoped lock. - Made the withdraw test load-bearing: it counted zero iterations and would have passed even with a wrong metric name. - Documented why the 020 index rebuild accepts a brief no-index window (its only consumer is the manual backfill command, which degrades to a seq scan) rather than adding a temp-name-and-rename dance. --- cmd/api/app.go | 14 +++++--- internal/observability/enrichment_backlog.go | 14 ++++---- .../observability/enrichment_backlog_test.go | 17 +++++++-- .../enrichment_status_repository.go | 12 +++++++ .../repository/feedback_records_repository.go | 6 ++-- migrations/020_add_emotions_classified_at.sql | 9 ++++- tests/enrichment_status_test.go | 35 +++++++++++++++++-- 7 files changed, 90 insertions(+), 17 deletions(-) diff --git a/cmd/api/app.go b/cmd/api/app.go index ab3a1032..8ee64f52 100644 --- a/cmd/api/app.go +++ b/cmd/api/app.go @@ -673,6 +673,12 @@ func runEnrichmentBacklogPoller( leader := repository.NewEnrichmentBacklogLeader(db) defer leader.Close(ctx) + // Withdraw the series whenever this process stops being the leader, including the most common + // case of all: shutdown. Close releases the advisory lock, so without this the meter provider's + // final collect-and-export on shutdown would publish one last reading for a backlog this + // process no longer owns. + defer backlog.ClearEnrichmentPending() + ticker := time.NewTicker(enrichmentBacklogInterval) defer ticker.Stop() @@ -697,13 +703,13 @@ func runEnrichmentBacklogPoller( // than leave it frozen at the last good reading while the new leader publishes its own. backlog.ClearEnrichmentPending() - // Always count the failure so a stale gauge is alertable, then escalate the log from - // warn to error once failures persist: a single blip is noise, a run of them means the - // gauge is frozen at its last value and silently lying about the backlog. + // Count the failure so the gap is alertable, then escalate the log from warn to error + // once failures persist: a single blip is noise, but a run of them means nobody is + // publishing the backlog at all. backlog.RecordPollError(ctx) if consecutiveFailures >= enrichmentBacklogFailuresBeforeError { - slog.ErrorContext(ctx, "enrichment backlog poll failing repeatedly; gauge is stale", + slog.ErrorContext(ctx, "enrichment backlog poll failing repeatedly; gauge is not being published", "error", err, "consecutive_failures", consecutiveFailures) } else { slog.WarnContext(ctx, "enrichment backlog poll failed", diff --git a/internal/observability/enrichment_backlog.go b/internal/observability/enrichment_backlog.go index 3509b1dd..23621450 100644 --- a/internal/observability/enrichment_backlog.go +++ b/internal/observability/enrichment_backlog.go @@ -31,9 +31,10 @@ type EnrichmentBacklogMetrics interface { // would look like a permanently stuck backlog. Callers must clear as soon as they are no longer // the leader. Exporting nothing is the honest state — absence is visible, a stale value is not. ClearEnrichmentPending() - // RecordPollError counts a failed refresh. The gauge holds its last value when a poll fails, - // which is indistinguishable from a healthy steady backlog on a dashboard — this counter is - // the signal to alert on (rate > 0 means the gauge is going stale). + // RecordPollError counts a failed refresh. A failed poll costs this process its leadership and + // withdraws the series (see ClearEnrichmentPending), so the symptom is a MISSING gauge rather + // than a stale value: alert on this counter, or on the absence of the gauge — a value-based + // staleness rule will not catch it. RecordPollError(ctx context.Context) } @@ -84,9 +85,10 @@ func NewEnrichmentBacklogMetrics(meter metric.Meter) (EnrichmentBacklogMetrics, pollErrors, err := meter.Int64Counter( MetricNameEnrichmentBacklogPollErrs, metric.WithDescription( - "Failed refreshes of the enrichment backlog gauge. A non-zero rate means the gauge is "+ - "stale (holding its last value), which on a dashboard is indistinguishable from a "+ - "steady backlog — alert on this rather than trusting a flat gauge.", + "Failed refreshes of the enrichment backlog gauge. A failed poll costs the process its "+ + "leadership and withdraws the series, so the symptom is a MISSING "+ + MetricNameEnrichmentPendingRecords+" rather than a stale value — alert on this "+ + "counter, or on absence of that gauge.", ), metric.WithUnit("1"), ) diff --git a/internal/observability/enrichment_backlog_test.go b/internal/observability/enrichment_backlog_test.go index 368b815a..db820c38 100644 --- a/internal/observability/enrichment_backlog_test.go +++ b/internal/observability/enrichment_backlog_test.go @@ -60,12 +60,25 @@ func TestEnrichmentBacklogMetricsClearWithdrawsSeries(t *testing.T) { var collected metricdata.ResourceMetrics require.NoError(t, reader.Collect(context.Background(), &collected)) + // Count rather than iterate-and-assert: after the clear the SDK drops the empty metric and then + // the empty scope, so a loop over ScopeMetrics runs zero times and would pass vacuously (it + // would also "pass" if the metric name were wrong, or if Collect returned nothing at all). + points := 0 + for _, scope := range collected.ScopeMetrics { for _, m := range scope.Metrics { - assert.NotEqual(t, MetricNameEnrichmentPendingRecords, m.Name, - "a cleared gauge must export no data points at all, not a stale value") + if m.Name != MetricNameEnrichmentPendingRecords { + continue + } + + gauge, ok := m.Data.(metricdata.Gauge[int64]) + require.True(t, ok, "expected Gauge[int64] for %s", MetricNameEnrichmentPendingRecords) + + points += len(gauge.DataPoints) } } + + assert.Zero(t, points, "a cleared gauge must export no data points at all, not a stale value") } // TestEnrichmentBacklogMetricsPollErrors verifies failed refreshes are counted, so a gauge frozen diff --git a/internal/repository/enrichment_status_repository.go b/internal/repository/enrichment_status_repository.go index 9c744328..e890913c 100644 --- a/internal/repository/enrichment_status_repository.go +++ b/internal/repository/enrichment_status_repository.go @@ -140,7 +140,13 @@ const ( // setLeaderIdleTimeoutSQL caps how long a stalled leader session can hold the lock. Applies to // this session only (SET, not ALTER ROLE); comfortably above enrichmentBacklogInterval so a // healthy leader, which queries every interval, is never affected. Requires PG14+. + // + // It cannot be SET LOCAL (the convention elsewhere in this package, e.g. embeddings_repository) + // because that is transaction-scoped and the leader deliberately holds no transaction -- which + // is exactly why release() must RESET it before the connection returns to the pool. setLeaderIdleTimeoutSQL = `SET idle_session_timeout = '30min'` + // resetLeaderIdleTimeoutSQL restores the server default so the pooled connection goes back clean. + resetLeaderIdleTimeoutSQL = `RESET idle_session_timeout` ) // CountEnrichmentStatus returns one tenant's eligible/done counts per enrichment. defaultLang is @@ -279,6 +285,12 @@ func (l *EnrichmentBacklogLeader) release(ctx context.Context) { // session lock survives being handed back to the pool. _, _ = l.conn.Exec(unlockCtx, sessionUnlockSQL, enrichmentBacklogLockKey) + // Undo the leader-only idle timeout for the same reason: SET is session state, and the pool + // hands this exact backend to unrelated callers afterwards (verified: same pid, setting still + // in effect). Leaving it behind would let Postgres terminate some other component's pooled + // connection after 30 idle minutes. + _, _ = l.conn.Exec(unlockCtx, resetLeaderIdleTimeoutSQL) + l.conn.Release() l.conn = nil } diff --git a/internal/repository/feedback_records_repository.go b/internal/repository/feedback_records_repository.go index 15b9d6af..0386b03b 100644 --- a/internal/repository/feedback_records_repository.go +++ b/internal/repository/feedback_records_repository.go @@ -545,8 +545,10 @@ const sentimentBackfillSelectSQL = classifyBackfillEligibleSQL + ` // Both NULL checks are required. emotions IS NULL alone would re-send every record whose // classification legitimately found no emotion (stored as NULL) to the provider on every run; // emotions_classified_at IS NULL alone would re-classify rows enriched before that column existed -// (migration 020). Together they select exactly the records never classified. The 016 partial index -// covers the emotions IS NULL arm; the marker is an additional filter on top. +// (migration 020). Together they select exactly the records never classified. Migration 020 also +// realigned the 016 partial index to this same pair, so BOTH checks are served by the index instead +// of the marker being re-checked as a post-filter -- and so a classified-empty row (which keeps +// emotions NULL forever) leaves the index instead of being retained in it permanently. const emotionsBackfillSelectSQL = classifyBackfillEligibleSQL + ` AND emotions IS NULL AND emotions_classified_at IS NULL diff --git a/migrations/020_add_emotions_classified_at.sql b/migrations/020_add_emotions_classified_at.sql index 04571901..ce5ca282 100644 --- a/migrations/020_add_emotions_classified_at.sql +++ b/migrations/020_add_emotions_classified_at.sql @@ -33,7 +33,14 @@ ALTER TABLE feedback_records -- directly instead of re-checking emotions_classified_at as a post-filter. -- -- Built CONCURRENTLY, and DROP-then-CREATE (not IF NOT EXISTS) so an interrupted build leaves an --- INVALID index that a re-run replaces, matching 016. +-- INVALID index that a re-run replaces, matching 016. Note the difference from 016 though: there the +-- DROP only ever removed an INVALID leftover, whereas here it removes an index that is currently in +-- service, so there is a window -- the length of the concurrent build, or until an operator re-runs +-- a persistently failing one -- with no emotions-backfill index. Accepted deliberately: the only +-- consumer is the manual classify-backfill command, which degrades to a sequential scan rather than +-- failing, and no serving path uses this index. The zero-window alternative (build under a temp +-- name, drop the old, then ALTER INDEX ... RENAME) adds two more statements and more partial-failure +-- states than that degradation is worth. DROP INDEX CONCURRENTLY IF EXISTS idx_feedback_records_emotions_backfill; CREATE INDEX CONCURRENTLY idx_feedback_records_emotions_backfill ON feedback_records (id) diff --git a/tests/enrichment_status_test.go b/tests/enrichment_status_test.go index 89d5607c..58895006 100644 --- a/tests/enrichment_status_test.go +++ b/tests/enrichment_status_test.go @@ -290,8 +290,8 @@ func TestCountEnrichmentBacklogAggregateIfLeader(t *testing.T) { cfg, err := config.Load() require.NoError(t, err) - // Two independent pools stand in for two API replicas: a transaction-scoped advisory lock is - // held per session, so the contention is only observable across separate connections. + // Two independent pools stand in for two API replicas: the advisory lock is session-scoped and + // held on one connection, so contention is only observable across separate connections. dbLeader, err := database.NewPostgresPool(ctx, cfg.Database.URL, database.WithPoolConfig(cfg.Database.PoolConfig())) require.NoError(t, err) @@ -346,3 +346,34 @@ func TestCountEnrichmentBacklogAggregateIfLeader(t *testing.T) { leaderTwo.Close(ctx) } + +// TestEnrichmentBacklogLeaderLeavesNoSessionState pins the release path: the leader sets a +// session-scoped idle_session_timeout on its connection, and pgxpool hands that same backend to +// unrelated callers afterwards. If release() did not RESET it, Postgres would eventually terminate +// some other component's pooled connection after 30 idle minutes. MaxConns=1 guarantees the +// connection reused below is the very one leadership was held on. +func TestEnrichmentBacklogLeaderLeavesNoSessionState(t *testing.T) { + ctx := context.Background() + + cfg, err := config.Load() + require.NoError(t, err) + + db, err := database.NewPostgresPool(ctx, cfg.Database.URL, + database.WithPoolConfig(database.PoolConfig{MaxConns: 1, MinConns: 1})) + require.NoError(t, err) + + defer db.Close() + + leader := repository.NewEnrichmentBacklogLeader(db) + + _, isLeader, err := leader.CountIfLeader(ctx, "") + require.NoError(t, err) + require.True(t, isLeader, "single-connection pool must win leadership") + + leader.Close(ctx) + + // Same backend, borrowed as any other caller would. + var timeout string + require.NoError(t, db.QueryRow(ctx, `SELECT current_setting('idle_session_timeout')`).Scan(&timeout)) + assert.Equal(t, "0", timeout, "the released connection must carry no leader-only session state") +} From 73a897dc4cb24875d86fdeeba439ef4fe622b529 Mon Sep 17 00:00:00 2001 From: Tiago Farto Date: Thu, 30 Jul 2026 11:13:43 +0000 Subject: [PATCH 2/3] fix(enrichment): never return an uncleaned leader connection to the pool (ENG-1670) Review follow-ups on #116: - release() shared one deadline between the advisory unlock and the idle-timeout RESET, so a slow unlock could consume the budget the reset needed, and both errors were discarded before handing the connection back -- exactly the leak the reset exists to prevent. Each statement now gets its own detached deadline, both errors are checked and logged, and if either fails the connection is taken out of the pool (Hijack + Close) instead of being reused: ending the session releases whatever the statements could not undo. Losing one pooled connection is cheaper than leaking a lock or a timeout onto a shared one. - The session-state test could have passed vacuously: tryAcquire sets the timeout best-effort and still grants leadership if the SET fails, so asserting a hard-coded "0" proved nothing on a server that ignores it. It now captures the original value, proves the server honours idle_session_timeout at all, and asserts release restores that captured value. - Register the leader's deferred Close before taking leadership so a failing assertion cannot leave the connection checked out and wedge db.Close(). The probe that verifies server support goes through the pool for the same reason, and uses set_config() because pgx's extended protocol rejects "SET ...; SELECT ...". --- .../enrichment_status_repository.go | 62 ++++++++++++++----- tests/enrichment_status_test.go | 32 ++++++++-- 2 files changed, 74 insertions(+), 20 deletions(-) diff --git a/internal/repository/enrichment_status_repository.go b/internal/repository/enrichment_status_repository.go index e890913c..981fe8c5 100644 --- a/internal/repository/enrichment_status_repository.go +++ b/internal/repository/enrichment_status_repository.go @@ -269,30 +269,60 @@ func (l *EnrichmentBacklogLeader) tryAcquire(ctx context.Context) (bool, error) return true, nil } -// release unlocks and returns the leader connection. The unlock uses a context detached from the -// caller's, so shutdown (whose context is already cancelled) still releases the lock rather than -// leaving it held until the backend session is reaped. +// release relinquishes leadership and disposes of the leader connection. +// +// Both cleanup statements undo SESSION state that would otherwise ride the connection back into the +// pool: the advisory lock (which survives Release and would wedge leadership for every replica) and +// the leader-only idle timeout (which would let Postgres terminate whichever component borrowed the +// connection next and left it idle). Each gets its OWN deadline, detached from the caller's context +// so shutdown -- whose context is already cancelled -- still cleans up, and so a slow unlock cannot +// eat the budget the reset needs. +// +// If either statement does not complete we cannot prove the session is clean, so the connection is +// taken out of the pool (Hijack + Close) rather than handed to an unrelated caller. Losing one +// pooled connection is strictly cheaper than leaking a lock or a timeout onto a shared one. func (l *EnrichmentBacklogLeader) release(ctx context.Context) { if l.conn == nil { return } - unlockCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), enrichmentBacklogUnlockTimeout) - defer cancel() + conn := l.conn + l.conn = nil - // Best effort: if this fails the connection is broken, and ending that session releases the - // lock anyway. Returning the connection WITHOUT unlocking would be the real leak, since a - // session lock survives being handed back to the pool. - _, _ = l.conn.Exec(unlockCtx, sessionUnlockSQL, enrichmentBacklogLockKey) + unlockErr := l.execDetached(ctx, conn, sessionUnlockSQL, enrichmentBacklogLockKey) + resetErr := l.execDetached(ctx, conn, resetLeaderIdleTimeoutSQL) - // Undo the leader-only idle timeout for the same reason: SET is session state, and the pool - // hands this exact backend to unrelated callers afterwards (verified: same pid, setting still - // in effect). Leaving it behind would let Postgres terminate some other component's pooled - // connection after 30 idle minutes. - _, _ = l.conn.Exec(unlockCtx, resetLeaderIdleTimeoutSQL) + if unlockErr != nil || resetErr != nil { + slog.WarnContext(ctx, "enrichment backlog: leader session cleanup incomplete, discarding connection", + "unlock_error", unlockErr, "reset_error", resetErr) - l.conn.Release() - l.conn = nil + // Hijack removes the connection from the pool and transfers ownership here; closing it ends + // the backend session, which releases anything the statements above failed to undo. + hijacked := conn.Hijack() + + closeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), enrichmentBacklogUnlockTimeout) + defer cancel() + + _ = hijacked.Close(closeCtx) + + return + } + + conn.Release() +} + +// execDetached runs one cleanup statement on its own deadline, independent of the caller's context. +func (l *EnrichmentBacklogLeader) execDetached( + ctx context.Context, conn *pgxpool.Conn, sql string, args ...any, +) error { + execCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), enrichmentBacklogUnlockTimeout) + defer cancel() + + if _, err := conn.Exec(execCtx, sql, args...); err != nil { + return fmt.Errorf("enrichment backlog leader cleanup: %w", err) + } + + return nil } // scanEnrichmentCounts reads the six-column count row; the scan order matches enrichmentCountSelect. diff --git a/tests/enrichment_status_test.go b/tests/enrichment_status_test.go index 58895006..381c369c 100644 --- a/tests/enrichment_status_test.go +++ b/tests/enrichment_status_test.go @@ -364,16 +364,40 @@ func TestEnrichmentBacklogLeaderLeavesNoSessionState(t *testing.T) { defer db.Close() + // Record the value the connection starts with, and prove this server/role actually honours the + // SET/RESET pair. tryAcquire applies the timeout best-effort and still grants leadership if the + // SET fails, so without this the assertion below could pass against an already-default value + // without ever exercising RESET. + var original string + require.NoError(t, db.QueryRow(ctx, `SELECT current_setting('idle_session_timeout')`).Scan(&original)) + + // set_config applies the value and returns it in one statement: pgx's extended protocol rejects + // a multi-statement "SET ...; SELECT ...", and going through the pool (rather than holding an + // acquired connection) means a failed assertion here cannot leave a connection checked out and + // wedge db.Close(). + var probeSet string + require.NoError(t, db.QueryRow(ctx, + `SELECT set_config('idle_session_timeout', '30min', false)`).Scan(&probeSet)) + require.Equal(t, "30min", probeSet, "server must honour idle_session_timeout for this test to mean anything") + + _, err = db.Exec(ctx, `RESET idle_session_timeout`) + require.NoError(t, err) + leader := repository.NewEnrichmentBacklogLeader(db) + // Registered before leadership is taken: if a require below fails, cleanup must still return the + // connection or db.Close() would block on it and wedge the suite. Close is idempotent. + defer leader.Close(ctx) + _, isLeader, err := leader.CountIfLeader(ctx, "") require.NoError(t, err) require.True(t, isLeader, "single-connection pool must win leadership") leader.Close(ctx) - // Same backend, borrowed as any other caller would. - var timeout string - require.NoError(t, db.QueryRow(ctx, `SELECT current_setting('idle_session_timeout')`).Scan(&timeout)) - assert.Equal(t, "0", timeout, "the released connection must carry no leader-only session state") + // Same backend, borrowed as any other caller would: it must be back to what it started as, not + // carrying the leader's 30min timeout. + var afterRelease string + require.NoError(t, db.QueryRow(ctx, `SELECT current_setting('idle_session_timeout')`).Scan(&afterRelease)) + assert.Equal(t, original, afterRelease, "the released connection must carry no leader-only session state") } From 0aba386c0ab5ee0a100472edc73b0bb0dab3bc06 Mon Sep 17 00:00:00 2001 From: Tiago Farto Date: Wed, 5 Aug 2026 11:55:36 +0000 Subject: [PATCH 3/3] fix(enrichment): join the backlog poller before shutting down metrics (ENG-1670) Run returns the moment its context is cancelled and never joins runEnrichmentBacklogPoller, so the poller's deferred ClearEnrichmentPending raced Shutdown. Losing that race defeats the point of the clear: the meter provider's final collect-and-export publishes one last backlog reading for a process that has already lost leadership, which is the stale series the clear exists to withdraw. Track the poller with a done channel and join it in Shutdown, registered after the observability defer so LIFO runs the join first. A defer rather than an inline call, so the early-return error paths are covered too. The join is bounded by the shutdown context: a poller stuck in its leader cleanup must not hold shutdown past its budget, and a stale gauge sample is cheaper than a wedged shutdown. Postgres releases the session lock when the socket closes regardless. Waiting also fixes a second, quieter race: the leader's advisory-lock release and idle-timeout reset are best-effort with nothing awaiting them, so process exit could cut them off mid-statement. Only this poller is tracked. runRiverQueueDepthPoller has no clear-on-exit semantics -- it just stops updating, and its last reading is meant to be exported -- so joining it would add shutdown latency for no benefit. TestShutdownWaitsForEnrichmentBacklogPoller pins the ordering; deleting the join defer makes it fail with [metrics-exported metrics-shutdown], the gauge exported before the clear runs at all. --- cmd/api/app.go | 61 ++++++++++++++--- cmd/api/app_test.go | 156 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 209 insertions(+), 8 deletions(-) diff --git a/cmd/api/app.go b/cmd/api/app.go index 8ee64f52..6bc693df 100644 --- a/cmd/api/app.go +++ b/cmd/api/app.go @@ -42,6 +42,9 @@ type App struct { tracerProvider *sdktrace.TracerProvider metrics *observability.Metrics taxonomyRepo *repository.TaxonomyRepository + // enrichmentBacklogDone closes when the enrichment-backlog poller goroutine has returned, + // including its deferred cleanup. Nil when the poller was never started (metrics disabled). + enrichmentBacklogDone chan struct{} } var ( @@ -609,14 +612,24 @@ func (a *App) Run(ctx context.Context) error { } if a.metrics != nil && a.metrics.EnrichmentBacklog != nil { - go runEnrichmentBacklogPoller(ctx, a.db, a.metrics.EnrichmentBacklog, enrichmentBacklogPollConfig{ - // Trim to stay consistent with NewEnrichmentStatusService (config already canonicalizes - // this, so it's defensive symmetry) — the endpoint and the gauge resolve the same target. - defaultLang: strings.TrimSpace(a.cfg.Translation.DefaultLanguage), - translationConfigured: a.cfg.Translation.Provider != "" && a.cfg.Translation.Model != "", - sentimentConfigured: a.cfg.Sentiment.Enabled(), - emotionsConfigured: a.cfg.Emotions.Enabled(), - }) + // Tracked, unlike the pollers above, because this one has cleanup that Shutdown must not + // race: it withdraws the backlog series and releases the leader's advisory lock on the way + // out. See App.awaitEnrichmentBacklogPoller. + done := make(chan struct{}) + a.enrichmentBacklogDone = done + + go func() { + defer close(done) + + runEnrichmentBacklogPoller(ctx, a.db, a.metrics.EnrichmentBacklog, enrichmentBacklogPollConfig{ + // Trim to stay consistent with NewEnrichmentStatusService (config already canonicalizes + // this, so it's defensive symmetry) — the endpoint and the gauge resolve the same target. + defaultLang: strings.TrimSpace(a.cfg.Translation.DefaultLanguage), + translationConfigured: a.cfg.Translation.Provider != "" && a.cfg.Translation.Model != "", + sentimentConfigured: a.cfg.Sentiment.Enabled(), + emotionsConfigured: a.cfg.Emotions.Enabled(), + }) + }() } // Reap taxonomy runs orphaned in a non-terminal state, but only when the taxonomy service is wired @@ -960,6 +973,13 @@ func (a *App) Shutdown(ctx context.Context) (err error) { } }() + // Registered AFTER the observability defer so LIFO runs it FIRST: the meter provider's final + // collect-and-export must not happen until the poller has withdrawn the backlog series, or + // shutdown publishes one last reading for a backlog this process no longer owns -- the exact + // thing the poller's ClearEnrichmentPending exists to prevent. A defer (not an inline call) + // so the early-return error paths below are covered too. + defer a.awaitEnrichmentBacklogPoller(ctx) + if err = a.server.Shutdown(ctx); err != nil && !errors.Is(err, http.ErrServerClosed) { if stopErr := a.river.Stop(ctx); stopErr != nil { slog.Error("river stop during server shutdown", "error", stopErr) @@ -974,3 +994,28 @@ func (a *App) Shutdown(ctx context.Context) (err error) { return nil } + +// awaitEnrichmentBacklogPoller blocks until the enrichment-backlog poller goroutine has returned, +// bounded by the shutdown context. +// +// Run returns as soon as its context is cancelled and does not join the poller, so without this the +// two race: Shutdown could reach the meter provider's final collect first and export a backlog +// reading for a process that is no longer the leader. Waiting also lets the leader's advisory-lock +// release and idle-timeout reset finish, which previously raced process exit with nothing awaiting +// them. +// +// On deadline this gives up rather than hanging shutdown past its budget -- the stale sample is a +// cosmetic gauge artifact, and a wedged shutdown is not. Postgres releases the session lock when the +// socket closes either way. +func (a *App) awaitEnrichmentBacklogPoller(ctx context.Context) { + if a.enrichmentBacklogDone == nil { + return + } + + select { + case <-a.enrichmentBacklogDone: + case <-ctx.Done(): + slog.Warn("enrichment backlog poller did not finish before the shutdown deadline; " + + "its final gauge reading may still be exported") + } +} diff --git a/cmd/api/app_test.go b/cmd/api/app_test.go index 53bc9183..c0afb502 100644 --- a/cmd/api/app_test.go +++ b/cmd/api/app_test.go @@ -7,13 +7,16 @@ import ( "net/http/httptest" "os" "path/filepath" + "slices" "strings" + "sync" "testing" "time" "github.com/riverqueue/river" "github.com/riverqueue/river/riverdriver/riverpgxv5" sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" sdktrace "go.opentelemetry.io/otel/sdk/trace" "github.com/formbricks/hub/internal/api/handlers" @@ -182,6 +185,159 @@ func TestShutdownObservabilityWithProviders(t *testing.T) { } } +// lifecycleRecorder orders the shutdown steps a test cares about. Both the exporter below and the +// test goroutine append to it, so it has to be mutex-guarded. +type lifecycleRecorder struct { + mu sync.Mutex + events []string +} + +func (r *lifecycleRecorder) record(event string) { + r.mu.Lock() + defer r.mu.Unlock() + + r.events = append(r.events, event) +} + +func (r *lifecycleRecorder) snapshot() []string { + r.mu.Lock() + defer r.mu.Unlock() + + return slices.Clone(r.events) +} + +// recordingExporter notes when the meter provider tears down. Its Shutdown is the step that must not +// happen before the backlog poller has withdrawn its gauge; Export is recorded too so the final +// collect is visible when instruments carry data. +type recordingExporter struct { + recorder *lifecycleRecorder +} + +func (e recordingExporter) Temporality(kind sdkmetric.InstrumentKind) metricdata.Temporality { + return sdkmetric.DefaultTemporalitySelector(kind) +} + +func (e recordingExporter) Aggregation(kind sdkmetric.InstrumentKind) sdkmetric.Aggregation { + return sdkmetric.DefaultAggregationSelector(kind) +} + +func (e recordingExporter) Export(_ context.Context, _ *metricdata.ResourceMetrics) error { + e.recorder.record("metrics-exported") + + return nil +} + +func (e recordingExporter) ForceFlush(context.Context) error { return nil } + +func (e recordingExporter) Shutdown(context.Context) error { + e.recorder.record("metrics-shutdown") + + return nil +} + +// TestShutdownWaitsForEnrichmentBacklogPoller pins the lifecycle ordering that makes the poller's +// deferred ClearEnrichmentPending observable at all. +// +// Run returns the moment its context is cancelled without joining the poller, so Shutdown used to +// race it: the meter provider's final collect-and-export could publish a backlog reading for a +// process that had already lost leadership -- the exact stale series the clear exists to prevent. +// The poller is slowed here so an unsynchronized Shutdown reliably loses the race. Verified by +// deleting the awaitEnrichmentBacklogPoller defer: the recorded events become +// [metrics-exported metrics-shutdown], i.e. the gauge is exported and the provider torn down before +// the clear runs at all. +func TestShutdownWaitsForEnrichmentBacklogPoller(t *testing.T) { + recorder := &lifecycleRecorder{} + + riverClient, err := river.NewClient(riverpgxv5.New(nil), &river.Config{}) + if err != nil { + t.Fatalf("river.NewClient() error = %v, want nil", err) + } + + done := make(chan struct{}) + + app := &App{ + cfg: &config.Config{Server: config.ServerConfig{Port: "0"}}, + server: &http.Server{Addr: "127.0.0.1:0", ReadHeaderTimeout: time.Second}, + river: riverClient, + message: service.NewMessagePublisherManager(1, time.Second, nil), + meterProvider: sdkmetric.NewMeterProvider( + sdkmetric.WithReader(sdkmetric.NewPeriodicReader(recordingExporter{recorder: recorder})), + ), + enrichmentBacklogDone: done, + } + + // Stands in for the poller returning: it withdraws the gauge, then signals. The delay is what + // makes the ordering assertion meaningful rather than incidentally true. + go func() { + time.Sleep(50 * time.Millisecond) + recorder.record("backlog-cleared") + close(done) + }() + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := app.Shutdown(shutdownCtx); err != nil { + t.Fatalf("Shutdown() error = %v, want nil", err) + } + + events := recorder.snapshot() + + clearedAt := slices.Index(events, "backlog-cleared") + if clearedAt < 0 { + t.Fatalf("Shutdown() events = %v, want the backlog clear to have happened", events) + } + + shutdownAt := slices.Index(events, "metrics-shutdown") + if shutdownAt < 0 { + t.Fatalf("Shutdown() events = %v, want the meter provider to have been shut down", events) + } + + if clearedAt > shutdownAt { + t.Fatalf("Shutdown() events = %v, want the backlog clear before the meter provider teardown", events) + } +} + +// TestShutdownWithoutEnrichmentBacklogPollerDoesNotBlock covers the metrics-disabled path, where the +// poller never started and there is nothing to join. +func TestShutdownWithoutEnrichmentBacklogPollerDoesNotBlock(t *testing.T) { + riverClient, err := river.NewClient(riverpgxv5.New(nil), &river.Config{}) + if err != nil { + t.Fatalf("river.NewClient() error = %v, want nil", err) + } + + app := &App{ + cfg: &config.Config{Server: config.ServerConfig{Port: "0"}}, + server: &http.Server{Addr: "127.0.0.1:0", ReadHeaderTimeout: time.Second}, + river: riverClient, + message: service.NewMessagePublisherManager(1, time.Second, nil), + } + + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if err := app.Shutdown(shutdownCtx); err != nil { + t.Fatalf("Shutdown() error = %v, want nil", err) + } +} + +// TestAwaitEnrichmentBacklogPollerGivesUpOnDeadline pins the bound: a poller stuck in its leader +// cleanup must not hold shutdown past its budget. +func TestAwaitEnrichmentBacklogPollerGivesUpOnDeadline(t *testing.T) { + app := &App{enrichmentBacklogDone: make(chan struct{})} // never closed + + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + + start := time.Now() + + app.awaitEnrichmentBacklogPoller(ctx) + + if elapsed := time.Since(start); elapsed > time.Second { + t.Fatalf("awaitEnrichmentBacklogPoller() blocked for %v, want it to give up on the deadline", elapsed) + } +} + func TestAppRunReturnsServerError(t *testing.T) { app := &App{ cfg: &config.Config{