Shut down background cache workers on file close

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:

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):

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):