Query Cancellation Design for ponylang/postgres

[SeanTAllen]

Problem

PostgreSQL query cancellation requires sending a CancelRequest on a separate TCP connection from the one executing the query. The current Session architecture uses a single actor with a single TCP connection, so cancellation requires a fundamentally different internal approach than other operations, which all go through the existing connection.

PostgreSQL CancelRequest Protocol

1. During startup, the server sends BackendKeyData with process_id and secret_key (already parsed and stored in _SessionLoggedIn.backend_pid / _SessionLoggedIn.backend_secret_key)
2. To cancel: open a new TCP connection to the same server, send CancelRequest(16, 80877102, process_id, secret_key), close the connection
3. No response on the cancel connection — it is fire-and-forget
4. If the server cancels the query, the original connection receives ErrorResponse with SQLSTATE 57014 (query_canceled), then ReadyForQuery
5. Cancellation is best-effort — the server may ignore it (e.g., if the query finishes before the cancel arrives)
6. The magic number 80877102 = 1234 << 16 | 5678

Design

New Actor: _CancelSender

A small, internal, fire-and-forget actor in a new file _cancel_sender.pony. It:

• Implements lori.TCPConnectionActor & lori.ClientLifecycleEventReceiver
• Opens a TCP connection to the PostgreSQL server
• On _on_connected(): sends the CancelRequest message then closes the connection
• On _on_connection_failure(): silently does nothing (fire-and-forget)
• Self-destructs (garbage collected after the connection closes and no references remain)

actor _CancelSender is (lori.TCPConnectionActor & lori.ClientLifecycleEventReceiver)
  var _tcp_connection: lori.TCPConnection = lori.TCPConnection.none()
  let _process_id: I32
  let _secret_key: I32

  new create(auth: lori.TCPConnectAuth, host: String, service: String,
    process_id: I32, secret_key: I32)
  =>
    _process_id = process_id
    _secret_key = secret_key
    _tcp_connection = lori.TCPConnection.client(
      auth, host, service, "", this, this)

  fun ref _connection(): lori.TCPConnection =>
    _tcp_connection

  fun ref _on_connected() =>
    _tcp_connection.send(_FrontendMessage.cancel_request(
      _process_id, _secret_key))
    _tcp_connection.close()

  fun ref _on_connection_failure() =>
    None

New Frontend Message: cancel_request

Added to the _FrontendMessage primitive:

fun cancel_request(process_id: I32, secret_key: I32): Array[U8] val =>
  """
  Build a CancelRequest message. Sent on a separate connection to request
  cancellation of a running query. No message type byte — same pattern as
  startup() and ssl_request().

  Format: Int32(16) Int32(80877102) Int32(process_id) Int32(secret_key)
  """
  // 16 bytes total: 4 (length) + 4 (cancel code) + 4 (pid) + 4 (key)

Session Changes

New fields on Session actor — immutable connection parameters already passed to the constructor, now stored for use by cancel:

let _auth: lori.TCPConnectAuth
let _host: String
let _service: String

New package-accessible methods on Session (no underscore prefix, so other types in the package can access them):

fun cancel_auth(): lori.TCPConnectAuth => _auth
fun cancel_host(): String => _host
fun cancel_service(): String => _service

New behavior on Session:

be cancel() =>
  """
  Request cancellation of the currently executing query. Opens a separate
  TCP connection to send a PostgreSQL CancelRequest. Cancellation is
  best-effort — the server may or may not honor it. If cancelled, the
  query's ResultReceiver receives `pg_query_failed` with an ErrorResponse
  (SQLSTATE 57014). Queued queries are not affected.

  Safe to call in any session state. No-op if no query is in flight.
  """
  state.cancel(this)

State Machine Changes

_SessionState interface — add:

fun ref cancel(s: Session ref)
  """
  The client requested query cancellation. Like `close`, this should never
  be an illegal state — it should be silently ignored when not applicable.
  """

