diff --git a/docs/filesystem.md b/docs/filesystem.md index 979d0cd64..e7570b229 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -203,6 +203,39 @@ 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 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. 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, 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. + +`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. 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 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 +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 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..392bf0711 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -2,8 +2,12 @@ from __future__ import annotations +import _thread +import atexit import contextlib +import contextvars import functools +import gc import hashlib import logging import mimetypes @@ -79,6 +83,15 @@ "client=connection.s3_client and retry_config=connection.retry_config." ) +# 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 ends of the closes that run in threads of the _thread module, which an +# 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. _LOOKUP_REQUEST_PARAMETERS = frozenset( @@ -4310,17 +4323,245 @@ 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 once the collection has ended (see :meth:`__del__`), and + 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 return + if not self.closed and S3File._in_collection(): + self._close_later() + return + 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.""" 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. 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. + 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. 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: + """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. 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: + complete = False + else: + 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 = [] + _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 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, + 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. + """ + if self._inherited(): + # Handed over in the parent, which forked before the collection + # ended. Its event, whose lock the fork may have copied held, is + # left alone. + self.closed = True + return + if sys.is_finalizing() or _handed_over_closes_waited: + self._close_collected(complete) + return + try: + _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._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. + + Args: + 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() + 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() + if getattr(self, "_discard_later", False): + self.discard() + except Exception: + _logger.exception(f"Failed to close s3://{self.bucket}/{self.key}.") + finally: + self._end_close_later() + + def _end_close_later(self) -> None: + """Mark the close that :meth:`_close_later` handed over as ended.""" + _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. Also waits for the closes that are + 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. + + 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. 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 has not ended and cannot be waited + for in this thread; otherwise True. + """ + closed_later = getattr(self, "_closed_later", None) + if closed_later is None or getattr(self, "_closer_ident", None) == threading.get_ident(): + return True + 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: + """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. @@ -4422,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( @@ -4575,17 +4821,21 @@ 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 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._raise_unless_closed_later() if self._session is None: # A file opened for reading has nothing to commit. return @@ -4617,13 +4867,18 @@ 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. + 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._raise_unless_closed_later() if self._session is not None: self._session.abort() @@ -4823,3 +5078,10 @@ def _get_ranges( else: ranges.append((start, end)) return 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/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..ea8a2c99f 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,13 +1,16 @@ +import _thread import array import asyncio import bz2 import contextlib +import contextvars import errno import functools import gc import gzip import io import itertools +import logging import lzma import os import pickle @@ -44,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 @@ -8947,6 +8955,545 @@ 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 + # given place, using the only thread of the request pool: + # - "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. + # 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"): + 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 + 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 not in _handed_over_closes: + # 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: + 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"): + 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) + assert not started_by_finalizer + return SimpleNamespace(closes=closes, finalizers=finalizers) + + @pytest.mark.parametrize( + ("where", "text"), + [ + ("last reference", False), + ("collection", False), + ("collection", True), + ("queued part", False), + ("own part", False), + ("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 + # 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: + 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 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(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. + 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 + + @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.usefixtures("handed_over_closes_end") + @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, 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 = [] + fs.core.complete_multipart_upload.side_effect = lambda *args, **kwargs: ( + committers.append(threading.current_thread()), + mock.DEFAULT, + )[1] + try: + with ( + 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=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" + ) + finally: + shutdown_bounded(fs._request_pool) + if finalizing: + assert self._uploaded_object(fs, b"") == b"abcde" + 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), ("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(), + # 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): + 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() + # 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() + assert waiting.wait(EVENT_TIMEOUT) + release.set() + waiter.join(EVENT_TIMEOUT) + assert not waiter.is_alive() + assert f.closed + 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"]) + @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 + # 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() + closed = [] + 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. + closed.append(file.closed) + try: + 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 closed == [False] + assert len(errors) == 1 + fs.core.complete_multipart_upload.assert_not_called() + + 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() + shutdown_bounded(fs._request_pool) + assert self._uploaded_object(fs, b"") == b"abcdefg" + + @pytest.mark.usefixtures("handed_over_closes_end") + @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, 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): + self.file = file + self.cycle = self + + def __del__(self): + # Hands the close over. + self.file.close() + + gc.disable() + try: + 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 != "commit", + ) + f.write(b"abcde") + f.write(b"fg") + Closer(f) + # 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