Fixes https://github.com/coder/internal/issues/1371
## Root causes
Two independent races cause this test to flake at ~2–3/1000:
### 1. Title-generation requests racing with the streaming request
counter
`maybeGenerateChatTitle` fires in a `context.WithoutCancel` goroutine
(line 2130) and makes a **non-streaming** request to the mock OpenAI
handler. The test handler was not filtering by request type, so these
title requests incremented the `requestCount` atomic — throwing off the
coordination logic that uses `requestCount == 1` to identify the first
streaming request and hold it open until shutdown.
**Fix:** Guard the test handler to return a canned response for
non-streaming requests before touching `requestCount`.
### 2. Phantom acquire: `AcquireChat` commits in Postgres but Go sees
`context.Canceled`
During `Close()`, the main loop's `select` can randomly pick
`acquireTicker.C` over `ctx.Done()` (Go spec: when multiple cases are
ready, one is chosen uniformly at random). This calls `processOnce(ctx)`
with an already-canceled context.
In the pq driver, `QueryContext` does **not** check `ctx.Err()` up
front. Instead it calls `watchCancel(ctx)` which spawns a goroutine
monitoring `ctx.Done()`, then sends the query on the existing
connection. When `ctx` is already canceled, a race ensues:
- **pq's watchCancel goroutine** immediately sees `<-done`, opens a
*new* TCP connection to Postgres, and sends a cancel request.
- **The query** is sent concurrently on the existing connection.
Because the `AcquireChat` UPDATE is fast (sub-millisecond, single row
with `SKIP LOCKED`), it often commits before the cancel arrives via the
second connection. Meanwhile in `database/sql`, `initContextClose`
spawns an `awaitDone` goroutine that fires immediately (context is
already canceled), stores `contextDone`, and calls `rs.close(ctx.Err())`
— which races with `Row.Scan` → `rows.Next()`. If `awaitDone` wins,
`Next()` sees `contextDone` is set and returns false, causing Scan to
return `context.Canceled` (or `ErrNoRows`).
**Result:** Postgres committed the UPDATE (chat is now `running` with
serverA's worker ID), but Go sees an error and never spawns a goroutine
to process it. The chat is stuck as `running` with no worker.
If the previous `processChat` cleanup already set the chat back to
`pending`, this phantom acquire flips it back to `running` — which is
exactly what the debug logs showed: after `Close()` returns, the DB
shows `status=running` with serverA's worker ID.
**Fix:** Three guards in `processOnce`:
1. Early `ctx.Err()` check — catches the common case where `select`
picked the ticker after cancellation.
2. `context.WithoutCancel(ctx)` for `AcquireChat` — prevents the pq
`watchCancel` race entirely, ensuring the driver sees the query
result if Postgres executed it.
3. Post-acquire `ctx.Err()` check — if the context was canceled while
`AcquireChat` ran (or between the early check and the call),
immediately release the chat back to `pending`.
## Verification
Passes 2000/2000 iterations (previously flaked at ~2–3/1000):
```
go test -run "TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica" \
-count=2000 -timeout 1800s -failfast ./coderd/chatd/
```
Kyle Carberry
fee5cc5e5b