Default implementations via trait hierarchy:

• _UnconnectedState: add cancel as no-op — covers _SessionUnopened, _SessionClosed
• _ConnectedState: add cancel as no-op — covers _SessionConnected (via _AuthenticableState) and provides default for _SessionLoggedIn (via _AuthenticatedState, which the class overrides)
• _SessionSSLNegotiating: add explicit cancel as no-op (standalone class, doesn't mix in the above)

_SessionLoggedIn.cancel() — the real implementation:

fun ref cancel(s: Session ref) =>
  match query_state
  | let _: _QueryReady => None
  | let _: _QueryNotReady => None
  else
    _CancelSender(s.cancel_auth(), s.cancel_host(), s.cancel_service(),
      backend_pid, backend_secret_key)
  end

The match checks whether a query is actually in flight. _QueryReady (idle) and _QueryNotReady (waiting for readiness) mean no query is executing on the server. Any other _QueryState (_SimpleQueryInFlight, _ExtendedQueryInFlight, _PrepareInFlight, _CloseStatementInFlight) means something is in flight.

Note on timing: due to Pony's actor model, cancel() is processed asynchronously. By the time it runs, the query may have already completed and the state may have transitioned. This is fine — a stale CancelRequest is harmless (the server ignores cancels for processes that aren't running a query). The check is a best-effort optimization to avoid unnecessary TCP connections.

What This Design Does NOT Change

• Error handling: Cancellation errors flow through the existing on_error_response path. The ResultReceiver gets pg_query_failed(session, query, ErrorResponseMessage) where the ErrorResponseMessage has SQLSTATE 57014. No new error types are needed — the server's ErrorResponse already fully describes the cancellation.
• Queue behavior: cancel() only targets the in-flight query on the server. Queued queries remain and execute normally after cancellation completes.
• State transitions: After cancellation, ReadyForQuery arrives normally on the original connection and triggers the usual transition to _QueryReady, which dequeues the next query.

SSL Considerations

Initial implementation: _CancelSender uses a plain TCP connection regardless of whether the original session uses SSL. CancelRequest contains no sensitive data (just pid + key that were already exchanged over the secure channel), so the security risk is minimal.

Limitation: If the PostgreSQL server is configured with hostssl entries that reject all non-SSL connections, the cancel connection will be silently refused. This is documented.

Future enhancement: Add SSL support to _CancelSender by:

