Query-in-Flight Sub-State Machine Design

[SeanTAllen]

Goal

Replace the boolean flags (_queryable and _query_in_flight) in _SessionLoggedIn with a proper sub-state machine that makes illegal query states unrepresentable and provides a natural extension point for the extended query protocol.

This is a prerequisite for the prepared statements work (discussion #62). With two query protocols (simple and extended), the booleans cannot distinguish which protocol is in flight, making it impossible to validate that the correct acknowledgment messages are arriving for the active protocol.

Current State

_SessionLoggedIn tracks query execution using two boolean flags and two data fields:

var _queryable: Bool = false
var _query_in_flight: Bool = false
var _data_rows: Array[Array[(String|None)] val] iso
var _row_description: (Array[(String, U32)] val | None)

These encode three logical states:

_queryable	_query_in_flight	Meaning
false	false	Server not ready (initial state, failed transaction)
true	false	Idle, ready to send a query
false	true	Query in flight, waiting for responses

The fourth combination (true, true) is never intentionally reached.

The data fields are reset in multiple places: on_command_complete resets both via destructive-read swap; on_empty_query_response resets both via destructive-read; on_error_response resets both explicitly; and on_ready_for_query resets both when dequeuing a completed query. This belt-and-suspenders approach prevents stale data but spreads the cleanup logic across four methods.

Problems

1. Repeated guards: Every query callback has an identical if _query_in_flight then ... else shutdown(s) end guard — five places today (on_command_complete, on_data_row, on_row_description, on_empty_query_response, on_error_response), growing to eleven with prepared statements.

2. No protocol distinction: With both simple and extended query protocols, the boolean cannot distinguish which protocol is in flight. ParseComplete arriving during a simple query should trigger shutdown, but the boolean-guarded approach cannot express this.

3. Double-delivery on failed transactions: When on_ready_for_query arrives with non-idle status during an in-flight query, the current code does not dequeue the query — but its error has already been delivered via on_error_response. If close() is called afterward, on_shutdown iterates the queue and sends SessionClosed to the same receiver, resulting in a double delivery of pg_query_failed.

4. Scattered cleanup: Per-query data (_data_rows, _row_description) lives on _SessionLoggedIn and must be explicitly reset in every command-terminating callback. With the data owned by the in-flight state instead, cleanup is structural — the data is destroyed when the state transitions out.

Design

The sub-state machine follows the same pattern as the outer session state machine: an interface defines the callbacks, a trait provides defaults for states that share behavior, and concrete classes implement state-specific logic. The sub-state is entirely internal to _SessionLoggedIn — no public API changes.

_QueryState Interface

Defines callbacks for all query-related protocol messages, plus try_run_query for attempting to start the next queued query.

interface _QueryState
  fun ref on_ready_for_query(s: Session ref, li: _SessionLoggedIn ref,
    msg: _ReadyForQueryMessage)
  fun ref on_command_complete(s: Session ref, li: _SessionLoggedIn ref,
    msg: _CommandCompleteMessage)
  fun ref on_data_row(s: Session ref, li: _SessionLoggedIn ref,
    msg: _DataRowMessage)
  fun ref on_row_description(s: Session ref, li: _SessionLoggedIn ref,
    msg: _RowDescriptionMessage)
  fun ref on_empty_query_response(s: Session ref, li: _SessionLoggedIn ref)
  fun ref on_error_response(s: Session ref, li: _SessionLoggedIn ref,
    msg: ErrorResponseMessage)
  fun ref try_run_query(s: Session ref, li: _SessionLoggedIn ref)

Query state methods receive both Session ref (for connection operations like send and shutdown) and _SessionLoggedIn ref (for queue access and state transitions). This mirrors the outer state machine where state methods receive Session ref and set s.state.

_QueryNoQueryInFlight Trait

Provides default shutdown behavior for states where no query is in flight. Query data callbacks and result callbacks trigger shutdown — consistent with the existing pattern where unexpected server messages during authenticated state cause a graceful shutdown rather than a panic.

trait _QueryNoQueryInFlight is _QueryState
  fun ref on_command_complete(s: Session ref, li: _SessionLoggedIn ref,
    msg: _CommandCompleteMessage) => li.shutdown(s)
  fun ref on_data_row(s: Session ref, li: _SessionLoggedIn ref,
    msg: _DataRowMessage) => li.shutdown(s)
  fun ref on_row_description(s: Session ref, li: _SessionLoggedIn ref,
    msg: _RowDescriptionMessage) => li.shutdown(s)
  fun ref on_empty_query_response(s: Session ref,
    li: _SessionLoggedIn ref) => li.shutdown(s)
  fun ref on_error_response(s: Session ref, li: _SessionLoggedIn ref,
    msg: ErrorResponseMessage) => li.shutdown(s)
  fun ref try_run_query(s: Session ref, li: _SessionLoggedIn ref) => None

The trait leaves on_ready_for_query abstract — each concrete no-query-in-flight state handles it differently.

Concrete States

_QueryNotReady

Server has not signaled readiness. This is the initial state after authentication and the state after a non-idle ReadyForQuery (failed transaction).

• on_ready_for_query(idle): transition to _QueryReady, call try_run_query
• on_ready_for_query(non-idle): remain in _QueryNotReady
• All other callbacks: shutdown(s) (via _QueryNoQueryInFlight)
• try_run_query: no-op (via _QueryNoQueryInFlight)

class _QueryNotReady is _QueryNoQueryInFlight
  fun ref on_ready_for_query(s: Session ref, li: _SessionLoggedIn ref,
    msg: _ReadyForQueryMessage)
  =>
    if msg.idle() then
      let ready = _QueryReady
      li._query_state = ready
      ready.try_run_query(s, li)
    end

_QueryReady

Server is idle and ready to accept a query. This is a transitional state — if the queue is non-empty, try_run_query immediately transitions to an in-flight state.

• on_ready_for_query: shutdown(s) — the server only sends ReadyForQuery in response to a query cycle or Sync, never unsolicited. Receiving it when no query is in flight indicates a protocol anomaly.
• All data/result callbacks: shutdown(s) (via _QueryNoQueryInFlight)
• try_run_query: if queue non-empty, send query, transition to _SimpleQueryInFlight

class _QueryReady is _QueryNoQueryInFlight
  fun ref on_ready_for_query(s: Session ref, li: _SessionLoggedIn ref,
    msg: _ReadyForQueryMessage)
  =>
    li.shutdown(s)

  fun ref try_run_query(s: Session ref, li: _SessionLoggedIn ref) =>
    try
      if li._query_queue.size() > 0 then
        (let query, _) = li._query_queue(0)?
        li._query_state = _SimpleQueryInFlight.create()
        s._connection().send(_FrontendMessage.query(query.string))
      end
    else
      _Unreachable()
    end

_SimpleQueryInFlight

Simple query protocol in progress. Owns the per-query accumulation data (_data_rows and _row_description), which are created fresh for each query and destroyed when the state transitions out.

• on_data_row: accumulate row data
• on_row_description: store column metadata
• on_command_complete: build and deliver result to receiver, reset accumulation data for potential multi-statement queries
• on_empty_query_response: validate no data accumulated, deliver SimpleResult to receiver
• on_error_response: deliver error to receiver
• on_ready_for_query(idle): dequeue completed query, transition to _QueryReady, call try_run_query
• on_ready_for_query(non-idle): dequeue completed query, transition to _QueryNotReady
• try_run_query: no-op (query already in flight)

class _SimpleQueryInFlight is _QueryState
  var _data_rows: Array[Array[(String|None)] val] iso
  var _row_description: (Array[(String, U32)] val | None)

  new create() =>
    _data_rows = recover iso Array[Array[(String|None)] val] end
    _row_description = None

  fun ref try_run_query(s: Session ref, li: _SessionLoggedIn ref) => None

  fun ref on_data_row(s: Session ref, li: _SessionLoggedIn ref,
    msg: _DataRowMessage)
  =>
    _data_rows.push(msg.columns)

  fun ref on_row_description(s: Session ref, li: _SessionLoggedIn ref,
    msg: _RowDescriptionMessage)
  =>
    _row_description = msg.columns

  fun ref on_ready_for_query(s: Session ref, li: _SessionLoggedIn ref,
    msg: _ReadyForQueryMessage)
  =>
    try
      li._query_queue.shift()?
    else
      _Unreachable()
    end
    if msg.idle() then
      let ready = _QueryReady
      li._query_state = ready
      ready.try_run_query(s, li)
    else
      li._query_state = _QueryNotReady
    end

  fun ref on_command_complete(s: Session ref, li: _SessionLoggedIn ref,
    msg: _CommandCompleteMessage)
  =>
    try
      (let query, let receiver) = li._query_queue(0)?
      let rows = _data_rows = recover iso
        Array[Array[(String|None)] val].create()
      end
      let rd = _row_description = None

      match rd
      | let desc: Array[(String, U32)] val =>
        try
          let rows_object = _RowsBuilder(consume rows, desc)?
          receiver.pg_query_result(ResultSet(query, rows_object, msg.id))
        else
          receiver.pg_query_failed(query, DataError)
        end
      | None =>
        if rows.size() > 0 then
          receiver.pg_query_failed(query, DataError)
        else
          receiver.pg_query_result(RowModifying(query, msg.id, msg.value))
        end
      end
    else
      _Unreachable()
    end

  fun ref on_empty_query_response(s: Session ref,
    li: _SessionLoggedIn ref)
  =>
    try
      (let query, let receiver) = li._query_queue(0)?
      let rows = _data_rows = recover iso
        Array[Array[(String|None)] val] end
      let rd = _row_description = None

      if (rows.size() > 0) or (rd isnt None) then
        receiver.pg_query_failed(query, DataError)
      else
        receiver.pg_query_result(SimpleResult(query))
      end
    else
      _Unreachable()
    end

  fun ref on_error_response(s: Session ref, li: _SessionLoggedIn ref,
    msg: ErrorResponseMessage)
  =>
    try
      (let query, let receiver) = li._query_queue(0)?
      _data_rows = recover iso Array[Array[(String|None)] val] end
      _row_description = None
      receiver.pg_query_failed(query, msg)
    else
      _Unreachable()
    end

The on_command_complete logic preserves the existing dispatch on _row_description: when a RowDescription was received (non-None), the result is a ResultSet — even if zero data rows arrived (a SELECT returning no matches). When no RowDescription was received (None), the result is RowModifying — unless data rows somehow accumulated without a description, which is a DataError.

The on_empty_query_response validation matches the current code: if any data rows or a row description accumulated before EmptyQueryResponse arrived, that's a protocol anomaly and the query fails with DataError.

The on_error_response resets _data_rows and _row_description before delivering the error. While the data would be abandoned on the next state transition anyway, the reset ensures the in-flight state is clean for potential subsequent messages in a multi-statement query (matching the current code's behavior).

State Transition Diagram

                      +------------------+
                      |  _QueryNotReady  |<---------------------+
                      |    (initial)     |                      |
                      +--------+---------+                      |
                               |                                |
                    ReadyForQuery(idle)             ReadyForQuery(non-idle)
                               |                                |
                               v                                |
                      +----------------+                        |
                 +--->|  _QueryReady   |                        |
                 |    +--------+-------+                        |
                 |             |                                |
                 |     queue non-empty                          |
                 |     (try_run_query)                          |
                 |             |                                |
                 |             v                                |
                 |    +------------------------+                |
                 |    | _SimpleQueryInFlight    |---------------+
                 |    |                        |
                 |    | owns: _data_rows       |
                 |    |       _row_description  |
                 |    +------------+-----------+
                 |                 |
                 |      ReadyForQuery(idle)
                 |      (dequeue completed query)
                 +-----------------+

Changes to _SessionLoggedIn

Remove:

• var _queryable: Bool
• var _query_in_flight: Bool
• var _data_rows: Array[Array[(String|None)] val] iso
• var _row_description: (Array[(String, U32)] val | None)

Add:

• var _query_state: _QueryState

Keep unchanged:

• _query_queue: Array[(SimpleQuery, ResultReceiver)]
• _notify, _readbuf
• on_shutdown (iterates queue, sends SessionClosed failures). This method does not need to interact with _query_state — the queue drain is sufficient for cleanup, and any per-query data held by the in-flight state is simply abandoned (garbage collected) when _SessionLoggedIn is replaced by _SessionClosed.

Constructor:

new ref create(notify': SessionStatusNotify, readbuf': Reader) =>
  _notify = notify'
  _readbuf = readbuf'
  _query_state = _QueryNotReady

Delegation: Each query callback on _SessionLoggedIn delegates to _query_state:

fun ref on_ready_for_query(s: Session ref, msg: _ReadyForQueryMessage) =>
  _query_state.on_ready_for_query(s, this, msg)

fun ref on_command_complete(s: Session ref, msg: _CommandCompleteMessage) =>
  _query_state.on_command_complete(s, this, msg)

fun ref on_data_row(s: Session ref, msg: _DataRowMessage) =>
  _query_state.on_data_row(s, this, msg)

fun ref on_row_description(s: Session ref, msg: _RowDescriptionMessage) =>
  _query_state.on_row_description(s, this, msg)

fun ref on_empty_query_response(s: Session ref) =>
  _query_state.on_empty_query_response(s, this)

fun ref on_error_response(s: Session ref, msg: ErrorResponseMessage) =>
  _query_state.on_error_response(s, this, msg)

execute: Pushes to queue and asks the current state to run:

fun ref execute(s: Session ref, query: SimpleQuery, receiver: ResultReceiver) =>
  _query_queue.push((query, receiver))
  _query_state.try_run_query(s, this)

_run_query is eliminated. Its logic moves into _QueryReady.try_run_query.

Data Ownership

Per-query accumulation data (_data_rows and _row_description) moves from _SessionLoggedIn into _SimpleQueryInFlight. Each new query creates a fresh _SimpleQueryInFlight with empty data, and when the query completes, the old state and its data are replaced.

The current code resets these fields in four places (on_command_complete, on_empty_query_response, on_error_response, and on_ready_for_query). The sub-state machine still resets within the command-terminating callbacks (for multi-statement query support), but the on_ready_for_query reset becomes unnecessary — transitioning to a new state inherently discards the old data.

Incidental Fix: Double-Delivery

The current code only dequeues in on_ready_for_query when msg.idle() is true. When ReadyForQuery(non-idle) arrives during an in-flight query, the query stays in the queue despite its error already having been delivered via on_error_response. If close() is called, on_shutdown sends SessionClosed to the same receiver — a double delivery of pg_query_failed.

The sub-state machine fixes this naturally: _SimpleQueryInFlight.on_ready_for_query always dequeues the completed query before transitioning, regardless of idle status. ReadyForQuery (any status) marks the end of the query cycle, so dequeuing is always correct.

The broader transaction issue — session becomes unusable in a failed transaction because _QueryNotReady prevents sending new queries — remains. That requires ROLLBACK support, which is out of scope.

External Interface

No changes. Session, SimpleQuery, Result, ResultReceiver, SessionStatusNotify, and all other public types remain unchanged. This is a purely internal refactoring.

Extension for Prepared Statements

The prepared statements work (discussion #62) will extend this sub-state machine by:

1. Adding new callbacks to _QueryState for extended query protocol messages (on_parse_complete, on_bind_complete, on_no_data, on_close_complete, on_parameter_description, on_portal_suspended).
2. Adding default shutdown implementations for these callbacks to _QueryNoQueryInFlight (unexpected when no query is in flight) and _SimpleQueryInFlight (unexpected during simple query protocol).
3. Adding _ExtendedQueryInFlight implementing real behavior for the extended query protocol's acknowledgment sequence.
4. Updating _QueryReady.try_run_query to match on query type (SimpleQuery vs PreparedQuery) and create the appropriate in-flight state.

This is what discussion #62 means by "distinguishing simple query in flight from extended query in flight at the type level."

Testing

This is a behavior-preserving refactoring (with the minor incidental fix to double-delivery on failed transactions, which has no existing test coverage). All existing unit and integration tests must continue to pass.

No new tests are required for the refactoring itself. The sub-state machine is an internal implementation detail whose correctness is validated by the existing end-to-end test suite.

Build and run:

make unit-tests ssl=$VERSION        # unit tests (no postgres needed)
make integration-tests ssl=$VERSION # integration tests (needs postgres)

Files Modified

• session.pony — Remove boolean flags and data fields from _SessionLoggedIn. Add _query_state field. Delegate query callbacks. Remove _run_query. Add _QueryState interface, _QueryNoQueryInFlight trait, _QueryNotReady, _QueryReady, and _SimpleQueryInFlight classes.

No other files are modified. The _SessionState interface, _ResponseMessageParser, and all public API types are unchanged.

[SeanTAllen]

this is mostly addressed via #67