Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/fsspec/caching.py: 19%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import asyncio
4import collections
5import functools
6import logging
7import math
8import os
9import threading
10from collections import OrderedDict
11from collections.abc import Callable
12from concurrent.futures import Future, ThreadPoolExecutor
13from itertools import groupby
14from operator import itemgetter
15from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar
17if TYPE_CHECKING:
18 import mmap
20 from typing_extensions import ParamSpec
22 P = ParamSpec("P")
23else:
24 P = TypeVar("P")
26T = TypeVar("T")
29logger = logging.getLogger("fsspec.caching")
31Fetcher = Callable[[int, int], bytes] # Maps (start, end) to bytes
32MultiFetcher = Callable[[list[int, int]], bytes] # Maps [(start, end)] to bytes
35class BaseCache:
36 """Pass-though cache: doesn't keep anything, calls every time
38 Acts as base class for other cachers
40 Parameters
41 ----------
42 blocksize: int
43 How far to read ahead in numbers of bytes
44 fetcher: func
45 Function of the form f(start, end) which gets bytes from remote as
46 specified
47 size: int
48 How big this file is
49 """
51 name: ClassVar[str] = "none"
53 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
54 self.blocksize = blocksize
55 self.nblocks = 0
56 self.fetcher = fetcher
57 self.size = size
58 self.hit_count = 0
59 self.miss_count = 0
60 # the bytes that we actually requested
61 self.total_requested_bytes = 0
63 def _fetch(self, start: int | None, stop: int | None) -> bytes:
64 if start is None:
65 start = 0
66 if stop is None:
67 stop = self.size
68 if start >= self.size or start >= stop:
69 return b""
70 return self.fetcher(start, stop)
72 def _reset_stats(self) -> None:
73 """Reset hit and miss counts for a more ganular report e.g. by file."""
74 self.hit_count = 0
75 self.miss_count = 0
76 self.total_requested_bytes = 0
78 def _log_stats(self) -> str:
79 """Return a formatted string of the cache statistics."""
80 if self.hit_count == 0 and self.miss_count == 0:
81 # a cache that does nothing, this is for logs only
82 return ""
83 return f" , {self.name}: {self.hit_count} hits, {self.miss_count} misses, {self.total_requested_bytes} total requested bytes"
85 def __repr__(self) -> str:
86 # TODO: use rich for better formatting
87 return f"""
88 <{self.__class__.__name__}:
89 block size : {self.blocksize}
90 block count : {self.nblocks}
91 file size : {self.size}
92 cache hits : {self.hit_count}
93 cache misses: {self.miss_count}
94 total requested bytes: {self.total_requested_bytes}>
95 """
98class MMapCache(BaseCache):
99 """memory-mapped sparse file cache
101 Opens temporary file, which is filled blocks-wise when data is requested.
102 Ensure there is enough disc space in the temporary location.
104 This cache method might only work on posix
106 Parameters
107 ----------
108 blocksize: int
109 How far to read ahead in numbers of bytes
110 fetcher: Fetcher
111 Function of the form f(start, end) which gets bytes from remote as
112 specified
113 size: int
114 How big this file is
115 location: str
116 Where to create the temporary file. If None, a temporary file is
117 created using tempfile.TemporaryFile().
118 blocks: set[int]
119 Set of block numbers that have already been fetched. If None, an empty
120 set is created.
121 multi_fetcher: MultiFetcher
122 Function of the form f([(start, end)]) which gets bytes from remote
123 as specified. This function is used to fetch multiple blocks at once.
124 If not specified, the fetcher function is used instead.
125 """
127 name = "mmap"
129 def __init__(
130 self,
131 blocksize: int,
132 fetcher: Fetcher,
133 size: int,
134 location: str | None = None,
135 blocks: set[int] | None = None,
136 multi_fetcher: MultiFetcher | None = None,
137 ) -> None:
138 super().__init__(blocksize, fetcher, size)
139 self.blocks = set() if blocks is None else blocks
140 self.location = location
141 self.multi_fetcher = multi_fetcher
142 self.cache = self._makefile()
144 def _makefile(self) -> mmap.mmap | bytearray:
145 import mmap
146 import tempfile
148 if self.size == 0:
149 return bytearray()
151 # posix version
152 if self.location is None or not os.path.exists(self.location):
153 self.blocks.clear()
154 if self.location is None:
155 fd = tempfile.TemporaryFile()
156 else:
157 fd = open(self.location, "wb+")
158 fd.seek(self.size - 1)
159 fd.write(b"1")
160 fd.flush()
161 else:
162 fd = open(self.location, "r+b")
164 return mmap.mmap(fd.fileno(), self.size)
166 def _fetch(self, start: int | None, end: int | None) -> bytes:
167 logger.debug(f"MMap cache fetching {start}-{end}")
168 if start is None:
169 start = 0
170 if end is None:
171 end = self.size
172 if start >= self.size or start >= end:
173 return b""
174 start_block = start // self.blocksize
175 end_block = end // self.blocksize
176 block_range = range(start_block, end_block + 1)
177 # Determine which blocks need to be fetched. This sequence is sorted by construction.
178 need = (i for i in block_range if i not in self.blocks)
179 # Count the number of blocks already cached
180 self.hit_count += sum(1 for i in block_range if i in self.blocks)
182 ranges = []
183 block_groups = []
185 # Consolidate needed blocks.
186 # Algorithm adapted from Python 2.x itertools documentation.
187 # We are grouping an enumerated sequence of blocks. By comparing when the difference
188 # between an ascending range (provided by enumerate) and the needed block numbers
189 # we can detect when the block number skips values. The key computes this difference.
190 # Whenever the difference changes, we know that we have previously cached block(s),
191 # and a new group is started. In other words, this algorithm neatly groups
192 # runs of consecutive block numbers so they can be fetched together.
193 for _, _blocks in groupby(enumerate(need), key=lambda x: x[0] - x[1]):
194 # Extract the blocks from the enumerated sequence
195 _blocks = tuple(map(itemgetter(1), _blocks))
196 # Compute start of first block
197 sstart = _blocks[0] * self.blocksize
198 # Compute the end of the last block. Last block may not be full size.
199 send = min(_blocks[-1] * self.blocksize + self.blocksize, self.size)
201 # Fetch bytes (could be multiple consecutive blocks)
202 self.total_requested_bytes += send - sstart
203 logger.debug(
204 f"MMap get blocks {_blocks[0]}-{_blocks[-1]} ({sstart}-{send})"
205 )
206 ranges.append((sstart, send))
207 block_groups.append(_blocks)
208 # Update cache statistics with number of blocks we had to cache
209 self.miss_count += len(_blocks)
211 if not ranges:
212 return self.cache[start:end]
214 if self.multi_fetcher:
215 logger.debug(f"MMap get blocks {ranges}")
216 for idx, r in enumerate(self.multi_fetcher(ranges)):
217 sstart, send = ranges[idx]
218 logger.debug(f"MMap copy block ({sstart}-{send}")
219 self.cache[sstart:send] = r
220 self.blocks.update(block_groups[idx])
221 else:
222 for idx, (sstart, send) in enumerate(ranges):
223 logger.debug(f"MMap get block ({sstart}-{send}")
224 self.cache[sstart:send] = self.fetcher(sstart, send)
225 self.blocks.update(block_groups[idx])
227 return self.cache[start:end]
229 def close(self) -> None:
230 if hasattr(self.cache, "close"):
231 self.cache.close()
233 def __getstate__(self) -> dict[str, Any]:
234 state = self.__dict__.copy()
235 # Remove the unpicklable entries.
236 del state["cache"]
237 return state
239 def __setstate__(self, state: dict[str, Any]) -> None:
240 # Restore instance attributes
241 self.__dict__.update(state)
242 self.cache = self._makefile()
245class ReadAheadCache(BaseCache):
246 """Cache which reads only when we get beyond a block of data
248 This is a much simpler version of BytesCache, and does not attempt to
249 fill holes in the cache or keep fragments alive. It is best suited to
250 many small reads in a sequential order (e.g., reading lines from a file).
251 """
253 name = "readahead"
255 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
256 super().__init__(blocksize, fetcher, size)
257 self.cache = b""
258 self.start = 0
259 self.end = 0
261 def _fetch(self, start: int | None, end: int | None) -> bytes:
262 if start is None:
263 start = 0
264 if end is None or end > self.size:
265 end = self.size
266 if start >= self.size or start >= end:
267 return b""
268 l = end - start
269 if start >= self.start and end <= self.end:
270 # cache hit
271 self.hit_count += 1
272 return self.cache[start - self.start : end - self.start]
273 elif self.start <= start < self.end:
274 # partial hit
275 self.miss_count += 1
276 part = self.cache[start - self.start :]
277 l -= len(part)
278 start = self.end
279 else:
280 # miss
281 self.miss_count += 1
282 part = b""
283 end = min(self.size, end + self.blocksize)
284 self.total_requested_bytes += end - start
285 self.cache = self.fetcher(start, end) # new block replaces old
286 self.start = start
287 self.end = self.start + len(self.cache)
288 return part + self.cache[:l]
291class AdaptiveReadaheadCache(BaseCache):
292 """Cache with adaptive asynchronous prefetching.
294 Optimized for sequential and near-sequential reads by dynamically
295 adjusting the amount of data prefetched in the background. The cache
296 uses the generic prefetch engine when async loop support is available,
297 and falls back to ``ReadAheadCache`` when it is not.
299 Parameters
300 ----------
301 blocksize: int
302 Nominal read size used by callers.
303 fetcher: Fetcher
304 Function of the form ``f(start, end)`` that returns bytes.
305 size: int
306 Total size of the file.
307 concurrency: int
308 Maximum number of concurrent background fetch tasks.
309 max_prefetch_size: int | None
310 Optional upper bound for adaptive prefetch size in bytes.
311 """
313 name = "adaptive"
315 def __init__(
316 self,
317 blocksize: int,
318 fetcher: Fetcher,
319 size: int,
320 concurrency: int = 4,
321 max_prefetch_size: int | None = None,
322 ) -> None:
323 super().__init__(blocksize, fetcher, size)
324 self._fallback = ReadAheadCache(blocksize, fetcher, size)
325 self._prefetcher = None
327 async def _default_fetcher_async(
328 start_offset: int,
329 total_size: int,
330 split_factor: int = 1,
331 ) -> bytes:
332 del split_factor
333 return await asyncio.to_thread(
334 self.fetcher, start_offset, start_offset + total_size
335 )
337 try:
338 from . import asyn as fsspec_asyn
339 except ImportError as e:
340 raise ImportError(
341 "AdaptiveReadaheadCache requires fsspec.asyn to be available"
342 ) from e
344 try:
345 from .prefetcher import BackgroundPrefetcher
347 try:
348 fs_loop = fetcher.__self__.fs.loop
349 except AttributeError:
350 fs_loop = None
351 self._prefetcher = BackgroundPrefetcher(
352 fetcher=_default_fetcher_async,
353 size=size,
354 concurrency=concurrency,
355 max_prefetch_size=max_prefetch_size,
356 loop=(
357 fs_loop
358 if isinstance(fs_loop, asyncio.AbstractEventLoop)
359 else fsspec_asyn.get_loop()
360 ),
361 )
362 logger.info(
363 "AdaptiveReadaheadCache enabled (blocksize=%d, size=%d, concurrency=%d, max_prefetch_size=%s)",
364 blocksize,
365 size,
366 concurrency,
367 max_prefetch_size,
368 )
369 except Exception as e:
370 logger.info(
371 "AdaptiveReadaheadCache fallback to ReadAheadCache: %s",
372 e,
373 exc_info=True,
374 )
375 self._prefetcher = None
377 def _fetch(self, start: int | None, end: int | None) -> bytes:
378 if self._prefetcher is None:
379 out = self._fallback._fetch(start, end)
380 self.hit_count = self._fallback.hit_count
381 self.miss_count = self._fallback.miss_count
382 self.total_requested_bytes = self._fallback.total_requested_bytes
383 return out
385 out = self._prefetcher.fetch(start, end)
386 self.miss_count += 1
387 self.total_requested_bytes += len(out)
388 return out
390 def close(self) -> None:
391 if self._prefetcher is not None:
392 self._prefetcher.close()
393 self._prefetcher = None
395 def __getstate__(self) -> dict[str, Any]:
396 # The prefetcher owns asyncio primitives that are not picklable.
397 self.close()
398 state = self.__dict__.copy()
399 state["_prefetcher"] = None
400 return state
402 def __setstate__(self, state: dict[str, Any]) -> None:
403 self.__dict__.update(state)
404 self._prefetcher = None
406 def __del__(self):
407 try:
408 self.close()
409 except Exception:
410 # Best-effort cleanup during GC.
411 pass
414class FirstChunkCache(BaseCache):
415 """Caches the first block of a file only
417 This may be useful for file types where the metadata is stored in the header,
418 but is randomly accessed.
419 """
421 name = "first"
423 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
424 if blocksize > size:
425 # this will buffer the whole thing
426 blocksize = size
427 super().__init__(blocksize, fetcher, size)
428 self.cache: bytes | None = None
430 def _fetch(self, start: int | None, end: int | None) -> bytes:
431 start = start or 0
432 if start > self.size:
433 logger.debug("FirstChunkCache: requested start > file size")
434 return b""
436 if end is None:
437 end = self.size
438 end = min(end, self.size)
440 if start < self.blocksize:
441 if self.cache is None:
442 self.miss_count += 1
443 if end > self.blocksize:
444 self.total_requested_bytes += end
445 data = self.fetcher(0, end)
446 self.cache = data[: self.blocksize]
447 return data[start:]
448 self.cache = self.fetcher(0, self.blocksize)
449 self.total_requested_bytes += self.blocksize
450 part = self.cache[start:end]
451 if end > self.blocksize:
452 self.total_requested_bytes += end - self.blocksize
453 part += self.fetcher(self.blocksize, end)
454 self.hit_count += 1
455 return part
456 else:
457 self.miss_count += 1
458 self.total_requested_bytes += end - start
459 return self.fetcher(start, end)
462class BlockCache(BaseCache):
463 """
464 Cache holding memory as a set of blocks.
466 Requests are only ever made ``blocksize`` at a time, and are
467 stored in an LRU cache. The least recently accessed block is
468 discarded when more than ``maxblocks`` are stored.
470 Parameters
471 ----------
472 blocksize : int
473 The number of bytes to store in each block.
474 Requests are only ever made for ``blocksize``, so this
475 should balance the overhead of making a request against
476 the granularity of the blocks.
477 fetcher : Callable
478 size : int
479 The total size of the file being cached.
480 maxblocks : int
481 The maximum number of blocks to cache for. The maximum memory
482 use for this cache is then ``blocksize * maxblocks``.
483 """
485 name = "blockcache"
487 def __init__(
488 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
489 ) -> None:
490 super().__init__(blocksize, fetcher, size)
491 self.nblocks = math.ceil(size / blocksize)
492 self.maxblocks = maxblocks
493 self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block)
495 def cache_info(self):
496 """
497 The statistics on the block cache.
499 Returns
500 -------
501 NamedTuple
502 Returned directly from the LRU Cache used internally.
503 """
504 return self._fetch_block_cached.cache_info()
506 def __getstate__(self) -> dict[str, Any]:
507 state = self.__dict__.copy()
508 del state["_fetch_block_cached"]
509 return state
511 def __setstate__(self, state: dict[str, Any]) -> None:
512 self.__dict__.update(state)
513 self._fetch_block_cached = functools.lru_cache(state["maxblocks"])(
514 self._fetch_block
515 )
517 def _fetch(self, start: int | None, end: int | None) -> bytes:
518 if start is None:
519 start = 0
520 if end is None or end > self.size:
521 end = self.size
522 if start >= self.size or start >= end:
523 return b""
525 return self._read_cache(
526 start, end, start // self.blocksize, (end - 1) // self.blocksize
527 )
529 def _fetch_block(self, block_number: int) -> bytes:
530 """
531 Fetch the block of data for `block_number`.
532 """
533 if block_number > self.nblocks:
534 raise ValueError(
535 f"'block_number={block_number}' is greater than "
536 f"the number of blocks ({self.nblocks})"
537 )
539 start = block_number * self.blocksize
540 if start >= self.size:
541 return b""
542 end = min(start + self.blocksize, self.size)
543 self.total_requested_bytes += end - start
544 self.miss_count += 1
545 logger.info("BlockCache fetching block %d", block_number)
546 block_contents = super()._fetch(start, end)
547 return block_contents
549 def _read_cache(
550 self, start: int, end: int, start_block_number: int, end_block_number: int
551 ) -> bytes:
552 """
553 Read from our block cache.
555 Parameters
556 ----------
557 start, end : int
558 The start and end byte positions.
559 start_block_number, end_block_number : int
560 The start and end block numbers.
561 """
562 start_pos = start % self.blocksize
563 end_pos = end % self.blocksize
564 if end_pos == 0:
565 end_pos = self.blocksize
567 self.hit_count += 1
568 if start_block_number == end_block_number:
569 block: bytes = self._fetch_block_cached(start_block_number)
570 return block[start_pos:end_pos]
572 else:
573 # read from the initial
574 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
576 # intermediate blocks
577 # Note: it'd be nice to combine these into one big request. However
578 # that doesn't play nicely with our LRU cache.
579 out.extend(
580 map(
581 self._fetch_block_cached,
582 range(start_block_number + 1, end_block_number),
583 )
584 )
586 # final block
587 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
589 return b"".join(out)
592class BytesCache(BaseCache):
593 """Cache which holds data in a in-memory bytes object
595 Implements read-ahead by the block size, for semi-random reads progressing
596 through the file.
598 Parameters
599 ----------
600 trim: bool
601 As we read more data, whether to discard the start of the buffer when
602 we are more than a blocksize ahead of it.
603 """
605 name: ClassVar[str] = "bytes"
607 def __init__(
608 self, blocksize: int, fetcher: Fetcher, size: int, trim: bool = True
609 ) -> None:
610 super().__init__(blocksize, fetcher, size)
611 self.cache = b""
612 self.start: int | None = None
613 self.end: int | None = None
614 self.trim = trim
616 def _fetch(self, start: int | None, end: int | None) -> bytes:
617 # TODO: only set start/end after fetch, in case it fails?
618 # is this where retry logic might go?
619 if start is None:
620 start = 0
621 if end is None:
622 end = self.size
623 if start >= self.size or start >= end:
624 return b""
625 if (
626 self.start is not None
627 and start >= self.start
628 and self.end is not None
629 and end <= self.end
630 ):
631 # cache hit: we have all the required data
632 offset = start - self.start
633 self.hit_count += 1
634 return self.cache[offset : offset + end - start]
636 if self.blocksize:
637 bend = min(self.size, end + self.blocksize)
638 else:
639 bend = end
641 if bend == start or start > self.size:
642 return b""
644 if (self.start is None or start < self.start) and (
645 self.end is None or end > self.end
646 ):
647 # First read, or extending both before and after
648 self.total_requested_bytes += bend - start
649 self.miss_count += 1
650 self.cache = self.fetcher(start, bend)
651 self.start = start
652 else:
653 assert self.start is not None
654 assert self.end is not None
655 self.miss_count += 1
657 if start < self.start:
658 if self.end is None or self.end - end > self.blocksize:
659 self.total_requested_bytes += bend - start
660 self.cache = self.fetcher(start, bend)
661 self.start = start
662 else:
663 self.total_requested_bytes += self.start - start
664 new = self.fetcher(start, self.start)
665 self.start = start
666 self.cache = new + self.cache
667 elif self.end is not None and bend > self.end:
668 if self.end > self.size:
669 pass
670 elif end - self.end > self.blocksize:
671 self.total_requested_bytes += bend - start
672 self.cache = self.fetcher(start, bend)
673 self.start = start
674 else:
675 self.total_requested_bytes += bend - self.end
676 new = self.fetcher(self.end, bend)
677 self.cache = self.cache + new
679 self.end = self.start + len(self.cache)
680 offset = start - self.start
681 out = self.cache[offset : offset + end - start]
682 if self.trim:
683 num = (self.end - self.start) // (self.blocksize + 1)
684 if num > 1:
685 self.start += self.blocksize * num
686 self.cache = self.cache[self.blocksize * num :]
687 return out
689 def __len__(self) -> int:
690 return len(self.cache)
693class AllBytes(BaseCache):
694 """Cache entire contents of the file"""
696 name: ClassVar[str] = "all"
698 def __init__(
699 self,
700 blocksize: int | None = None,
701 fetcher: Fetcher | None = None,
702 size: int | None = None,
703 data: bytes | None = None,
704 ) -> None:
705 super().__init__(blocksize, fetcher, size) # type: ignore[arg-type]
706 if data is None:
707 self.miss_count += 1
708 self.total_requested_bytes += self.size
709 data = self.fetcher(0, self.size)
710 self.data = data
712 def _fetch(self, start: int | None, stop: int | None) -> bytes:
713 self.hit_count += 1
714 return self.data[start:stop]
717class KnownPartsOfAFile(BaseCache):
718 """
719 Cache holding known file parts.
721 Parameters
722 ----------
723 blocksize: int
724 How far to read ahead in numbers of bytes
725 fetcher: func
726 Function of the form f(start, end) which gets bytes from remote as
727 specified
728 size: int
729 How big this file is
730 data: dict
731 A dictionary mapping explicit `(start, stop)` file-offset tuples
732 with known bytes.
733 strict: bool, default True
734 Whether to fetch reads that go beyond a known byte-range boundary.
735 If `False`, any read that ends outside a known part will be zero
736 padded. Note that zero padding will not be used for reads that
737 begin outside a known byte-range.
738 """
740 name: ClassVar[str] = "parts"
742 def __init__(
743 self,
744 blocksize: int,
745 fetcher: Fetcher,
746 size: int,
747 data: dict[tuple[int, int], bytes] | None = None,
748 strict: bool = False,
749 **_: Any,
750 ):
751 super().__init__(blocksize, fetcher, size)
752 self.strict = strict
754 # simple consolidation of contiguous blocks
755 if data:
756 old_offsets = sorted(data.keys())
757 offsets = [old_offsets[0]]
758 blocks = [data.pop(old_offsets[0])]
759 for start, stop in old_offsets[1:]:
760 start0, stop0 = offsets[-1]
761 if start == stop0:
762 offsets[-1] = (start0, stop)
763 blocks[-1] += data.pop((start, stop))
764 else:
765 offsets.append((start, stop))
766 blocks.append(data.pop((start, stop)))
768 self.data = dict(zip(offsets, blocks))
769 else:
770 self.data = {}
772 @property
773 def size(self):
774 return sum(_[1] - _[0] for _ in self.data)
776 @size.setter
777 def size(self, value):
778 pass
780 @property
781 def nblocks(self):
782 return len(self.data)
784 @nblocks.setter
785 def nblocks(self, value):
786 pass
788 def _fetch(self, start: int | None, stop: int | None) -> bytes:
789 logger.debug("Known parts request %s %s", start, stop)
790 if start is None:
791 start = 0
792 if stop is None:
793 stop = self.size
794 self.total_requested_bytes += stop - start
795 out = b""
796 started = False
797 loc_old = 0
798 for loc0, loc1 in sorted(self.data):
799 if (loc0 <= start < loc1) and (loc0 <= stop <= loc1):
800 # entirely within the block
801 off = start - loc0
802 self.hit_count += 1
803 return self.data[(loc0, loc1)][off : off + stop - start]
804 if stop <= loc0:
805 break
806 if started and loc0 > loc_old:
807 # a gap where we need data
808 self.miss_count += 1
809 if self.strict:
810 raise ValueError
811 out += b"\x00" * (loc0 - loc_old)
812 if loc0 <= start < loc1:
813 # found the start
814 self.hit_count += 1
815 off = start - loc0
816 out = self.data[(loc0, loc1)][off : off + stop - start]
817 started = True
818 elif start < loc0 and stop > loc1:
819 # the whole block
820 self.hit_count += 1
821 out += self.data[(loc0, loc1)]
822 elif started and loc0 <= stop <= loc1:
823 # end block. The start of the request must have been found in
824 # an earlier part, else these bytes come from the wrong offsets
825 self.hit_count += 1
826 out = out + self.data[(loc0, loc1)][: stop - loc0]
827 return out
828 loc_old = loc1
829 self.miss_count += 1
830 if started and not self.strict:
831 out = out + b"\x00" * (stop - loc_old)
832 return out
833 raise ValueError
836class UpdatableLRU(Generic[P, T]):
837 """
838 Custom implementation of LRU cache that allows updating keys
840 Used by BackgroundBlockCache
841 """
843 class CacheInfo(NamedTuple):
844 hits: int
845 misses: int
846 maxsize: int
847 currsize: int
849 def __init__(self, func: Callable[P, T], max_size: int = 128) -> None:
850 self._cache: OrderedDict[Any, T] = collections.OrderedDict()
851 self._func = func
852 self._max_size = max_size
853 self._hits = 0
854 self._misses = 0
855 self._lock = threading.Lock()
857 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
858 if kwargs:
859 raise TypeError(f"Got unexpected keyword argument {kwargs.keys()}")
860 with self._lock:
861 if args in self._cache:
862 self._cache.move_to_end(args)
863 self._hits += 1
864 return self._cache[args]
866 result = self._func(*args, **kwargs)
868 with self._lock:
869 self._cache[args] = result
870 self._misses += 1
871 if len(self._cache) > self._max_size:
872 self._cache.popitem(last=False)
874 return result
876 def is_key_cached(self, *args: Any) -> bool:
877 with self._lock:
878 return args in self._cache
880 def add_key(self, result: T, *args: Any) -> None:
881 with self._lock:
882 self._cache[args] = result
883 if len(self._cache) > self._max_size:
884 self._cache.popitem(last=False)
886 def cache_info(self) -> UpdatableLRU.CacheInfo:
887 with self._lock:
888 return self.CacheInfo(
889 maxsize=self._max_size,
890 currsize=len(self._cache),
891 hits=self._hits,
892 misses=self._misses,
893 )
896class BackgroundBlockCache(BaseCache):
897 """
898 Cache holding memory as a set of blocks with pre-loading of
899 the next block in the background.
901 Requests are only ever made ``blocksize`` at a time, and are
902 stored in an LRU cache. The least recently accessed block is
903 discarded when more than ``maxblocks`` are stored. If the
904 next block is not in cache, it is loaded in a separate thread
905 in non-blocking way.
907 Parameters
908 ----------
909 blocksize : int
910 The number of bytes to store in each block.
911 Requests are only ever made for ``blocksize``, so this
912 should balance the overhead of making a request against
913 the granularity of the blocks.
914 fetcher : Callable
915 size : int
916 The total size of the file being cached.
917 maxblocks : int
918 The maximum number of blocks to cache for. The maximum memory
919 use for this cache is then ``blocksize * maxblocks``.
920 """
922 name: ClassVar[str] = "background"
924 def __init__(
925 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
926 ) -> None:
927 super().__init__(blocksize, fetcher, size)
928 self.nblocks = math.ceil(size / blocksize)
929 self.maxblocks = maxblocks
930 self._fetch_block_cached = UpdatableLRU(self._fetch_block, maxblocks)
932 self._thread_executor = ThreadPoolExecutor(max_workers=1)
933 self._fetch_future_block_number: int | None = None
934 self._fetch_future: Future[bytes] | None = None
935 self._fetch_future_lock = threading.Lock()
936 self._closed = False
938 def cache_info(self) -> UpdatableLRU.CacheInfo:
939 """
940 The statistics on the block cache.
942 Returns
943 -------
944 NamedTuple
945 Returned directly from the LRU Cache used internally.
946 """
947 return self._fetch_block_cached.cache_info()
949 def close(self) -> None:
950 """Cancel pending work and shut down the background worker."""
951 with self._fetch_future_lock:
952 if self._closed:
953 return
954 self._closed = True
955 future = self._fetch_future
956 self._fetch_future = None
957 self._fetch_future_block_number = None
959 if future is not None:
960 future.cancel()
961 self._thread_executor.shutdown(wait=True, cancel_futures=True)
963 # UpdatableLRU stores a bound method and otherwise forms a reference cycle.
964 del self._fetch_block_cached
966 def __getstate__(self) -> dict[str, Any]:
967 state = self.__dict__.copy()
968 del state["_fetch_block_cached"]
969 del state["_thread_executor"]
970 del state["_fetch_future_block_number"]
971 del state["_fetch_future"]
972 del state["_fetch_future_lock"]
973 return state
975 def __setstate__(self, state) -> None:
976 self.__dict__.update(state)
977 self._fetch_block_cached = UpdatableLRU(self._fetch_block, state["maxblocks"])
978 self._thread_executor = ThreadPoolExecutor(max_workers=1)
979 self._fetch_future_block_number = None
980 self._fetch_future = None
981 self._fetch_future_lock = threading.Lock()
982 self._closed = False
984 def _fetch(self, start: int | None, end: int | None) -> bytes:
985 if start is None:
986 start = 0
987 if end is None or end > self.size:
988 end = self.size
989 if start >= self.size or start >= end:
990 return b""
992 # byte position -> block numbers; ``end`` is exclusive
993 start_block_number = start // self.blocksize
994 end_block_number = (end - 1) // self.blocksize
996 fetch_future_block_number = None
997 fetch_future = None
998 with self._fetch_future_lock:
999 # Background thread is running. Check we we can or must join it.
1000 if self._fetch_future is not None:
1001 assert self._fetch_future_block_number is not None
1002 if self._fetch_future.done():
1003 logger.info("BlockCache joined background fetch without waiting.")
1004 self._fetch_block_cached.add_key(
1005 self._fetch_future.result(), self._fetch_future_block_number
1006 )
1007 # Cleanup the fetch variables. Done with fetching the block.
1008 self._fetch_future_block_number = None
1009 self._fetch_future = None
1010 else:
1011 # Must join if we need the block for the current fetch
1012 must_join = bool(
1013 start_block_number
1014 <= self._fetch_future_block_number
1015 <= end_block_number
1016 )
1017 if must_join:
1018 # Copy to the local variables to release lock
1019 # before waiting for result
1020 fetch_future_block_number = self._fetch_future_block_number
1021 fetch_future = self._fetch_future
1023 # Cleanup the fetch variables. Have a local copy.
1024 self._fetch_future_block_number = None
1025 self._fetch_future = None
1027 # Need to wait for the future for the current read
1028 if fetch_future is not None:
1029 logger.info("BlockCache waiting for background fetch.")
1030 # Wait until result and put it in cache
1031 self._fetch_block_cached.add_key(
1032 fetch_future.result(), fetch_future_block_number
1033 )
1035 # fetch next block in the background if nothing is running in the background,
1036 # the block is within file and it is not already cached
1037 end_block_plus_1 = end_block_number + 1
1038 with self._fetch_future_lock:
1039 if (
1040 self._fetch_future is None
1041 and end_block_plus_1 <= self.nblocks
1042 and not self._fetch_block_cached.is_key_cached(end_block_plus_1)
1043 ):
1044 self._fetch_future_block_number = end_block_plus_1
1045 self._fetch_future = self._thread_executor.submit(
1046 self._fetch_block, end_block_plus_1, "async"
1047 )
1049 return self._read_cache(
1050 start,
1051 end,
1052 start_block_number=start_block_number,
1053 end_block_number=end_block_number,
1054 )
1056 def _fetch_block(self, block_number: int, log_info: str = "sync") -> bytes:
1057 """
1058 Fetch the block of data for `block_number`.
1059 """
1060 if block_number > self.nblocks:
1061 raise ValueError(
1062 f"'block_number={block_number}' is greater than "
1063 f"the number of blocks ({self.nblocks})"
1064 )
1066 start = block_number * self.blocksize
1067 if start >= self.size:
1068 return b""
1069 end = min(start + self.blocksize, self.size)
1070 logger.info("BlockCache fetching block (%s) %d", log_info, block_number)
1071 self.total_requested_bytes += end - start
1072 self.miss_count += 1
1073 block_contents = super()._fetch(start, end)
1074 return block_contents
1076 def _read_cache(
1077 self, start: int, end: int, start_block_number: int, end_block_number: int
1078 ) -> bytes:
1079 """
1080 Read from our block cache.
1082 Parameters
1083 ----------
1084 start, end : int
1085 The start and end byte positions.
1086 start_block_number, end_block_number : int
1087 The start and end block numbers.
1088 """
1089 start_pos = start % self.blocksize
1090 end_pos = end % self.blocksize
1091 if end_pos == 0:
1092 end_pos = self.blocksize
1094 # kind of pointless to count this as a hit, but it is
1095 self.hit_count += 1
1097 if start_block_number == end_block_number:
1098 block = self._fetch_block_cached(start_block_number)
1099 return block[start_pos:end_pos]
1101 else:
1102 # read from the initial
1103 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
1105 # intermediate blocks
1106 # Note: it'd be nice to combine these into one big request. However
1107 # that doesn't play nicely with our LRU cache.
1108 out.extend(
1109 map(
1110 self._fetch_block_cached,
1111 range(start_block_number + 1, end_block_number),
1112 )
1113 )
1115 # final block
1116 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
1118 return b"".join(out)
1121caches: dict[str | None, type[BaseCache]] = {
1122 # one custom case
1123 None: BaseCache,
1124}
1127def register_cache(cls: type[BaseCache], clobber: bool = False) -> None:
1128 """'Register' cache implementation.
1130 Parameters
1131 ----------
1132 clobber: bool, optional
1133 If set to True (default is False) - allow to overwrite existing
1134 entry.
1136 Raises
1137 ------
1138 ValueError
1139 """
1140 name = cls.name
1141 if not clobber and name in caches:
1142 raise ValueError(f"Cache with name {name!r} is already known: {caches[name]}")
1143 caches[name] = cls
1146for c in (
1147 BaseCache,
1148 MMapCache,
1149 BytesCache,
1150 ReadAheadCache,
1151 AdaptiveReadaheadCache,
1152 BlockCache,
1153 FirstChunkCache,
1154 AllBytes,
1155 KnownPartsOfAFile,
1156 BackgroundBlockCache,
1157):
1158 register_cache(c)