Lock backend to frontend when in progress cancel requests are in flight to prevent cancel race between two frontends

[KennanHunter]

The issue

Fixes #1505, a cancel request can be made by the frontend, then before pgdog's cancel packet actually lands at Postgres, the physical backend that frontend was using can be returned to the pool and handed to a different frontend. The cancel then lands on the new frontend's query instead of the original one.

Concretely, in Pool::cancel:

let key = self.lock().cancel_key(id).cloned();   // phase 1: snapshot under lock
if let Some(key) = key {
    Server::cancel(self.addr(), key).await?;      // phase 2: TCP send, no lock
}

The fix

Taken picks up a backend_cancels_in_flight map (backend pid → count). Pool::cancel takes a CancelLease, a small RAII guard in pool/cancel.rs, that bumps the count going in and drops it on the way out whether the cancel succeeded or not. Pool::checkin looks at the map: if there's a pending cancel for that backend, the check-in waits until the count hits zero, otherwise it's the exact same sync path as today. I went with a count instead of a bool so a couple of cancels landing on the same backend at once just work. Pool::move_conns_to bails with a retryable Error::CancelInFlight if anything's on the wire.

Server::cancel now hangs on until Postgres sends EOF. Postgres only does that after it's validated the request and poked the target backend, so that's our signal the cancel is really done. There's a 5s cap so a wedged server can't pin backends forever.

Cost-wise, only waiters on that one backend notice if a cancel is in flight, the rest of the pool's connections are unblocked.

Tests: test_cancel_request_cross_frontend_race (A starts a pg_sleep, cancel fires, B grabs the backend and runs its own pg_sleep, make sure B doesn't get 57014), cancel_leases_stack_on_same_backend, and the repro script from #1505 as integration/python/test_cancel_race.py.

Next steps

Style

Theres some duplication in logic between locking due to cancel requests, and locking due to advisory locks. I think for right now it's fine, but if we see more locking behaviors it might be a good idea to merge those in some way.

Could look something like this, with a single queue for locked connections on a pool side.

#[derive(Default)]
struct Pin {
    client_locked: bool,
    cancels_in_flight: u32,
}

pub(super) struct Taken {
    frontend_to_cancel: HashMap<FrontendPid, Checkout>,
    backend_to_frontend: HashMap<BackendPid, FrontendPid>,
    pins: HashMap<BackendPid, Pin>,
}

Scoping

An interesting idea longer-term is that the whole lease mechanism could disappear if pgdog tagged every forwarded statement with a /* pgdog:<FrontendPid> */ comment and routed cancels through a helper connection running SELECT pg_cancel_backend(pid) FROM pg_stat_activity WHERE pid = $1 AND state = 'active' AND query LIKE '/* pgdog:F<n> */%'

The problem is it gets pretty complicated injecting the tag cleanly across every query, especially when we consider mechanisms like PREPARE, EXECUTE, and COPY.

Benchmarking

I think we should probably benchmark this on a cancel heavy workflow with actual network delay between frontend/backend, to verify that nothing goes haywire. Lmk if there's any changes you'd like to see.

[codecov]

Codecov Report

❌ Patch coverage is 98.32536% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
pgdog/src/backend/pool/cancel.rs	92.50%	3 Missing ⚠️
pgdog/src/backend/server.rs	96.05%	3 Missing ⚠️
pgdog-stats/src/pool.rs	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[KennanHunter]

Might want CANCEL_ACK_TIMEOUT to be configurable, not sure.