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 collections
4import functools
5import logging
6import math
7import os
8import threading
9from collections import OrderedDict
10from collections.abc import Callable
11from concurrent.futures import Future, ThreadPoolExecutor
12from itertools import groupby
13from operator import itemgetter
14from typing import TYPE_CHECKING, Any, ClassVar, Generic, NamedTuple, TypeVar
16if TYPE_CHECKING:
17 import mmap
19 from typing_extensions import ParamSpec
21 P = ParamSpec("P")
22else:
23 P = TypeVar("P")
25T = TypeVar("T")
28logger = logging.getLogger("fsspec.caching")
30Fetcher = Callable[[int, int], bytes] # Maps (start, end) to bytes
31MultiFetcher = Callable[[list[int, int]], bytes] # Maps [(start, end)] to bytes
34class BaseCache:
35 """Pass-though cache: doesn't keep anything, calls every time
37 Acts as base class for other cachers
39 Parameters
40 ----------
41 blocksize: int
42 How far to read ahead in numbers of bytes
43 fetcher: func
44 Function of the form f(start, end) which gets bytes from remote as
45 specified
46 size: int
47 How big this file is
48 """
50 name: ClassVar[str] = "none"
52 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
53 self.blocksize = blocksize
54 self.nblocks = 0
55 self.fetcher = fetcher
56 self.size = size
57 self.hit_count = 0
58 self.miss_count = 0
59 # the bytes that we actually requested
60 self.total_requested_bytes = 0
62 def _fetch(self, start: int | None, stop: int | None) -> bytes:
63 if start is None:
64 start = 0
65 if stop is None:
66 stop = self.size
67 if start >= self.size or start >= stop:
68 return b""
69 return self.fetcher(start, stop)
71 def _reset_stats(self) -> None:
72 """Reset hit and miss counts for a more ganular report e.g. by file."""
73 self.hit_count = 0
74 self.miss_count = 0
75 self.total_requested_bytes = 0
77 def _log_stats(self) -> str:
78 """Return a formatted string of the cache statistics."""
79 if self.hit_count == 0 and self.miss_count == 0:
80 # a cache that does nothing, this is for logs only
81 return ""
82 return f" , {self.name}: {self.hit_count} hits, {self.miss_count} misses, {self.total_requested_bytes} total requested bytes"
84 def __repr__(self) -> str:
85 # TODO: use rich for better formatting
86 return f"""
87 <{self.__class__.__name__}:
88 block size : {self.blocksize}
89 block count : {self.nblocks}
90 file size : {self.size}
91 cache hits : {self.hit_count}
92 cache misses: {self.miss_count}
93 total requested bytes: {self.total_requested_bytes}>
94 """
97class MMapCache(BaseCache):
98 """memory-mapped sparse file cache
100 Opens temporary file, which is filled blocks-wise when data is requested.
101 Ensure there is enough disc space in the temporary location.
103 This cache method might only work on posix
105 Parameters
106 ----------
107 blocksize: int
108 How far to read ahead in numbers of bytes
109 fetcher: Fetcher
110 Function of the form f(start, end) which gets bytes from remote as
111 specified
112 size: int
113 How big this file is
114 location: str
115 Where to create the temporary file. If None, a temporary file is
116 created using tempfile.TemporaryFile().
117 blocks: set[int]
118 Set of block numbers that have already been fetched. If None, an empty
119 set is created.
120 multi_fetcher: MultiFetcher
121 Function of the form f([(start, end)]) which gets bytes from remote
122 as specified. This function is used to fetch multiple blocks at once.
123 If not specified, the fetcher function is used instead.
124 """
126 name = "mmap"
128 def __init__(
129 self,
130 blocksize: int,
131 fetcher: Fetcher,
132 size: int,
133 location: str | None = None,
134 blocks: set[int] | None = None,
135 multi_fetcher: MultiFetcher | None = None,
136 ) -> None:
137 super().__init__(blocksize, fetcher, size)
138 self.blocks = set() if blocks is None else blocks
139 self.location = location
140 self.multi_fetcher = multi_fetcher
141 self.cache = self._makefile()
143 def _makefile(self) -> mmap.mmap | bytearray:
144 import mmap
145 import tempfile
147 if self.size == 0:
148 return bytearray()
150 # posix version
151 if self.location is None or not os.path.exists(self.location):
152 if self.location is None:
153 fd = tempfile.TemporaryFile()
154 self.blocks = set()
155 else:
156 fd = open(self.location, "wb+")
157 fd.seek(self.size - 1)
158 fd.write(b"1")
159 fd.flush()
160 else:
161 fd = open(self.location, "r+b")
163 return mmap.mmap(fd.fileno(), self.size)
165 def _fetch(self, start: int | None, end: int | None) -> bytes:
166 logger.debug(f"MMap cache fetching {start}-{end}")
167 if start is None:
168 start = 0
169 if end is None:
170 end = self.size
171 if start >= self.size or start >= end:
172 return b""
173 start_block = start // self.blocksize
174 end_block = end // self.blocksize
175 block_range = range(start_block, end_block + 1)
176 # Determine which blocks need to be fetched. This sequence is sorted by construction.
177 need = (i for i in block_range if i not in self.blocks)
178 # Count the number of blocks already cached
179 self.hit_count += sum(1 for i in block_range if i in self.blocks)
181 ranges = []
183 # Consolidate needed blocks.
184 # Algorithm adapted from Python 2.x itertools documentation.
185 # We are grouping an enumerated sequence of blocks. By comparing when the difference
186 # between an ascending range (provided by enumerate) and the needed block numbers
187 # we can detect when the block number skips values. The key computes this difference.
188 # Whenever the difference changes, we know that we have previously cached block(s),
189 # and a new group is started. In other words, this algorithm neatly groups
190 # runs of consecutive block numbers so they can be fetched together.
191 for _, _blocks in groupby(enumerate(need), key=lambda x: x[0] - x[1]):
192 # Extract the blocks from the enumerated sequence
193 _blocks = tuple(map(itemgetter(1), _blocks))
194 # Compute start of first block
195 sstart = _blocks[0] * self.blocksize
196 # Compute the end of the last block. Last block may not be full size.
197 send = min(_blocks[-1] * self.blocksize + self.blocksize, self.size)
199 # Fetch bytes (could be multiple consecutive blocks)
200 self.total_requested_bytes += send - sstart
201 logger.debug(
202 f"MMap get blocks {_blocks[0]}-{_blocks[-1]} ({sstart}-{send})"
203 )
204 ranges.append((sstart, send))
206 # Update set of cached blocks
207 self.blocks.update(_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 else:
221 for sstart, send in ranges:
222 logger.debug(f"MMap get block ({sstart}-{send}")
223 self.cache[sstart:send] = self.fetcher(sstart, send)
225 return self.cache[start:end]
227 def __getstate__(self) -> dict[str, Any]:
228 state = self.__dict__.copy()
229 # Remove the unpicklable entries.
230 del state["cache"]
231 return state
233 def __setstate__(self, state: dict[str, Any]) -> None:
234 # Restore instance attributes
235 self.__dict__.update(state)
236 self.cache = self._makefile()
239class ReadAheadCache(BaseCache):
240 """Cache which reads only when we get beyond a block of data
242 This is a much simpler version of BytesCache, and does not attempt to
243 fill holes in the cache or keep fragments alive. It is best suited to
244 many small reads in a sequential order (e.g., reading lines from a file).
245 """
247 name = "readahead"
249 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
250 super().__init__(blocksize, fetcher, size)
251 self.cache = b""
252 self.start = 0
253 self.end = 0
255 def _fetch(self, start: int | None, end: int | None) -> bytes:
256 if start is None:
257 start = 0
258 if end is None or end > self.size:
259 end = self.size
260 if start >= self.size or start >= end:
261 return b""
262 l = end - start
263 if start >= self.start and end <= self.end:
264 # cache hit
265 self.hit_count += 1
266 return self.cache[start - self.start : end - self.start]
267 elif self.start <= start < self.end:
268 # partial hit
269 self.miss_count += 1
270 part = self.cache[start - self.start :]
271 l -= len(part)
272 start = self.end
273 else:
274 # miss
275 self.miss_count += 1
276 part = b""
277 end = min(self.size, end + self.blocksize)
278 self.total_requested_bytes += end - start
279 self.cache = self.fetcher(start, end) # new block replaces old
280 self.start = start
281 self.end = self.start + len(self.cache)
282 return part + self.cache[:l]
285class FirstChunkCache(BaseCache):
286 """Caches the first block of a file only
288 This may be useful for file types where the metadata is stored in the header,
289 but is randomly accessed.
290 """
292 name = "first"
294 def __init__(self, blocksize: int, fetcher: Fetcher, size: int) -> None:
295 if blocksize > size:
296 # this will buffer the whole thing
297 blocksize = size
298 super().__init__(blocksize, fetcher, size)
299 self.cache: bytes | None = None
301 def _fetch(self, start: int | None, end: int | None) -> bytes:
302 start = start or 0
303 if start > self.size:
304 logger.debug("FirstChunkCache: requested start > file size")
305 return b""
307 if end is None:
308 end = self.size
309 end = min(end, self.size)
311 if start < self.blocksize:
312 if self.cache is None:
313 self.miss_count += 1
314 if end > self.blocksize:
315 self.total_requested_bytes += end
316 data = self.fetcher(0, end)
317 self.cache = data[: self.blocksize]
318 return data[start:]
319 self.cache = self.fetcher(0, self.blocksize)
320 self.total_requested_bytes += self.blocksize
321 part = self.cache[start:end]
322 if end > self.blocksize:
323 self.total_requested_bytes += end - self.blocksize
324 part += self.fetcher(self.blocksize, end)
325 self.hit_count += 1
326 return part
327 else:
328 self.miss_count += 1
329 self.total_requested_bytes += end - start
330 return self.fetcher(start, end)
333class BlockCache(BaseCache):
334 """
335 Cache holding memory as a set of blocks.
337 Requests are only ever made ``blocksize`` at a time, and are
338 stored in an LRU cache. The least recently accessed block is
339 discarded when more than ``maxblocks`` are stored.
341 Parameters
342 ----------
343 blocksize : int
344 The number of bytes to store in each block.
345 Requests are only ever made for ``blocksize``, so this
346 should balance the overhead of making a request against
347 the granularity of the blocks.
348 fetcher : Callable
349 size : int
350 The total size of the file being cached.
351 maxblocks : int
352 The maximum number of blocks to cache for. The maximum memory
353 use for this cache is then ``blocksize * maxblocks``.
354 """
356 name = "blockcache"
358 def __init__(
359 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
360 ) -> None:
361 super().__init__(blocksize, fetcher, size)
362 self.nblocks = math.ceil(size / blocksize)
363 self.maxblocks = maxblocks
364 self._fetch_block_cached = functools.lru_cache(maxblocks)(self._fetch_block)
366 def cache_info(self):
367 """
368 The statistics on the block cache.
370 Returns
371 -------
372 NamedTuple
373 Returned directly from the LRU Cache used internally.
374 """
375 return self._fetch_block_cached.cache_info()
377 def __getstate__(self) -> dict[str, Any]:
378 state = self.__dict__
379 del state["_fetch_block_cached"]
380 return state
382 def __setstate__(self, state: dict[str, Any]) -> None:
383 self.__dict__.update(state)
384 self._fetch_block_cached = functools.lru_cache(state["maxblocks"])(
385 self._fetch_block
386 )
388 def _fetch(self, start: int | None, end: int | None) -> bytes:
389 if start is None:
390 start = 0
391 if end is None or end > self.size:
392 end = self.size
393 if start >= self.size or start >= end:
394 return b""
396 return self._read_cache(
397 start, end, start // self.blocksize, (end - 1) // self.blocksize
398 )
400 def _fetch_block(self, block_number: int) -> bytes:
401 """
402 Fetch the block of data for `block_number`.
403 """
404 if block_number > self.nblocks:
405 raise ValueError(
406 f"'block_number={block_number}' is greater than "
407 f"the number of blocks ({self.nblocks})"
408 )
410 start = block_number * self.blocksize
411 end = start + self.blocksize
412 self.total_requested_bytes += end - start
413 self.miss_count += 1
414 logger.info("BlockCache fetching block %d", block_number)
415 block_contents = super()._fetch(start, end)
416 return block_contents
418 def _read_cache(
419 self, start: int, end: int, start_block_number: int, end_block_number: int
420 ) -> bytes:
421 """
422 Read from our block cache.
424 Parameters
425 ----------
426 start, end : int
427 The start and end byte positions.
428 start_block_number, end_block_number : int
429 The start and end block numbers.
430 """
431 start_pos = start % self.blocksize
432 end_pos = end % self.blocksize
433 if end_pos == 0:
434 end_pos = self.blocksize
436 self.hit_count += 1
437 if start_block_number == end_block_number:
438 block: bytes = self._fetch_block_cached(start_block_number)
439 return block[start_pos:end_pos]
441 else:
442 # read from the initial
443 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
445 # intermediate blocks
446 # Note: it'd be nice to combine these into one big request. However
447 # that doesn't play nicely with our LRU cache.
448 out.extend(
449 map(
450 self._fetch_block_cached,
451 range(start_block_number + 1, end_block_number),
452 )
453 )
455 # final block
456 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
458 return b"".join(out)
461class BytesCache(BaseCache):
462 """Cache which holds data in a in-memory bytes object
464 Implements read-ahead by the block size, for semi-random reads progressing
465 through the file.
467 Parameters
468 ----------
469 trim: bool
470 As we read more data, whether to discard the start of the buffer when
471 we are more than a blocksize ahead of it.
472 """
474 name: ClassVar[str] = "bytes"
476 def __init__(
477 self, blocksize: int, fetcher: Fetcher, size: int, trim: bool = True
478 ) -> None:
479 super().__init__(blocksize, fetcher, size)
480 self.cache = b""
481 self.start: int | None = None
482 self.end: int | None = None
483 self.trim = trim
485 def _fetch(self, start: int | None, end: int | None) -> bytes:
486 # TODO: only set start/end after fetch, in case it fails?
487 # is this where retry logic might go?
488 if start is None:
489 start = 0
490 if end is None:
491 end = self.size
492 if start >= self.size or start >= end:
493 return b""
494 if (
495 self.start is not None
496 and start >= self.start
497 and self.end is not None
498 and end < self.end
499 ):
500 # cache hit: we have all the required data
501 offset = start - self.start
502 self.hit_count += 1
503 return self.cache[offset : offset + end - start]
505 if self.blocksize:
506 bend = min(self.size, end + self.blocksize)
507 else:
508 bend = end
510 if bend == start or start > self.size:
511 return b""
513 if (self.start is None or start < self.start) and (
514 self.end is None or end > self.end
515 ):
516 # First read, or extending both before and after
517 self.total_requested_bytes += bend - start
518 self.miss_count += 1
519 self.cache = self.fetcher(start, bend)
520 self.start = start
521 else:
522 assert self.start is not None
523 assert self.end is not None
524 self.miss_count += 1
526 if start < self.start:
527 if self.end is None or self.end - end > self.blocksize:
528 self.total_requested_bytes += bend - start
529 self.cache = self.fetcher(start, bend)
530 self.start = start
531 else:
532 self.total_requested_bytes += self.start - start
533 new = self.fetcher(start, self.start)
534 self.start = start
535 self.cache = new + self.cache
536 elif self.end is not None and bend > self.end:
537 if self.end > self.size:
538 pass
539 elif end - self.end > self.blocksize:
540 self.total_requested_bytes += bend - start
541 self.cache = self.fetcher(start, bend)
542 self.start = start
543 else:
544 self.total_requested_bytes += bend - self.end
545 new = self.fetcher(self.end, bend)
546 self.cache = self.cache + new
548 self.end = self.start + len(self.cache)
549 offset = start - self.start
550 out = self.cache[offset : offset + end - start]
551 if self.trim:
552 num = (self.end - self.start) // (self.blocksize + 1)
553 if num > 1:
554 self.start += self.blocksize * num
555 self.cache = self.cache[self.blocksize * num :]
556 return out
558 def __len__(self) -> int:
559 return len(self.cache)
562class AllBytes(BaseCache):
563 """Cache entire contents of the file"""
565 name: ClassVar[str] = "all"
567 def __init__(
568 self,
569 blocksize: int | None = None,
570 fetcher: Fetcher | None = None,
571 size: int | None = None,
572 data: bytes | None = None,
573 ) -> None:
574 super().__init__(blocksize, fetcher, size) # type: ignore[arg-type]
575 if data is None:
576 self.miss_count += 1
577 self.total_requested_bytes += self.size
578 data = self.fetcher(0, self.size)
579 self.data = data
581 def _fetch(self, start: int | None, stop: int | None) -> bytes:
582 self.hit_count += 1
583 return self.data[start:stop]
586class KnownPartsOfAFile(BaseCache):
587 """
588 Cache holding known file parts.
590 Parameters
591 ----------
592 blocksize: int
593 How far to read ahead in numbers of bytes
594 fetcher: func
595 Function of the form f(start, end) which gets bytes from remote as
596 specified
597 size: int
598 How big this file is
599 data: dict
600 A dictionary mapping explicit `(start, stop)` file-offset tuples
601 with known bytes.
602 strict: bool, default True
603 Whether to fetch reads that go beyond a known byte-range boundary.
604 If `False`, any read that ends outside a known part will be zero
605 padded. Note that zero padding will not be used for reads that
606 begin outside a known byte-range.
607 """
609 name: ClassVar[str] = "parts"
611 def __init__(
612 self,
613 blocksize: int,
614 fetcher: Fetcher,
615 size: int,
616 data: dict[tuple[int, int], bytes] | None = None,
617 strict: bool = False,
618 **_: Any,
619 ):
620 super().__init__(blocksize, fetcher, size)
621 self.strict = strict
623 # simple consolidation of contiguous blocks
624 if data:
625 old_offsets = sorted(data.keys())
626 offsets = [old_offsets[0]]
627 blocks = [data.pop(old_offsets[0])]
628 for start, stop in old_offsets[1:]:
629 start0, stop0 = offsets[-1]
630 if start == stop0:
631 offsets[-1] = (start0, stop)
632 blocks[-1] += data.pop((start, stop))
633 else:
634 offsets.append((start, stop))
635 blocks.append(data.pop((start, stop)))
637 self.data = dict(zip(offsets, blocks))
638 else:
639 self.data = {}
641 @property
642 def size(self):
643 return sum(_[1] - _[0] for _ in self.data)
645 @size.setter
646 def size(self, value):
647 pass
649 @property
650 def nblocks(self):
651 return len(self.data)
653 @nblocks.setter
654 def nblocks(self, value):
655 pass
657 def _fetch(self, start: int | None, stop: int | None) -> bytes:
658 logger.debug("Known parts request %s %s", start, stop)
659 if start is None:
660 start = 0
661 if stop is None:
662 stop = self.size
663 self.total_requested_bytes += stop - start
664 out = b""
665 started = False
666 loc_old = 0
667 for loc0, loc1 in sorted(self.data):
668 if (loc0 <= start < loc1) and (loc0 <= stop <= loc1):
669 # entirely within the block
670 off = start - loc0
671 self.hit_count += 1
672 return self.data[(loc0, loc1)][off : off + stop - start]
673 if stop <= loc0:
674 break
675 if started and loc0 > loc_old:
676 # a gap where we need data
677 self.miss_count += 1
678 if self.strict:
679 raise ValueError
680 out += b"\x00" * (loc0 - loc_old)
681 if loc0 <= start < loc1:
682 # found the start
683 self.hit_count += 1
684 off = start - loc0
685 out = self.data[(loc0, loc1)][off : off + stop - start]
686 started = True
687 elif start < loc0 and stop > loc1:
688 # the whole block
689 self.hit_count += 1
690 out += self.data[(loc0, loc1)]
691 elif loc0 <= stop <= loc1:
692 # end block
693 self.hit_count += 1
694 out = out + self.data[(loc0, loc1)][: stop - loc0]
695 return out
696 loc_old = loc1
697 self.miss_count += 1
698 if started and not self.strict:
699 out = out + b"\x00" * (stop - loc_old)
700 return out
701 raise ValueError
704class UpdatableLRU(Generic[P, T]):
705 """
706 Custom implementation of LRU cache that allows updating keys
708 Used by BackgroundBlockCache
709 """
711 class CacheInfo(NamedTuple):
712 hits: int
713 misses: int
714 maxsize: int
715 currsize: int
717 def __init__(self, func: Callable[P, T], max_size: int = 128) -> None:
718 self._cache: OrderedDict[Any, T] = collections.OrderedDict()
719 self._func = func
720 self._max_size = max_size
721 self._hits = 0
722 self._misses = 0
723 self._lock = threading.Lock()
725 def __call__(self, *args: P.args, **kwargs: P.kwargs) -> T:
726 if kwargs:
727 raise TypeError(f"Got unexpected keyword argument {kwargs.keys()}")
728 with self._lock:
729 if args in self._cache:
730 self._cache.move_to_end(args)
731 self._hits += 1
732 return self._cache[args]
734 result = self._func(*args, **kwargs)
736 with self._lock:
737 self._cache[args] = result
738 self._misses += 1
739 if len(self._cache) > self._max_size:
740 self._cache.popitem(last=False)
742 return result
744 def is_key_cached(self, *args: Any) -> bool:
745 with self._lock:
746 return args in self._cache
748 def add_key(self, result: T, *args: Any) -> None:
749 with self._lock:
750 self._cache[args] = result
751 if len(self._cache) > self._max_size:
752 self._cache.popitem(last=False)
754 def cache_info(self) -> UpdatableLRU.CacheInfo:
755 with self._lock:
756 return self.CacheInfo(
757 maxsize=self._max_size,
758 currsize=len(self._cache),
759 hits=self._hits,
760 misses=self._misses,
761 )
764class BackgroundBlockCache(BaseCache):
765 """
766 Cache holding memory as a set of blocks with pre-loading of
767 the next block in the background.
769 Requests are only ever made ``blocksize`` at a time, and are
770 stored in an LRU cache. The least recently accessed block is
771 discarded when more than ``maxblocks`` are stored. If the
772 next block is not in cache, it is loaded in a separate thread
773 in non-blocking way.
775 Parameters
776 ----------
777 blocksize : int
778 The number of bytes to store in each block.
779 Requests are only ever made for ``blocksize``, so this
780 should balance the overhead of making a request against
781 the granularity of the blocks.
782 fetcher : Callable
783 size : int
784 The total size of the file being cached.
785 maxblocks : int
786 The maximum number of blocks to cache for. The maximum memory
787 use for this cache is then ``blocksize * maxblocks``.
788 """
790 name: ClassVar[str] = "background"
792 def __init__(
793 self, blocksize: int, fetcher: Fetcher, size: int, maxblocks: int = 32
794 ) -> None:
795 super().__init__(blocksize, fetcher, size)
796 self.nblocks = math.ceil(size / blocksize)
797 self.maxblocks = maxblocks
798 self._fetch_block_cached = UpdatableLRU(self._fetch_block, maxblocks)
800 self._thread_executor = ThreadPoolExecutor(max_workers=1)
801 self._fetch_future_block_number: int | None = None
802 self._fetch_future: Future[bytes] | None = None
803 self._fetch_future_lock = threading.Lock()
804 self._closed = False
806 def cache_info(self) -> UpdatableLRU.CacheInfo:
807 """
808 The statistics on the block cache.
810 Returns
811 -------
812 NamedTuple
813 Returned directly from the LRU Cache used internally.
814 """
815 return self._fetch_block_cached.cache_info()
817 def close(self) -> None:
818 """Cancel pending work and shut down the background worker."""
819 with self._fetch_future_lock:
820 if self._closed:
821 return
822 self._closed = True
823 future = self._fetch_future
824 self._fetch_future = None
825 self._fetch_future_block_number = None
827 if future is not None:
828 future.cancel()
829 self._thread_executor.shutdown(wait=True, cancel_futures=True)
831 # UpdatableLRU stores a bound method and otherwise forms a reference cycle.
832 del self._fetch_block_cached
834 def __getstate__(self) -> dict[str, Any]:
835 state = self.__dict__
836 del state["_fetch_block_cached"]
837 del state["_thread_executor"]
838 del state["_fetch_future_block_number"]
839 del state["_fetch_future"]
840 del state["_fetch_future_lock"]
841 return state
843 def __setstate__(self, state) -> None:
844 self.__dict__.update(state)
845 self._fetch_block_cached = UpdatableLRU(self._fetch_block, state["maxblocks"])
846 self._thread_executor = ThreadPoolExecutor(max_workers=1)
847 self._fetch_future_block_number = None
848 self._fetch_future = None
849 self._fetch_future_lock = threading.Lock()
850 self._closed = False
852 def _fetch(self, start: int | None, end: int | None) -> bytes:
853 if start is None:
854 start = 0
855 if end is None or end > self.size:
856 end = self.size
857 if start >= self.size or start >= end:
858 return b""
860 # byte position -> block numbers
861 start_block_number = start // self.blocksize
862 end_block_number = end // self.blocksize
864 fetch_future_block_number = None
865 fetch_future = None
866 with self._fetch_future_lock:
867 # Background thread is running. Check we we can or must join it.
868 if self._fetch_future is not None:
869 assert self._fetch_future_block_number is not None
870 if self._fetch_future.done():
871 logger.info("BlockCache joined background fetch without waiting.")
872 self._fetch_block_cached.add_key(
873 self._fetch_future.result(), self._fetch_future_block_number
874 )
875 # Cleanup the fetch variables. Done with fetching the block.
876 self._fetch_future_block_number = None
877 self._fetch_future = None
878 else:
879 # Must join if we need the block for the current fetch
880 must_join = bool(
881 start_block_number
882 <= self._fetch_future_block_number
883 <= end_block_number
884 )
885 if must_join:
886 # Copy to the local variables to release lock
887 # before waiting for result
888 fetch_future_block_number = self._fetch_future_block_number
889 fetch_future = self._fetch_future
891 # Cleanup the fetch variables. Have a local copy.
892 self._fetch_future_block_number = None
893 self._fetch_future = None
895 # Need to wait for the future for the current read
896 if fetch_future is not None:
897 logger.info("BlockCache waiting for background fetch.")
898 # Wait until result and put it in cache
899 self._fetch_block_cached.add_key(
900 fetch_future.result(), fetch_future_block_number
901 )
903 # these are cached, so safe to do multiple calls for the same start and end.
904 for block_number in range(start_block_number, end_block_number + 1):
905 self._fetch_block_cached(block_number)
907 # fetch next block in the background if nothing is running in the background,
908 # the block is within file and it is not already cached
909 end_block_plus_1 = end_block_number + 1
910 with self._fetch_future_lock:
911 if (
912 self._fetch_future is None
913 and end_block_plus_1 <= self.nblocks
914 and not self._fetch_block_cached.is_key_cached(end_block_plus_1)
915 ):
916 self._fetch_future_block_number = end_block_plus_1
917 self._fetch_future = self._thread_executor.submit(
918 self._fetch_block, end_block_plus_1, "async"
919 )
921 return self._read_cache(
922 start,
923 end,
924 start_block_number=start_block_number,
925 end_block_number=end_block_number,
926 )
928 def _fetch_block(self, block_number: int, log_info: str = "sync") -> bytes:
929 """
930 Fetch the block of data for `block_number`.
931 """
932 if block_number > self.nblocks:
933 raise ValueError(
934 f"'block_number={block_number}' is greater than "
935 f"the number of blocks ({self.nblocks})"
936 )
938 start = block_number * self.blocksize
939 end = start + self.blocksize
940 logger.info("BlockCache fetching block (%s) %d", log_info, block_number)
941 self.total_requested_bytes += end - start
942 self.miss_count += 1
943 block_contents = super()._fetch(start, end)
944 return block_contents
946 def _read_cache(
947 self, start: int, end: int, start_block_number: int, end_block_number: int
948 ) -> bytes:
949 """
950 Read from our block cache.
952 Parameters
953 ----------
954 start, end : int
955 The start and end byte positions.
956 start_block_number, end_block_number : int
957 The start and end block numbers.
958 """
959 start_pos = start % self.blocksize
960 end_pos = end % self.blocksize
962 # kind of pointless to count this as a hit, but it is
963 self.hit_count += 1
965 if start_block_number == end_block_number:
966 block = self._fetch_block_cached(start_block_number)
967 return block[start_pos:end_pos]
969 else:
970 # read from the initial
971 out = [self._fetch_block_cached(start_block_number)[start_pos:]]
973 # intermediate blocks
974 # Note: it'd be nice to combine these into one big request. However
975 # that doesn't play nicely with our LRU cache.
976 out.extend(
977 map(
978 self._fetch_block_cached,
979 range(start_block_number + 1, end_block_number),
980 )
981 )
983 # final block
984 out.append(self._fetch_block_cached(end_block_number)[:end_pos])
986 return b"".join(out)
989caches: dict[str | None, type[BaseCache]] = {
990 # one custom case
991 None: BaseCache,
992}
995def register_cache(cls: type[BaseCache], clobber: bool = False) -> None:
996 """'Register' cache implementation.
998 Parameters
999 ----------
1000 clobber: bool, optional
1001 If set to True (default is False) - allow to overwrite existing
1002 entry.
1004 Raises
1005 ------
1006 ValueError
1007 """
1008 name = cls.name
1009 if not clobber and name in caches:
1010 raise ValueError(f"Cache with name {name!r} is already known: {caches[name]}")
1011 caches[name] = cls
1014for c in (
1015 BaseCache,
1016 MMapCache,
1017 BytesCache,
1018 ReadAheadCache,
1019 BlockCache,
1020 FirstChunkCache,
1021 AllBytes,
1022 KnownPartsOfAFile,
1023 BackgroundBlockCache,
1024):
1025 register_cache(c)