Design: a lifecycle model for in-flight work under close, cancellation, and interrupts #1205

Problem

PyAthena has no common model for how in-flight work ends when its owner is closed, cancelled, or interrupted. In-flight work here means API retries, polling, S3 reads and writes, and threads started for asyncio callers. Each owner (connection, cursor, result set, Spark session, filesystem, file) handles these events on its own, and many cases are not handled at all.

Recent fixes each covered a single symptom of this gap:

• Take buffered rows in bulk in fetchmany() and fetchall() #1187: a close-safe bulk take
• Defer closing a chunk reader while another thread reads it #1182: a deferred reader close in the pandas and Polars result sets
• Give the chunk readers the state that close() clears when they are created #1190: reader state captured when the reader is created
• Discard the page of a request that close() interrupts #1197: an is_closed check after the page is stored

#1202 and #1203 were the next two. A static audit of master b2d7aa2, followed by AWS-free stub reproductions, found the failures below. They fall under five root causes.

This issue proposes a lifecycle model, together with an order of implementation, for agreement before any of it is implemented. #1202 and #1203 become parts of it.

Verified failures

Unless stated otherwise, each item was reproduced with stubbed boto3 clients and local files, with no AWS calls. Paths are relative to pyathena/.

R1. Retries and abandoned sibling work cannot be stopped

retry_api_call() (util.py:268) has no stop signal. With the default RetryConfig, a throttled request runs 8 attempts with about 127 s of waits plus jitter. That happens even when the work it serves has been abandoned.

1. A result set's close() during a throttled GetQueryResults retry is followed by 7 more requests, then OperationalError. This happens for page fetches (result_set.py:461, aio/result_set.py:145) and for whole reads through _iter_all_row_pages() (result_set.py:745), which the pandas, Arrow, Polars, and S3FS result sets use when the output location is unavailable (Retries continue after a result set is closed or an asyncio API call is cancelled #1202).

