diff --git a/fsspec/caching.py b/fsspec/caching.py index 3cfb343a6..8b150e818 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import collections import functools import logging @@ -282,6 +283,121 @@ def _fetch(self, start: int | None, end: int | None) -> bytes: return part + self.cache[:l] +class AdaptiveReadaheadCache(BaseCache): + """Cache with adaptive asynchronous prefetching. + + Optimized for sequential and near-sequential reads by dynamically + adjusting the amount of data prefetched in the background. The cache + uses the generic prefetch engine when async loop support is available, + and falls back to ``ReadAheadCache`` when it is not. + + Parameters + ---------- + blocksize: int + Nominal read size used by callers. + fetcher: Fetcher + Function of the form ``f(start, end)`` that returns bytes. + size: int + Total size of the file. + concurrency: int + Maximum number of concurrent background fetch tasks. + max_prefetch_size: int | None + Optional upper bound for adaptive prefetch size in bytes. + """ + + name = "adaptive" + + 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: + from . import asyn as fsspec_asyn + except ImportError as e: + raise ImportError( + "AdaptiveReadaheadCache requires fsspec.asyn to be available" + ) from e + + try: + from .prefetcher import BackgroundPrefetcher + + self._prefetcher = BackgroundPrefetcher( + fetcher=_default_fetcher_async, + size=size, + concurrency=concurrency, + max_prefetch_size=max_prefetch_size, + loop=fsspec_asyn.get_loop(), + ) + logger.info( + "AdaptiveReadaheadCache enabled (blocksize=%d, size=%d, concurrency=%d, max_prefetch_size=%s)", + blocksize, + size, + concurrency, + max_prefetch_size, + ) + except Exception as e: + logger.info( + "AdaptiveReadaheadCache fallback 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 @@ -1016,6 +1132,7 @@ def register_cache(cls: type[BaseCache], clobber: bool = False) -> None: MMapCache, BytesCache, ReadAheadCache, + AdaptiveReadaheadCache, BlockCache, FirstChunkCache, AllBytes, diff --git a/fsspec/implementations/local.py b/fsspec/implementations/local.py index fd7531e7e..4623df4d9 100644 --- a/fsspec/implementations/local.py +++ b/fsspec/implementations/local.py @@ -8,11 +8,12 @@ import tempfile from functools import lru_cache -from fsspec import AbstractFileSystem from fsspec.compression import compr from fsspec.core import get_compression from fsspec.utils import isfilelike, stringify_path +from ..spec import AbstractFileSystem + logger = logging.getLogger("fsspec.local") diff --git a/fsspec/prefetcher.py b/fsspec/prefetcher.py new file mode 100644 index 000000000..e73b40e64 --- /dev/null +++ b/fsspec/prefetcher.py @@ -0,0 +1,878 @@ +import asyncio +import ctypes +import logging +import weakref +from collections import deque + +from . import asyn as fsspec_asyn + +logger = logging.getLogger(__name__) + +try: + PyBytes_FromStringAndSize = ctypes.pythonapi.PyBytes_FromStringAndSize + PyBytes_FromStringAndSize.argtypes = (ctypes.c_void_p, ctypes.c_ssize_t) + PyBytes_FromStringAndSize.restype = ctypes.py_object + + PyBytes_AsString = ctypes.pythonapi.PyBytes_AsString + PyBytes_AsString.argtypes = (ctypes.py_object,) + PyBytes_AsString.restype = ctypes.c_void_p + HAS_CPYTHON_API = True +except Exception: + PyBytes_FromStringAndSize = None + PyBytes_AsString = None + HAS_CPYTHON_API = False + + +# Please refer to following discussion to understand why this is required at this point +# Discussion = https://github.com/fsspec/gcsfs/pull/795#discussion_r3032749881 +def _fast_slice(src_bytes, offset, read_size): + if read_size == 0: + return b"" + if offset < 0 or offset + read_size > len(src_bytes): + raise ValueError("Slice indices out of bounds") + + if HAS_CPYTHON_API: + dest_bytes = PyBytes_FromStringAndSize(None, read_size) + src_ptr = PyBytes_AsString(src_bytes) + dest_ptr = PyBytes_AsString(dest_bytes) + # Releases the GIL + ctypes.memmove(dest_ptr, src_ptr + offset, read_size) + return dest_bytes + else: + # Standard fallback for PyPy/non-CPython + return src_bytes[offset : offset + read_size] + + +class RunningAverageTracker: + """Tracks a running average of values over a sliding window. + + This is used to monitor read sizes and adaptively scale the + prefetching strategy based on recent user behavior. + """ + + def __init__(self, maxlen=10): + """Initializes the tracker with a specific window size. + + Args: + maxlen (int): The maximum number of historical values to keep. + """ + logger.debug("Initializing RunningAverageTracker with maxlen: %d", maxlen) + self._history = deque(maxlen=maxlen) + self._sum = 0 + + def add(self, value: int): + """Adds a new value to the sliding window and updates the rolling sum. + + Args: + value (int): The integer value to add to the history. + """ + if value <= 0: + raise ValueError( + "Internal error, RunningAverageTracker tried inserting negative value" + ) + if len(self._history) == self._history.maxlen: + self._sum -= self._history[0] + + self._history.append(value) + self._sum += value + logger.debug( + "RunningAverageTracker added value: %d, new sum: %d", value, self._sum + ) + + @property + def average(self) -> int: + """Calculates and returns the current running average. + + Returns: + int: The integer average of the current history. + """ + count = len(self._history) + if count == 0: + return 1024 * 1024 # 1MB + return self._sum // count + + @property + def is_variable(self) -> bool: + """Determines if the history contains distinct chunk sizes.""" + count = len(self._history) + if count < 2: + return False + + return len(set(self._history)) > 1 + + @property + def last_value(self) -> int: + """Returns the most recent entry in the history.""" + if not self._history: + raise RuntimeError("No entry found in history") + + return self._history[-1] + + def clear(self): + """Clears the history and resets the sum to zero.""" + logger.debug("Clearing RunningAverageTracker history.") + self._history.clear() + self._sum = 0 + + +class PrefetchProducer: + """Background worker that fetches sequential blocks of data. + + This class handles the network requests. It spawns asynchronous tasks + to fetch data ahead of the user's current reading position and + places those task promises into a queue for the consumer. + """ + + # If the request is too small, and prefetch window is expanded till 5MB + # we then make request in 5MB blocks. + MIN_CHUNK_SIZE = 5 * 1024 * 1024 + + # If user doesn't specify any max_prefetch_size, the prefetcher defaults + # to maximum of 2 * io_size and 128MB + MIN_PREFETCH_SIZE = 128 * 1024 * 1024 + + # The prefetching starts on the third read. + MIN_STREAKS_FOR_PREFETCHING = 3 + + # Threshold for disabling proactive prefetching on large, variable reads. + # + # If the average read size exceeds this value and patterns are variable, + # prefetching shifts from an I/O bottleneck to a memory(CPU) bottleneck. When a user + # requests random massive sizes (e.g., jumping between 64MB and INF), the + # producer still fetches chunks based on the rolling average. The consumer + # then has to pick up multiple chunks and stitch them together to match the + # exact requested size. + # + # For small average read sizes, this byte assembly is fast and the bottleneck + # remains the network I/O. However, for massive reads (>= 64MB), the extra + # step of copying and assembling huge byte strings in memory severely slows + # down the operation. + VARIABLE_IO_THRESHOLD = 64 * 1024 * 1024 + + def __init__( + self, + fetcher, + size: int, + concurrency: int, + queue: asyncio.Queue, + wakeup_event: asyncio.Event, + consumer: "PrefetchConsumer", + tracker: RunningAverageTracker, + orchestrator: "BackgroundPrefetcher", + user_max_prefetch_size=None, + ): + """Initializes the background producer. + + Args: + fetcher (Callable): A coroutine function to fetch bytes from a remote source. + size (int): Total size of the file being fetched. + concurrency (int): Maximum number of concurrent fetch tasks. + queue (asyncio.Queue): The shared queue to push download tasks into. + wakeup_event (asyncio.Event): Event used to wake the producer from an idle state. + consumer (PrefetchConsumer): The consumer reading the prefetched chunks. + tracker (RunningAverageTracker): Tracker for history of read sizes. + orchestrator (BackgroundPrefetcher): The parent object managing the operation. + user_max_prefetch_size (int, optional): A hard limit for prefetch size overrides. + """ + logger.debug( + "Initializing PrefetchProducer: size=%d, concurrency=%d, user_max_prefetch_size=%s", + size, + concurrency, + user_max_prefetch_size, + ) + self.fetcher = fetcher + self.size = size + self.concurrency = concurrency + self.queue = queue + self.wakeup_event = wakeup_event + + self.consumer = consumer + self.tracker = tracker + self.orchestrator = weakref.proxy(orchestrator) + self._user_max_prefetch_size = user_max_prefetch_size + + self.current_offset = 0 + self.is_stopped = False + self._active_tasks = set() + self._producer_task = None + + @property + def max_prefetch_size(self) -> int: + """Calculates the maximum prefetch size based on user intent or io size. + + Returns: + int: The maximum number of bytes to prefetch ahead. + """ + if self._user_max_prefetch_size is not None: + return min( + self._user_max_prefetch_size, + max(2 * self.tracker.average, self.MIN_PREFETCH_SIZE), + ) + return max(2 * self.tracker.average, self.MIN_PREFETCH_SIZE) + + def start(self): + """Starts the background producer loop. + + This clears any previous wakeup events and spawns the main loop task. + """ + logger.debug("Starting PrefetchProducer loop.") + self.is_stopped = False + self.wakeup_event.clear() + self._producer_task = asyncio.create_task(self._loop()) + + async def stop(self): + """Cancels all active fetch tasks and shuts down the producer loop. + + This method ensures the queue is flushed and waits for cancelled + tasks to finish cleaning up. + """ + logger.debug( + "Stopping PrefetchProducer. Active fetch tasks: %d", len(self._active_tasks) + ) + self.is_stopped = True + self.wakeup_event.set() + + tasks_to_wait = [] + if self._producer_task and not self._producer_task.done(): + self._producer_task.cancel() + tasks_to_wait.append(self._producer_task) + + tasks_to_wait.extend(task for task in self._active_tasks if not task.done()) + + # We do not cancel the network task, instead we wait on them. + # This is intentionally done to avoid MRD stream disruption. + self._active_tasks.clear() + + # Clear out any leftover items in the queue + cleared_items = 0 + while not self.queue.empty(): + try: + item = self.queue.get_nowait() + if ( + isinstance(item, asyncio.Task) + and item.done() + and not item.cancelled() + ): + item.exception() + cleared_items += 1 + except asyncio.QueueEmpty: + break + + if cleared_items > 0: + logger.debug( + "Cleared %d leftover items from the queue during stop.", cleared_items + ) + + if tasks_to_wait: + logger.debug( + "Waiting for %d cancelled tasks to finish their teardown.", + len(tasks_to_wait), + ) + await asyncio.gather(*tasks_to_wait, return_exceptions=True) + + self.wakeup_event.clear() + + async def restart(self, new_offset: int): + """Stops current tasks and restarts the background loop at a new byte offset. + + Args: + new_offset (int): The new byte position to start prefetching from. + """ + logger.debug("Restarting PrefetchProducer at new offset: %d", new_offset) + await self.stop() + self.current_offset = new_offset + self.start() + + async def _loop(self): + """The main background loop that delegates calculations and spawns tasks.""" + logger.debug("PrefetchProducer internal loop is now running.") + try: + while not self.is_stopped: + await self.wakeup_event.wait() + self.wakeup_event.clear() + + if self.is_stopped: + break + + await self._process_prefetch_cycle() + + except asyncio.CancelledError: + logger.debug("PrefetchProducer loop was cancelled.") + except Exception as e: + logger.exception( + "PrefetchProducer loop encountered an unexpected error: %s", + e, + ) + self.is_stopped = True + self.orchestrator.set_error(e) + await self.queue.put(e) + + def _calculate_prefetch_params(self) -> tuple[int, int, int]: + """ + Evaluates current trackers and state to determine sizes. + + Returns: + tuple: (prefetch_size, io_size, effective_prefetch_size) + """ + avg_io_size = self.tracker.average + streak = self.consumer.sequential_streak + is_variable = self.tracker.is_variable + last_read_size = self.tracker.last_value + + exceeds_user_max = ( + self._user_max_prefetch_size is not None + and avg_io_size > self._user_max_prefetch_size + ) + + # Disable prefetching ahead if variable AND average > 64MB, or if it exceeds user max + if ( + is_variable and avg_io_size > self.VARIABLE_IO_THRESHOLD + ) or exceeds_user_max: + logger.debug( + "Large IO detected (variable > 64MB or > user max). Disabling background prefetching." + ) + prefetch_multiplier = 1 + elif streak < self.MIN_STREAKS_FOR_PREFETCHING: + prefetch_multiplier = 1 + else: + prefetch_multiplier = streak - self.MIN_STREAKS_FOR_PREFETCHING + 1 + + if self.queue.empty() or prefetch_multiplier == 1: + io_size = last_read_size + else: + io_size = avg_io_size + + prefetch_size = min(prefetch_multiplier * io_size, self.max_prefetch_size) + if self.consumer.offset + prefetch_size < self.consumer.target_offset: + prefetch_size = self.consumer.target_offset - self.consumer.offset + + if is_variable: + effective_prefetch_size = prefetch_size + else: + effective_prefetch_size = (prefetch_size // io_size) * io_size + if effective_prefetch_size == 0: + effective_prefetch_size = prefetch_size + + return prefetch_size, io_size, effective_prefetch_size + + async def _process_prefetch_cycle(self): + """Executes a single cycle of enqueuing fetch tasks.""" + prefetch_size, io_size, effective_prefetch_size = ( + self._calculate_prefetch_params() + ) + + logger.debug( + "Producer awake. Current offset: %d, User offset: %d, Prefetch size: %d", + self.current_offset, + self.consumer.offset, + prefetch_size, + ) + + while ( + not self.is_stopped + and (self.current_offset - self.consumer.offset) < prefetch_size + and self.current_offset < self.size + ): + user_offset = self.consumer.offset + space_remaining = self.size - self.current_offset + prefetch_space_available = prefetch_size - ( + self.current_offset - user_offset + ) + + if prefetch_size >= self.MIN_CHUNK_SIZE: + if prefetch_space_available >= self.MIN_CHUNK_SIZE: + actual_size = min( + max(self.MIN_CHUNK_SIZE, io_size), space_remaining + ) + else: + break + else: + actual_size = min(io_size, space_remaining) + + if prefetch_space_available < actual_size: + if ( + self.tracker.is_variable + or prefetch_space_available == prefetch_size + ): + actual_size = prefetch_space_available + else: + break + + streak = self.consumer.sequential_streak + if streak < self.MIN_STREAKS_FOR_PREFETCHING: + sfactor = self.concurrency + else: + sfactor = min( + self.concurrency, + max( + 1, + actual_size * self.concurrency // effective_prefetch_size, + ), + ) + + logger.debug( + "Spawning fetch task. Offset: %d, Size: %d, Split Factor: %d", + self.current_offset, + actual_size, + sfactor, + ) + + download_task = asyncio.create_task( + self.fetcher(self.current_offset, actual_size, split_factor=sfactor) + ) + self._active_tasks.add(download_task) + download_task.add_done_callback(self._active_tasks.discard) + + await self.queue.put(download_task) + self.current_offset += actual_size + + if self.current_offset >= self.size: + logger.debug("Producer reached EOF. Exiting background loop.") + self.is_stopped = True + + +class PrefetchConsumer: + """Consumes prefetched chunks from the queue and manages byte slicing. + + This class pulls data out of the shared queue and slices it into the + exact byte sizes requested by the user. It also manages the local block buffer. + """ + + def __init__( + self, + queue: asyncio.Queue, + wakeup_event: asyncio.Event, + tracker: RunningAverageTracker, + orchestrator: "BackgroundPrefetcher", + ): + """Initializes the consumer. + + Args: + queue (asyncio.Queue): The shared queue containing fetch tasks. + wakeup_event (asyncio.Event): Event used to wake the producer when more data is needed. + tracker (RunningAverageTracker): Tracker for history of read sizes. + orchestrator (BackgroundPrefetcher): The parent object managing the operation. + """ + logger.debug("Initializing PrefetchConsumer.") + self.queue = queue + self.wakeup_event = wakeup_event + self.tracker = tracker + self.orchestrator = weakref.proxy(orchestrator) + self.sequential_streak = 0 + self.offset = 0 + self.target_offset = 0 + self._current_block = b"" + self._current_block_idx = 0 + + def seek(self, new_offset: int): + """Clears the buffer and resets the internal offset for a hard seek. + + Args: + new_offset (int): The byte position the consumer is jumping to. + """ + logger.debug( + "Consumer executing hard seek to offset %d. Clearing internal buffer.", + new_offset, + ) + self.offset = new_offset + self.target_offset = new_offset + self.sequential_streak = 0 + self._current_block = b"" + self._current_block_idx = 0 + + def clear_buffer(self): + """Discards the local byte buffer. Useful during shutdown or resets.""" + logger.debug("Consumer local block buffer cleared.") + self._current_block = b"" + self._current_block_idx = 0 + + async def _advance(self, size: int, save_data: bool) -> list[bytes]: + """Internal method to advance the offset and optionally extract data. + + Handles queue exhaustion, producer wakeups, and streak tracking. + """ + if size <= 0: + return [] + + chunks = [] + processed = 0 + self.target_offset = self.offset + size + + while processed < size: + available = len(self._current_block) - self._current_block_idx + trigger_wakeup = False + + if not available: + is_producer_stopped = ( + self.orchestrator.producer is None + or self.orchestrator.producer.is_stopped + ) + if is_producer_stopped and self.queue.empty(): + logger.debug("Consumer reached EOF.") + break + + if self.queue.empty(): + logger.debug("Queue is empty. Waking up producer.") + self.wakeup_event.set() + + task = await self.queue.get() + + if isinstance(task, Exception): + logger.error("Consumer retrieved an exception: %s", task) + self.orchestrator.set_error(task) + raise task + + try: + block = await task + + self.sequential_streak += 1 + if ( + self.sequential_streak + >= PrefetchProducer.MIN_STREAKS_FOR_PREFETCHING + ): + exceeds_user_max = ( + self.orchestrator.max_prefetch_size is not None + and self.tracker.average + > self.orchestrator.max_prefetch_size + ) + is_massive_variable = ( + self.tracker.is_variable + and self.tracker.average + > PrefetchProducer.VARIABLE_IO_THRESHOLD + ) + + # Suppress proactive wakeups to prevent large CPU assembly + # on erratic large reads or exceeding max + if not (is_massive_variable or exceeds_user_max): + trigger_wakeup = True + else: + logger.debug( + "Suppressing proactive producer wakeup due to massive variable" + " workload or exceeding user max prefetch." + ) + + self._current_block = block + self._current_block_idx = 0 + available = len(self._current_block) + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception("Consumer caught an error: %s", e) + self.orchestrator.set_error(e) + raise e + + if not self._current_block: + break + + needed = size - processed + take = min(needed, available) + + if save_data: + if take == len(self._current_block) and self._current_block_idx == 0: + chunk = self._current_block + else: + # Native Python slicing was GIL bound in my experiments. + chunk = await asyncio.to_thread( + _fast_slice, self._current_block, self._current_block_idx, take + ) + chunks.append(chunk) + + self._current_block_idx += take + processed += take + self.offset += take + if trigger_wakeup: + self.wakeup_event.set() + + return chunks + + async def consume(self, size: int) -> bytes: + """Pulls exactly 'size' bytes from the local block or the task queue. + + If the local block is exhausted, this will wait on the queue for the next + available chunk of data. + + Args: + size (int): The exact number of bytes to retrieve. + + Returns: + bytes: The requested bytes. This may be shorter than 'size' if EOF is reached. + + Raises: + Exception: Re-raises any exceptions encountered by the producer fetch tasks. + """ + if size <= 0: + return b"" + + chunks = await self._advance(size, save_data=True) + + if not chunks: + return b"" + + if len(chunks) == 1: + return chunks[0] + + return await asyncio.to_thread(b"".join, chunks) + + async def skip(self, size: int) -> None: + """Advances the consumer offset without allocating memory.""" + await self._advance(size, save_data=False) + + +class BackgroundPrefetcher: + """Orchestrator that manages reading behavior and coordinates background work. + + This acts as the main public interface for the file reader. It tracks the + user's reading history, routes seek operations, and links the producer's + network tasks with the consumer's data slicing logic. + """ + + producer = None + + def __init__( + self, fetcher, size: int, concurrency: int, max_prefetch_size=None, loop=None + ): + """Initializes the background prefetcher. + + Args: + fetcher (Callable): A coroutine of the form `f(start, end)` which gets bytes from the remote. + size (int): Total byte size of the file being read. + concurrency (int): Number of concurrent network requests to use for large chunks. + max_prefetch_size (int, optional): Maximum bytes to prefetch ahead of the current user offset. + loop (asyncio.AbstractEventLoop, optional): The event loop to attach the prefetcher to. + If executing synchronously, this should be the fsspec background loop. If executing + asynchronously (asynchronous=True), this should be None so it can automatically + inherit the user's currently running event loop. + + Raises: + ValueError: If max_prefetch_size is provided but is not a positive integer. + """ + logger.debug( + "Starting BackgroundPrefetcher. Size: %d, Concurrency: %d, Max Prefetch: %s", + size, + concurrency, + max_prefetch_size, + ) + self.size = size + self.concurrency = concurrency + self.max_prefetch_size = max_prefetch_size + + if max_prefetch_size is not None and max_prefetch_size <= 0: + logger.error("Invalid max_prefetch_size provided: %s", max_prefetch_size) + raise ValueError( + "max_prefetch_size should be a positive integer to use adaptive prefetching!" + ) + + self.loop = loop + self._error = None + self.is_stopped = False + self.user_offset = 0 + self.read_tracker = RunningAverageTracker(maxlen=10) + + self.queue = None + self.wakeup_event = None + self._async_lock = None + self.consumer = None + self.producer = None + + def _start(): + # Ensures all primitives bind directly to `self.loop` + self.queue = asyncio.Queue() + self.wakeup_event = asyncio.Event() + self._async_lock = asyncio.Lock() + + self.consumer = PrefetchConsumer( + queue=self.queue, + wakeup_event=self.wakeup_event, + tracker=self.read_tracker, + orchestrator=self, + ) + + self.producer = PrefetchProducer( + fetcher=fetcher, + size=self.size, + concurrency=self.concurrency, + queue=self.queue, + wakeup_event=self.wakeup_event, + consumer=self.consumer, + tracker=self.read_tracker, + orchestrator=self, + user_max_prefetch_size=max_prefetch_size, + ) + self.producer.start() + + try: + current_loop = asyncio.get_running_loop() + except RuntimeError: + current_loop = None + + if current_loop is self.loop and self.loop is not None: + # We are already safely running inside the fsspec background loop + _start() + elif self.loop is not None: + # We are on the main thread; schedule setup on the fsspec background loop + async def _start_wrapper(): + _start() + + fsspec_asyn.sync(self.loop, _start_wrapper) + elif current_loop is not None: + # asynchronous=True: use the user's active event loop + self.loop = current_loop + _start() + else: + # asynchronous=True but called completely outside of an async context + raise RuntimeError("No event loop found") + + logger.debug("BackgroundPrefetcher initialization complete.") + + def __enter__(self): + """Context manager entry point.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit point. Ensures the prefetcher is cleanly closed.""" + self.close() + + async def __aenter__(self): + """Async context manager entry point.""" + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """Async context manager exit point. Ensures the prefetcher is cleanly closed.""" + await self.aclose() + + def set_error(self, e: Exception): + logger.error("Global error state set in BackgroundPrefetcher: %s", e) + self._error = e + + async def _restart_producer(self, new_offset: int): + logger.debug( + "Handling seek request. Restarting producer at offset: %d", new_offset + ) + self._error = None + await self.producer.restart(new_offset) + self.consumer.seek(new_offset) + self.read_tracker.clear() + + async def _async_fetch(self, start, end): + """Core internal async fetching logic, protected safely by the async lock.""" + async with self._async_lock: + try: + if self.is_stopped: + raise RuntimeError("The file instance has been closed.") + + logger.debug("Executing _async_fetch for range %d - %d.", start, end) + + # If the prefetcher is in error state, let's do a hard seek to start offset. + if self._error: + logger.info( + "Recovering from error state. Restarting producer at offset: %d", + start, + ) + self.user_offset = start + await self._restart_producer(start) + elif start != self.user_offset: + block_offset = ( + self.consumer.offset - self.consumer._current_block_idx + ) + if self.user_offset < start <= self.producer.current_offset: + logger.debug( + "Soft seek detected. Skipping ahead from %d to %d.", + self.user_offset, + start, + ) + skip_amount = start - self.user_offset + await self.consumer.skip(skip_amount) + self.user_offset = start + elif block_offset <= start < self.consumer.offset: + logger.debug( + "Local seek performed. User offset moved from %d to %d. " + "Adjusting buffer index from %d to %d.", + self.user_offset, + start, + self.consumer._current_block_idx, + start - block_offset, + ) + self.consumer._current_block_idx = start - block_offset + self.consumer.offset = start + self.consumer.target_offset = start + self.user_offset = start + else: + logger.debug( + "Hard seek detected. Moving user offset from %d to %d.", + self.user_offset, + start, + ) + self.user_offset = start + await self._restart_producer(start) + + requested_size = end - start + self.read_tracker.add(requested_size) + + chunk = await self.consumer.consume(requested_size) + self.user_offset += len(chunk) + + logger.debug("Completed _async_fetch. Returned %d bytes.", len(chunk)) + return chunk + except asyncio.CancelledError as e: + self._error = e + raise + except Exception as e: + logger.exception("Exception raised during asynchronous fetch: %s", e) + self._error = e + if self.producer and not self.producer.is_stopped: + await self.producer.stop() + raise + + async def _async_close(self): + """Asynchronous teardown logic protected by the async lock.""" + async with self._async_lock: + if self.is_stopped: + return + + self.is_stopped = True + logger.debug("Acquired async lock. Tearing down producer and buffers.") + + if self.producer: + await self.producer.stop() + + self.consumer.clear_buffer() + logger.debug("BackgroundPrefetcher closed successfully.") + + async def afetch(self, start: int | None, end: int | None) -> bytes: + """Asynchronous API counterpart to `_fetch`.""" + if start is None: + start = 0 + if end is None: + end = self.size + + end = min(end, self.size) + logger.debug( + "Asynchronous afetch called for bounds start=%s, end=%s.", start, end + ) + + if start >= self.size or start >= end: + return b"" + + if self.is_stopped: + logger.error( + "Cannot fetch data: BackgroundPrefetcher is stopped or closed." + ) + raise RuntimeError( + "The file instance has been closed. This can occur if a close operation " + "is executed concurrently while a read operation is still in progress." + ) + + return await self._async_fetch(start, end) + + def fetch(self, start: int | None, end: int | None) -> bytes: + """Synchronous API wrapper delegating to `afetch`.""" + # Delegates all boundaries, checking, and fetching to the async event loop perfectly + return fsspec_asyn.sync(self.loop, self.afetch, start, end) + + async def aclose(self): + """Safely shuts down the prefetcher from an asynchronous context.""" + await self._async_close() + + def close(self): + """Safely shuts down the prefetcher from a synchronous context.""" + fsspec_asyn.sync(self.loop, self._async_close) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index a1977f307..f486f9a76 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -359,3 +359,70 @@ def test_cache_kwargs(mocker): # It is a random location that cannot be predicted. # The important thing is the 'overwrite' kwarg fs.fs.put.assert_called_with(fs.fs.put.call_args[0][0], ["/test"], overwrite=True) + + +def test_adaptive_cache_with_async_fetcher(): + data = string.ascii_letters.encode() + + cache = caches["adaptive"]( + 8, + letters_fetcher, + len(data), + # Keep this below the default prefetch minimum so behavior is deterministic + # and bounded for unit testing. + concurrency=2, + max_prefetch_size=64, + ) + try: + assert cache._fetch(0, 0) == b"" + assert cache._fetch(0, 5) == data[0:5] + assert cache._fetch(5, 12) == data[5:12] + assert cache._fetch(12, 20) == data[12:20] + finally: + cache.close() + + +def test_adaptive_cache_registered(): + assert "adaptive" in caches + + +def test_adaptive_cache_fallback_when_prefetcher_init_fails(monkeypatch): + import fsspec.prefetcher as prefetcher_mod + + class FailingPrefetcher: + def __init__(self, *args, **kwargs): + raise RuntimeError("prefetch init failed") + + monkeypatch.setattr(prefetcher_mod, "BackgroundPrefetcher", FailingPrefetcher) + + data = string.ascii_letters.encode() + cache = caches["adaptive"](8, letters_fetcher, len(data)) + + # On prefetch setup failure, adaptive must still serve reads through readahead. + assert cache._fetch(0, 10) == data[0:10] + assert cache._fetch(10, 17) == data[10:17] + assert cache._prefetcher is None + + +def test_adaptive_cache_selected_in_open_flow(): + from fsspec.spec import AbstractBufferedFile + + data = string.ascii_letters.encode() + + class TestFile(AbstractBufferedFile): + DEFAULT_BLOCK_SIZE = 8 + + def _fetch_range(self, start, end): + return data[start:end] + + with TestFile(None, "afile", mode="rb", cache_type="adaptive", size=len(data)) as f: + assert f.cache.name == "adaptive" + assert f.read(12) == data[:12] + + +def test_adaptive_cache_fallback_without_loop(): + data = string.ascii_letters.encode() + cache = caches["adaptive"](8, letters_fetcher, len(data)) + + assert cache._fetch(0, 10) == data[0:10] + assert cache._fetch(10, 17) == data[10:17] diff --git a/fsspec/tests/test_prefetcher.py b/fsspec/tests/test_prefetcher.py new file mode 100644 index 000000000..6a03ad83c --- /dev/null +++ b/fsspec/tests/test_prefetcher.py @@ -0,0 +1,727 @@ +import asyncio +from unittest import mock + +import pytest + +import fsspec.asyn +from fsspec.prefetcher import BackgroundPrefetcher, RunningAverageTracker, _fast_slice + + +@pytest.fixture +def prefetcher_factory(): + prefetchers = [] + + def _make_prefetcher(**kwargs): + if "loop" not in kwargs: + kwargs["loop"] = fsspec.asyn.get_loop() + + bp = BackgroundPrefetcher(**kwargs) + prefetchers.append(bp) + return bp + + yield _make_prefetcher + + for bp in prefetchers: + bp.is_stopped = False + bp.close() + + +class MockFetcher: + def __init__(self, data, fail_at_call=None, hang_at_call=None): + self.data = data + self.calls = [] + self.fail_at_call = fail_at_call + self.hang_at_call = hang_at_call + self.call_count = 0 + + async def __call__(self, start, size, split_factor=1): + self.call_count += 1 + self.calls.append({"start": start, "size": size, "split_factor": split_factor}) + + await asyncio.sleep(0.001) + + if self.hang_at_call is not None and self.call_count >= self.hang_at_call: + await asyncio.sleep(1000) + + if self.fail_at_call is not None and self.call_count >= self.fail_at_call: + raise OSError("Simulated Network Timeout") + + return self.data[start : start + size] + + +def test_fast_slice_direct(): + src = b"0123456789" + assert _fast_slice(src, 2, 4) == b"2345" + assert _fast_slice(src, 5, 0) == b"" + assert _fast_slice(src, 0, 10) == b"0123456789" + + +def test_running_average_tracker(): + tracker = RunningAverageTracker(maxlen=3) + assert tracker.average == 1024 * 1024 # Default 1MB fallback + + tracker.add(512) + tracker.add(512) + assert tracker.average == 512 + + tracker.add(2048) + assert tracker.average == 1024 # (512 + 512 + 2048) // 3 + + tracker.clear() + assert tracker.average == 1024 * 1024 + + +def test_max_prefetch_size_property(prefetcher_factory): + bp1 = prefetcher_factory(fetcher=MockFetcher(b""), size=10000, concurrency=4) + assert bp1.producer.max_prefetch_size == bp1.producer.MIN_PREFETCH_SIZE + + bp2 = prefetcher_factory(fetcher=MockFetcher(b""), size=1000000000, concurrency=4) + # Give it a history so it calculates 2x the io_size + bp2.read_tracker.add(100 * 1024 * 1024) + assert bp2.producer.max_prefetch_size == 200 * 1024 * 1024 + + +def test_sequential_read_spanning_blocks(prefetcher_factory): + data = b"A" * 100 + b"B" * 100 + b"C" * 100 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=300, concurrency=4) + bp.read_tracker.add(100) # Seed the adaptive tracker + + assert bp.fetch(0, 100) == b"A" * 100 + assert bp.fetch(100, 150) == b"B" * 50 + assert bp.consumer._current_block_idx == 50 + assert bp.fetch(150, 250) == b"B" * 50 + b"C" * 50 + assert bp.fetch(250, 300) == b"C" * 50 + assert bp.fetch(300, 310) == b"" + + +def test_fetch_default_args_and_out_of_bounds(prefetcher_factory): + fetcher = MockFetcher(b"12345") + bp = prefetcher_factory(fetcher=fetcher, size=5, concurrency=4) + + assert bp.fetch(None, None) == b"12345" + assert bp.fetch(None, 2) == b"12" + assert bp.fetch(5, 10) == b"" + assert bp.fetch(10, 20) == b"" + assert bp.fetch(2, 2) == b"" + assert bp.fetch(4, 2) == b"" + + +def test_seek_logic(prefetcher_factory): + data = b"0123456789" * 10 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + + assert bp.fetch(0, 10) == data[0:10] + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + assert bp.fetch(50, 60) == data[50:60] + assert bp.user_offset == 60 + assert bp.fetch(10, 20) == data[10:20] + assert bp.user_offset == 20 + + +def test_exception_placed_in_queue(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject_error(): + await bp.queue.put(ValueError("Injected Producer Error")) + + fsspec.asyn.sync(bp.loop, inject_error) + + with pytest.raises(ValueError, match="Injected Producer Error"): + bp.fetch(0, 50) + + assert isinstance(bp._error, ValueError) + + +def test_producer_concurrency_streak_and_min_chunk(prefetcher_factory): + data = b"X" * 1000 + fetcher = MockFetcher(data) + + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + bp.read_tracker.add(50) + + # Temporarily lower chunk limit for test + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 10 + + # Do 6 reads to push the streak well past the MIN_STREAKS threshold + # Update these values as BackgroundPrefetcher constant changes. + target_streak = bp.producer.MIN_STREAKS_FOR_PREFETCHING + 3 + for i in range(target_streak): + bp.fetch(i * 50, (i + 1) * 50) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + split_factors = [call["split_factor"] for call in fetcher.calls] + assert split_factors[0] == 4 + assert max(split_factors) > 1 + assert max(split_factors) <= 4 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_loop_space_constraints(prefetcher_factory): + data = b"Y" * 100 + fetcher = MockFetcher(data) + + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + bp.read_tracker.add(60) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + assert bp.fetch(0, 10) == b"Y" * 10 + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + sizes = [call["size"] for call in fetcher.calls] + assert all(s <= 100 for s in sizes) + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_error_propagation_and_recovery(prefetcher_factory): + fetcher = MockFetcher(b"A" * 2000, fail_at_call=3) + bp = prefetcher_factory(fetcher=fetcher, size=2000, concurrency=4) + + for i in range(2): + bp.fetch(i * 100, (i + 1) * 100) + + # 3rd read triggers the network timeout + with pytest.raises(OSError, match="Simulated Network Timeout"): + bp.fetch(400, 500) + + # The prefetcher is now in an error state + assert isinstance(bp._error, OSError) + + # Disable the mock failure so it can succeed on retry + fetcher.fail_at_call = None + + # The next fetch should seamlessly recover, wiping the error and returning data + data = bp.fetch(400, 500) + assert data == b"A" * 100 + assert bp._error is None + + +def test_read_after_close(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + bp.close() + + assert bp.is_stopped is True + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 10) + + +def test_read_recovers_after_error(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + # Simulate an error state from a previous failed read + bp._error = ValueError("Pre-existing error") + + # The new error-recovery logic allows a subsequent read to clear the error and succeed + assert bp.fetch(0, 10) == b"X" * 10 + assert bp._error is None + + +def test_empty_queue_when_stopped(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 500), size=500, concurrency=4) + bp.is_stopped = True + + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 100) + + +def test_cancel_all_tasks_cleans_queue_with_exceptions(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject_task(): + async def dummy_exception_task(): + raise ValueError("Hidden error") + + task = asyncio.create_task(dummy_exception_task()) + await bp.queue.put(task) + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, inject_task) + bp.close() + assert bp.queue.empty() + + +def test_cleanup_cancels_active_tasks(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"Z" * 1000), size=1000, concurrency=4) + + async def inject_task(): + async def dummy_task(): + await asyncio.sleep(3) + + task = asyncio.create_task(dummy_task()) + bp.producer._active_tasks.add(task) + + fsspec.asyn.sync(bp.loop, inject_task) + + assert len(bp.producer._active_tasks) > 0 + assert bp.is_stopped is False + + bp.close() + + assert bp.is_stopped is True + assert len(bp.producer._active_tasks) == 0 + + +def test_read_task_cancellation(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 1000), size=1000, concurrency=4) + + async def inject_and_read(): + bp.is_stopped = True + while not bp.queue.empty(): + bp.queue.get_nowait() + + cancel_task = asyncio.create_task(asyncio.sleep(10)) + cancel_task.cancel() + await bp.queue.put(cancel_task) + + with pytest.raises(asyncio.CancelledError): + await bp.consumer.consume(10) + + fsspec.asyn.sync(bp.loop, inject_and_read) + + +def test_async_fetch_exception_trapping(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def bad_consume(*args, **kwargs): + raise RuntimeError("Simulated async crash") + + bp.consumer.consume = bad_consume + + with pytest.raises(RuntimeError, match="Simulated async crash"): + bp.fetch(0, 10) + + # Orchestrator should capture the error internally and halt producer processing correctly + assert isinstance(bp._error, RuntimeError) + assert bp.producer.is_stopped is True + + +def test_read_past_eof_internal(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 50), size=50, concurrency=4) + bp.user_offset = 50 + res = bp.fetch(50, 60) + assert res == b"" + + +def test_fetch_with_exact_block_matches(prefetcher_factory): + data = b"X" * 100 + bp = prefetcher_factory(fetcher=MockFetcher(data), size=100, concurrency=4) + bp.read_tracker.add(50) + + assert bp.fetch(0, 50) == b"X" * 50 + assert bp.consumer._current_block_idx == 50 + assert bp.fetch(50, 100) == b"X" * 50 + + +def test_queue_empty_race_condition(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def inject(): + bp.queue.put_nowait(asyncio.create_task(asyncio.sleep(0))) + with mock.patch.object(bp.queue, "get_nowait", side_effect=asyncio.QueueEmpty): + await bp.producer.stop() + + fsspec.asyn.sync(bp.loop, inject) + + +def test_producer_space_remaining_break(prefetcher_factory): + bp = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=150, + ) + bp.fetch(0, 10) + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + +def test_producer_min_chunk_logic(prefetcher_factory): + bp1 = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=300, + ) + bp1.producer.MIN_CHUNK_SIZE = 100 + + fsspec.asyn.sync(bp1.loop, asyncio.sleep, 0.1) + + bp2 = prefetcher_factory( + fetcher=MockFetcher(b"X" * 1000), + size=1000, + concurrency=4, + max_prefetch_size=150, + ) + bp2.producer.MIN_CHUNK_SIZE = 100 + fsspec.asyn.sync(bp2.loop, asyncio.sleep, 0.1) + + +def test_producer_loop_exception(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"A" * 100), size=100, concurrency=4) + error_object = ValueError("Producer crash") + + with mock.patch( + "fsspec.prefetcher.RunningAverageTracker.average", + new_callable=mock.PropertyMock, + ) as mocked_avg: + mocked_avg.side_effect = error_object + with pytest.raises(ValueError, match="Producer crash"): + bp.fetch(0, 10) + + assert bp.is_stopped is False + assert bp._error == error_object + + +def test_seek_same_offset(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b""), size=100, concurrency=4) + bp.fetch(0, 10) + + +def test_read_history_maxlen(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 2000), size=2000, concurrency=4) + for i in range(12): + bp.fetch(i * 10, (i + 1) * 10) + assert len(bp.read_tracker._history) == 10 + + +def test_fast_slice_branch(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 200), size=200, concurrency=4) + assert bp.fetch(0, 10) == b"X" * 10 + assert bp.fetch(10, 20) == b"X" * 10 + + +def test_async_fetch_not_block_break(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b""), size=100, concurrency=4) + + async def fake_consume(size): + return b"" + + bp.consumer.consume = fake_consume + bp.user_offset = 0 + + res = bp.fetch(0, 50) + assert res == b"" + + +def test_fetch_stopped_before_execution(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + bp.is_stopped = True + bp._error = None + + with pytest.raises(RuntimeError, match="The file instance has been closed"): + bp.fetch(0, 10) + + +def test_async_fetch_zero_copy_remainder(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X"), size=100, concurrency=4) + bp.consumer._current_block = b"ABCDE" + bp.consumer._current_block_idx = 0 + bp.user_offset = 0 + res = bp.fetch(0, 5) + assert res == b"ABCDE" + assert bp.consumer._current_block_idx == 5 + + +def test_read_runtime_error_on_stopped_empty(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X"), size=100, concurrency=4) + bp.is_stopped = True + bp.producer.is_stopped = True + + while not bp.queue.empty(): + bp.queue.get_nowait() + + res = fsspec.asyn.sync(bp.loop, bp.consumer.consume, 10) + assert res == b"" + + +def test_init_invalid_max_prefetch_size(): + with pytest.raises( + ValueError, + match=r"max_prefetch_size should be a positive integer", + ): + BackgroundPrefetcher( + fetcher=MockFetcher(b""), size=1000, concurrency=4, max_prefetch_size=0 + ) + + +def test_init_valid_max_prefetch_size_edge_case(prefetcher_factory): + bp = prefetcher_factory( + fetcher=MockFetcher(b""), size=1000, concurrency=4, max_prefetch_size=100 + ) + assert bp.producer._user_max_prefetch_size == 100 + + +def test_consumer_zero_size_checks(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + # 1. Test consume size <= 0 + res_consume_zero = fsspec.asyn.sync(bp.loop, bp.consumer.consume, 0) + assert res_consume_zero == b"" + res_consume_neg = fsspec.asyn.sync(bp.loop, bp.consumer.consume, -5) + assert res_consume_neg == b"" + + # 2. Test _advance size <= 0 directly + # (consume catches it early, so we call _advance directly to hit its internal check) + res_advance_zero = fsspec.asyn.sync( + bp.loop, bp.consumer._advance, 0, save_data=True + ) + assert res_advance_zero == [] + res_advance_neg = fsspec.asyn.sync( + bp.loop, bp.consumer._advance, -10, save_data=False + ) + assert res_advance_neg == [] + + +def test_producer_min_chunk_inner_break(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=400 + ) + + bp.read_tracker.add(100) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + async def trigger_loop(): + bp.producer.current_offset = 250 + bp.consumer.offset = 0 + bp.consumer.target_offset = 0 + # streak=6 makes prefetch_multiplier = 4 (6 - 3 + 1) + # prefetch_size = 4 * 100 = 400 + bp.consumer.sequential_streak = 6 + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_loop) + + assert fetcher.call_count == 0 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_producer_loop_break_on_stopped_after_wakeup(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + + async def trigger_stop_and_wake(): + bp.producer.is_stopped = True + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_stop_and_wake) + + # Verify the producer gracefully exited without doing work + assert fetcher.call_count == 0 + + +def test_massive_read_disables_proactive_prefetching(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + + # max_prefetch_size = 40 + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=40 + ) + + # Do enough reads to build a sequential streak and trigger large averages + # Reading 60 bytes at a time. Average = 60. Threshold = 50. + for i in range(4): + bp.fetch(i * 60, (i + 1) * 60) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # Because average (60) > threshold (40), prefetch_multiplier is pinned to 1. + # The producer should only fetch what the user specifically read (4 * 60 = 240) + # and should NOT have pre-fetched any additional data ahead into the queue. + assert bp.producer.current_offset == 240 + + +def test_normal_read_allows_proactive_prefetching(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + + # max_prefetch_size = 200 makes dynamic threshold = 100 + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=200 + ) + + # Reading 60 bytes at a time. Average = 60. Threshold = 100. + for i in range(4): + bp.fetch(i * 60, (i + 1) * 60) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # Because average (60) <= threshold (100), the producer allows prefetching. + # It calculates a normal prefetch_multiplier > 1 and pre-fetches data ahead. + assert bp.producer.current_offset > 240 + + +def test_target_offset_expands_prefetch(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory(fetcher=fetcher, size=1000, concurrency=4) + + # Seed tracker to keep the default `max_prefetch_size` calculation small + bp.read_tracker.add(10) + + # The consumer requests a massive chunk (500 bytes), far exceeding normal prefetch windows + bp.fetch(0, 500) + + fsspec.asyn.sync(bp.loop, asyncio.sleep, 0.1) + + # The new target_offset logic should explicitly tell the producer to expand its + # boundary to cover the requested 500 bytes, overriding the tiny multiplier logic. + assert bp.consumer.target_offset == 500 + assert bp.producer.current_offset >= 500 + + +def test_producer_min_chunk_inner_empty_queue_shrink(prefetcher_factory): + fetcher = MockFetcher(b"X" * 1000) + bp = prefetcher_factory( + fetcher=fetcher, size=1000, concurrency=4, max_prefetch_size=400 + ) + + bp.read_tracker.add(100) + + original_min_chunk = bp.producer.MIN_CHUNK_SIZE + bp.producer.MIN_CHUNK_SIZE = 200 + + async def trigger_loop(): + # Setup conditions where the queue is empty and the user is waiting + # This makes prefetch_space_available exactly equal to prefetch_size + bp.producer.current_offset = 0 + bp.consumer.offset = 0 + bp.consumer.target_offset = 0 + bp.consumer.sequential_streak = 6 + bp.wakeup_event.set() + await asyncio.sleep(0.05) + + fsspec.asyn.sync(bp.loop, trigger_loop) + + # Because space_available == prefetch_size, it triggers the shrink condition + # instead of breaking, ensuring the blocked consumer gets its data. + assert fetcher.call_count > 0 + + bp.producer.MIN_CHUNK_SIZE = original_min_chunk + + +def test_async_context_manager_and_afetch(prefetcher_factory): + bp = prefetcher_factory(fetcher=MockFetcher(b"X" * 100), size=100, concurrency=4) + + async def run_async(): + async with bp as ctx: + res = await ctx.afetch(0, 10) + assert res == b"X" * 10 + # Test default bounds -> (0, 100). Validates a backwards hard seek internally. + res_all = await ctx.afetch(None, None) + assert len(res_all) == 100 + + assert bp.is_stopped is True + assert bp.consumer._current_block == b"" # Buffer cleanly cleared on async exit + + fsspec.asyn.sync(bp.loop, run_async) + + +def test_init_with_explicit_loop(prefetcher_factory): + """Verify that passing an explicit loop assigns it correctly.""" + loop = fsspec.asyn.get_loop() + bp = prefetcher_factory( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=loop + ) + assert bp.loop is loop + + +@pytest.mark.asyncio +async def test_init_with_running_loop(): + """Verify asynchronous=True behavior where it inherits the user's running loop.""" + current = asyncio.get_running_loop() + # explicitly passing loop=None simulates asynchronous=True behavior + bp = BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=None + ) + assert bp.loop is current + await bp.aclose() + + +def test_init_within_fsspec_loop(prefetcher_factory): + """Verify the edge case where the prefetcher is initialized while already running inside the target loop.""" + loop = fsspec.asyn.get_loop() + + async def init_inside_loop(): + bp = BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=loop + ) + assert bp.loop is loop + await bp.aclose() + + # Submit the initialization directly onto the background thread + fsspec.asyn.sync(loop, init_inside_loop) + + +def test_init_no_loop_raises_error(): + """Verify synchronous execution strictly fails if no explicit loop or active loop is found.""" + with pytest.raises(RuntimeError, match="No event loop found"): + BackgroundPrefetcher( + fetcher=MockFetcher(b"X"), size=100, concurrency=1, loop=None + ) + + +def test_local_seek_optimization(prefetcher_factory): + data = b"0123456789" * 10 + fetcher = MockFetcher(data) + bp = prefetcher_factory(fetcher=fetcher, size=100, concurrency=4) + + # First fetch: read first 50 bytes (0-50) + # This will trigger a fetch. The block size fetched will be 50 bytes (covering [0, 50]). + assert bp.fetch(0, 50) == data[0:50] + assert len(bp.consumer._current_block) == 50 + initial_fetch_calls = fetcher.call_count + assert initial_fetch_calls > 0 + + # 1. Perform a backward seek *within* the currently buffered block [0, 50]. + # Seek back to 5 and read 10 bytes (5 to 15). + # This should be a zero-cost local seek and NOT increment the fetcher call count. + assert bp.fetch(5, 15) == data[5:15] + assert fetcher.call_count == initial_fetch_calls + assert bp.user_offset == 15 + + # 2. Perform a forward seek *within* the currently buffered block [0, 50]. + # Seek forward to 30 and read 10 bytes (30 to 40). + # This should also be resolved locally and NOT increment the fetcher call count. + assert bp.fetch(30, 40) == data[30:40] + assert fetcher.call_count == initial_fetch_calls + assert bp.user_offset == 40 + + # 3. Read further to trigger a hard seek and load the next block [60, 100]. + assert bp.fetch(60, 80) == data[60:80] + calls_after_next = fetcher.call_count + assert calls_after_next > initial_fetch_calls + + # The currently buffered block in the consumer is now [60, 100]. + # Seek backward to 10 (which is outside [60, 100]). + # This should trigger a hard seek and increment the fetcher call count. + assert bp.fetch(10, 20) == data[10:20] + assert fetcher.call_count > calls_after_next + + +@mock.patch("fsspec.prefetcher.HAS_CPYTHON_API", False) +def test_fast_slice_pypy_fallback(): + """ + Tests that when HAS_CPYTHON_API is False (e.g., on PyPy), _fast_slice + correctly falls back to standard Python slicing and respects boundaries. + """ + src = b"0123456789_pypy_fallback_test" + + # 1. Verify standard extraction works perfectly + assert _fast_slice(src, 11, 13) == b"pypy_fallback" + + # 2. Verify zero-length read + assert _fast_slice(src, 5, 0) == b"" + + # 3. Verify exact bounds + assert _fast_slice(src, 0, len(src)) == src