feat(adaptive-readahead): introducing new cache adaptive-readahead

Author: raj-prince

fsspec/caching.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import collections
 import functools
 import logging
@@ -282,6 +283,95 @@ def _fetch(self, start: int | None, end: int | None) -> bytes:
         return part + self.cache[:l]
 
 
+class AdaptiveReadaheadCache(BaseCache):
+    """Adaptive readahead cache using the generic prefetch engine.
+
+    This cache remains API-compatible with existing fsspec cache wiring and
+    falls back to classic ``ReadAheadCache`` when async prefetching is not
+    available.
+    """
+
+    name = "adaptive_readahead"
+
+    def __init__(
+        self,
+        blocksize: int,
+        fetcher: Fetcher,
+        size: int,
+        concurrency: int = 4,
+        max_prefetch_size: int | None = None,
+    ) -> None:
+        super().__init__(blocksize, fetcher, size)
+        self._fallback = ReadAheadCache(blocksize, fetcher, size)
+        self._prefetcher = None
+
+        async def _default_fetcher_async(
+            start_offset: int,
+            total_size: int,
+            split_factor: int = 1,
+        ) -> bytes:
+            del split_factor
+            return await asyncio.to_thread(
+                self.fetcher, start_offset, start_offset + total_size
+            )
+
+        try:
+            import fsspec.asyn
+
+            from .prefetch import BackgroundPrefetcher
+
+            self._prefetcher = BackgroundPrefetcher(
+                fetcher=_default_fetcher_async,
+                size=size,
+                concurrency=concurrency,
+                max_prefetch_size=max_prefetch_size,
+                loop=fsspec.asyn.get_loop(),
+            )
+        except Exception as e:
+            logger.debug(
+                "AdaptiveReadaheadCache falling back to ReadAheadCache: %s",
+                e,
+                exc_info=True,
+            )
+            self._prefetcher = None
+
+    def _fetch(self, start: int | None, end: int | None) -> bytes:
+        if self._prefetcher is None:
+            out = self._fallback._fetch(start, end)
+            self.hit_count = self._fallback.hit_count
+            self.miss_count = self._fallback.miss_count
+            self.total_requested_bytes = self._fallback.total_requested_bytes
+            return out
+
+        out = self._prefetcher.fetch(start, end)
+        self.miss_count += 1
+        self.total_requested_bytes += len(out)
+        return out
+
+    def close(self) -> None:
+        if self._prefetcher is not None:
+            self._prefetcher.close()
+            self._prefetcher = None
+
+    def __getstate__(self) -> dict[str, Any]:
+        # The prefetcher owns asyncio primitives that are not picklable.
+        self.close()
+        state = self.__dict__.copy()
+        state["_prefetcher"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self._prefetcher = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            # Best-effort cleanup during GC.
+            pass
+
+
 class FirstChunkCache(BaseCache):
     """Caches the first block of a file only
 
@@ -997,6 +1087,7 @@ def register_cache(cls: type[BaseCache], clobber: bool = False) -> None:
     MMapCache,
     BytesCache,
     ReadAheadCache,
+    AdaptiveReadaheadCache,
     BlockCache,
     FirstChunkCache,
     AllBytes,