fix(zonal): clean up failed ZonalFile opens; support size= and pickling of read handles

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,24 @@ 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."""
+    with gcs_bucket_mocks(json_data, bucket_type_val=layouts.ZONAL_HIERARCHICAL):
+        with extended_gcsfs.open(file_path, "rb") as f:
+            assert isinstance(f, ZonalFile)
+            with pickle.loads(pickle.dumps(f)) as f2:
+                assert f2.read() == json_data
+
+
 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:

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,
 )
@@ -1074,6 +1075,150 @@ 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
+
+
+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__
+
+
 @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."""

gcsfs/zonal_file.py

@@ -1,6 +1,7 @@
 import logging
 
 from fsspec import asyn
+from fsspec.caching import caches
 from google.cloud.storage.asyncio.async_appendable_object_writer import (
     _DEFAULT_FLUSH_INTERVAL_BYTES,
 )
@@ -76,17 +77,23 @@ def __init__(
         # 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)
+        kwargs.pop("_supports_append", None)  # replayed by fsspec.spec.reopen
         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)
             )
+            # Reject an unknown cache type before checking out a pool and stream.
+            if resolved_cache_type not in caches:
+                raise KeyError(resolved_cache_type)
             self.mrd_pool = asyn.sync(
                 self.gcsfs.loop,
                 self.gcsfs._mrd_pool_cache.get,
@@ -97,45 +104,75 @@ 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:
+        elif self.mode in ("wb", "ab", "xb"):
             pass
         else:
             raise NotImplementedError(
                 "Only read, write and append operations are currently supported for Zonal buckets."
             )
 
-        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,
-        )
+        # Everything after the pool checkout runs under the cleanup below.
+        try:
+            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:
+            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
 
     async def _init_mrd(self, bucket_name, object_name, generation=None):
         """