fix: reset leader session state and withdraw backlog gauge on shutdown (ENG-1670) - #116
Conversation
WalkthroughThe enrichment backlog poller now withdraws its pending gauge when it exits, and related logs, documentation, and tests describe the missing-series behavior. Leader release resets PostgreSQL session timeout state before returning connections to the pool, with a test covering reuse. Migration 020 updates the emotions backfill partial index to require both null classification fields and restores the previous predicate on rollback. Repository comments were aligned with the revised index behavior. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/repository/enrichment_status_repository.go`:
- Around line 288-294: Update the connection cleanup in release to use a fresh
deadline/context for resetLeaderIdleTimeoutSQL rather than reusing unlockCtx.
Check errors from both the unlock and RESET executions; if either fails or
cannot complete, close/discard the connection instead of calling l.conn.Release.
When cleanup succeeds, propagate any unlock failure as required by the existing
release contract.
In `@tests/enrichment_status_test.go`:
- Around line 367-373: In the test around NewEnrichmentBacklogLeader, register
defer leader.Close(ctx) immediately after constructing the leader so cleanup
runs if a later assertion fails. Retain the existing explicit leader.Close(ctx)
before the reuse assertion.
- Around line 369-378: Update the leadership test around tryAcquire and
leader.CountIfLeader to query and store the connection’s original
idle_session_timeout before acquisition. Assert acquisition changes it to
"30min", then close the leader and assert the borrowed connection restores the
captured original value instead of hard-coding "0".
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d7dd9ccf-3335-44ae-b660-05fbdd082088
📒 Files selected for processing (7)
cmd/api/app.gointernal/observability/enrichment_backlog.gointernal/observability/enrichment_backlog_test.gointernal/repository/enrichment_status_repository.gointernal/repository/feedback_records_repository.gomigrations/020_add_emotions_classified_at.sqltests/enrichment_status_test.go
…ool (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 ...".
BhagyaAmarasinghe
left a comment
There was a problem hiding this comment.
Requesting changes for the shutdown-ordering issue below. The remaining reviewed paths and validation were clean.
…(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.
…ool (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 ...".
… (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.
cfdcd76 to
0aba386
Compare
What does this PR do?
Follow-up to #111 (ENG-1670), which merged before a last round of review comments were addressed. Six fixes, one with real (if bounded) production impact.
The one that matters — the backlog poller's leader leaked a session setting onto a pooled connection. To bound how long a leader lost to a node failure or network partition can hold its advisory lock, the leader runs
SET idle_session_timeout = '30min'on its connection. ButSETis session state, and pgxpool hands that exact backend to unrelated callers once the leader releases it — verified with aMaxConns=1pool: same backend pid, setting still in effect,reset_val = 0. Postgres would then terminate whichever component next borrowed that connection and left it idle for 30 minutes.release()nowRESETs it before returning the connection. Pinned by a test that fails withactual: "30min"if the reset is removed.Impact today is bounded rather than severe:
pgxpooldestroys connections pastMaxConnLifetime(1h) and pings anything idle >1s on acquire, so the realistic worst case is one extra connection churn — and there is currently no long-idle pooled consumer to inherit the setting, because the API's River client is constructed but neverStarted. Starting it (or adding any consumer that holds a pooled connection idle >30min outside a transaction) would make this bite, so it is worth fixing now rather than later.The rest:
Closereleased the advisory lock while the values stayed, so the meter provider's final collect-and-export published a backlog this process no longer owned. Added the matchingdefer.feedback_records_repository.goclaiming the marker is a post-filter on the 016 index (migration 020 realigned the index predicate, so it isn't), and one in the leader test left over from the transaction-scoped lock.Collect.Note on the migration file: this touches
migrations/020_add_emotions_classified_at.sql, which is already applied — the change is comments only, no DDL, so no re-run or checksum concern. Verified: no non-comment line differs.No API, response-shape, OpenAPI or SDK change.
Linear: https://linear.app/formbricks/issue/ENG-1670
How should this be tested?
Automated
make tests—TestEnrichmentBacklogLeaderLeavesNoSessionStateis the load-bearing one: it takes leadership on aMaxConns=1pool, releases, then borrows the same connection as any other caller would and assertscurrent_setting('idle_session_timeout')is back to0. Temporarily removing theRESETmakes it fail withactual: "30min"(checked).TestEnrichmentBacklogMetricsClearWithdrawsSeriesnow counts data points instead of iterating an empty collection.go test ./internal/... ./cmd/...,make migrate-validate,make lint-openapi,golangci-lint— all clean on currentmain(which has moved on since feat: expose per-tenant enrichment status and backlog metric (ENG-1670) #111).Manual
The behaviour was smoke-tested end to end before #111 merged and re-verified after these changes: three API replicas against one database hold exactly one advisory lock (
pg_locks), killing non-leaders leaves it at 1, and killing the leader drops it to 0 immediately via the gracefulCloserelease — no leaked lock. Emotion completion was verified with real providers: records that classify to no emotions are storedemotions=NULL, emotions_classified_at=setand correctly count as done, so the backlog drains to zero.Checklist
Required
make buildmake tests(integration tests intests/)make fmtandmake lint; no new warningsgit pull origin main(branched off currentmain)migrations/020is touched for comments only (verified no DDL diff);make migrate-validatepasses and the up/down round-trip was re-checkedAppreciated
docs/— n/a; the hand-authoredhub-api-docspages for ENG-1670 (metrics reference + core-concepts) remain a separate follow-upmake tests-coveragefor meaningful logic changes