mirror of
https://github.com/coder/coder.git
synced 2026-09-01 14:53:15 +08:00
J. Scott Miller
866e676320
## Summary Closes PLAT-305. When a provisioner key is deleted, the associated daemon kept operating on its existing WebSocket connection, because authentication was only checked at connection establishment and deletion was a bare `DELETE` with no session invalidation. This adds four layers of defense so a deleted key promptly stops doing work: 1. **Publish on delete.** `deleteProvisionerKey` publishes to a new per-key pubsub channel (`coderd/pubsub.ProvisionerKeyDeletedChannel`) after a successful delete. Publish errors are logged but still return `204`, since layer 3 is the durable backstop. 2. **Subscribe and tear down.** The daemon serve handler subscribes to its key's channel and terminates the DRPC session on a deletion event. Termination is deferred while a job claimed by the session is active: the daemon may finish and report the in-flight job (`UpdateJob`/`CompleteJob` have no key check), and the last active job's completion performs the cancellation. Because Postgres `LISTEN`/`NOTIFY` does not buffer for non-listeners, the handler also performs a synchronous key-existence re-check immediately after subscribing to close the race between auth and subscription. The subscription uses `SubscribeWithErr` so that an `ErrDroppedMessages` signal (emitted when the pubsub listener reconnects) triggers the same key re-check, closing the listener-outage window in which a deletion notification could be missed. 3. **Backstop on acquire.** `AcquireJob` and `AcquireJobWithCancel` verify the key still exists before waiting for a job, and the `Acquirer` claims jobs in a transaction that first locks the worker's deletable key (`LockProvisionerKeyByIDForShare`, a `FOR KEY SHARE` row lock held until commit) before running the `AcquireProvisionerJob` claim, so a claim cannot commit after the key's deletion. This guards against a missed pubsub message. A missing key row surfaces as its own result rather than overloading the claim query's no-rows response: the acquire terminates with `ErrProvisionerKeyDeleted` (terminating the session, with the same active-job deferral) and hands the consumed wakeup to another waiting daemon in the same domain, rather than silently re-parking and starving peers of job postings. 4. **Heartbeat watchdog.** The per-session heartbeat loop (1m interval) also re-checks the key, so even a session whose deletion notification was silently lost terminates within one heartbeat interval instead of living until the connection breaks (same active-job deferral as layer 2). Reserved keys skip the check. A job that is claimed but never delivered (the session or connection dies between the database claim and the stream send) is marked failed immediately on a fresh context, instead of staying assigned to the worker until the job reaper. Reserved keys (built-in, user-auth, PSK) are exempt throughout, since they are not deletable rows. The acquire-time lookup runs as `dbauthz.AsSystemReadProvisionerDaemons`, because the provisionerd role cannot read provisioner keys and a provisioner key's RBAC object is a provisioner daemon. A single key can back many daemons (and span HA replicas), so the per-key channel fans out to invalidate all of them at once. Per-key channels keep the `LISTEN` count proportional to distinct keys rather than waking every daemon on unrelated deletions. ### Known limitations - **`UpdateJob`/`CompleteJob` intentionally have no key check.** By the time those RPCs arrive the work has already run; rejecting completion would strand a build in "running" (until the job reaper fails it) with real infrastructure left unreconciled. Session termination is deferred while a job is active so the completion can be reported; the daemon may not receive the final RPC response when the deferred termination fires, but the job's outcome is already persisted. - **After termination, the daemon process redials and receives 401s until restarted.** The dial-time exit logic only triggers on 403, and the auth middleware returns 401 for an invalid key; this dial behavior predates this PR and is tracked as a follow-up in [PLAT-452](https://linear.app/codercom/issue/PLAT-452) (return 403 for invalid provisioner keys). ## Tests - `coderd/provisionerdserver`: `TestAcquireJob_ProvisionerKeyDeleted` (both RPC variants), `TestAcquireJob_ReservedProvisionerKey`, `TestHeartbeat_ProvisionerKeyDeleted` (heartbeat watchdog cancels the session after key deletion), `TestAcquirer_ProvisionerKeyDeleted` (a dead-key acquiree exits terminally and its clearance is promoted to a peer in the same domain), and `TestTerminateSession_Deferral` (termination is immediate when idle and deferred until the last active job finishes). - `coderd/database`: `TestAcquireProvisionerJob/ProvisionerKeyLock` covers the lock query against real Postgres: it returns the key ID while the row exists and no rows once it is deleted. The lock-then-claim composition is pinned by `TestAcquirer_ProvisionerKeyDeleted`. - `enterprise/coderd`: `TestProvisionerDaemonServe/KeyDeletionClosesSession` asserts an active session closes after its key is deleted. `KeyDeletedDuringSetupClosesSession` covers the post-subscribe re-check when a key is deleted between auth and subscription, and `DroppedMessageClosesSession` covers the `ErrDroppedMessages` re-check when a deletion is missed during a listener outage. ## Validation - `make` pre-commit (gen/fmt/lint/build) passed via git hooks. - Targeted tests pass; existing acquire tests pass with no regression. - Manual: brought up a dev deployment (coder-in-coder) with a Premium license, created a deletable provisioner key, and started an external daemon with `coder provisionerd start`. Confirmed it authenticated via the key and connected, appearing as `idle` in both `coder provisioner list` (with the key name) and the organization Provisioners UI. - Manual, idle teardown: deleted the key while the daemon was idle. The server logged `provisioner key deleted, terminating session`, the daemon's session closed immediately, and it dropped from `coder provisioner list` (then entered the known 401 redial loop, PLAT-452). - Manual, deferred termination: ran a workspace build (tagged template, `sleep 45` in `local-exec`) pinned to the external daemon and deleted the key mid-build. The server logged `deferring session cancellation until active jobs finish`; the heartbeat watchdog re-checked mid-build and re-deferred rather than force-killing. The build ran to completion (`Apply complete`, workspace `Started`) and only then did `canceling session after job completion` fire. The documented caveat reproduced: the daemon lost the final `CompleteJob` ack, and the build outcome was still persisted correctly. <details> <summary>Implementation plan and design decisions</summary> ### Design - **Per-key vs global channel:** chose per-key (`provisioner_key_deleted:<keyID>`) so daemons do not wake on unrelated deletions. The cost is one `LISTEN` per distinct key per replica on the shared listener connection, which is negligible against Coder's existing channels. - **Missing-key behavior on acquire:** returns an error that tears down the acquire rather than silently returning an empty job. - **Subscribe-startup race:** ordering is `authorize -> UpsertProvisionerDaemon -> Subscribe -> GetProvisionerKeyByID`. The post-subscribe re-check handles a deletion that committed before the `LISTEN` registered (Postgres does not buffer notifications for non-listeners; the in-process buffer only smooths bursts and drops on overflow). - **`NewServer` change:** `KeyID` was added to `provisionerdserver.Options` to avoid a positional signature change across call sites. The in-memory (built-in) daemon leaves it unset and is therefore exempt. ### Files - `coderd/pubsub/provisionerkeydeleted.go` (new) — channel helper. - `enterprise/coderd/provisionerkeys.go` — publish on delete. - `enterprise/coderd/provisionerdaemons.go` — subscribe, re-check, cancel session; pass `KeyID`. - `coderd/provisionerdserver/provisionerdserver.go` — `KeyID` option and acquire-time existence check. </details> --- This pull request was created by Coder Agents on behalf of @jscottmiller.