1. Passing SSLMode from the original session to _CancelSender
2. If SSLRequired: send SSLRequest, wait for 'S', TLS handshake, then CancelRequest
3. This would make _CancelSender a mini state machine rather than purely fire-and-forget
4. The same SSLContext val from the original session can be reused (it's val, freely shareable)

This SSL enhancement could be a follow-up PR. The non-SSL cancel covers the common deployment case (mixed SSL/non-SSL acceptance, or non-SSL deployments).

File Changes

File	Change
_cancel_sender.pony	New file: _CancelSender actor
_frontend_message.pony	Add cancel_request(process_id, secret_key) method
session.pony	Store _auth, _host, _service; add cancel() behavior; add cancel to _SessionState interface and trait defaults; implement _SessionLoggedIn.cancel()
_test_frontend_message.pony	Unit test: verify cancel_request() wire format
_test.pony	Unit tests: (a) mock server verifies CancelRequest is sent on separate connection when query in flight; (b) cancel when logged in but idle (no query in flight) does not open a connection; (c) cancel in non-applicable states (unopened, closed, etc.) is no-op
_test_query.pony	Integration test: cancel a pg_sleep() query, verify pg_query_failed with SQLSTATE 57014
examples/cancel/cancel-example.pony	New example: execute a long-running query and cancel it, showing the cancellation error flow
CLAUDE.md	Update: (1) add cancel() to Session behaviors in Query Execution Flow section; (2) add _CancelSender and cancel_request() to Architecture/Protocol sections; (3) add _cancel_sender.pony and examples/cancel/ to File Layout; (4) add cancel-related tests to Test Organization

Building and Running Tests

make unit-tests ssl=3.0.x             # wire format + mock server cancel tests
make integration-tests ssl=3.0.x      # pg_sleep cancel test (requires containers)
make build-examples ssl=3.0.x         # verify cancel example compiles

Design Decisions and Rationale

1. Cancel on Session, not standalone: Cancel is a method on Session because Session already holds the backend key data and connection parameters. A standalone API would require exposing backend_pid/backend_secret_key publicly, leaking internal protocol details.

2. No new error types: The server already provides a fully descriptive ErrorResponse for cancellation (SQLSTATE 57014). Adding a client-side QueryCancelled type would be redundant.

3. No queue draining: cancel() only targets the server-side in-flight query. Cancelling queued-but-not-sent queries is a different concern (queue management, not PostgreSQL cancellation protocol). If needed, a separate clear_queue() API could be added independently.

4. Fire-and-forget _CancelSender: No callback or notification about whether the cancel connection succeeded. This matches the PostgreSQL protocol semantics — the result is always delivered via the original connection's ErrorResponse/CommandComplete flow.

5. Check before sending: Although a stale CancelRequest is harmless, we check query_state to avoid creating an unnecessary TCP connection + actor when we already know the server is idle.

[SeanTAllen]

Two design decisions made during review:

1. Accessor naming: The Session accessor methods should be named for what they are, not what they're used for. Use auth(), host(), service() instead of cancel_auth(), cancel_host(), cancel_service(). These are the session's general connection parameters — naming them after their current consumer (cancel) would create false coupling.

2. SSL support included, phased implementation: The initial design deferred SSL on the cancel connection. Decision: SSL support will be included in the same PR, but implemented in phases — get non-SSL cancel working first as a checkpoint, then add SSL support to _CancelSender. The cancel connection should use the same SSLMode as the original session.

[SeanTAllen]

Missing _on_received on _CancelSender: every actor implementing lori.TCPConnectionActor in the codebase (Session, _JunkSendingTestServer, _DoesntAnswerTestServer, _TerminateSentTestServer, etc.) implements _on_received — lori requires it. The _CancelSender code snippet should include it as a no-op:

fun ref _on_received(data: Array[U8] iso) =>
  None

[SeanTAllen]

SSL support needs to be part of the main design: the comment decision says SSL is included in the same PR, but the main design body doesn't account for this — SSLMode isn't stored on Session, and _CancelSender doesn't receive it. Since this is now an in-scope requirement, the design should show:

• _ssl_mode: SSLMode field on Session (stored from constructor)
• ssl_mode() accessor on Session
• _CancelSender receiving SSLMode and, for SSLRequired, performing the SSLRequest/'S'/TLS handshake before sending CancelRequest
• This makes _CancelSender a mini state machine (two states: waiting for SSL response vs sending cancel) — worth sketching even if the implementation is phased

Related: the SSL-capable _CancelSender would also need _on_tls_ready/_on_tls_failure callbacks from lori (similar to _SessionSSLNegotiating). The non-SSL version doesn't need them, but since SSL is in scope, the design should address which lori callbacks are needed and how they're handled.

[SeanTAllen]

Implementation Plan

Phased implementation plan for query cancellation, incorporating the two design decisions from previous comments (generic accessor names, SSL included). This is "two PRs in one" — Phase 1 (non-SSL cancel) is implemented and reviewed at a checkpoint before Phase 2 (SSL cancel) is added on top. Final PR is a single squashed commit.

Prerequisites

Create feature branch cancel-query from main. This is an added feature (changelog - added label).

Phase 1: Non-SSL Cancel

Step 1 — _FrontendMessage.cancel_request(): Add to _frontend_message.pony. Format: Int32(16) Int32(80877102) Int32(process_id) Int32(secret_key) — no message type byte, same pattern as startup() and ssl_request(). 16 bytes total.

Step 2 — Session fields and accessors: Store _auth: lori.TCPConnectAuth, _host: String, _service: String on the Session actor from constructor params. Add package-accessible accessors: auth(), host(), service().

Step 3 — _CancelSender actor (non-SSL): New file _cancel_sender.pony. Implements lori.TCPConnectionActor & lori.ClientLifecycleEventReceiver. Fields: _tcp_connection, _process_id, _secret_key, _host (stored now for Phase 2 SSL use). Explicitly implements: _connection(), _on_connected() (calls _send_cancel_and_close()), _on_received() (no-op — not expected in non-SSL path but every TCPConnectionActor in the codebase implements it explicitly). All remaining lori callbacks use trait-provided no-op defaults. Helper _send_cancel_and_close() sends cancel_request then closes.

Step 4 — State machine: Add cancel(s: Session ref) to _SessionState interface with docstring noting the "never illegal" contract (like close). Default no-op on _UnconnectedState, _ConnectedState, _SessionSSLNegotiating. _SessionLoggedIn.cancel() matches on query_state — no-op for _QueryReady/_QueryNotReady, creates _CancelSender for any in-flight state (including _PrepareInFlight and _CloseStatementInFlight — intentional, server may be doing real work).

Step 5 — Session cancel() behavior: New be cancel() on Session with docstring. Delegates to state.cancel(this).

Step 6 — Tests:

• _TestFrontendMessageCancelRequest — wire format verification (in _test_frontend_message.pony)
• _TestCancelQueryInFlight (port 7675) — mock server accepts TWO connections. First: AuthOk + BackendKeyData(pid=12345, key=67890) + ReadyForQuery, receives query, holds. Client executes query then calls cancel. Second connection (cancel sender): verify 16-byte CancelRequest with correct magic/pid/key. Uses _IncomingBackendKeyDataTestMessage builder.
• _TestCancelPgSleep (integration, integration/Cancel/Query) — execute SELECT pg_sleep(30), cancel, verify pg_query_failed with SQLSTATE 57014

Step 7 — Example: examples/cancel/cancel-example.pony — execute pg_sleep(10), cancel, show error flow.

Step 8 — CLAUDE.md: Update architecture, protocol, file layout, test organization sections.

Checkpoint: Squash, push, run reviewer subagent. Fix findings. Iterate until clean.

---

Phase 2: SSL Cancel

Step 9 — SSL mode on Session: Add _ssl_mode: SSLMode field, store from constructor. Add ssl_mode() accessor.

Step 10 — Upgrade _CancelSender for SSL: Add _ssl_mode field (_host already stored from Phase 1). Updated callbacks:

• _on_connected(): match on _ssl_mode — SSLDisabled calls _send_cancel_and_close(), SSLRequired sets expect(1) (CVE-2021-23222) and sends SSLRequest
• _on_received(): handles SSL response byte — 'S' triggers start_tls(), anything else closes (fire-and-forget)
• _on_tls_ready(): resets expect(0), calls _send_cancel_and_close()
• _on_tls_failure(): no-op (fire-and-forget, lori follows with _on_closed default)

Update _SessionLoggedIn.cancel() to pass s.ssl_mode().

Step 11 — SSL cancel tests:

• _TestCancelSSLQueryInFlight (port 7676) — same mock pattern as Phase 1 but with SSL on both connections, using test certificates
• _TestCancelSSLPgSleep (integration, integration/SSL/Cancel) — cancel over SSL connection

Step 12 — CLAUDE.md: Document SSL support in _CancelSender.

---

Final Steps

Reviewer subagent against all changes. Fix findings. Squash all commits into single commit. Open PR with changelog - added label.

Verification

make unit-tests ssl=3.0.x
make start-pg-containers
make integration-tests ssl=3.0.x
make build-examples ssl=3.0.x

[SeanTAllen]

Implemented in PR #89 (merged). Both non-SSL and SSL cancel paths are complete.