From 543b88c48e387aade97b03fda3acee7b45b0d3ec Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 12 Nov 2026 13:38:29 +0900 Subject: [PATCH 1/8] Close an unclosed S3File outside the pools when it is finalized (GH-1364) A file finalized without having been closed waited for its parts in the finalizing thread. In a thread of the request pool, those parts could be queued behind that thread, or include the request in which the collection ran, so the finalizer waited forever. A file whose last reference is dropped outside the threads of an S3ThreadPool is still closed by its finalizer, and committed before the drop returns. A file that the cyclic garbage collector finalizes, or that is finalized in a pool thread, is closed in a new thread with an empty context. close() called by another finalizer during a collection, such as that of a TextIOWrapper, is handed over the same way, once. The write buffer is copied first, since the collector closes it when it finalizes it; if it is already closed, the write is discarded instead of committed with incomplete data. If no thread can be started, the file is left unclosed rather than closed where it could wait forever. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 17 +++ pyathena/filesystem/s3.py | 112 ++++++++++++++++- pyathena/filesystem/s3_executor.py | 10 ++ tests/pyathena/filesystem/test_s3.py | 174 +++++++++++++++++++++++++++ 4 files changed, 312 insertions(+), 1 deletion(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 979d0cd64..bc09f141d 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -203,6 +203,23 @@ In a forked child process, create a new filesystem. A file that the parent opene `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, and closing it only marks it closed, without uploading or aborting anything. +Close a file returned by `open` explicitly or with a `with` block. A file that is not +closed is closed when it is finalized, which commits its written data unless +`autocommit` is False: + +- When its last reference is dropped outside the threads of an `S3ThreadPool`, such as + the request pool of a filesystem, it is closed before that returns. +- When the cyclic garbage collector finalizes it, or it is finalized in a thread of an + `S3ThreadPool`, it is closed in a new thread after the finalization returns. If the + collector has finalized its write buffer first, its data is not committed, and its + multipart upload is aborted. If no thread can be started, as during interpreter + finalization, the file is left unclosed and its multipart upload incomplete. +- During interpreter shutdown, a file whose last part can no longer be submitted is + closed without committing it, and its multipart upload is aborted. + +A file opened in text mode writes its remaining text to the file in the thread that +finalizes it, before the file is closed. + An interrupt, such as Ctrl+C, or a failure while a file reads ranges in parallel cancels the range reads that have not started, and the running ones make no further attempt. The interrupt or the error is raised once they have ended, and their data is diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index da6dfcb60..cc4589013 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3,7 +3,9 @@ from __future__ import annotations import contextlib +import contextvars import functools +import gc import hashlib import logging import mimetypes @@ -78,6 +80,8 @@ "major version. Use pyathena.filesystem.s3_filesystem(connection), or pass " "client=connection.s3_client and retry_config=connection.retry_config." ) +# Whether the cyclic garbage collector runs in a thread (S3File._track_collection). +_collecting = threading.local() # The request parameters on which the authorization of a lookup (HeadObject, # HeadBucket, or ListObjectsV2) depends. @@ -4310,17 +4314,120 @@ def close(self) -> None: In a forked child process, a file that the parent opened is only marked closed: its data is neither flushed nor committed, and its - upload is not aborted, since they belong to the parent. + upload is not aborted, since they belong to the parent. Called by + another finalizer while the cyclic garbage collector runs in this + thread, such as that of a ``TextIOWrapper``, the file is closed in a + new thread instead (see :meth:`__del__`). """ if self._inherited(): self.closed = True return + if not self.closed and S3File._in_collection(): + self._close_later() + return try: super().close() finally: # The executor is shut down even if the final flush fails. self._executor.shutdown() + def __del__(self) -> None: + """Close the file when it is finalized without having been closed. + + A file whose last reference is dropped in a user thread is closed + here. A file that the cyclic garbage collector finalizes, or that is + finalized in a thread of an :class:`S3ThreadPool`, is closed in a new + thread instead, since a wait here could last forever (T1): the + requests of the file may be queued behind this thread or include the + one that runs the collection, and this thread may hold a lock that + the close takes. + """ + if getattr(self, "closed", True) or self._inherited(): + # The upload of a file inherited across a fork is the parent's. + return + if S3ThreadPool.in_pool_thread(): + self._close_later() + else: + # Hands the close over itself during a collection. + self.close() + + @staticmethod + def _track_collection(phase: str, info: dict[str, int]) -> None: + """Record whether the cyclic garbage collector runs in this thread. + + Registered in ``gc.callbacks``, which the collector calls in the + thread that runs it. + + Args: + phase: ``"start"`` or ``"stop"``. + info: The details of the collection, which are not used. + """ + _collecting.active = phase == "start" + + @staticmethod + def _in_collection() -> bool: + """Whether the cyclic garbage collector runs in this thread. + + Returns: + True while a collection runs in the calling thread. + """ + return getattr(_collecting, "active", False) + + def _close_later(self) -> None: + """Close the file in a new thread, once. + + The written data is copied first: the collector finalizes every + object that it found unreachable before it checks whether a + finalizer has made it reachable again, and finalizing the buffer + closes it. If the buffer is already closed, the data is incomplete, + so the write is discarded instead of committed. The thread runs with + an empty context, so that the stop scopes of the collecting thread do + not stop the requests of the file. If no thread can be started, as + during interpreter finalization, the file is left unclosed rather + than closed in this thread. + """ + if getattr(self, "_closing_later", False): + return + self._closing_later = True + complete = True + if self.mode != "rb" and self.buffer is not None: + if self.buffer.closed: + complete = False + else: + buffer = BytesIO(self.buffer.getvalue()) + buffer.seek(self.buffer.tell()) + self.buffer = buffer + try: + threading.Thread( + target=contextvars.Context().run, + args=(self._close_collected, complete), + name=f"pyathena-close-s3://{self.bucket}/{self.key}", + ).start() + except RuntimeError: + _logger.warning( + f"s3://{self.bucket}/{self.key} was not closed, and no thread could be " + "started to close it when it was finalized; its data was not committed." + ) + + def _close_collected(self, complete: bool) -> None: + """Close a finalized file, logging a failure instead of raising it. + + Args: + complete: Whether the written data is complete. If not, the + write is discarded instead of committed. + """ + try: + if complete: + self.close() + else: + _logger.warning( + f"s3://{self.bucket}/{self.key} was not closed, and its buffer was " + "finalized before it; its data was not committed." + ) + self._close_without_commit() + except Exception: + _logger.exception(f"Failed to close s3://{self.bucket}/{self.key}.") + def read(self, length: int = -1) -> bytes: """Read bytes from the file. @@ -4823,3 +4930,6 @@ def _get_ranges( else: ranges.append((start, end)) return ranges + + +gc.callbacks.append(S3File._track_collection) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 95e8760f1..cd4146e9a 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -930,6 +930,16 @@ def owns_current_thread(self) -> bool: pool = getattr(_pool_thread, "pool", None) return pool is not None and pool() is self + @staticmethod + def in_pool_thread() -> bool: + """Whether the calling thread is a thread of any pool of this type. + + Returns: + True in a thread that a pool of this type started, also once the + pool has been collected. + """ + return getattr(_pool_thread, "pool", None) is not None + @classmethod def _reset_after_fork(cls) -> None: """Drop the threads and locks that a forked child copied from its parent. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index f653496e8..94bad35b3 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -8,6 +8,7 @@ import gzip import io import itertools +import logging import lzma import os import pickle @@ -8947,6 +8948,179 @@ def interrupted(timeout=None): assert operation.calls == 1 assert file.closed + @staticmethod + def _finalize_unclosed_file(fs, where, *, text=False): + # Write a multipart file without closing it, and finalize it in the + # given place, using the only thread of the request pool: + # - "last reference": its last reference is dropped in this thread; + # - "collection": the cyclic garbage collector runs in this thread; + # - "queued part": the collector runs in the thread of the pool, + # while the first part of the file is queued behind it; + # - "own part": the collector runs inside that first part; + # - "last reference in pool": its last reference is dropped in the + # thread of the pool, while the first part is queued behind it. + # A close handed over to another thread starts only once the + # finalization has returned, after the collector has finalized the + # buffer of the file. Returns the files whose close was handed over. + holder = [] + ready = threading.Event() + finalized = threading.Event() + + def finalize(): + assert ready.wait(EVENT_TIMEOUT) + if where.startswith("last reference"): + holder.clear() + else: + gc.collect() + finalized.set() + + if where in ("queued part", "last reference in pool"): + upload = fs.core.create_multipart_upload.return_value + + def create(*args, **kwargs): + # Queued ahead of the first part of the file. + fs._request_pool.submit(finalize) + return upload + + fs.core.create_multipart_upload.side_effect = create + elif where == "own part": + part = fs.core.upload_part.side_effect + + def upload_part(**kwargs): + if kwargs["part_number"] == 1: + finalize() + return part(**kwargs) + + fs.core.upload_part.side_effect = upload_part + + close_collected = S3File._close_collected + closes = [] + + def close_after_finalization(file, complete): + closes.append(file) + assert finalized.wait(EVENT_TIMEOUT) + close_collected(file, complete) + + gc.disable() + try: + with mock.patch.object(S3File, "_close_collected", close_after_finalization): + file = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) + if text: + holder.append(io.TextIOWrapper(file, encoding="utf-8")) + holder[0].write("abcde") + else: + holder.append(file) + file.write(b"abcde") + if not where.startswith("last reference"): + # A reference cycle, as the read caches of fsspec make, + # leaves the file to the cyclic garbage collector. + holder[0].cycle = holder[0] + holder.clear() + del file + ready.set() + if where in ("last reference", "collection"): + finalize() + assert finalized.wait(EVENT_TIMEOUT) + finally: + gc.enable() + return closes + + @pytest.mark.parametrize( + ("where", "text"), + [ + ("last reference", False), + ("collection", False), + ("collection", True), + ("queued part", False), + ("own part", False), + ("last reference in pool", False), + ], + ) + def test_finalized_unclosed_file(self, where, text): + # GH-1364: an unclosed file is committed when it is finalized. Only a + # file whose last reference is dropped in a user thread is committed + # by the finalizer itself. One that the cyclic garbage collector + # finalizes, or that is finalized in a thread of the request pool, is + # committed from a new thread, with the data that it had buffered, + # although the collector closes its buffer, so that the finalizer + # does not wait forever for a part queued behind it, or for the part + # in which it runs. + fs = self._make_append_fs(b"") + fs._request_pool = S3ThreadPool(1) + completed = threading.Event() + committers = [] + + def complete(*args, **kwargs): + committers.append((threading.current_thread(), fs._request_pool.owns_current_thread())) + completed.set() + return mock.DEFAULT + + fs.core.complete_multipart_upload.side_effect = complete + try: + closes = self._finalize_unclosed_file(fs, where, text=text) + assert completed.wait(EVENT_TIMEOUT) + finally: + shutdown_bounded(fs._request_pool) + assert self._uploaded_object(fs, b"") == b"abcde" + # A sync file commits in the thread that closes it. + ((committer, in_pool),) = committers + assert (committer is threading.current_thread()) == (where == "last reference") + assert not in_pool + # The close is handed over once, also when the finalizer of a text + # wrapper closes the file before the finalizer of the file runs. + assert len(closes) == (0 if where == "last reference" else 1) + + def test_finalized_unclosed_file_with_finalized_buffer(self, caplog): + # GH-1364: an unclosed file whose buffer the collector finalized + # first is discarded, instead of committed without the buffered data. + fs = self._make_append_fs(b"") + aborted = threading.Event() + fs._call.side_effect = lambda *args, **kwargs: aborted.set() + gc.disable() + try: + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) + f.write(b"abcde") + # As the collector closes a buffer that it finalizes first. + f.buffer.close() + f.cycle = f + del f + with caplog.at_level(logging.WARNING, logger="pyathena.filesystem.s3"): + gc.collect() + assert aborted.wait(EVENT_TIMEOUT) + finally: + gc.enable() + shutdown_bounded(fs._request_pool) + fs._call.assert_called_once_with( + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid" + ) + fs.core.complete_multipart_upload.assert_not_called() + assert "buffer was finalized" in caplog.text + + def test_finalized_unclosed_file_without_thread(self, caplog): + # GH-1364: if no thread can be started for the close of a file + # finalized in a thread of the request pool, as during interpreter + # finalization, the file is left unclosed, and its data uncommitted, + # rather than closed in the pool, where it would wait forever. + fs = self._make_append_fs(b"") + fs._request_pool = S3ThreadPool(1) + start = threading.Thread.start + + def start_unless_closing(thread): + if thread.name == "pyathena-close-s3://bucket/key.txt": + raise RuntimeError("can't create new thread at interpreter shutdown") + start(thread) + + try: + with ( + mock.patch.object(threading.Thread, "start", start_unless_closing), + caplog.at_level(logging.WARNING, logger="pyathena.filesystem.s3"), + ): + self._finalize_unclosed_file(fs, "queued part") + finally: + shutdown_bounded(fs._request_pool) + fs.core.complete_multipart_upload.assert_not_called() + assert "no thread could be started" in caplog.text + @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): # A file larger than the block size is uploaded in parts and completed From fcdd2022907715be2f37af0935d2e80d6701218a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 00:02:32 +0900 Subject: [PATCH 2/8] Close a finalized S3File inline during interpreter finalization (GH-1364) During interpreter finalization no thread is started for the close of a finalized file, and the file is closed in the finalizing thread, as it was before this change: Python 3.11 waits forever in Thread.start() then, and later versions raise, which left the file unclosed. No other thread runs requests then, and the request pools take no more calls. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 7 +++--- pyathena/filesystem/s3.py | 19 ++++++++++++---- tests/pyathena/filesystem/test_s3.py | 34 ++++++++++++++++++++-------- 3 files changed, 42 insertions(+), 18 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index bc09f141d..4b7229ec0 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -212,10 +212,9 @@ closed is closed when it is finalized, which commits its written data unless - When the cyclic garbage collector finalizes it, or it is finalized in a thread of an `S3ThreadPool`, it is closed in a new thread after the finalization returns. If the collector has finalized its write buffer first, its data is not committed, and its - multipart upload is aborted. If no thread can be started, as during interpreter - finalization, the file is left unclosed and its multipart upload incomplete. -- During interpreter shutdown, a file whose last part can no longer be submitted is - closed without committing it, and its multipart upload is aborted. + multipart upload is aborted. During interpreter finalization, the file is closed in + the finalizing thread instead. If a thread cannot be started otherwise, the file is + left unclosed and its multipart upload incomplete. A file opened in text mode writes its remaining text to the file in the thread that finalizes it, before the file is closed. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index cc4589013..0f88b122d 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -4325,6 +4325,10 @@ def close(self) -> None: if not self.closed and S3File._in_collection(): self._close_later() return + self._close() + + def _close(self) -> None: + """Close the file in this thread, flushing any written data, and shut down its executor.""" try: super().close() finally: @@ -4382,9 +4386,12 @@ def _close_later(self) -> None: closes it. If the buffer is already closed, the data is incomplete, so the write is discarded instead of committed. The thread runs with an empty context, so that the stop scopes of the collecting thread do - not stop the requests of the file. If no thread can be started, as - during interpreter finalization, the file is left unclosed rather - than closed in this thread. + not stop the requests of the file. During interpreter finalization, + when a started thread never runs, the file is closed in this thread, + as no other thread runs requests then and the pools take no more + calls. If no thread can be started otherwise, the file is left + unclosed rather than closed in this thread, where it could wait + forever. """ if getattr(self, "_closing_later", False): return @@ -4397,6 +4404,10 @@ def _close_later(self) -> None: buffer = BytesIO(self.buffer.getvalue()) buffer.seek(self.buffer.tell()) self.buffer = buffer + if sys.is_finalizing(): + # Python 3.11 waits forever in Thread.start() here. + self._close_collected(complete) + return try: threading.Thread( target=contextvars.Context().run, @@ -4418,7 +4429,7 @@ def _close_collected(self, complete: bool) -> None: """ try: if complete: - self.close() + self._close() else: _logger.warning( f"s3://{self.bucket}/{self.key} was not closed, and its buffer was " diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 94bad35b3..2cd7f620c 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -8961,7 +8961,7 @@ def _finalize_unclosed_file(fs, where, *, text=False): # thread of the pool, while the first part is queued behind it. # A close handed over to another thread starts only once the # finalization has returned, after the collector has finalized the - # buffer of the file. Returns the files whose close was handed over. + # buffer of the file. Returns the files closed by `_close_collected`. holder = [] ready = threading.Event() finalized = threading.Event() @@ -8998,7 +8998,8 @@ def upload_part(**kwargs): def close_after_finalization(file, complete): closes.append(file) - assert finalized.wait(EVENT_TIMEOUT) + if threading.current_thread().name.startswith("pyathena-close-"): + assert finalized.wait(EVENT_TIMEOUT) close_collected(file, complete) gc.disable() @@ -9096,13 +9097,21 @@ def test_finalized_unclosed_file_with_finalized_buffer(self, caplog): fs.core.complete_multipart_upload.assert_not_called() assert "buffer was finalized" in caplog.text - def test_finalized_unclosed_file_without_thread(self, caplog): - # GH-1364: if no thread can be started for the close of a file - # finalized in a thread of the request pool, as during interpreter - # finalization, the file is left unclosed, and its data uncommitted, - # rather than closed in the pool, where it would wait forever. + @pytest.mark.parametrize("finalizing", [False, True]) + def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): + # GH-1364: if no thread can be started for the close of a finalized + # file, it is left unclosed, and its data uncommitted, rather than + # closed where it would wait forever, here in the thread of the + # request pool with its first part queued behind it. During + # interpreter finalization, when no other thread runs requests, it is + # closed in the finalizing thread instead, without starting a thread. fs = self._make_append_fs(b"") fs._request_pool = S3ThreadPool(1) + committers = [] + fs.core.complete_multipart_upload.side_effect = lambda *args, **kwargs: ( + committers.append(threading.current_thread()), + mock.DEFAULT, + )[1] start = threading.Thread.start def start_unless_closing(thread): @@ -9113,13 +9122,18 @@ def start_unless_closing(thread): try: with ( mock.patch.object(threading.Thread, "start", start_unless_closing), + mock.patch.object(sys, "is_finalizing", return_value=finalizing), caplog.at_level(logging.WARNING, logger="pyathena.filesystem.s3"), ): - self._finalize_unclosed_file(fs, "queued part") + self._finalize_unclosed_file(fs, "collection" if finalizing else "queued part") finally: shutdown_bounded(fs._request_pool) - fs.core.complete_multipart_upload.assert_not_called() - assert "no thread could be started" in caplog.text + if finalizing: + assert self._uploaded_object(fs, b"") == b"abcde" + assert committers == [threading.current_thread()] + else: + fs.core.complete_multipart_upload.assert_not_called() + assert "no thread could be started" in caplog.text @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): From 35db97735a654eaba45cdf132c0f57d7332cbd9b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 00:18:59 +0900 Subject: [PATCH 3/8] Document how a file is closed at interpreter finalization (GH-1364) A file finalized during a pass of the cyclic garbage collector is handed over also when another finalizer drops its last reference, and during interpreter finalization the request pools take no more calls, so a file with data left to upload as a part is not committed and its multipart upload is aborted. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 19 +++++++++++-------- pyathena/filesystem/s3.py | 12 ++++++------ 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 4b7229ec0..84e7cf063 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -207,14 +207,17 @@ Close a file returned by `open` explicitly or with a `with` block. A file that i closed is closed when it is finalized, which commits its written data unless `autocommit` is False: -- When its last reference is dropped outside the threads of an `S3ThreadPool`, such as - the request pool of a filesystem, it is closed before that returns. -- When the cyclic garbage collector finalizes it, or it is finalized in a thread of an - `S3ThreadPool`, it is closed in a new thread after the finalization returns. If the - collector has finalized its write buffer first, its data is not committed, and its - multipart upload is aborted. During interpreter finalization, the file is closed in - the finalizing thread instead. If a thread cannot be started otherwise, the file is - left unclosed and its multipart upload incomplete. +- When it is finalized during a pass of the cyclic garbage collector, or in a thread of + an `S3ThreadPool`, such as the request pool of a filesystem, it is closed in a new + thread after the finalization returns. If no thread can be started, the file is left + unclosed and its multipart upload incomplete. +- When its last reference is dropped elsewhere, it is closed before that returns. +- During interpreter finalization, it is closed in the finalizing thread instead of a + new thread. The request pools take no more calls then, so a file with data left to + upload as a part is not committed, and its multipart upload is aborted. + +If the garbage collector finalizes the write buffer of a file before the file, the data +of the file is not committed, and its multipart upload is aborted. A file opened in text mode writes its remaining text to the file in the thread that finalizes it, before the file is closed. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 0f88b122d..e4ac8f7b1 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -4387,11 +4387,12 @@ def _close_later(self) -> None: so the write is discarded instead of committed. The thread runs with an empty context, so that the stop scopes of the collecting thread do not stop the requests of the file. During interpreter finalization, - when a started thread never runs, the file is closed in this thread, - as no other thread runs requests then and the pools take no more - calls. If no thread can be started otherwise, the file is left - unclosed rather than closed in this thread, where it could wait - forever. + the file is closed in this thread without starting one: Python 3.11 + waits forever for a thread started then, and later versions raise. + No other thread runs requests then, and the pools take no more calls, + so the close does not wait for a part queued behind it. If no thread + can be started otherwise, the file is left unclosed rather than + closed in this thread, where it could wait forever. """ if getattr(self, "_closing_later", False): return @@ -4405,7 +4406,6 @@ def _close_later(self) -> None: buffer.seek(self.buffer.tell()) self.buffer = buffer if sys.is_finalizing(): - # Python 3.11 waits forever in Thread.start() here. self._close_collected(complete) return try: From 6a4a5bbb8d207d06b0ff753d740ca680d684083e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 01:04:58 +0900 Subject: [PATCH 4/8] Start a handed-over close once the collection has ended (GH-1364) A close handed over during a collection starts once the collection has ended, so that the other finalizers of the collection, such as that of the TextIOWrapper of a text-mode file, which the collector may run after the finalizer of the file, still write to it first. A later close(), commit(), or discard() of the file outside a collection waits for the handed-over close to end, so that a transaction does not complete the upload without the data that the close uploads, and an explicit close does not race with it. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 20 ++++-- pyathena/filesystem/s3.py | 92 +++++++++++++++++++----- tests/pyathena/filesystem/test_s3.py | 100 +++++++++++++++++++++++++-- 3 files changed, 181 insertions(+), 31 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 84e7cf063..e17181615 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -207,10 +207,11 @@ Close a file returned by `open` explicitly or with a `with` block. A file that i closed is closed when it is finalized, which commits its written data unless `autocommit` is False: -- When it is finalized during a pass of the cyclic garbage collector, or in a thread of - an `S3ThreadPool`, such as the request pool of a filesystem, it is closed in a new - thread after the finalization returns. If no thread can be started, the file is left - unclosed and its multipart upload incomplete. +- When it is finalized during a pass of the cyclic garbage collector, it is closed in a + new thread that starts once the collection has ended. When it is finalized in a + thread of an `S3ThreadPool`, such as the request pool of a filesystem, it is closed + in a new thread. If no thread can be started, the file is left unclosed and its + multipart upload incomplete. - When its last reference is dropped elsewhere, it is closed before that returns. - During interpreter finalization, it is closed in the finalizing thread instead of a new thread. The request pools take no more calls then, so a file with data left to @@ -219,8 +220,15 @@ closed is closed when it is finalized, which commits its written data unless If the garbage collector finalizes the write buffer of a file before the file, the data of the file is not committed, and its multipart upload is aborted. -A file opened in text mode writes its remaining text to the file in the thread that -finalizes it, before the file is closed. +`close()` called during a pass of the cyclic garbage collector, such as by another +finalizer, closes the file in the same way, in a new thread that starts once the +collection has ended, and returns before the file is closed. A later `close()`, +`commit()`, or `discard()` of the file waits for that close to end, unless it is +called during a collection. + +A file opened in text mode is closed by its text wrapper when the wrapper is finalized. +The wrapper writes its remaining text to the file in the thread that finalizes it, and +then closes the file there, or as above during a collection. An interrupt, such as Ctrl+C, or a failure while a file reads ranges in parallel cancels the range reads that have not started, and the running ones make no further diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index e4ac8f7b1..c795fb398 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -80,7 +80,9 @@ "major version. Use pyathena.filesystem.s3_filesystem(connection), or pass " "client=connection.s3_client and retry_config=connection.retry_config." ) -# Whether the cyclic garbage collector runs in a thread (S3File._track_collection). + +# Whether the cyclic garbage collector runs in a thread, and the closes of +# files that it handed over (S3File._track_collection). _collecting = threading.local() # The request parameters on which the authorization of a lookup (HeadObject, @@ -4317,7 +4319,9 @@ def close(self) -> None: upload is not aborted, since they belong to the parent. Called by another finalizer while the cyclic garbage collector runs in this thread, such as that of a ``TextIOWrapper``, the file is closed in a - new thread instead (see :meth:`__del__`). + new thread once the collection has ended (see :meth:`__del__`), and + this returns before that. Otherwise, a close handed over to another + thread is waited for. """ if self._inherited(): self.closed = True @@ -4325,6 +4329,7 @@ def close(self) -> None: if not self.closed and S3File._in_collection(): self._close_later() return + self._wait_for_close_later() self._close() def _close(self) -> None: @@ -4344,7 +4349,9 @@ def __del__(self) -> None: thread instead, since a wait here could last forever (T1): the requests of the file may be queued behind this thread or include the one that runs the collection, and this thread may hold a lock that - the close takes. + the close takes. During a collection, the thread starts once the + collection has ended, after the other finalizers of the collection, + such as that of a ``TextIOWrapper``, have written to the file. """ if getattr(self, "closed", True) or self._inherited(): # The upload of a file inherited across a fork is the parent's. @@ -4360,13 +4367,19 @@ def _track_collection(phase: str, info: dict[str, int]) -> None: """Record whether the cyclic garbage collector runs in this thread. Registered in ``gc.callbacks``, which the collector calls in the - thread that runs it. + thread that runs it. When the collection stops, the closes that it + handed over start. Args: phase: ``"start"`` or ``"stop"``. info: The details of the collection, which are not used. """ _collecting.active = phase == "start" + pending = getattr(_collecting, "pending", None) + if not _collecting.active and pending: + _collecting.pending = [] + for file, complete in pending: + file._start_close(complete) @staticmethod def _in_collection() -> bool: @@ -4384,19 +4397,14 @@ def _close_later(self) -> None: object that it found unreachable before it checks whether a finalizer has made it reachable again, and finalizing the buffer closes it. If the buffer is already closed, the data is incomplete, - so the write is discarded instead of committed. The thread runs with - an empty context, so that the stop scopes of the collecting thread do - not stop the requests of the file. During interpreter finalization, - the file is closed in this thread without starting one: Python 3.11 - waits forever for a thread started then, and later versions raise. - No other thread runs requests then, and the pools take no more calls, - so the close does not wait for a part queued behind it. If no thread - can be started otherwise, the file is left unclosed rather than - closed in this thread, where it could wait forever. + so the write is discarded instead of committed. During a collection, + the thread starts once the collection has ended, so that the other + finalizers of the collection can still write to the file. """ if getattr(self, "_closing_later", False): return self._closing_later = True + self._closed_later = threading.Event() complete = True if self.mode != "rb" and self.buffer is not None: if self.buffer.closed: @@ -4405,20 +4413,47 @@ def _close_later(self) -> None: buffer = BytesIO(self.buffer.getvalue()) buffer.seek(self.buffer.tell()) self.buffer = buffer + if S3File._in_collection(): + if not hasattr(_collecting, "pending"): + _collecting.pending = [] + _collecting.pending.append((self, complete)) + else: + self._start_close(complete) + + def _start_close(self, complete: bool) -> None: + """Start the close that :meth:`_close_later` handed over. + + The thread runs with an empty context, so that the stop scopes of + the finalizing thread do not stop the requests of the file. During + interpreter finalization, the file is closed in this thread without + starting one: Python 3.11 waits forever for a thread started then, + and later versions raise. No other thread runs requests then, and the + pools take no more calls, so the close does not wait for a part + queued behind it. If no thread can be started otherwise, the file is + left unclosed rather than closed in this thread, where it could wait + forever. + + Args: + complete: Whether the written data is complete. + """ if sys.is_finalizing(): + self._closer = threading.current_thread() self._close_collected(complete) return try: - threading.Thread( + self._closer = threading.Thread( target=contextvars.Context().run, args=(self._close_collected, complete), name=f"pyathena-close-s3://{self.bucket}/{self.key}", - ).start() - except RuntimeError: + ) + self._closer.start() + except Exception: + # Such as RuntimeError, when no more threads can be started. _logger.warning( f"s3://{self.bucket}/{self.key} was not closed, and no thread could be " "started to close it when it was finalized; its data was not committed." ) + self._closed_later.set() def _close_collected(self, complete: bool) -> None: """Close a finalized file, logging a failure instead of raising it. @@ -4438,6 +4473,22 @@ def _close_collected(self, complete: bool) -> None: self._close_without_commit() except Exception: _logger.exception(f"Failed to close s3://{self.bucket}/{self.key}.") + finally: + self._closed_later.set() + + def _wait_for_close_later(self) -> None: + """Wait for a close that :meth:`_close_later` handed over to end. + + Not waited for in the thread that closes the file, or during a + collection, whose end starts the close. + """ + closed_later = getattr(self, "_closed_later", None) + if ( + closed_later is not None + and getattr(self, "_closer", None) is not threading.current_thread() + and not S3File._in_collection() + ): + closed_later.wait() def read(self, length: int = -1) -> bytes: """Read bytes from the file. @@ -4693,7 +4744,8 @@ def commit(self) -> None: if the abort also fails, the upload is kept so that :meth:`discard` can abort it. Committing again after a successful commit sends no request. Invalidates the cache of the path afterwards. Does nothing - for a file opened for reading. + for a file opened for reading. Waits first for a close of the file + that was handed over to another thread (see :meth:`close`). Raises: FileExistsError: If an object was created at the path after the @@ -4704,6 +4756,7 @@ def commit(self) -> None: or an earlier commit failed. """ self._check_process() + self._wait_for_close_later() if self._session is None: # A file opened for reading has nothing to commit. return @@ -4735,13 +4788,16 @@ def discard(self) -> None: interrupted, the upload is kept, and calling this again retries the abort. Afterwards, :meth:`commit` raises :class:`~pyathena.filesystem.s3_writer.S3MultipartStateError`. Does - nothing for a committed file or a file opened for reading. + nothing for a committed file or a file opened for reading. Waits + first for a close of the file that was handed over to another thread + (see :meth:`close`). Raises: RuntimeError: In a forked child process of the process that opened the file. """ self._check_process() + self._wait_for_close_later() if self._session is not None: self._session.abort() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 2cd7f620c..ee6125aee 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -8959,9 +8959,9 @@ def _finalize_unclosed_file(fs, where, *, text=False): # - "own part": the collector runs inside that first part; # - "last reference in pool": its last reference is dropped in the # thread of the pool, while the first part is queued behind it. - # A close handed over to another thread starts only once the - # finalization has returned, after the collector has finalized the - # buffer of the file. Returns the files closed by `_close_collected`. + # Checks that a close handed over during a collection starts once the + # collection has ended. Returns the files closed by + # `_close_collected`. holder = [] ready = threading.Event() finalized = threading.Event() @@ -8994,17 +8994,24 @@ def upload_part(**kwargs): fs.core.upload_part.side_effect = upload_part close_collected = S3File._close_collected + start_close = S3File._start_close closes = [] + started_in_collection = [] - def close_after_finalization(file, complete): + def count_close(file, complete): closes.append(file) - if threading.current_thread().name.startswith("pyathena-close-"): - assert finalized.wait(EVENT_TIMEOUT) close_collected(file, complete) + def record_start(file, complete): + started_in_collection.append(S3File._in_collection()) + start_close(file, complete) + gc.disable() try: - with mock.patch.object(S3File, "_close_collected", close_after_finalization): + with ( + mock.patch.object(S3File, "_close_collected", count_close), + mock.patch.object(S3File, "_start_close", record_start), + ): file = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) if text: holder.append(io.TextIOWrapper(file, encoding="utf-8")) @@ -9024,6 +9031,7 @@ def close_after_finalization(file, complete): assert finalized.wait(EVENT_TIMEOUT) finally: gc.enable() + assert not any(started_in_collection) return closes @pytest.mark.parametrize( @@ -9135,6 +9143,84 @@ def start_unless_closing(thread): fs.core.complete_multipart_upload.assert_not_called() assert "no thread could be started" in caplog.text + @pytest.mark.parametrize( + ("method", "autocommit"), [("close", True), ("close", False), ("commit", False)] + ) + def test_close_handed_over_is_waited_for(self, method, autocommit): + # GH-1364: a close of a file that the finalizer of its text wrapper + # hands over during a collection is waited for by a later close() or + # commit() of the file, as at the end of a transaction, which would + # otherwise complete the upload without the text that the wrapper + # wrote when it was finalized, or race with the close. + fs = self._make_append_fs(b"") + release = threading.Event() + close_collected = S3File._close_collected + + def close_when_released(file, complete): + assert release.wait(EVENT_TIMEOUT) + close_collected(file, complete) + + gc.disable() + try: + with mock.patch.object(S3File, "_close_collected", close_when_released): + f = S3File( + fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=autocommit + ) + wrapper = io.TextIOWrapper(f, encoding="utf-8") + wrapper.write("abcde") + wrapper.flush() + wrapper.write("fg") + wrapper.cycle = wrapper + del wrapper + gc.collect() + waiter = threading.Thread(target=getattr(f, method)) + waiter.start() + waiter.join(0.2) + assert waiter.is_alive() + release.set() + waiter.join(EVENT_TIMEOUT) + assert not waiter.is_alive() + assert f.closed + if not autocommit and method == "close": + f.commit() + finally: + gc.enable() + release.set() + shutdown_bounded(fs._request_pool) + assert self._uploaded_object(fs, b"") == b"abcdefg" + + def test_close_handed_over_is_not_waited_for_in_collection(self): + # GH-1364: commit() called during the collection that handed over + # the close of the file does not wait for that close, which starts + # once the collection has ended. + fs = self._make_append_fs(b"") + + class Committer: + def __init__(self, file): + self.file = file + self.cycle = self + + def __del__(self): + self.file.close() + self.file.commit() + + gc.disable() + try: + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=False) + f.write(b"ab") + Committer(f) + del f + # A daemon thread, so that a collection that waits forever does + # not keep the process alive. + collector = threading.Thread(target=gc.collect, daemon=True) + collector.start() + collector.join(EVENT_TIMEOUT) + assert not collector.is_alive() + finally: + gc.enable() + shutdown_bounded(fs._request_pool) + assert self._uploaded_object(fs, b"") == b"ab" + @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): # A file larger than the block size is uploaded in parts and completed From 01b49c3093bd61d7d46e4a15892139357e3ea596 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 02:32:19 +0900 Subject: [PATCH 5/8] Close handed-over files in _thread threads and never wait where unsafe (GH-1364) A handed-over close runs in a thread of the _thread module, since Thread.start() takes locks of the threading module that the finalizing or collecting thread may hold, and waits for the new thread, which may take another (_shutdown_locks_lock on 3.11 and 3.12, _active_limbo_lock on every version). An exit waits for these closes with atexit, since the interpreter does not join such threads; a fork clears them in the child, and a close queued in the parent is not run in a child. During a collection or in a thread of an S3ThreadPool, where a wait for a handed-over close could last forever, close() returns without waiting, and commit() and discard() raise RuntimeError instead of completing the upload without the data that the close uploads. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 10 +- pyathena/filesystem/s3.py | 122 +++++++++++---- tests/pyathena/filesystem/test_s3.py | 225 +++++++++++++++++++++------ 3 files changed, 277 insertions(+), 80 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index e17181615..33fc7039d 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -210,8 +210,8 @@ closed is closed when it is finalized, which commits its written data unless - When it is finalized during a pass of the cyclic garbage collector, it is closed in a new thread that starts once the collection has ended. When it is finalized in a thread of an `S3ThreadPool`, such as the request pool of a filesystem, it is closed - in a new thread. If no thread can be started, the file is left unclosed and its - multipart upload incomplete. + in a new thread. The interpreter waits for these closes when it exits. If no thread + can be started, the file is left unclosed and its multipart upload incomplete. - When its last reference is dropped elsewhere, it is closed before that returns. - During interpreter finalization, it is closed in the finalizing thread instead of a new thread. The request pools take no more calls then, so a file with data left to @@ -223,8 +223,10 @@ of the file is not committed, and its multipart upload is aborted. `close()` called during a pass of the cyclic garbage collector, such as by another finalizer, closes the file in the same way, in a new thread that starts once the collection has ended, and returns before the file is closed. A later `close()`, -`commit()`, or `discard()` of the file waits for that close to end, unless it is -called during a collection. +`commit()`, or `discard()` of the file waits for that close to end. During a collection +or in a thread of an `S3ThreadPool`, `close()` returns without waiting instead, and +`commit()` and `discard()` raise `RuntimeError`. If no thread could be started for that +close, a later `close()` closes the file. A file opened in text mode is closed by its text wrapper when the wrapper is finalized. The wrapper writes its remaining text to the file in the thread that finalizes it, and diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index c795fb398..90f593de9 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -2,6 +2,8 @@ from __future__ import annotations +import _thread +import atexit import contextlib import contextvars import functools @@ -85,6 +87,10 @@ # files that it handed over (S3File._track_collection). _collecting = threading.local() +# The ends of the closes that run in threads of the _thread module, which an +# exit waits for (S3File._wait_for_handed_over_closes). +_handed_over_closes: set[threading.Event] = set() + # The request parameters on which the authorization of a lookup (HeadObject, # HeadBucket, or ListObjectsV2) depends. _LOOKUP_REQUEST_PARAMETERS = frozenset( @@ -4320,8 +4326,9 @@ def close(self) -> None: another finalizer while the cyclic garbage collector runs in this thread, such as that of a ``TextIOWrapper``, the file is closed in a new thread once the collection has ended (see :meth:`__del__`), and - this returns before that. Otherwise, a close handed over to another - thread is waited for. + this returns before that. A close handed over to another thread is + waited for, except during a collection or in a thread of an + :class:`S3ThreadPool`, where this returns without waiting. """ if self._inherited(): self.closed = True @@ -4329,8 +4336,8 @@ def close(self) -> None: if not self.closed and S3File._in_collection(): self._close_later() return - self._wait_for_close_later() - self._close() + if self._wait_for_close_later(): + self._close() def _close(self) -> None: """Close the file in this thread, flushing any written data, and shut down its executor.""" @@ -4413,6 +4420,7 @@ def _close_later(self) -> None: buffer = BytesIO(self.buffer.getvalue()) buffer.seek(self.buffer.tell()) self.buffer = buffer + _handed_over_closes.add(self._closed_later) if S3File._in_collection(): if not hasattr(_collecting, "pending"): _collecting.pending = [] @@ -4423,12 +4431,17 @@ def _close_later(self) -> None: def _start_close(self, complete: bool) -> None: """Start the close that :meth:`_close_later` handed over. - The thread runs with an empty context, so that the stop scopes of - the finalizing thread do not stop the requests of the file. During - interpreter finalization, the file is closed in this thread without - starting one: Python 3.11 waits forever for a thread started then, - and later versions raise. No other thread runs requests then, and the - pools take no more calls, so the close does not wait for a part + The file is closed in a thread of the ``_thread`` module, since + ``Thread.start()`` takes locks of the ``threading`` module, which this + thread, in which a finalizer or the collector runs, may hold, and + waits for the new thread, which may take another. The interpreter + does not join such a thread, so an exit waits for the close instead + (see :meth:`_wait_for_handed_over_closes`). The thread runs with an + empty context, so that the stop scopes of the finalizing thread do + not stop the requests of the file. During interpreter finalization, + the file is closed in this thread without starting one, since no + thread started then runs. No other thread runs requests then, and + the pools take no more calls, so the close does not wait for a part queued behind it. If no thread can be started otherwise, the file is left unclosed rather than closed in this thread, where it could wait forever. @@ -4436,24 +4449,24 @@ def _start_close(self, complete: bool) -> None: Args: complete: Whether the written data is complete. """ + if self._inherited(): + # Handed over in the parent, which forked before the collection + # ended. + self.closed = True + self._end_close_later() + return if sys.is_finalizing(): - self._closer = threading.current_thread() self._close_collected(complete) return try: - self._closer = threading.Thread( - target=contextvars.Context().run, - args=(self._close_collected, complete), - name=f"pyathena-close-s3://{self.bucket}/{self.key}", - ) - self._closer.start() + _thread.start_new_thread(contextvars.Context().run, (self._close_collected, complete)) except Exception: # Such as RuntimeError, when no more threads can be started. _logger.warning( f"s3://{self.bucket}/{self.key} was not closed, and no thread could be " "started to close it when it was finalized; its data was not committed." ) - self._closed_later.set() + self._end_close_later() def _close_collected(self, complete: bool) -> None: """Close a finalized file, logging a failure instead of raising it. @@ -4462,6 +4475,7 @@ def _close_collected(self, complete: bool) -> None: complete: Whether the written data is complete. If not, the write is discarded instead of committed. """ + self._closer_ident = threading.get_ident() try: if complete: self._close() @@ -4474,21 +4488,61 @@ def _close_collected(self, complete: bool) -> None: except Exception: _logger.exception(f"Failed to close s3://{self.bucket}/{self.key}.") finally: - self._closed_later.set() + self._end_close_later() + + def _end_close_later(self) -> None: + """Mark the close that :meth:`_close_later` handed over as ended.""" + self._closed_later.set() + _handed_over_closes.discard(self._closed_later) - def _wait_for_close_later(self) -> None: + @staticmethod + def _wait_for_handed_over_closes() -> None: + """Wait for the closes that run in threads of the ``_thread`` module. + + Registered with :mod:`atexit`, since the interpreter does not join + these threads when it exits. + """ + for closed in list(_handed_over_closes): + closed.wait() + + def _wait_for_close_later(self) -> bool: """Wait for a close that :meth:`_close_later` handed over to end. - Not waited for in the thread that closes the file, or during a - collection, whose end starts the close. + It is not waited for in the thread that runs it, during a + collection, whose end starts it, or in a thread of an + :class:`S3ThreadPool`, since it may wait for a request queued behind + that thread. + + Returns: + False if a handed-over close runs and cannot be waited for in + this thread; otherwise True. """ closed_later = getattr(self, "_closed_later", None) if ( - closed_later is not None - and getattr(self, "_closer", None) is not threading.current_thread() - and not S3File._in_collection() + closed_later is None + or closed_later.is_set() + or getattr(self, "_closer_ident", None) == threading.get_ident() ): - closed_later.wait() + return True + if S3File._in_collection() or S3ThreadPool.in_pool_thread(): + return False + closed_later.wait() + return True + + def _raise_unless_closed_later(self) -> None: + """Wait for a handed-over close to end, or raise if it cannot be waited for here. + + Raises: + RuntimeError: If a close of the file was handed over to another + thread and runs, and this is called during a collection or + in a thread of an :class:`S3ThreadPool`. + """ + if not self._wait_for_close_later(): + raise RuntimeError( + f"s3://{self.bucket}/{self.key} is being closed in another thread, which " + "cannot be waited for during a garbage collection or in a thread of an " + "S3ThreadPool." + ) def read(self, length: int = -1) -> bytes: """Read bytes from the file. @@ -4751,12 +4805,14 @@ def commit(self) -> None: FileExistsError: If an object was created at the path after the file was opened in exclusive-create mode. RuntimeError: In a forked child process of the process that - opened the file. + opened the file, or if a close of the file that was handed + over to another thread runs and this is called during a + garbage collection or in a thread of an :class:`S3ThreadPool`. S3MultipartStateError: If the write was discarded, or the write or an earlier commit failed. """ self._check_process() - self._wait_for_close_later() + self._raise_unless_closed_later() if self._session is None: # A file opened for reading has nothing to commit. return @@ -4794,10 +4850,12 @@ def discard(self) -> None: Raises: RuntimeError: In a forked child process of the process that - opened the file. + opened the file, or if a close of the file that was handed + over to another thread runs and this is called during a + garbage collection or in a thread of an :class:`S3ThreadPool`. """ self._check_process() - self._wait_for_close_later() + self._raise_unless_closed_later() if self._session is not None: self._session.abort() @@ -5000,3 +5058,7 @@ def _get_ranges( gc.callbacks.append(S3File._track_collection) +atexit.register(S3File._wait_for_handed_over_closes) +if hasattr(os, "register_at_fork"): + # The closes of the parent do not run in a child. + os.register_at_fork(after_in_child=_handed_over_closes.clear) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index ee6125aee..9b0fce0cb 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,7 +1,9 @@ +import _thread import array import asyncio import bz2 import contextlib +import contextvars import errno import functools import gc @@ -8952,26 +8954,35 @@ def interrupted(timeout=None): def _finalize_unclosed_file(fs, where, *, text=False): # Write a multipart file without closing it, and finalize it in the # given place, using the only thread of the request pool: - # - "last reference": its last reference is dropped in this thread; - # - "collection": the cyclic garbage collector runs in this thread; + # - "last reference": its last reference is dropped in a user thread; + # - "collection": the cyclic garbage collector runs in a user thread; # - "queued part": the collector runs in the thread of the pool, # while the first part of the file is queued behind it; # - "own part": the collector runs inside that first part; # - "last reference in pool": its last reference is dropped in the # thread of the pool, while the first part is queued behind it. - # Checks that a close handed over during a collection starts once the - # collection has ended. Returns the files closed by - # `_close_collected`. + # The user thread is a daemon, so that a finalizer that waits forever + # fails the test instead of hanging it. Checks that a close handed + # over during a collection starts once the collection has ended, and + # that the finalization starts no thread through `threading`, whose + # locks the finalizing thread may hold. Returns once each handed-over + # close has started or failed to start: the files closed by + # `_close_collected`, and the threads that finalized the file. holder = [] ready = threading.Event() finalized = threading.Event() + finalizers = [] + finalizing = set() def finalize(): assert ready.wait(EVENT_TIMEOUT) + finalizers.append(threading.current_thread()) + finalizing.add(threading.current_thread()) if where.startswith("last reference"): holder.clear() else: gc.collect() + finalizing.discard(threading.current_thread()) finalized.set() if where in ("queued part", "last reference in pool"): @@ -8995,22 +9006,39 @@ def upload_part(**kwargs): close_collected = S3File._close_collected start_close = S3File._start_close + start = threading.Thread.start closes = [] started_in_collection = [] + started_by_finalizer = [] + resolved = threading.Semaphore(0) def count_close(file, complete): closes.append(file) + resolved.release() close_collected(file, complete) def record_start(file, complete): started_in_collection.append(S3File._in_collection()) start_close(file, complete) + if file not in closes and file._closed_later.is_set(): + # No thread could be started. + resolved.release() + + def record_thread_start(thread): + # Except a worker of the request pool, which the wrapper of a + # text-mode file may start when it writes to the file. + if threading.current_thread() in finalizing and not thread.name.startswith( + "ThreadPoolExecutor" + ): + started_by_finalizer.append(thread) + start(thread) gc.disable() try: with ( mock.patch.object(S3File, "_close_collected", count_close), mock.patch.object(S3File, "_start_close", record_start), + mock.patch.object(threading.Thread, "start", record_thread_start), ): file = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) if text: @@ -9027,12 +9055,15 @@ def record_start(file, complete): del file ready.set() if where in ("last reference", "collection"): - finalize() + threading.Thread(target=finalize, daemon=True).start() assert finalized.wait(EVENT_TIMEOUT) + for _ in started_in_collection: + assert resolved.acquire(timeout=EVENT_TIMEOUT) finally: gc.enable() assert not any(started_in_collection) - return closes + assert not started_by_finalizer + return SimpleNamespace(closes=closes, finalizers=finalizers) @pytest.mark.parametrize( ("where", "text"), @@ -9066,18 +9097,18 @@ def complete(*args, **kwargs): fs.core.complete_multipart_upload.side_effect = complete try: - closes = self._finalize_unclosed_file(fs, where, text=text) + finalization = self._finalize_unclosed_file(fs, where, text=text) assert completed.wait(EVENT_TIMEOUT) finally: shutdown_bounded(fs._request_pool) assert self._uploaded_object(fs, b"") == b"abcde" # A sync file commits in the thread that closes it. ((committer, in_pool),) = committers - assert (committer is threading.current_thread()) == (where == "last reference") + assert (committer is finalization.finalizers[0]) == (where == "last reference") assert not in_pool # The close is handed over once, also when the finalizer of a text # wrapper closes the file before the finalizer of the file runs. - assert len(closes) == (0 if where == "last reference" else 1) + assert len(finalization.closes) == (0 if where == "last reference" else 1) def test_finalized_unclosed_file_with_finalized_buffer(self, caplog): # GH-1364: an unclosed file whose buffer the collector finalized @@ -9105,6 +9136,17 @@ def test_finalized_unclosed_file_with_finalized_buffer(self, caplog): fs.core.complete_multipart_upload.assert_not_called() assert "buffer was finalized" in caplog.text + @staticmethod + def _start_new_thread_unless_closing(start_new_thread): + # Fails to start the thread of a handed-over close, as when no more + # threads can be started. + def start(function, args): + if isinstance(getattr(function, "__self__", None), contextvars.Context): + raise RuntimeError("can't start new thread") + return start_new_thread(function, args) + + return start + @pytest.mark.parametrize("finalizing", [False, True]) def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): # GH-1364: if no thread can be started for the close of a finalized @@ -9120,40 +9162,41 @@ def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): committers.append(threading.current_thread()), mock.DEFAULT, )[1] - start = threading.Thread.start - - def start_unless_closing(thread): - if thread.name == "pyathena-close-s3://bucket/key.txt": - raise RuntimeError("can't create new thread at interpreter shutdown") - start(thread) - try: with ( - mock.patch.object(threading.Thread, "start", start_unless_closing), + mock.patch.object( + _thread, + "start_new_thread", + self._start_new_thread_unless_closing(_thread.start_new_thread), + ), mock.patch.object(sys, "is_finalizing", return_value=finalizing), caplog.at_level(logging.WARNING, logger="pyathena.filesystem.s3"), ): - self._finalize_unclosed_file(fs, "collection" if finalizing else "queued part") + finalization = self._finalize_unclosed_file( + fs, "collection" if finalizing else "queued part" + ) finally: shutdown_bounded(fs._request_pool) if finalizing: assert self._uploaded_object(fs, b"") == b"abcde" - assert committers == [threading.current_thread()] + assert committers == finalization.finalizers else: fs.core.complete_multipart_upload.assert_not_called() assert "no thread could be started" in caplog.text @pytest.mark.parametrize( - ("method", "autocommit"), [("close", True), ("close", False), ("commit", False)] + ("method", "autocommit"), + [("close", True), ("close", False), ("commit", False), ("discard", False)], ) def test_close_handed_over_is_waited_for(self, method, autocommit): # GH-1364: a close of a file that the finalizer of its text wrapper - # hands over during a collection is waited for by a later close() or - # commit() of the file, as at the end of a transaction, which would - # otherwise complete the upload without the text that the wrapper - # wrote when it was finalized, or race with the close. + # hands over during a collection is waited for by a later close(), + # commit(), or discard() of the file, as at the end of a transaction, + # which would otherwise complete the upload without the text that + # the wrapper wrote when it was finalized, or race with the close. fs = self._make_append_fs(b"") release = threading.Event() + waiting = threading.Event() close_collected = S3File._close_collected def close_when_released(file, complete): @@ -9173,53 +9216,143 @@ def close_when_released(file, complete): wrapper.cycle = wrapper del wrapper gc.collect() - waiter = threading.Thread(target=getattr(f, method)) + # Observes the wait for the handed-over close, which ends when + # it is released. + closed_later = f._closed_later + wait = closed_later.wait + + def observed_wait(timeout=None): + waiting.set() + return wait(timeout) + + closed_later.wait = observed_wait + waiter = threading.Thread(target=getattr(f, method), daemon=True) waiter.start() - waiter.join(0.2) - assert waiter.is_alive() + assert waiting.wait(EVENT_TIMEOUT) release.set() waiter.join(EVENT_TIMEOUT) assert not waiter.is_alive() assert f.closed - if not autocommit and method == "close": + if method == "close" and not autocommit: f.commit() finally: gc.enable() release.set() shutdown_bounded(fs._request_pool) + if method == "discard": + fs._call.assert_called_once_with( + S3_CLIENT.abort_multipart_upload, + Bucket="bucket", + Key="key.txt", + UploadId="uploadid", + ) + fs.core.complete_multipart_upload.assert_not_called() + else: + assert self._uploaded_object(fs, b"") == b"abcdefg" + + @pytest.mark.parametrize("where", ["collection", "pool"]) + def test_close_handed_over_is_not_waited_for(self, where): + # GH-1364: where a handed-over close of a file cannot be waited for, + # during a collection, whose end starts it, or in the thread of the + # request pool that one of its parts may be queued behind, close() + # returns and commit() raises, instead of waiting forever or + # completing the upload without the data that the close uploads. A + # commit() once the close has ended commits all of the data. + fs = self._make_append_fs(b"") + fs._request_pool = S3ThreadPool(1) + release = threading.Event() + errors = [] + close_collected = S3File._close_collected + + def close_when_released(file, complete): + assert release.wait(EVENT_TIMEOUT) + close_collected(file, complete) + + def close_and_commit(file): + try: + file.close() + # The close is left to the thread that runs it. + assert not file.closed + file.commit() + except RuntimeError as e: + errors.append(e) + finally: + release.set() + + class Closer: + def __init__(self, file): + self.file = file + self.cycle = self + + def __del__(self): + # Hands the close over. + self.file.close() + if where == "collection": + close_and_commit(self.file) + + gc.disable() + try: + with mock.patch.object(S3File, "_close_collected", close_when_released): + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=False) + f.write(b"abcde") + f.write(b"fg") + Closer(f) + # A daemon thread, so that a collection that waits forever + # does not keep the process alive. + collector = threading.Thread(target=gc.collect, daemon=True) + collector.start() + collector.join(EVENT_TIMEOUT) + assert not collector.is_alive() + if where == "pool": + fs._request_pool.submit(close_and_commit, f).result(EVENT_TIMEOUT) + assert len(errors) == 1 + fs.core.complete_multipart_upload.assert_not_called() + # Waits for the handed-over close. + f.close() + f.commit() + finally: + gc.enable() + release.set() + shutdown_bounded(fs._request_pool) assert self._uploaded_object(fs, b"") == b"abcdefg" - def test_close_handed_over_is_not_waited_for_in_collection(self): - # GH-1364: commit() called during the collection that handed over - # the close of the file does not wait for that close, which starts - # once the collection has ended. + def test_close_handed_over_not_started_is_not_waited_for(self): + # GH-1364: when no thread could be started for a handed-over close, + # a later close() of the file closes it instead of waiting forever. fs = self._make_append_fs(b"") - class Committer: + class Closer: def __init__(self, file): self.file = file self.cycle = self def __del__(self): + # Hands the close over. self.file.close() - self.file.commit() gc.disable() try: - f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=False) - f.write(b"ab") - Committer(f) - del f - # A daemon thread, so that a collection that waits forever does - # not keep the process alive. - collector = threading.Thread(target=gc.collect, daemon=True) - collector.start() - collector.join(EVENT_TIMEOUT) - assert not collector.is_alive() + with mock.patch.object( + _thread, + "start_new_thread", + self._start_new_thread_unless_closing(_thread.start_new_thread), + ): + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) + f.write(b"abcde") + Closer(f) + gc.collect() + assert not f.closed + # A daemon thread, so that a close that waits forever does not + # keep the process alive. + closer = threading.Thread(target=f.close, daemon=True) + closer.start() + closer.join(EVENT_TIMEOUT) + assert not closer.is_alive() finally: gc.enable() shutdown_bounded(fs._request_pool) - assert self._uploaded_object(fs, b"") == b"ab" + assert f.closed + assert self._uploaded_object(fs, b"") == b"abcde" @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): From aba5b654ff0e1cfdd1d6621d3ff7ef6b50875702 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 03:21:32 +0900 Subject: [PATCH 6/8] Run a handed-over close that could not start on the next call (GH-1364) A close for which no thread could be started is no longer marked as ended, so that a later commit() does not complete the upload without the data that the close uploads: the next close(), commit(), or discard() outside a collection and the pool threads runs it first, and an exit does not wait for it. In a forked child, an inherited file whose close was queued is only marked closed, and its event is left alone. The exit waits until no handed-over close is left, also for those handed over while it waits. The tests that hand over a close fail when one of their closes does not end, instead of leaving the exit to wait for it. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 2 +- pyathena/filesystem/s3.py | 31 +++++++++++------- tests/pyathena/filesystem/test_s3.py | 48 +++++++++++++++++++++------- 3 files changed, 58 insertions(+), 23 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 33fc7039d..7abe05388 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -226,7 +226,7 @@ collection has ended, and returns before the file is closed. A later `close()`, `commit()`, or `discard()` of the file waits for that close to end. During a collection or in a thread of an `S3ThreadPool`, `close()` returns without waiting instead, and `commit()` and `discard()` raise `RuntimeError`. If no thread could be started for that -close, a later `close()` closes the file. +close, a later `close()`, `commit()`, or `discard()` of the file closes it first. A file opened in text mode is closed by its text wrapper when the wrapper is finalized. The wrapper writes its remaining text to the file in the thread that finalizes it, and diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 90f593de9..18a0b4df2 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -4444,16 +4444,17 @@ def _start_close(self, complete: bool) -> None: the pools take no more calls, so the close does not wait for a part queued behind it. If no thread can be started otherwise, the file is left unclosed rather than closed in this thread, where it could wait - forever. + forever, until a later :meth:`close`, :meth:`commit`, or + :meth:`discard` closes it (see :meth:`_wait_for_close_later`). Args: complete: Whether the written data is complete. """ if self._inherited(): # Handed over in the parent, which forked before the collection - # ended. + # ended. Its event, whose lock the fork may have copied held, is + # left alone. self.closed = True - self._end_close_later() return if sys.is_finalizing(): self._close_collected(complete) @@ -4466,7 +4467,8 @@ def _start_close(self, complete: bool) -> None: f"s3://{self.bucket}/{self.key} was not closed, and no thread could be " "started to close it when it was finalized; its data was not committed." ) - self._end_close_later() + self._unstarted_close: bool | None = complete + _handed_over_closes.discard(self._closed_later) def _close_collected(self, complete: bool) -> None: """Close a finalized file, logging a failure instead of raising it. @@ -4492,18 +4494,20 @@ def _close_collected(self, complete: bool) -> None: def _end_close_later(self) -> None: """Mark the close that :meth:`_close_later` handed over as ended.""" - self._closed_later.set() _handed_over_closes.discard(self._closed_later) + self._closed_later.set() @staticmethod def _wait_for_handed_over_closes() -> None: """Wait for the closes that run in threads of the ``_thread`` module. Registered with :mod:`atexit`, since the interpreter does not join - these threads when it exits. + these threads when it exits. Also waits for the closes that are + handed over while it waits. """ - for closed in list(_handed_over_closes): - closed.wait() + while _handed_over_closes: + for closed in list(_handed_over_closes): + closed.wait() def _wait_for_close_later(self) -> bool: """Wait for a close that :meth:`_close_later` handed over to end. @@ -4511,11 +4515,12 @@ def _wait_for_close_later(self) -> bool: It is not waited for in the thread that runs it, during a collection, whose end starts it, or in a thread of an :class:`S3ThreadPool`, since it may wait for a request queued behind - that thread. + that thread. A close for which no thread could be started is run + here instead, where it can be waited for. Returns: - False if a handed-over close runs and cannot be waited for in - this thread; otherwise True. + False if a handed-over close has not ended and cannot be waited + for in this thread; otherwise True. """ closed_later = getattr(self, "_closed_later", None) if ( @@ -4526,6 +4531,10 @@ def _wait_for_close_later(self) -> bool: return True if S3File._in_collection() or S3ThreadPool.in_pool_thread(): return False + unstarted = getattr(self, "_unstarted_close", None) + if unstarted is not None: + self._unstarted_close = None + self._close_collected(unstarted) closed_later.wait() return True diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 9b0fce0cb..aac6e022d 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -47,7 +47,12 @@ import pyathena from pyathena.filesystem import register_s3_filesystem, s3_filesystem -from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem +from pyathena.filesystem.s3 import ( + CompressedBuffer, + S3File, + S3FileSystem, + _handed_over_closes, +) from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3Batch, _S3PoolScope @@ -8950,6 +8955,15 @@ def interrupted(timeout=None): assert operation.calls == 1 assert file.closed + @pytest.fixture + def handed_over_closes_end(self): + # Each close that the test hands over ends; one that does not fails + # the test instead of leaving an exit to wait for it forever. + yield + pending = list(_handed_over_closes) + _handed_over_closes.clear() + assert all(closed.wait(EVENT_TIMEOUT) for closed in pending) + @staticmethod def _finalize_unclosed_file(fs, where, *, text=False): # Write a multipart file without closing it, and finalize it in the @@ -9020,7 +9034,7 @@ def count_close(file, complete): def record_start(file, complete): started_in_collection.append(S3File._in_collection()) start_close(file, complete) - if file not in closes and file._closed_later.is_set(): + if file not in closes and file._closed_later not in _handed_over_closes: # No thread could be started. resolved.release() @@ -9076,6 +9090,7 @@ def record_thread_start(thread): ("last reference in pool", False), ], ) + @pytest.mark.usefixtures("handed_over_closes_end") def test_finalized_unclosed_file(self, where, text): # GH-1364: an unclosed file is committed when it is finalized. Only a # file whose last reference is dropped in a user thread is committed @@ -9110,6 +9125,7 @@ def complete(*args, **kwargs): # wrapper closes the file before the finalizer of the file runs. assert len(finalization.closes) == (0 if where == "last reference" else 1) + @pytest.mark.usefixtures("handed_over_closes_end") def test_finalized_unclosed_file_with_finalized_buffer(self, caplog): # GH-1364: an unclosed file whose buffer the collector finalized # first is discarded, instead of committed without the buffered data. @@ -9148,6 +9164,7 @@ def start(function, args): return start @pytest.mark.parametrize("finalizing", [False, True]) + @pytest.mark.usefixtures("handed_over_closes_end") def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): # GH-1364: if no thread can be started for the close of a finalized # file, it is left unclosed, and its data uncommitted, rather than @@ -9188,6 +9205,7 @@ def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): ("method", "autocommit"), [("close", True), ("close", False), ("commit", False), ("discard", False)], ) + @pytest.mark.usefixtures("handed_over_closes_end") def test_close_handed_over_is_waited_for(self, method, autocommit): # GH-1364: a close of a file that the finalizer of its text wrapper # hands over during a collection is waited for by a later close(), @@ -9251,6 +9269,7 @@ def observed_wait(timeout=None): assert self._uploaded_object(fs, b"") == b"abcdefg" @pytest.mark.parametrize("where", ["collection", "pool"]) + @pytest.mark.usefixtures("handed_over_closes_end") def test_close_handed_over_is_not_waited_for(self, where): # GH-1364: where a handed-over close of a file cannot be waited for, # during a collection, whose end starts it, or in the thread of the @@ -9316,9 +9335,13 @@ def __del__(self): shutdown_bounded(fs._request_pool) assert self._uploaded_object(fs, b"") == b"abcdefg" - def test_close_handed_over_not_started_is_not_waited_for(self): + @pytest.mark.usefixtures("handed_over_closes_end") + @pytest.mark.parametrize("method", ["close", "commit"]) + def test_close_handed_over_not_started(self, method): # GH-1364: when no thread could be started for a handed-over close, - # a later close() of the file closes it instead of waiting forever. + # a later close() or commit() of the file closes it first, instead of + # waiting forever, or completing the upload, as at the end of a + # transaction, without the data that the close uploads. fs = self._make_append_fs(b"") class Closer: @@ -9337,22 +9360,25 @@ def __del__(self): "start_new_thread", self._start_new_thread_unless_closing(_thread.start_new_thread), ): - f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) + f = S3File( + fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=method == "close" + ) f.write(b"abcde") + f.write(b"fg") Closer(f) gc.collect() assert not f.closed - # A daemon thread, so that a close that waits forever does not + # A daemon thread, so that a call that waits forever does not # keep the process alive. - closer = threading.Thread(target=f.close, daemon=True) - closer.start() - closer.join(EVENT_TIMEOUT) - assert not closer.is_alive() + caller = threading.Thread(target=getattr(f, method), daemon=True) + caller.start() + caller.join(EVENT_TIMEOUT) + assert not caller.is_alive() finally: gc.enable() shutdown_bounded(fs._request_pool) assert f.closed - assert self._uploaded_object(fs, b"") == b"abcde" + assert self._uploaded_object(fs, b"") == b"abcdefg" @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): From 8ac2ef11f22b163bd114384de46909ad7ba947cb Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 04:25:06 +0900 Subject: [PATCH 7/8] Wake waiters on a failed start and close after the exit wait inline (GH-1364) A failed start of a handed-over close now wakes the callers that wait for it, which then run the close where they can wait, as a later call does. A write that fails during the collection that handed over the close of its file leaves the discard of the write to that close, which it cannot wait for, instead of logging the refused discard and leaving the upload. A close handed over after the exit has waited for the handed-over closes, such as in an exit handler that runs later, is run in the finalizing thread, since nothing would wait for a new thread. The test that waits for a handed-over close from the test thread now does so from a daemon thread with a timeout. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 10 +- pyathena/filesystem/s3.py | 60 ++++++---- tests/pyathena/filesystem/test_s3.py | 161 ++++++++++++++++++++++----- 3 files changed, 179 insertions(+), 52 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 7abe05388..39664eef9 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -213,9 +213,10 @@ closed is closed when it is finalized, which commits its written data unless in a new thread. The interpreter waits for these closes when it exits. If no thread can be started, the file is left unclosed and its multipart upload incomplete. - When its last reference is dropped elsewhere, it is closed before that returns. -- During interpreter finalization, it is closed in the finalizing thread instead of a - new thread. The request pools take no more calls then, so a file with data left to - upload as a part is not committed, and its multipart upload is aborted. +- During interpreter finalization, or after the exit has waited for the closes in new + threads, it is closed in the finalizing thread instead of a new thread. The request + pools take no more calls then, so a file with data left to upload as a part is not + committed, and its multipart upload is aborted. If the garbage collector finalizes the write buffer of a file before the file, the data of the file is not committed, and its multipart upload is aborted. @@ -226,7 +227,8 @@ collection has ended, and returns before the file is closed. A later `close()`, `commit()`, or `discard()` of the file waits for that close to end. During a collection or in a thread of an `S3ThreadPool`, `close()` returns without waiting instead, and `commit()` and `discard()` raise `RuntimeError`. If no thread could be started for that -close, a later `close()`, `commit()`, or `discard()` of the file closes it first. +close, a later `close()`, `commit()`, or `discard()` of the file elsewhere closes it +first. A file opened in text mode is closed by its text wrapper when the wrapper is finalized. The wrapper writes its remaining text to the file in the thread that finalizes it, and diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 18a0b4df2..392bf0711 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -88,8 +88,9 @@ _collecting = threading.local() # The ends of the closes that run in threads of the _thread module, which an -# exit waits for (S3File._wait_for_handed_over_closes). +# exit waits for (S3File._wait_for_handed_over_closes), and whether it has. _handed_over_closes: set[threading.Event] = set() +_handed_over_closes_waited = False # The request parameters on which the authorization of a lookup (HeadObject, # HeadBucket, or ListObjectsV2) depends. @@ -4439,13 +4440,15 @@ def _start_close(self, complete: bool) -> None: (see :meth:`_wait_for_handed_over_closes`). The thread runs with an empty context, so that the stop scopes of the finalizing thread do not stop the requests of the file. During interpreter finalization, - the file is closed in this thread without starting one, since no - thread started then runs. No other thread runs requests then, and - the pools take no more calls, so the close does not wait for a part - queued behind it. If no thread can be started otherwise, the file is - left unclosed rather than closed in this thread, where it could wait - forever, until a later :meth:`close`, :meth:`commit`, or - :meth:`discard` closes it (see :meth:`_wait_for_close_later`). + and once the exit has waited for the handed-over closes, the file is + closed in this thread without starting one, since no thread started + during finalization runs, and none started afterwards is waited for. + No other thread runs requests then, and the pools take no more + calls, so the close does not wait for a part queued behind it. If no + thread can be started otherwise, the file is left unclosed rather + than closed in this thread, where it could wait forever, until a + later :meth:`close`, :meth:`commit`, or :meth:`discard` closes it + (see :meth:`_wait_for_close_later`). Args: complete: Whether the written data is complete. @@ -4456,7 +4459,7 @@ def _start_close(self, complete: bool) -> None: # left alone. self.closed = True return - if sys.is_finalizing(): + if sys.is_finalizing() or _handed_over_closes_waited: self._close_collected(complete) return try: @@ -4469,6 +4472,8 @@ def _start_close(self, complete: bool) -> None: ) self._unstarted_close: bool | None = complete _handed_over_closes.discard(self._closed_later) + # Wakes the callers that wait for it, which then close the file. + self._closed_later.set() def _close_collected(self, complete: bool) -> None: """Close a finalized file, logging a failure instead of raising it. @@ -4487,6 +4492,8 @@ def _close_collected(self, complete: bool) -> None: "finalized before it; its data was not committed." ) self._close_without_commit() + if getattr(self, "_discard_later", False): + self.discard() except Exception: _logger.exception(f"Failed to close s3://{self.bucket}/{self.key}.") finally: @@ -4503,11 +4510,14 @@ def _wait_for_handed_over_closes() -> None: Registered with :mod:`atexit`, since the interpreter does not join these threads when it exits. Also waits for the closes that are - handed over while it waits. + handed over while it waits; later ones, such as in an exit handler + that runs after it, are closed in the finalizing thread. """ + global _handed_over_closes_waited while _handed_over_closes: for closed in list(_handed_over_closes): closed.wait() + _handed_over_closes_waited = True def _wait_for_close_later(self) -> bool: """Wait for a close that :meth:`_close_later` handed over to end. @@ -4523,19 +4533,18 @@ def _wait_for_close_later(self) -> bool: for in this thread; otherwise True. """ closed_later = getattr(self, "_closed_later", None) - if ( - closed_later is None - or closed_later.is_set() - or getattr(self, "_closer_ident", None) == threading.get_ident() - ): + if closed_later is None or getattr(self, "_closer_ident", None) == threading.get_ident(): return True - if S3File._in_collection() or S3ThreadPool.in_pool_thread(): - return False - unstarted = getattr(self, "_unstarted_close", None) - if unstarted is not None: - self._unstarted_close = None - self._close_collected(unstarted) - closed_later.wait() + unsafe = S3File._in_collection() or S3ThreadPool.in_pool_thread() + if not closed_later.is_set(): + if unsafe: + return False + closed_later.wait() + if getattr(self, "_unstarted_close", None) is not None: + if unsafe: + return False + unstarted, self._unstarted_close = self._unstarted_close, None + self._close_collected(cast(bool, unstarted)) return True def _raise_unless_closed_later(self) -> None: @@ -4654,7 +4663,12 @@ def _close_without_commit(self) -> None: self.buffer = None self.closed = True try: - self.discard() + if self._wait_for_close_later(): + self.discard() + else: + # The handed-over close, which cannot be waited for here, + # discards the write once it ends. + self._discard_later = True except Exception: upload_id = cast(S3MultipartUpload, self.multipart_upload).upload_id _logger.exception( diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index aac6e022d..cea8e75d7 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -9163,15 +9163,18 @@ def start(function, args): return start - @pytest.mark.parametrize("finalizing", [False, True]) @pytest.mark.usefixtures("handed_over_closes_end") - def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): + @pytest.mark.parametrize("when", ["running", "finalizing", "exiting"]) + def test_finalized_unclosed_file_without_thread(self, caplog, when): # GH-1364: if no thread can be started for the close of a finalized # file, it is left unclosed, and its data uncommitted, rather than # closed where it would wait forever, here in the thread of the # request pool with its first part queued behind it. During - # interpreter finalization, when no other thread runs requests, it is - # closed in the finalizing thread instead, without starting a thread. + # interpreter finalization, when no other thread runs requests, and + # once the exit has waited for the handed-over closes, which it would + # not wait for again, it is closed in the finalizing thread instead, + # without starting a thread. + finalizing = when != "running" fs = self._make_append_fs(b"") fs._request_pool = S3ThreadPool(1) committers = [] @@ -9186,9 +9189,12 @@ def test_finalized_unclosed_file_without_thread(self, caplog, finalizing): "start_new_thread", self._start_new_thread_unless_closing(_thread.start_new_thread), ), - mock.patch.object(sys, "is_finalizing", return_value=finalizing), + mock.patch.object(sys, "is_finalizing", return_value=when == "finalizing"), + mock.patch("pyathena.filesystem.s3._handed_over_closes_waited", False), caplog.at_level(logging.WARNING, logger="pyathena.filesystem.s3"), ): + if when == "exiting": + S3File._wait_for_handed_over_closes() finalization = self._finalize_unclosed_file( fs, "collection" if finalizing else "queued part" ) @@ -9326,9 +9332,18 @@ def __del__(self): fs._request_pool.submit(close_and_commit, f).result(EVENT_TIMEOUT) assert len(errors) == 1 fs.core.complete_multipart_upload.assert_not_called() - # Waits for the handed-over close. - f.close() - f.commit() + + def close_and_commit_later(): + # Waits for the handed-over close. + f.close() + f.commit() + + # A daemon thread, so that a close that waits forever does + # not keep the process alive. + committer = threading.Thread(target=close_and_commit_later, daemon=True) + committer.start() + committer.join(EVENT_TIMEOUT) + assert not committer.is_alive() finally: gc.enable() release.set() @@ -9336,13 +9351,40 @@ def __del__(self): assert self._uploaded_object(fs, b"") == b"abcdefg" @pytest.mark.usefixtures("handed_over_closes_end") - @pytest.mark.parametrize("method", ["close", "commit"]) - def test_close_handed_over_not_started(self, method): + @pytest.mark.parametrize( + ("method", "waiting"), + [("close", False), ("commit", False), ("commit", True), ("close in pool", False)], + ) + def test_close_handed_over_not_started(self, method, waiting): # GH-1364: when no thread could be started for a handed-over close, - # a later close() or commit() of the file closes it first, instead of - # waiting forever, or completing the upload, as at the end of a - # transaction, without the data that the close uploads. + # a later close() or commit() of the file closes it first, also one + # that already waits for it, instead of waiting forever, or + # completing the upload, as at the end of a transaction, without the + # data that the close uploads. In a thread of the request pool, where + # the close could wait for a part queued behind it, close() leaves + # the file open and commit() raises. fs = self._make_append_fs(b"") + fs._request_pool = S3ThreadPool(1) + errors = [] + + def close_and_commit_in_pool(file): + file.close() + assert not file.closed + try: + file.commit() + except RuntimeError as e: + errors.append(e) + + handed_over = threading.Event() + waited = threading.Event() + start_new_thread = self._start_new_thread_unless_closing(_thread.start_new_thread) + + def start_new_thread_once_waited(function, args): + if isinstance(getattr(function, "__self__", None), contextvars.Context): + handed_over.set() + if waiting: + assert waited.wait(EVENT_TIMEOUT) + return start_new_thread(function, args) class Closer: def __init__(self, file): @@ -9355,31 +9397,100 @@ def __del__(self): gc.disable() try: - with mock.patch.object( - _thread, - "start_new_thread", - self._start_new_thread_unless_closing(_thread.start_new_thread), - ): + with mock.patch.object(_thread, "start_new_thread", start_new_thread_once_waited): f = S3File( - fs, "s3://bucket/key.txt", mode="wb", block_size=4, autocommit=method == "close" + fs, + "s3://bucket/key.txt", + mode="wb", + block_size=4, + autocommit=method != "commit", ) f.write(b"abcde") f.write(b"fg") Closer(f) - gc.collect() - assert not f.closed - # A daemon thread, so that a call that waits forever does not - # keep the process alive. - caller = threading.Thread(target=getattr(f, method), daemon=True) - caller.start() + # Daemon threads, so that a call that waits forever does not + # keep the process alive. + collector = threading.Thread(target=gc.collect, daemon=True) + collector.start() + assert handed_over.wait(EVENT_TIMEOUT) + caller = threading.Thread(target=getattr(f, method.split()[0]), daemon=True) + if waiting: + # The caller waits for the close before its start fails. + closed_later = f._closed_later + wait = closed_later.wait + + def observed_wait(timeout=None): + waited.set() + return wait(timeout) + + closed_later.wait = observed_wait + caller.start() + collector.join(EVENT_TIMEOUT) + assert not collector.is_alive() + if method == "close in pool": + fs._request_pool.submit(close_and_commit_in_pool, f).result(EVENT_TIMEOUT) + assert len(errors) == 1 + if not waiting: + assert not f.closed + caller.start() caller.join(EVENT_TIMEOUT) assert not caller.is_alive() finally: gc.enable() + waited.set() shutdown_bounded(fs._request_pool) assert f.closed assert self._uploaded_object(fs, b"") == b"abcdefg" + @pytest.mark.usefixtures("handed_over_closes_end") + def test_close_handed_over_failed_write_is_discarded(self): + # GH-1364: when a write fails during the collection that handed over + # the close of its file, as a text wrapper writes when it is + # finalized, the handed-over close discards the write, which the + # failing write cannot wait for. + fs = self._make_append_fs(b"") + failed = threading.Event() + aborted = threading.Event() + errors = [] + + def fail_part(**kwargs): + failed.set() + raise OSError("part failed") + + fs.core.upload_part.side_effect = fail_part + fs._call.side_effect = lambda *args, **kwargs: aborted.set() + + class Writer: + def __init__(self, file): + self.file = file + self.cycle = self + + def __del__(self): + # Hands the close over, then writes another part. + self.file.close() + try: + self.file.write(b"fghij") + except OSError as e: + errors.append(e) + + gc.disable() + try: + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4) + f.write(b"abcde") + assert failed.wait(EVENT_TIMEOUT) + Writer(f) + del f + gc.collect() + assert aborted.wait(EVENT_TIMEOUT) + finally: + gc.enable() + shutdown_bounded(fs._request_pool) + assert len(errors) == 1 + fs._call.assert_called_once_with( + S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid" + ) + fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.parametrize("autocommit", [True, False]) def test_upload_chunk_multipart(self, autocommit): # A file larger than the block size is uploaded in parts and completed From e5eabc4e69c5b00c5c3e3b7f8209867276322b26 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 13 Nov 2026 04:41:04 +0900 Subject: [PATCH 8/8] Check close() apart from commit() and document the wrapper write wait (GH-1364) The test of the calls that cannot wait for a handed-over close checks that close() returns with the file open separately from the RuntimeError of commit(), so that close() raising is not taken for that error. The docs state that the write of the remaining text of a text wrapper waits as any write does, so in a thread of an S3ThreadPool it can wait forever for a part queued behind that thread, as it could before. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 4 +++- tests/pyathena/filesystem/test_s3.py | 11 +++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 39664eef9..e7570b229 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -232,7 +232,9 @@ first. A file opened in text mode is closed by its text wrapper when the wrapper is finalized. The wrapper writes its remaining text to the file in the thread that finalizes it, and -then closes the file there, or as above during a collection. +then closes the file there, or as above during a collection. That write waits as any +write does, so in a thread of an `S3ThreadPool` it can wait forever for a part of the +file that is queued behind that thread. An interrupt, such as Ctrl+C, or a failure while a file reads ranges in parallel cancels the range reads that have not started, and the running ones make no further diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index cea8e75d7..ea8a2c99f 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -9286,6 +9286,7 @@ def test_close_handed_over_is_not_waited_for(self, where): fs = self._make_append_fs(b"") fs._request_pool = S3ThreadPool(1) release = threading.Event() + closed = [] errors = [] close_collected = S3File._close_collected @@ -9297,10 +9298,11 @@ def close_and_commit(file): try: file.close() # The close is left to the thread that runs it. - assert not file.closed - file.commit() - except RuntimeError as e: - errors.append(e) + closed.append(file.closed) + try: + file.commit() + except RuntimeError as e: + errors.append(e) finally: release.set() @@ -9330,6 +9332,7 @@ def __del__(self): assert not collector.is_alive() if where == "pool": fs._request_pool.submit(close_and_commit, f).result(EVENT_TIMEOUT) + assert closed == [False] assert len(errors) == 1 fs.core.complete_multipart_upload.assert_not_called()