pgconn: add CancelAndDrainContextWatcherHandler

Add a context watcher handler that deterministically drains stale cancel
signals instead of relying on a fixed 100ms sleep. When a context is
cancelled, HandleCancel sends the cancel request immediately. After the
query returns, HandleUnwatchAfterCancel polls SELECT 1 to absorb any
in-flight SQLSTATE 57014 before releasing the connection.

The drain runs inside Unwatch via HandleUnwatchAfterCancel, so every
pgconn operation that watches a context (Exec, ExecParams, Prepare,
Deallocate, CopyTo, CopyFrom, Pipeline, WaitForNotification) gets
drain-after-cancel automatically with no per-call-site wiring.

Author: ecordell

pgconn/cancel_and_drain.go

@@ -0,0 +1,138 @@
+package pgconn
+
+import (
+	"context"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
+	"github.com/jackc/pgx/v5/pgproto3"
+)
+
+// CancelAndDrainContextWatcherHandler handles cancelled contexts by sending a cancel request to the server and then
+// draining any in-flight SQLSTATE 57014 (query_canceled) by polling SELECT 1. Unlike [CancelRequestContextWatcherHandler],
+// no fixed sleep is used; the drain is deterministic.
+type CancelAndDrainContextWatcherHandler struct {
+	Conn *PgConn
+
+	// DeadlineDelay is the network deadline set on the connection when the context
+	// is cancelled, used as a fallback to unblock any blocked read. Defaults to 1s.
+	DeadlineDelay time.Duration
+
+	// DrainTimeout is the maximum time to spend draining a cancelled query's
+	// in-flight results via SELECT 1 polling. Defaults to 5s.
+	DrainTimeout time.Duration
+
+	cancelFinishedChan chan struct{}
+	stopFn             context.CancelFunc
+}
+
+var _ ctxwatch.Handler = (*CancelAndDrainContextWatcherHandler)(nil)
+
+func (h *CancelAndDrainContextWatcherHandler) deadlineDelay() time.Duration {
+	if h.DeadlineDelay == 0 {
+		return time.Second
+	}
+	return h.DeadlineDelay
+}
+
+func (h *CancelAndDrainContextWatcherHandler) drainTimeout() time.Duration {
+	if h.DrainTimeout == 0 {
+		return 5 * time.Second
+	}
+	return h.DrainTimeout
+}
+
+// HandleCancel is called when the context is cancelled. It sets a net.Conn deadline
+// as a fallback and sends a PostgreSQL cancel request in a goroutine.
+func (h *CancelAndDrainContextWatcherHandler) HandleCancel(_ context.Context) {
+	h.cancelFinishedChan = make(chan struct{})
+	cancelCtx, stop := context.WithCancel(context.Background())
+	h.stopFn = stop
+
+	deadline := time.Now().Add(h.deadlineDelay())
+	h.Conn.conn.SetDeadline(deadline)
+
+	doneCh := h.cancelFinishedChan
+	go func() {
+		defer close(doneCh)
+		reqCtx, cancel := context.WithDeadline(cancelCtx, deadline)
+		defer cancel()
+		h.Conn.CancelRequest(reqCtx)
+	}()
+}
+
+// HandleUnwatchAfterCancel is called after the cancelled query returns. It stops the cancel goroutine (if still
+// running), clears the net.Conn deadline, and drains any in-flight cancel by polling SELECT 1.
+func (h *CancelAndDrainContextWatcherHandler) HandleUnwatchAfterCancel() {
+	if h.stopFn != nil {
+		h.stopFn()
+	}
+	if h.cancelFinishedChan != nil {
+		<-h.cancelFinishedChan
+	}
+	h.Conn.conn.SetDeadline(time.Time{})
+	h.cancelFinishedChan = nil
+	h.stopFn = nil
+
+	if !h.Conn.IsClosed() {
+		ctx, cancel := context.WithTimeout(context.Background(), h.drainTimeout())
+		defer cancel()
+		h.Conn.execInternalForDrain(ctx)
+	}
+}
+
+// queryCanceledSQLStateCode is SQLSTATE 57014 (query_canceled).
+const queryCanceledSQLStateCode = "57014"
+
+// execInternalForDrain sends an empty query (";") in a loop, absorbing any
+// SQLSTATE 57014 responses, until the connection is confirmed clean or a
+// non-57014 error occurs. On any failure the connection is asyncClosed.
+//
+// Called while the connection is still logically "busy" from pgconn's perspective
+// (lock is held and contextWatcher.Unwatch has been called) but idle from the
+// PostgreSQL server's perspective (ReadyForQuery was just received). This means
+// it bypasses the normal lock/unlock and contextWatcher.Watch paths.
+//
+// The deadline from ctx is applied directly to the net.Conn.
+func (pgConn *PgConn) execInternalForDrain(ctx context.Context) {
+	if deadline, ok := ctx.Deadline(); ok {
+		pgConn.conn.SetDeadline(deadline)
+		defer pgConn.conn.SetDeadline(time.Time{})
+	}
+
+outer:
+	for {
+		pgConn.frontend.Send(&pgproto3.Query{String: ";"})
+		if err := pgConn.frontend.Flush(); err != nil {
+			pgConn.asyncClose()
+			return
+		}
+		clean := true
+		for {
+			msg, err := pgConn.receiveMessage()
+			if err != nil {
+				pgConn.asyncClose()
+				return
+			}
+
+			switch msg := msg.(type) {
+			case *pgproto3.ReadyForQuery:
+				if !clean {
+					clean = true
+					continue outer // absorbed 57014; send another to confirm clean
+				}
+				return // clean ReadyForQuery — done
+			case *pgproto3.ErrorResponse:
+				pgErr := ErrorResponseToPgError(msg)
+				if pgErr.Code == queryCanceledSQLStateCode {
+					clean = false // cancel hit this query; will confirm after ReadyForQuery
+				} else {
+					pgConn.asyncClose()
+					return
+				}
+			case *pgproto3.EmptyQueryResponse:
+				// Expected response for ";".
+			}
+		}
+	}
+}