diff --git a/gcsfs/core.py b/gcsfs/core.py index 918c74d76..58165e048 100644 --- a/gcsfs/core.py +++ b/gcsfs/core.py @@ -2665,6 +2665,16 @@ def __init__( ) else: kwargs.pop("max_prefetch_size", None) + # Kept so that a pickled read handle reopens with the same cache options. + self._reopen_cache_options = dict(cache_options) + + # _supports_append is an internal argument not meant to be used directly. + # If True, allows opening file in append mode. This is generally not supported + # by GCS, but may be supported by subclasses (e.g. ZonalFile). This flag should + # be set by subclasses that support append operations. Otherwise, the mode + # will be overwritten to "wb" mode with a warning. Pop it before + # super().__init__() so that it stays out of self.kwargs. + _supports_append = kwargs.pop("_supports_append", False) fs_loop = getattr(gcsfs, "loop", None) has_usable_fs_loop = ( @@ -2732,12 +2742,6 @@ def __init__( # The prefetcher will instead be closed cleanly inside GCSFile._close_impl(). cache.close = lambda: None - # _supports_append is an internal argument not meant to be used directly. - # If True, allows opening file in append mode. This is generally not supported - # by GCS, but may be supported by subclasses (e.g. ZonalFile). This flag should - # be set by subclasses that support append operations. Otherwise, the mode - # will be overwritten to "wb" mode with a warning. - _supports_append = kwargs.pop("_supports_append", False) if "a" in self.mode and not _supports_append: warnings.warn( "Append mode 'a' is not supported in GCS. Using overwrite mode instead." @@ -2766,6 +2770,38 @@ def __init__( self.blocksize = GCS_MIN_BLOCK_SIZE self.location = None + def _reopen_kwargs(self): + """Keyword arguments that reopen this file with the options it was opened with.""" + kwargs = dict(self.kwargs) + if self.generation is not None: + kwargs["generation"] = self.generation + if self._reopen_cache_options: + kwargs["cache_options"] = dict(self._reopen_cache_options) + kwargs["consistency"] = self.consistency + if self.timeout is not None: + kwargs["timeout"] = self.timeout + return kwargs + + def __reduce__(self): + # AbstractBufferedFile.__reduce__ replays only self.kwargs, which misses the + # named arguments above (so the reopened file could read a newer + # generation), and replays the cache name, which makes a default cache + # explicit. + if self.mode != "rb": + raise RuntimeError("Pickling a writeable file is not supported") + cache_type = None if self.cache_source == "default" else self.cache_type + return fsspec.spec.reopen, ( + self.fs, + self.path, + self.mode, + self.blocksize, + self.loc, + self.size, + self.autocommit, + cache_type, + self._reopen_kwargs(), + ) + @property def details(self): if self._details is None: diff --git a/gcsfs/tests/test_zonal.py b/gcsfs/tests/test_zonal.py index a2a90dd3c..b449a0e7e 100644 --- a/gcsfs/tests/test_zonal.py +++ b/gcsfs/tests/test_zonal.py @@ -7,6 +7,7 @@ import io import multiprocessing import os +import pickle import random import threading import uuid @@ -38,6 +39,7 @@ from gcsfs.tests.settings import TEST_BUCKET, TEST_ZONAL_BUCKET from gcsfs.tests.utils import is_real_gcs, tempdir, tmpfile from gcsfs.zb_hns_utils import MRDPool +from gcsfs.zonal_file import ZonalFile file = "test/accounts.1.json" file_path = f"{TEST_ZONAL_BUCKET}/{file}" @@ -243,6 +245,35 @@ def test_readline_blocksize_zb(extended_gcsfs, gcs_bucket_mocks): assert result == expected +def test_open_with_size_zb(extended_gcsfs, gcs_bucket_mocks): + """A caller-supplied size= is accepted, and the MRD's persisted_size takes precedence.""" + with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + # A wrong size must not truncate the read. + with extended_gcsfs.open(file_path, "rb", size=1) as f: + assert isinstance(f, ZonalFile) + assert f.read() == json_data + + +def test_pickle_read_handle_zb(extended_gcsfs, gcs_bucket_mocks): + """A zonal read handle can be pickled and reopened, as DataLoader and Dask do. + + The reopened handle keeps its position and the options it was opened with. + """ + generation = extended_gcsfs.info(file_path)["generation"] + with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL): + with extended_gcsfs.open( + file_path, "rb", generation=generation, pool_size=2 + ) as f: + assert isinstance(f, ZonalFile) + assert f.read(10) == json_data[:10] + with pickle.loads(pickle.dumps(f)) as f2: + assert f2.tell() == 10 + assert f2.generation == generation + assert f2.pool_size == 2 + assert f2.cache_source == "default" + assert f2.read() == json_data[10:] + + def test_read_unfinalized_file_using_mrd(extended_gcsfs, file_path): "Tests that mrd can read from an unfinalized file successfully" if not extended_gcsfs.on_google: diff --git a/gcsfs/tests/test_zonal_file.py b/gcsfs/tests/test_zonal_file.py index 965e80b40..5cc323228 100644 --- a/gcsfs/tests/test_zonal_file.py +++ b/gcsfs/tests/test_zonal_file.py @@ -16,6 +16,7 @@ import fsspec.asyn import pytest from fsspec.asyn import FSTimeoutError +from fsspec.prefetcher import BackgroundPrefetcher from google.cloud.storage.asyncio.async_appendable_object_writer import ( _DEFAULT_FLUSH_INTERVAL_BYTES, ) @@ -118,20 +119,18 @@ def mock_gcsfs(): return fs -@mock.patch("gcsfs.zonal_file.asyn.sync") -def test_zonal_file_generation_kwarg_handling(mock_sync, mock_gcsfs): +def test_zonal_file_generation_kwarg_handling(mock_gcsfs): """Test that ZonalFile explicitly provided generation kwarg is not dropped.""" from gcsfs.zonal_file import ZonalFile + mock_gcsfs.loop = fsspec.asyn.get_loop() # split_path returns None for generation when it's not in the path mock_gcsfs.split_path.return_value = ("test-bucket", "test-key", None) zf = ZonalFile( gcsfs=mock_gcsfs, path="gs://test-bucket/test-key", mode="rb", generation="456" ) - expected_call = mock.call( - mock_gcsfs.loop, - mock_gcsfs._mrd_pool_cache.get, + mock_gcsfs._mrd_pool_cache.get.assert_awaited_once_with( "test-bucket", "test-key", "456", @@ -139,21 +138,18 @@ def test_zonal_file_generation_kwarg_handling(mock_sync, mock_gcsfs): cache_type=mock.ANY, cache_source=mock.ANY, ) - assert expected_call in mock_sync.call_args_list zf.close() -@mock.patch("gcsfs.zonal_file.asyn.sync") -def test_zonal_file_generation_path_handling(mock_sync, mock_gcsfs): +def test_zonal_file_generation_path_handling(mock_gcsfs): """Test that ZonalFile parses generation from the path when provided.""" from gcsfs.zonal_file import ZonalFile + mock_gcsfs.loop = fsspec.asyn.get_loop() mock_gcsfs.split_path.return_value = ("test-bucket", "test-key", "123") zf = ZonalFile(gcsfs=mock_gcsfs, path="gs://test-bucket/test-key#123", mode="rb") - expected_call = mock.call( - mock_gcsfs.loop, - mock_gcsfs._mrd_pool_cache.get, + mock_gcsfs._mrd_pool_cache.get.assert_awaited_once_with( "test-bucket", "test-key", "123", @@ -161,7 +157,6 @@ def test_zonal_file_generation_path_handling(mock_sync, mock_gcsfs): cache_type=mock.ANY, cache_source=mock.ANY, ) - assert expected_call in mock_sync.call_args_list zf.close() @@ -1074,6 +1069,216 @@ async def coro_func(): assert result == ["done"] +def test_zonal_file_failed_init_closes_mrd_pool(mock_gcsfs): + """If GCSFile.__init__ fails after the MRD pool is checked out, the pool and prefetcher are closed.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + pool = mock_gcsfs._mrd_pool_cache.get.return_value + pool.close = mock.AsyncMock() + zf = ZonalFile.__new__(ZonalFile) + real_close = BackgroundPrefetcher.close + + with mock.patch.object( + BackgroundPrefetcher, "close", autospec=True, side_effect=real_close + ) as prefetcher_close: + # An invalid fixed_key_metadata fails in GCSFile.__init__ after the cache is built. + with pytest.raises(TypeError): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + fixed_key_metadata=123, + ) + + mock_gcsfs._mrd_pool_cache.get.assert_awaited_once() + pool.close.assert_awaited_once() + prefetcher_close.assert_called_once() + assert zf.mrd_pool is None + assert zf.closed is True + assert zf.cache._prefetcher is None + + +def test_zonal_file_failed_init_preserves_exception_when_pool_close_fails( + mock_gcsfs, caplog +): + """If MRD pool teardown fails during failed __init__, log a warning and propagate the init error.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + pool = mock_gcsfs._mrd_pool_cache.get.return_value + pool.close = mock.AsyncMock(side_effect=RuntimeError("pool close boom")) + zf = ZonalFile.__new__(ZonalFile) + + with caplog.at_level(logging.WARNING, logger="gcsfs.zonal_file"): + with pytest.raises(TypeError): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + fixed_key_metadata=123, + ) + + pool.close.assert_awaited_once() + assert zf.mrd_pool is None + assert zf.closed is True + assert any( + r.levelno == logging.WARNING and "pool close boom" in r.message + for r in caplog.records + ) + + +def test_zonal_file_interrupt_after_pool_checkout_closes_mrd_pool(mock_gcsfs): + """A KeyboardInterrupt right after the MRD pool is checked out still releases the pool.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + pool = mock_gcsfs._mrd_pool_cache.get.return_value + pool.close = mock.AsyncMock() + # Ctrl-C arriving while the pool's size is read, before GCSFile.__init__ runs. + type(pool).persisted_size = mock.PropertyMock(side_effect=KeyboardInterrupt) + zf = ZonalFile.__new__(ZonalFile) + + with pytest.raises(KeyboardInterrupt): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + ) + + pool.close.assert_awaited_once() + assert zf.mrd_pool is None + + +@pytest.mark.parametrize( + "checkout_finished", [False, True], ids=["during_checkout", "after_checkout"] +) +def test_zonal_file_interrupt_while_waiting_for_pool_checkout_closes_mrd_pool( + mock_gcsfs, checkout_finished +): + """A Ctrl-C while asyn.sync() waits for the pool checkout still releases the pool. + + asyn.sync() does not cancel the coroutine it scheduled, so the checkout can + finish after the caller has already given up. + """ + loop = fsspec.asyn.get_loop() + mock_gcsfs.loop = loop + pool = mock_gcsfs._mrd_pool_cache.get.return_value + pool.close = mock.AsyncMock() + release = asyncio.Event() + + async def slow_get(*args, **kwargs): + await release.wait() + return pool + + mock_gcsfs._mrd_pool_cache.get.side_effect = slow_get + checkouts = [] + + def interrupted_sync(loop, func, *args, **kwargs): + # Like asyn.sync(), schedule the checkout, then get Ctrl-C while waiting. + checkouts.append(asyncio.run_coroutine_threadsafe(func(*args, **kwargs), loop)) + if checkout_finished: + loop.call_soon_threadsafe(release.set) + checkouts[0].result(timeout=5) + raise KeyboardInterrupt + + zf = ZonalFile.__new__(ZonalFile) + with mock.patch("gcsfs.zonal_file.asyn.sync", side_effect=interrupted_sync): + with pytest.raises(KeyboardInterrupt): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + ) + + # Let a checkout that is still running finish after the caller gave up. + loop.call_soon_threadsafe(release.set) + checkouts[0].result(timeout=5) + pool.close.assert_awaited_once() + assert zf.mrd_pool is None + + +def test_zonal_file_failed_write_init_is_marked_closed(mock_gcsfs): + """A failed write-mode open is marked closed so __del__ never flushes (creates) the object.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + zf = ZonalFile.__new__(ZonalFile) + + with pytest.raises(TypeError): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="wb", + fixed_key_metadata=123, + ) + + mock_gcsfs._mrd_pool_cache.get.assert_not_called() + # AbstractBufferedFile.__init__ already set _closed = False, so this is not the default. + assert zf.closed is True + + +def test_zonal_file_unknown_cache_type_fails_before_pool_checkout(mock_gcsfs): + """An unknown cache_type fails before checking out an MRD pool or arming __del__.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + zf = ZonalFile.__new__(ZonalFile) + + with pytest.raises(KeyError, match="no-such-cache"): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode="rb", + cache_type="no-such-cache", + ) + + mock_gcsfs._mrd_pool_cache.get.assert_not_called() + assert "_closed" not in zf.__dict__ + + +@pytest.mark.parametrize( + "mode", + [ + "r+", + "rb+", + "r+b", + "w+", + "wb+", + "w+b", + "a+", + "ab+", + "a+b", + "x+", + "xb+", + "x+b", + "rwb", + ], +) +def test_zonal_file_plus_modes_rejected_before_pool_or_super_init(mock_gcsfs, mode): + """Unsupported modes ('+' modes or 'rwb') fail before checking out an MRD pool or arming __del__.""" + mock_gcsfs.loop = fsspec.asyn.get_loop() + zf = ZonalFile.__new__(ZonalFile) + + with pytest.raises(NotImplementedError): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode=mode, + ) + + mock_gcsfs._mrd_pool_cache.get.assert_not_called() + assert "_closed" not in zf.__dict__ + + +@pytest.mark.parametrize( + "mode, if_generation_match", [("r+b", 5), ("w+b", -1), ("x+b", 5)] +) +def test_zonal_file_unsupported_mode_checked_before_if_generation_match( + mock_gcsfs, mode, if_generation_match +): + """An unsupported mode raises NotImplementedError, not an if_generation_match error.""" + zf = ZonalFile.__new__(ZonalFile) + + with pytest.raises(NotImplementedError): + zf.__init__( + gcsfs=mock_gcsfs, + path="gs://test-bucket/test-key", + mode=mode, + if_generation_match=if_generation_match, + ) + + @mock.patch("gcsfs.zonal_file.asyn.sync") def test_close_impl_mrd_pool_failure_does_not_skip_aaow(mock_sync, mock_gcsfs): """Verify that if mrd_pool.close fails, aaow teardown is still executed.""" diff --git a/gcsfs/zonal_file.py b/gcsfs/zonal_file.py index e4cfc128e..142c9ca02 100644 --- a/gcsfs/zonal_file.py +++ b/gcsfs/zonal_file.py @@ -1,6 +1,8 @@ import logging +import threading from fsspec import asyn +from fsspec.caching import caches from google.cloud.storage.asyncio.async_appendable_object_writer import ( _DEFAULT_FLUSH_INTERVAL_BYTES, ) @@ -71,25 +73,39 @@ def __init__( generation = _coalesce_generation(generation, path_generation) if not key: raise OSError("Attempt to open a bucket") + if mode not in ("rb", "wb", "ab", "xb"): + raise NotImplementedError( + "Only read, write and append operations are currently supported for Zonal buckets." + ) # GCSFile.__init__ resolves the precondition as well, but only after the # MRD pool below has been opened; validate first so an invalid value # fails without a lookup. _resolve_open_generation_match(mode, if_generation_match) self.aaow = None + self.mrd_pool = None self.finalize_on_close = finalize_on_close self.finalized = False self.mode = mode self.flush_interval_bytes = flush_interval_bytes self.gcsfs = gcsfs self.pool_size = pool_size + caller_size = kwargs.pop("size", None) object_size = None - if "r" in self.mode: + if self.mode == "rb": resolved_cache_type, resolved_cache_source = ( _get_prefetcher_and_cache_config(cache_type) ) - self.mrd_pool = asyn.sync( - self.gcsfs.loop, - self.gcsfs._mrd_pool_cache.get, + # Reject an unknown cache type before checking out a pool and stream. + if resolved_cache_type not in caches: + raise KeyError(resolved_cache_type) + + # asyn.sync() does not cancel the checkout when the caller is interrupted + # (e.g. by Ctrl-C), so whichever side sees the abort closes the pool. + handoff_lock = threading.Lock() + handoff = {"aborted": False, "pool": None} + + async def _checkout_pool(): + pool = await self.gcsfs._mrd_pool_cache.get( bucket, key, generation, @@ -97,45 +113,83 @@ def __init__( cache_type=resolved_cache_type, cache_source=resolved_cache_source, ) - if getattr(self.mrd_pool, "details", None) is not None: - self._details = self.mrd_pool.details - object_size = self.mrd_pool.persisted_size - - if object_size is None: - logger.warning( - "AsyncMultiRangeDownloader (MRD) exists but has no 'persisted_size'. " - "This may result in incorrect behavior for unfinalized objects." - ) - elif "w" in self.mode or "a" in self.mode or "x" in self.mode: - pass - else: - raise NotImplementedError( - "Only read, write and append operations are currently supported for Zonal buckets." + with handoff_lock: + if not handoff["aborted"]: + handoff["pool"] = pool + return pool + await pool.close() # The caller already gave up. + + # The pool checkout and everything after it run under the cleanup below. + try: + if self.mode == "rb": + self.mrd_pool = asyn.sync(self.gcsfs.loop, _checkout_pool) + if self.mrd_pool is not None: + if getattr(self.mrd_pool, "details", None) is not None: + self._details = self.mrd_pool.details + object_size = self.mrd_pool.persisted_size + + if object_size is None: + logger.warning( + "AsyncMultiRangeDownloader (MRD) exists but has no 'persisted_size'. " + "This may result in incorrect behavior for unfinalized objects." + ) + + super().__init__( + gcsfs, + path, + mode, + block_size, + autocommit, + cache_type, + cache_options, + acl, + consistency, + metadata, + content_type, + timeout, + fixed_key_metadata, + generation, + kms_key_name, + if_generation_match=if_generation_match, + # Zonal buckets support append; this prevents GCSFile from forcing 'w' mode + _supports_append="a" in mode, + # pass persisted_size here so that Cache is initialized with correct object size + size=object_size if object_size is not None else caller_size, + **kwargs, ) + except BaseException: + with handoff_lock: + handoff["aborted"] = True + if self.mrd_pool is None: + self.mrd_pool = handoff["pool"] + self.closed = True + # GCSFile.__init__ leaves the prefetcher to _close_impl(), which a + # file marked closed never reaches, so stop it here. Like + # _close_impl(), rely on BackgroundPrefetcher.close() to log its + # own teardown errors. + cache = getattr(self, "cache", None) + prefetcher = getattr(cache, "_prefetcher", None) + if prefetcher is not None: + cache._prefetcher = None + prefetcher.close() + if self.mrd_pool is not None: + pool, self.mrd_pool = self.mrd_pool, None + try: + sync_teardown( + self.gcsfs.loop, + pool.close, + timeout=timeout or DEFAULT_TEARDOWN_TIMEOUT_SECONDS, + description=f"closing mrd_pool for {path} after failed open", + ) + except Exception as e: + logger.warning( + "Failed to close MRD pool during failed initialization cleanup: %s", + e, + ) + raise - super().__init__( - gcsfs, - path, - mode, - block_size, - autocommit, - cache_type, - cache_options, - acl, - consistency, - metadata, - content_type, - timeout, - fixed_key_metadata, - generation, - kms_key_name, - if_generation_match=if_generation_match, - # Zonal buckets support append; this prevents GCSFile from forcing 'w' mode - _supports_append="a" in mode, - # pass persisted_size here so that Cache is initialized with correct object size - size=object_size, - **kwargs, - ) + def _reopen_kwargs(self): + return {**super()._reopen_kwargs(), "pool_size": self.pool_size} async def _init_mrd(self, bucket_name, object_name, generation=None): """