2. A cancelled async_retry_api_call() (aio/util.py:53) keeps retrying in its thread, with 7 more requests. This covers every asyncio Athena and Spark API call (Retries continue after a result set is closed or an asyncio API call is cancelled #1202).

3. A cancelled create_in_thread() (aio/util.py:72) keeps running construction, including GetQueryResults retries (7 more requests) and the pandas Parquet footer reads (pandas/result_set.py:1631). The aio pandas, Arrow, Polars, and S3FS cursors use it.

4. A cancelled AioS3FileSystem operation keeps its S3 retries running in S3AioThreadPool (7 more requests) (Cancelled AioS3FileSystem operations keep retrying in executor threads #1203).

5. S3File.discard(), the abort paths, and AioS3StreamedFile._discard() wait for running parts to finish their retry chain. In the reproduction (attempt=3, multiplier=1), the abort went out 4.5 s (sync) and 3.7 s (async) after the first part failed. With the defaults that wait is about 127 s. Locations: filesystem/s3.py:3348, 1871, 1722; filesystem/s3_async.py:1604.

6. After a failure, sibling work is not cancelled:

   • S3File range reads: queued ranges start after one range fails. The sync close() waits for them; the async close() returns, but the reads are still sent (filesystem/s3.py:3435-3442, 3038; filesystem/s3_async.py:1400).
   • Recursive chmod: all 8 of 8 ACL calls ran after the first failed (filesystem/s3.py:2298-2304).
   • Cancelled async multipart copy: queued parts were sent before the abort when the copy's max_workers exceeds the shared pool (filesystem/s3_async.py:616-728).

   By contrast, ReadAheadStream cancels queued reads (filesystem/reader.py:147).

R2. Result set operations are not coordinated with close() or with each other

A result set's fetches and close() can run at the same time. Each backend guards this differently, or not at all.

1. Concurrent fetches. Cancelling an asyncio fetch that runs in asyncio.to_thread() (aio/common.py:807, 831, 851) leaves the thread fetching, and the next fetch runs alongside it (2 concurrent fetches observed). For pandas and Polars, the second fetch raises ValueError: generator already executing (pandas/result_set.py:982, 996; polars/result_set.py:637, 651).
2. S3FS close during a read. close() from another thread during a fetch makes the reading thread fail with TypeError / AttributeError / ValueError: I/O operation on closed file (s3fs/result_set.py:220-262, s3fs/reader.py:163-165). Alternatively, the fetch returns a row of the closed result set while dropping the buffered one. The pandas and Polars result sets defer the close in the same situation (pandas/result_set.py:338-416, polars/result_set.py:101-171).
3. Arrow and S3FS refill after close. Their _fetch() overrides lack the base check after the page is stored (result_set.py:478-482). Rows read during close() stay in the buffer and are returned by later fetches on the closed result set:
   • Arrow returns them without conversion, for example JSON as text (arrow/result_set.py:199-217).
   • S3FS returns () (s3fs/result_set.py:220-254).
4. Closed semantics differ by path.
   • Page fetches treat close as the end of the data.
   • _iter_all_row_pages() raises ProgrammingError("AthenaResultSet is closed.").
   • _get_content_length() and _read_data_manifest() raise OutputLocation is none or empty. / DataManifestLocation is none or empty., because close() cleared the query execution.
   • A close() during HeadObject still returns the length.
5. A conversion error drops a page. _fetch() advances _next_token before converting the rows (result_set.py:475-477). Any exception during conversion, including KeyboardInterrupt, therefore loses that page; the next fetch continues from the following page.

R3. Cursors have no lifecycle state, and the start and cancel paths have gaps

1. No closed state. close() (result_set.py:1328-1332) closes the result set but neither marks the cursor closed nor stops an execute() in progress. A close() from another thread during polling is followed by execute() attaching a new result set; the cursor then works as if it had never been closed. Fetches after close() return None / [] instead of raising. The sync and asyncio cursors behave the same here.
2. Interrupt before polling. A KeyboardInterrupt in the sync on_start_query_execution callback (cursor.py:187) is raised before _poll(). As a result, kill_on_interrupt never sends StopQueryExecution, and the query keeps running.
3. cancel() races. cancel() from another thread, which the docs describe (docs/usage.md:400-402), raises ProgrammingError until query_id is assigned, and the query then runs on. After a cache hit, it stops the earlier query that the cache returned.
4. Spark session leaks. StartSession catches only Exception (spark/common.py:256-265), so a KeyboardInterrupt during or after the request leaves the session running (TerminateSession = 0). The calculation start, in contrast, is interrupt-safe.

R4. Construction under cancellation leaves orphans

create_in_thread() closes an object built after its caller was cancelled, but only the result sets of the aio pandas, Arrow, Polars, and S3FS cursors use it.

1. AioConnection.create() (aio/connection.py:62), cancelled: the connection is built and never closed.
2. The documented await asyncio.to_thread(conn.cursor) for AioSparkCursor (aio/spark/cursor.py:36-41), cancelled: the Spark session goes IDLE and is never terminated.
3. open_async() (filesystem/s3_async.py:902), cancelled: the file is built and left to the garbage collector. In xb mode, collection emits a misleading "was not committed" ResourceWarning.
4. AsyncCursor and AsyncPandasCursor: a result set built for a running future that the caller cancels is never closed (async_cursor.py:274, pandas/async_cursor.py:274). AsyncCursor._description (async_cursor.py:139) builds a result set and never closes it (source inspection only).
5. A KeyboardInterrupt during S3FS result set construction leaves the CSV file open until garbage collection, because cleanup runs only on Exception (s3fs/result_set.py:210-217). pandas closes it through an ExitStack (pandas/result_set.py:1054).

R5. Blocking or unowned resources in asyncio paths

1. AioS3FileSystem._thread_pool (filesystem/s3_async.py:145) is never shut down. A cached instance keeps its threads for the life of the process.
2. AioPandasCursor.as_pandas() returns the synchronous PandasDataFrameIterator (aio/pandas/cursor.py:239-248), so its chunk reads, and with automatic chunks the whole read, run on the event loop thread. AioPolarsCursor.as_polars() / as_arrow() follow the same pattern (source inspection only).

Checked and not defects

• The asyncio StartQueryExecution with kill_on_interrupt: its retries run in a task shielded on purpose so that the started query can be stopped.
• The sync cursor's status polling during cancel().
• Glue requests in asyncio.to_thread(): botocore retries only.
• The Parquet prefetch thread (arrow/reader.py:196-206): one file at most, joined at close.
• submit_to() on a shut-down pool: a future can stay pending, but no S3 path reaches a bare wait on it.
• Cancelled sync and async deletes: both leave running batches. The difference reported in the audit came from a batch failure, where both wait.
• Breaking out of iter_chunks(): both cursors close the iterator, the Polars one through yield from.

Proposed model

1. Cooperative stop for in-flight work

This addresses R1.

• retry_api_call() honors a stop signal. Once the signal is set, it makes no further attempt, ends the backoff wait early, and raises the last error.
• The asyncio side carries the signal in a ContextVar holding a per-call threading.Event. The variable is set only in the context copied into the worker thread, never in the task's context, so shielded work does not inherit it.
• Every asyncio entry point that hands sync work to a thread sets the event on cancellation. The entry points are async_retry_api_call(), create_in_thread(), the to_thread() fetches, and S3AioThreadPool.run().
• S3ThreadPoolExecutor.submit() copies the submitter's context; it does not today, while S3AioThreadPool and S3AioExecutor already propagate it. As a result, part uploads and range reads started from a cancelled call stop too.
• A result set passes stop_when=lambda: self.is_closed for its own requests. Withdrawn in Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206: a close() from another thread is unsupported (see the contract below).
• Protected cleanup masks the signal. This covers multipart abort and discard, the close in create_in_thread(), run_shielded(), and the shielded query and calculation starts, so cleanup always completes.
• When a part or range fails, the owner cancels the queued sibling work and signals the running work, as ReadAheadStream already does. This covers range reads, chmod, and multipart copy.

The guarantee is cooperative: a request already sent, botocore's own retries, and native I/O can still finish.

2. One operation guard per result set

Withdrawn (2026-10-17). PyAthena declares threadsafety = 2, so a cursor or result set used from several threads, including a close() from another thread, is unsupported, and no guard is added for it. The guards from #1182, #1190, #1197, and #1187 are removed. #1206 keeps only the single-task cases: a failed or cancelled fetch breaks the result set (#1225 option (c)). See #1206 (comment).

This addresses R2 and replaces the per-backend guards from #1182, #1190, #1197, and #1187.

• The base result set wraps each fetch and read in an operation scope that serializes fetches. A fetch started while another runs waits for it, which covers the cancelled asyncio fetch whose thread is still running.
• close() marks the result set closed and signals the stop. Releasing resources is deferred until the running operation leaves its scope. pandas and Polars do this today for their chunk iterators.
• An operation that starts or completes after close() returns no rows. Rows read during the close are discarded.
• Paths that cannot return "nothing", such as whole reads and HeadObject lookups, raise one ProgrammingError for the closed result set.
• Page state advances only after its rows are converted, so a failed conversion can be retried.

3. Cursor lifecycle state

This addresses R3.

• Cursors get a closed state. execute() and fetches on a closed cursor raise ProgrammingError, and an execute() running at close() closes the result set it creates.
• The protected region of kill_on_interrupt starts once the execution ID exists, which takes in the start callback.
• cancel() stops only the execution that the current execute() started, never a cached one. Before the ID exists, it waits for it or is recorded and applied once the ID arrives. Withdrawn: a cancel() from another thread is unsupported.
• Spark StartSession gets the same interrupt-safe start as calculations.

4. One construction rule

This addresses R4.

• Every asyncio construction of a closable object goes through create_in_thread(). This covers connections, Spark cursors, files from open_async(), and result sets.
• Sync constructors release what they opened on BaseException.
• AsyncCursor owns the result sets of its futures and closes the ones the caller abandons when the cursor closes.

5. Resource ownership in asyncio paths

This addresses R5.

• AioS3FileSystem shuts down its pool, either through a close() or at garbage collection; how is to be decided with Free the data of failed S3 requests and clean up async copies at loop shutdown #1192.
• The aio cursors' as_pandas(), as_polars(), and as_arrow() read through threads, or document that the returned iterator blocks.

Contract to document

• threadsafety = 2 stays: threads may share connections but not cursors.
• cancel() from another thread on the same cursor is not supported either (decided in Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206); a query can be stopped from elsewhere with its ID, for example with AsyncCursor.cancel(query_id). The documentation example is rewritten in the Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206 PR.
• close() from another thread is not supported, like any other use of one cursor or result set from several threads (decided in Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206). Misuse is the caller's bug.
• Asyncio cancellation stops the PyAthena retries and queued work of the cancelled call, but not a request already sent.

Decisions

• PEP 249. Cursors follow PEP 249. A closed cursor raises ProgrammingError from execute(), executemany(), and the fetch methods instead of returning None / []. This is a breaking change and will be called out in the release notes.
• Tracking. This issue is the parent. The sub-issues below are implemented in order, each with its own PR, and the reproductions above become unit tests.

Sub-issues

Implemented in this order, one PR per step unless noted:

1. Cooperative stop (model 1): Retries continue after a result set is closed or an asyncio API call is cancelled #1202 for the Athena API calls and asyncio entry points, and Cancelled AioS3FileSystem operations keep retrying in executor threads #1203 for the S3 filesystem.
2. Failed and cancelled fetches break the result set, without an operation guard: Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206 (with Define the behavior of result set fetches after an error #1225)
3. Cursor lifecycle per PEP 249 (model 3): Give cursors a PEP 249 closed state and make query and session starts interrupt-safe #1207
4. Construction under cancellation (model 4): Close objects built for cancelled asyncio constructions #1208
5. Asyncio resource ownership (model 5): Own the AioS3FileSystem pool and keep aio DataFrame reads off the event loop #1209, coordinated with Free the data of failed S3 requests and clean up async copies at loop shutdown #1192

The contract above remains open for comment until step 1 starts.