diff --git a/fsspec/caching.py b/fsspec/caching.py index 3499b4d26..9373460a6 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -799,6 +799,7 @@ def __init__( self._fetch_future_block_number: int | None = None self._fetch_future: Future[bytes] | None = None self._fetch_future_lock = threading.Lock() + self._closed = False def cache_info(self) -> UpdatableLRU.CacheInfo: """ @@ -811,6 +812,23 @@ def cache_info(self) -> UpdatableLRU.CacheInfo: """ return self._fetch_block_cached.cache_info() + def close(self) -> None: + """Cancel pending work and shut down the background worker.""" + with self._fetch_future_lock: + if self._closed: + return + self._closed = True + future = self._fetch_future + self._fetch_future = None + self._fetch_future_block_number = None + + if future is not None: + future.cancel() + self._thread_executor.shutdown(wait=True, cancel_futures=True) + + # UpdatableLRU stores a bound method and otherwise forms a reference cycle. + del self._fetch_block_cached + def __getstate__(self) -> dict[str, Any]: state = self.__dict__ del state["_fetch_block_cached"] @@ -827,6 +845,7 @@ def __setstate__(self, state) -> None: self._fetch_future_block_number = None self._fetch_future = None self._fetch_future_lock = threading.Lock() + self._closed = False def _fetch(self, start: int | None, end: int | None) -> bytes: if start is None: diff --git a/fsspec/spec.py b/fsspec/spec.py index a390bf060..25cb9c4d7 100644 --- a/fsspec/spec.py +++ b/fsspec/spec.py @@ -2220,6 +2220,11 @@ def close(self): return try: if self.mode == "rb": + cache = getattr(self, "cache", None) + if cache is not None: + close = getattr(cache, "close", None) + if callable(close): + close() self.cache = None else: if not getattr(self, "forced", True): diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index b9c5293c0..e0362534e 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -4,6 +4,7 @@ import pytest from fsspec.caching import ( + BackgroundBlockCache, BlockCache, FirstChunkCache, MMapCache, @@ -292,6 +293,40 @@ def wrapped(*a, **kw): assert len(thread_ids) == 2 +def test_background_shutdown_on_close(): + import weakref + + from fsspec.spec import AbstractBufferedFile + + data = b"abcdefgh" + + class TestFile(AbstractBufferedFile): + DEFAULT_BLOCK_SIZE = 4 + + def _fetch_range(self, start, end): + return data[start:end] + + f = TestFile( + None, + "test", + mode="rb", + cache_type="background", + size=len(data), + ) + f.read(1) + cache = f.cache + assert isinstance(cache, BackgroundBlockCache) + cache_ref = weakref.ref(cache) + executor = cache._thread_executor + del cache + + f.close() + + with pytest.raises(RuntimeError, match="cannot schedule new futures"): + executor.submit(lambda: None) + assert cache_ref() is None + + def test_register_cache(): # just test that we have them populated and fail to re-add again unless overload with pytest.raises(ValueError):