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 if self.location is None:
154 fd = tempfile.TemporaryFile()
155 self.blocks = set()
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 = []
184 # Consolidate needed blocks.
185 # Algorithm adapted from Python 2.x itertools documentation.
186 # We are grouping an enumerated sequence of blocks. By comparing when the difference
187 # between an ascending range (provided by enumerate) and the needed block numbers
188 # we can detect when the block number skips values. The key computes this difference.
189 # Whenever the difference changes, we know that we have previously cached block(s),
190 # and a new group is started. In other words, this algorithm neatly groups
191 # runs of consecutive block numbers so they can be fetched together.
192 for _, _blocks in groupby(enumerate(need), key=lambda x: x[0] - x[1]):
193 # Extract the blocks from the enumerated sequence
194 _blocks = tuple(map(itemgetter(1), _blocks))
195 # Compute start of first block
196 sstart = _blocks[0] * self.blocksize
197 # Compute the end of the last block. Last block may not be full size.
198 send = min(_blocks[-1] * self.blocksize + self.blocksize, self.size)
200 # Fetch bytes (could be multiple consecutive blocks)
201 self.total_requested_bytes += send - sstart
202 logger.debug(
203 f"MMap get blocks {_blocks[0]}-{_blocks[-1]} ({sstart}-{send})"
204 )
205 ranges.append((sstart, send))
207 # Update set of cached blocks
208 self.blocks.update(_blocks)
209 # Update cache statistics with number of blocks we had to cache
210 self.miss_count += len(_blocks)
212 if not ranges:
213 return self.cache[start:end]
215 if self.multi_fetcher:
216 logger.debug(f"MMap get blocks {ranges}")
217 for idx, r in enumerate(self.multi_fetcher(ranges)):
218 sstart, send = ranges[idx]
219 logger.debug(f"MMap copy block ({sstart}-{send}")
220 self.cache[sstart:send] = r
221 else:
222 for sstart, send in ranges:
223 logger.debug(f"MMap get block ({sstart}-{send}")
224 self.cache[sstart:send] = self.fetcher(sstart, send)
226 return self.cache[start:end]
228 def __getstate__(self) -> dict[str, Any]:
229 state = self.__dict__.copy()
230 # Remove the unpicklable entries.
231 del state["cache"]
232 return state
234 def __setstate__(self, state: dict[str, Any]) -> None:
235 # Restore instance attributes
236 self.__dict__.update(state)
237 self.cache = self._makefile()
240class ReadAheadCache(BaseCache):
241 """Cache which reads only when we get beyond a block of data
243 This is a much simpler version of BytesCache, and does not attempt to
244 fill holes in the cache or keep fragments alive. It is best suited to
245 many small reads in a sequential order (e.g., reading lines from a file).
246 """
248 name = "readahead"
250 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
251 super().__init__(blocksize, fetcher, size)
252 self.cache = b""
253 self.start = 0
254 self.end = 0
256 def _fetch(self, start: int | None, end: int | None) -> bytes:
257 if start is None:
258 start = 0
259 if end is None or end > self.size:
260 end = self.size
261 if start >= self.size or start >= end:
262 return b""
263 l = end - start
264 if start >= self.start and end <= self.end:
265 # cache hit
266 self.hit_count += 1
267 return self.cache[start - self.start : end - self.start]
268 elif self.start <= start < self.end:
269 # partial hit
270 self.miss_count += 1
271 part = self.cache[start - self.start :]
272 l -= len(part)
273 start = self.end
274 else:
275 # miss
276 self.miss_count += 1
277 part = b""
278 end = min(self.size, end + self.blocksize)
279 self.total_requested_bytes += end - start
280 self.cache = self.fetcher(start, end) # new block replaces old
281 self.start = start
282 self.end = self.start + len(self.cache)
283 return part + self.cache[:l]
286class AdaptiveReadaheadCache(BaseCache):
287 """Cache with adaptive asynchronous prefetching.
289 Optimized for sequential and near-sequential reads by dynamically
290 adjusting the amount of data prefetched in the background. The cache
291 uses the generic prefetch engine when async loop support is available,
292 and falls back to ``ReadAheadCache`` when it is not.
294 Parameters
295 ----------
296 blocksize: int
297 Nominal read size used by callers.
298 fetcher: Fetcher
299 Function of the form ``f(start, end)`` that returns bytes.
300 size: int
301 Total size of the file.
302 concurrency: int
303 Maximum number of concurrent background fetch tasks.
304 max_prefetch_size: int | None
305 Optional upper bound for adaptive prefetch size in bytes.
306 """
308 name = "adaptive"
310 def __init__(
311 self,
312 blocksize: int,
313 fetcher: Fetcher,
314 size: int,
315 concurrency: int = 4,
316 max_prefetch_size: int | None = None,
317 ) -> None:
318 super().__init__(blocksize, fetcher, size)
319 self._fallback = ReadAheadCache(blocksize, fetcher, size)
320 self._prefetcher = None
322 async def _default_fetcher_async(
323 start_offset: int,
324 total_size: int,
325 split_factor: int = 1,
326 ) -> bytes:
327 del split_factor
328 return await asyncio.to_thread(
329 self.fetcher, start_offset, start_offset + total_size
330 )
332 try:
333 from . import asyn as fsspec_asyn
334 except ImportError as e:
335 raise ImportError(
336 "AdaptiveReadaheadCache requires fsspec.asyn to be available"
337 ) from e
339 try:
340 from .prefetcher import BackgroundPrefetcher
342 self._prefetcher = BackgroundPrefetcher(
343 fetcher=_default_fetcher_async,
344 size=size,
345 concurrency=concurrency,
346 max_prefetch_size=max_prefetch_size,
347 loop=fsspec_asyn.get_loop(),
348 )
349 logger.info(
350 "AdaptiveReadaheadCache enabled (blocksize=%d, size=%d, concurrency=%d, max_prefetch_size=%s)",
351 blocksize,
352 size,
353 concurrency,
354 max_prefetch_size,
355 )
356 except Exception as e:
357 logger.info(
358 "AdaptiveReadaheadCache fallback to ReadAheadCache: %s",
359 e,
360 exc_info=True,
361 )
362 self._prefetcher = None
364 def _fetch(self, start: int | None, end: int | None) -> bytes:
365 if self._prefetcher is None:
366 out = self._fallback._fetch(start, end)
367 self.hit_count = self._fallback.hit_count
368 self.miss_count = self._fallback.miss_count
369 self.total_requested_bytes = self._fallback.total_requested_bytes
370 return out
372 out = self._prefetcher.fetch(start, end)
373 self.miss_count += 1
374 self.total_requested_bytes += len(out)
375 return out
377 def close(self) -> None:
378 if self._prefetcher is not None:
379 self._prefetcher.close()
380 self._prefetcher = None
382 def __getstate__(self) -> dict[str, Any]:
383 # The prefetcher owns asyncio primitives that are not picklable.
384 self.close()
385 state = self.__dict__.copy()
386 state["_prefetcher"] = None
387 return state
389 def __setstate__(self, state: dict[str, Any]) -> None:
390 self.__dict__.update(state)
391 self._prefetcher = None
393 def __del__(self):
394 try:
395 self.close()
396 except Exception:
397 # Best-effort cleanup during GC.
398 pass
401class FirstChunkCache(BaseCache):
402 """Caches the first block of a file only
404 This may be useful for file types where the metadata is stored in the header,
405 but is randomly accessed.
406 """
408 name = "first"
410 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
411 if blocksize > size:
412 # this will buffer the whole thing
413 blocksize = size
414 super().__init__(blocksize, fetcher, size)
415 self.cache: bytes | None = None
417 def _fetch(self, start: int | None, end: int | None) -> bytes:
418 start = start or 0
419 if start > self.size:
420 logger.debug("FirstChunkCache: requested start > file size")
421 return b""
423 if end is None:
424 end = self.size
425 end = min(end, self.size)
427 if start < self.blocksize:
428 if self.cache is None:
429 self.miss_count += 1
430 if end > self.blocksize:
431 self.total_requested_bytes += end
432 data = self.fetcher(0, end)
433 self.cache = data[: self.blocksize]
434 return data[start:]
435 self.cache = self.fetcher(0, self.blocksize)
436 self.total_requested_bytes += self.blocksize
437 part = self.cache[start:end]
438 if end > self.blocksize:
439 self.total_requested_bytes += end - self.blocksize
440 part += self.fetcher(self.blocksize, end)
441 self.hit_count += 1
442 return part
443 else:
444 self.miss_count += 1
445 self.total_requested_bytes += end - start
446 return self.fetcher(start, end)
449class BlockCache(BaseCache):
450 """
451 Cache holding memory as a set of blocks.
453 Requests are only ever made ``blocksize`` at a time, and are
454 stored in an LRU cache. The least recently accessed block is
455 discarded when more than ``maxblocks`` are stored.
457 Parameters
458 ----------
459 blocksize : int
460 The number of bytes to store in each block.
461 Requests are only ever made for ``blocksize``, so this
462 should balance the overhead of making a request against
463 the granularity of the blocks.
464 fetcher : Callable
465 size : int
466 The total size of the file being cached.
467 maxblocks : int
468 The maximum number of blocks to cache for. The maximum memory
469 use for this cache is then ``blocksize * maxblocks``.
470 """
472 name = "blockcache"
474 def __init__(
475 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
476 ) -> None:
477 super().__init__(blocksize, fetcher, size)
478 self.nblocks = math.ceil(size / blocksize)
479 self.maxblocks = maxblocks
480 self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block)
482 def cache_info(self):
483 """
484 The statistics on the block cache.
486 Returns
487 -------
488 NamedTuple
489 Returned directly from the LRU Cache used internally.
490 """
491 return self._fetch_block_cached.cache_info()
493 def __getstate__(self) -> dict[str, Any]:
494 state = self.__dict__.copy()
495 del state["_fetch_block_cached"]
496 return state
498 def __setstate__(self, state: dict[str, Any]) -> None:
499 self.__dict__.update(state)
500 self._fetch_block_cached = functools.lru_cache(state["maxblocks"])(
501 self._fetch_block
502 )
504 def _fetch(self, start: int | None, end: int | None) -> bytes:
505 if start is None:
506 start = 0
507 if end is None or end > self.size:
508 end = self.size
509 if start >= self.size or start >= end:
510 return b""
512 return self._read_cache(
513 start, end, start // self.blocksize, (end - 1) // self.blocksize
514 )
516 def _fetch_block(self, block_number: int) -> bytes:
517 """
518 Fetch the block of data for `block_number`.
519 """
520 if block_number > self.nblocks:
521 raise ValueError(
522 f"'block_number={block_number}' is greater than "
523 f"the number of blocks ({self.nblocks})"
524 )
526 start = block_number * self.blocksize
527 end = start + self.blocksize
528 self.total_requested_bytes += end - start
529 self.miss_count += 1
530 logger.info("BlockCache fetching block %d", block_number)
531 block_contents = super()._fetch(start, end)
532 return block_contents
534 def _read_cache(
535 self, start: int, end: int, start_block_number: int, end_block_number: int
536 ) -> bytes:
537 """
538 Read from our block cache.
540 Parameters
541 ----------
542 start, end : int
543 The start and end byte positions.
544 start_block_number, end_block_number : int
545 The start and end block numbers.
546 """
547 start_pos = start % self.blocksize
548 end_pos = end % self.blocksize
549 if end_pos == 0:
550 end_pos = self.blocksize
552 self.hit_count += 1
553 if start_block_number == end_block_number:
554 block: bytes = self._fetch_block_cached(start_block_number)
555 return block[start_pos:end_pos]
557 else:
558 # read from the initial
559 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
561 # intermediate blocks
562 # Note: it'd be nice to combine these into one big request. However
563 # that doesn't play nicely with our LRU cache.
564 out.extend(
565 map(
566 self._fetch_block_cached,
567 range(start_block_number + 1, end_block_number),
568 )
569 )
571 # final block
572 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
574 return b"".join(out)
577class BytesCache(BaseCache):
578 """Cache which holds data in a in-memory bytes object
580 Implements read-ahead by the block size, for semi-random reads progressing
581 through the file.
583 Parameters
584 ----------
585 trim: bool
586 As we read more data, whether to discard the start of the buffer when
587 we are more than a blocksize ahead of it.
588 """
590 name: ClassVar[str] = "bytes"
592 def __init__(
593 self, blocksize: int, fetcher: Fetcher, size: int, trim: bool = True
594 ) -> None:
595 super().__init__(blocksize, fetcher, size)
596 self.cache = b""
597 self.start: int | None = None
598 self.end: int | None = None
599 self.trim = trim
601 def _fetch(self, start: int | None, end: int | None) -> bytes:
602 # TODO: only set start/end after fetch, in case it fails?
603 # is this where retry logic might go?
604 if start is None:
605 start = 0
606 if end is None:
607 end = self.size
608 if start >= self.size or start >= end:
609 return b""
610 if (
611 self.start is not None
612 and start >= self.start
613 and self.end is not None
614 and end <= self.end
615 ):
616 # cache hit: we have all the required data
617 offset = start - self.start
618 self.hit_count += 1
619 return self.cache[offset : offset + end - start]
621 if self.blocksize:
622 bend = min(self.size, end + self.blocksize)
623 else:
624 bend = end
626 if bend == start or start > self.size:
627 return b""
629 if (self.start is None or start < self.start) and (
630 self.end is None or end > self.end
631 ):
632 # First read, or extending both before and after
633 self.total_requested_bytes += bend - start
634 self.miss_count += 1
635 self.cache = self.fetcher(start, bend)
636 self.start = start
637 else:
638 assert self.start is not None
639 assert self.end is not None
640 self.miss_count += 1
642 if start < self.start:
643 if self.end is None or self.end - end > self.blocksize:
644 self.total_requested_bytes += bend - start
645 self.cache = self.fetcher(start, bend)
646 self.start = start
647 else:
648 self.total_requested_bytes += self.start - start
649 new = self.fetcher(start, self.start)
650 self.start = start
651 self.cache = new + self.cache
652 elif self.end is not None and bend > self.end:
653 if self.end > self.size:
654 pass
655 elif end - self.end > self.blocksize:
656 self.total_requested_bytes += bend - start
657 self.cache = self.fetcher(start, bend)
658 self.start = start
659 else:
660 self.total_requested_bytes += bend - self.end
661 new = self.fetcher(self.end, bend)
662 self.cache = self.cache + new
664 self.end = self.start + len(self.cache)
665 offset = start - self.start
666 out = self.cache[offset : offset + end - start]
667 if self.trim:
668 num = (self.end - self.start) // (self.blocksize + 1)
669 if num > 1:
670 self.start += self.blocksize * num
671 self.cache = self.cache[self.blocksize * num :]
672 return out
674 def __len__(self) -> int:
675 return len(self.cache)
678class AllBytes(BaseCache):
679 """Cache entire contents of the file"""
681 name: ClassVar[str] = "all"
683 def __init__(
684 self,
685 blocksize: int | None = None,
686 fetcher: Fetcher | None = None,
687 size: int | None = None,
688 data: bytes | None = None,
689 ) -> None:
690 super().__init__(blocksize, fetcher, size) # type: ignore[arg-type]
691 if data is None:
692 self.miss_count += 1
693 self.total_requested_bytes += self.size
694 data = self.fetcher(0, self.size)
695 self.data = data
697 def _fetch(self, start: int | None, stop: int | None) -> bytes:
698 self.hit_count += 1
699 return self.data[start:stop]
702class KnownPartsOfAFile(BaseCache):
703 """
704 Cache holding known file parts.
706 Parameters
707 ----------
708 blocksize: int
709 How far to read ahead in numbers of bytes
710 fetcher: func
711 Function of the form f(start, end) which gets bytes from remote as
712 specified
713 size: int
714 How big this file is
715 data: dict
716 A dictionary mapping explicit `(start, stop)` file-offset tuples
717 with known bytes.
718 strict: bool, default True
719 Whether to fetch reads that go beyond a known byte-range boundary.
720 If `False`, any read that ends outside a known part will be zero
721 padded. Note that zero padding will not be used for reads that
722 begin outside a known byte-range.
723 """
725 name: ClassVar[str] = "parts"
727 def __init__(
728 self,
729 blocksize: int,
730 fetcher: Fetcher,
731 size: int,
732 data: dict[tuple[int, int], bytes] | None = None,
733 strict: bool = False,
734 **_: Any,
735 ):
736 super().__init__(blocksize, fetcher, size)
737 self.strict = strict
739 # simple consolidation of contiguous blocks
740 if data:
741 old_offsets = sorted(data.keys())
742 offsets = [old_offsets[0]]
743 blocks = [data.pop(old_offsets[0])]
744 for start, stop in old_offsets[1:]:
745 start0, stop0 = offsets[-1]
746 if start == stop0:
747 offsets[-1] = (start0, stop)
748 blocks[-1] += data.pop((start, stop))
749 else:
750 offsets.append((start, stop))
751 blocks.append(data.pop((start, stop)))
753 self.data = dict(zip(offsets, blocks))
754 else:
755 self.data = {}
757 @property
758 def size(self):
759 return sum(_[1] - _[0] for _ in self.data)
761 @size.setter
762 def size(self, value):
763 pass
765 @property
766 def nblocks(self):
767 return len(self.data)
769 @nblocks.setter
770 def nblocks(self, value):
771 pass
773 def _fetch(self, start: int | None, stop: int | None) -> bytes:
774 logger.debug("Known parts request %s %s", start, stop)
775 if start is None:
776 start = 0
777 if stop is None:
778 stop = self.size
779 self.total_requested_bytes += stop - start
780 out = b""
781 started = False
782 loc_old = 0
783 for loc0, loc1 in sorted(self.data):
784 if (loc0 <= start < loc1) and (loc0 <= stop <= loc1):
785 # entirely within the block
786 off = start - loc0
787 self.hit_count += 1
788 return self.data[(loc0, loc1)][off : off + stop - start]
789 if stop <= loc0:
790 break
791 if started and loc0 > loc_old:
792 # a gap where we need data
793 self.miss_count += 1
794 if self.strict:
795 raise ValueError
796 out += b"\x00" * (loc0 - loc_old)
797 if loc0 <= start < loc1:
798 # found the start
799 self.hit_count += 1
800 off = start - loc0
801 out = self.data[(loc0, loc1)][off : off + stop - start]
802 started = True
803 elif start < loc0 and stop > loc1:
804 # the whole block
805 self.hit_count += 1
806 out += self.data[(loc0, loc1)]
807 elif loc0 <= stop <= loc1:
808 # end block
809 self.hit_count += 1
810 out = out + self.data[(loc0, loc1)][: stop - loc0]
811 return out
812 loc_old = loc1
813 self.miss_count += 1
814 if started and not self.strict:
815 out = out + b"\x00" * (stop - loc_old)
816 return out
817 raise ValueError
820class UpdatableLRU(Generic[P, T]):
821 """
822 Custom implementation of LRU cache that allows updating keys
824 Used by BackgroundBlockCache
825 """
827 class CacheInfo(NamedTuple):
828 hits: int
829 misses: int
830 maxsize: int
831 currsize: int
833 def __init__(self, func: Callable[P, T], max_size: int = 128) -> None:
834 self._cache: OrderedDict[Any, T] = collections.OrderedDict()
835 self._func = func
836 self._max_size = max_size
837 self._hits = 0
838 self._misses = 0
839 self._lock = threading.Lock()
841 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
842 if kwargs:
843 raise TypeError(f"Got unexpected keyword argument {kwargs.keys()}")
844 with self._lock:
845 if args in self._cache:
846 self._cache.move_to_end(args)
847 self._hits += 1
848 return self._cache[args]
850 result = self._func(*args, **kwargs)
852 with self._lock:
853 self._cache[args] = result
854 self._misses += 1
855 if len(self._cache) > self._max_size:
856 self._cache.popitem(last=False)
858 return result
860 def is_key_cached(self, *args: Any) -> bool:
861 with self._lock:
862 return args in self._cache
864 def add_key(self, result: T, *args: Any) -> None:
865 with self._lock:
866 self._cache[args] = result
867 if len(self._cache) > self._max_size:
868 self._cache.popitem(last=False)
870 def cache_info(self) -> UpdatableLRU.CacheInfo:
871 with self._lock:
872 return self.CacheInfo(
873 maxsize=self._max_size,
874 currsize=len(self._cache),
875 hits=self._hits,
876 misses=self._misses,
877 )
880class BackgroundBlockCache(BaseCache):
881 """
882 Cache holding memory as a set of blocks with pre-loading of
883 the next block in the background.
885 Requests are only ever made ``blocksize`` at a time, and are
886 stored in an LRU cache. The least recently accessed block is
887 discarded when more than ``maxblocks`` are stored. If the
888 next block is not in cache, it is loaded in a separate thread
889 in non-blocking way.
891 Parameters
892 ----------
893 blocksize : int
894 The number of bytes to store in each block.
895 Requests are only ever made for ``blocksize``, so this
896 should balance the overhead of making a request against
897 the granularity of the blocks.
898 fetcher : Callable
899 size : int
900 The total size of the file being cached.
901 maxblocks : int
902 The maximum number of blocks to cache for. The maximum memory
903 use for this cache is then ``blocksize * maxblocks``.
904 """
906 name: ClassVar[str] = "background"
908 def __init__(
909 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
910 ) -> None:
911 super().__init__(blocksize, fetcher, size)
912 self.nblocks = math.ceil(size / blocksize)
913 self.maxblocks = maxblocks
914 self._fetch_block_cached = UpdatableLRU(self._fetch_block, maxblocks)
916 self._thread_executor = ThreadPoolExecutor(max_workers=1)
917 self._fetch_future_block_number: int | None = None
918 self._fetch_future: Future[bytes] | None = None
919 self._fetch_future_lock = threading.Lock()
920 self._closed = False
922 def cache_info(self) -> UpdatableLRU.CacheInfo:
923 """
924 The statistics on the block cache.
926 Returns
927 -------
928 NamedTuple
929 Returned directly from the LRU Cache used internally.
930 """
931 return self._fetch_block_cached.cache_info()
933 def close(self) -> None:
934 """Cancel pending work and shut down the background worker."""
935 with self._fetch_future_lock:
936 if self._closed:
937 return
938 self._closed = True
939 future = self._fetch_future
940 self._fetch_future = None
941 self._fetch_future_block_number = None
943 if future is not None:
944 future.cancel()
945 self._thread_executor.shutdown(wait=True, cancel_futures=True)
947 # UpdatableLRU stores a bound method and otherwise forms a reference cycle.
948 del self._fetch_block_cached
950 def __getstate__(self) -> dict[str, Any]:
951 state = self.__dict__.copy()
952 del state["_fetch_block_cached"]
953 del state["_thread_executor"]
954 del state["_fetch_future_block_number"]
955 del state["_fetch_future"]
956 del state["_fetch_future_lock"]
957 return state
959 def __setstate__(self, state) -> None:
960 self.__dict__.update(state)
961 self._fetch_block_cached = UpdatableLRU(self._fetch_block, state["maxblocks"])
962 self._thread_executor = ThreadPoolExecutor(max_workers=1)
963 self._fetch_future_block_number = None
964 self._fetch_future = None
965 self._fetch_future_lock = threading.Lock()
966 self._closed = False
968 def _fetch(self, start: int | None, end: int | None) -> bytes:
969 if start is None:
970 start = 0
971 if end is None or end > self.size:
972 end = self.size
973 if start >= self.size or start >= end:
974 return b""
976 # byte position -> block numbers; ``end`` is exclusive
977 start_block_number = start // self.blocksize
978 end_block_number = (end - 1) // self.blocksize
980 fetch_future_block_number = None
981 fetch_future = None
982 with self._fetch_future_lock:
983 # Background thread is running. Check we we can or must join it.
984 if self._fetch_future is not None:
985 assert self._fetch_future_block_number is not None
986 if self._fetch_future.done():
987 logger.info("BlockCache joined background fetch without waiting.")
988 self._fetch_block_cached.add_key(
989 self._fetch_future.result(), self._fetch_future_block_number
990 )
991 # Cleanup the fetch variables. Done with fetching the block.
992 self._fetch_future_block_number = None
993 self._fetch_future = None
994 else:
995 # Must join if we need the block for the current fetch
996 must_join = bool(
997 start_block_number
998 <= self._fetch_future_block_number
999 <= end_block_number
1000 )
1001 if must_join:
1002 # Copy to the local variables to release lock
1003 # before waiting for result
1004 fetch_future_block_number = self._fetch_future_block_number
1005 fetch_future = self._fetch_future
1007 # Cleanup the fetch variables. Have a local copy.
1008 self._fetch_future_block_number = None
1009 self._fetch_future = None
1011 # Need to wait for the future for the current read
1012 if fetch_future is not None:
1013 logger.info("BlockCache waiting for background fetch.")
1014 # Wait until result and put it in cache
1015 self._fetch_block_cached.add_key(
1016 fetch_future.result(), fetch_future_block_number
1017 )
1019 # fetch next block in the background if nothing is running in the background,
1020 # the block is within file and it is not already cached
1021 end_block_plus_1 = end_block_number + 1
1022 with self._fetch_future_lock:
1023 if (
1024 self._fetch_future is None
1025 and end_block_plus_1 <= self.nblocks
1026 and not self._fetch_block_cached.is_key_cached(end_block_plus_1)
1027 ):
1028 self._fetch_future_block_number = end_block_plus_1
1029 self._fetch_future = self._thread_executor.submit(
1030 self._fetch_block, end_block_plus_1, "async"
1031 )
1033 return self._read_cache(
1034 start,
1035 end,
1036 start_block_number=start_block_number,
1037 end_block_number=end_block_number,
1038 )
1040 def _fetch_block(self, block_number: int, log_info: str = "sync") -> bytes:
1041 """
1042 Fetch the block of data for `block_number`.
1043 """
1044 if block_number > self.nblocks:
1045 raise ValueError(
1046 f"'block_number={block_number}' is greater than "
1047 f"the number of blocks ({self.nblocks})"
1048 )
1050 start = block_number * self.blocksize
1051 end = start + self.blocksize
1052 logger.info("BlockCache fetching block (%s) %d", log_info, block_number)
1053 self.total_requested_bytes += end - start
1054 self.miss_count += 1
1055 block_contents = super()._fetch(start, end)
1056 return block_contents
1058 def _read_cache(
1059 self, start: int, end: int, start_block_number: int, end_block_number: int
1060 ) -> bytes:
1061 """
1062 Read from our block cache.
1064 Parameters
1065 ----------
1066 start, end : int
1067 The start and end byte positions.
1068 start_block_number, end_block_number : int
1069 The start and end block numbers.
1070 """
1071 start_pos = start % self.blocksize
1072 end_pos = end % self.blocksize
1073 if end_pos == 0:
1074 end_pos = self.blocksize
1076 # kind of pointless to count this as a hit, but it is
1077 self.hit_count += 1
1079 if start_block_number == end_block_number:
1080 block = self._fetch_block_cached(start_block_number)
1081 return block[start_pos:end_pos]
1083 else:
1084 # read from the initial
1085 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
1087 # intermediate blocks
1088 # Note: it'd be nice to combine these into one big request. However
1089 # that doesn't play nicely with our LRU cache.
1090 out.extend(
1091 map(
1092 self._fetch_block_cached,
1093 range(start_block_number + 1, end_block_number),
1094 )
1095 )
1097 # final block
1098 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
1100 return b"".join(out)
1103caches: dict[str | None, type[BaseCache]] = {
1104 # one custom case
1105 None: BaseCache,
1106}
1109def register_cache(cls: type[BaseCache], clobber: bool = False) -> None:
1110 """'Register' cache implementation.
1112 Parameters
1113 ----------
1114 clobber: bool, optional
1115 If set to True (default is False) - allow to overwrite existing
1116 entry.
1118 Raises
1119 ------
1120 ValueError
1121 """
1122 name = cls.name
1123 if not clobber and name in caches:
1124 raise ValueError(f"Cache with name {name!r} is already known: {caches[name]}")
1125 caches[name] = cls
1128for c in (
1129 BaseCache,
1130 MMapCache,
1131 BytesCache,
1132 ReadAheadCache,
1133 AdaptiveReadaheadCache,
1134 BlockCache,
1135 FirstChunkCache,
1136 AllBytes,
1137 KnownPartsOfAFile,
1138 BackgroundBlockCache,
1139):
1140 register_cache(c)