Merge pull request #9836 from mr-raj12/pack-files-get-many-grouping

repository: cache whole packs in get_many, replacing adjacent-read grouping

src/borg/repository.py

@@ -18,6 +18,7 @@ from .helpers import Location
 from .helpers import bin_to_hex, hex_to_bin
 from .helpers import get_cache_dir
 from .helpers import ProgressIndicatorPercent
+from .helpers.lrucache import LRUCache
 from .storelocking import Lock
 from .logger import create_logger
 from .manifest import NoManifestError
@@ -206,7 +207,7 @@ class PackReader:
         self.key = "packs/" + bin_to_hex(pack_id) if pack_id is not None else None
         self.pack_contents = pack_contents
 
-    def _read(self, offset, size):
+    def read(self, offset, size):
         # read from the in-memory pack if we have it, else range-read from the store
         if self.pack_contents is not None:
             return self.pack_contents[offset : offset + size]
@@ -221,7 +222,7 @@ class PackReader:
         hdr_size = RepoObj.obj_header.size
         offset = 0
         while True:
-            hdr_data = self._read(offset, hdr_size)
+            hdr_data = self.read(offset, hdr_size)
             if len(hdr_data) < hdr_size:
                 break  # clean EOF, or trailing partial bytes
             hdr = RepoObj.ObjHeader(*RepoObj.obj_header.unpack(hdr_data))
@@ -303,6 +304,10 @@ class Repository:
 
         exit_mcode = 21
 
+    # Whole packs kept in memory for reads; the least recently used is evicted first.
+    # Memory use is this count times the pack size.
+    PACK_READER_CACHE_SIZE = 3
+
     def __init__(
         self,
         path_or_location,
@@ -387,6 +392,8 @@ class Repository:
         self.exclusive = exclusive
         self._pack_writer = None
         self._chunks = None  # ChunkIndex; loaded lazily on first access to .chunks
+        # pack_id -> PackReader holding the whole pack; get_many loads into it, get() reuses it
+        self._pack_cache = LRUCache(capacity=self.PACK_READER_CACHE_SIZE)
 
     def __repr__(self):
         return f"<{self.__class__.__name__} {self._location}>"
@@ -594,6 +601,7 @@ class Repository:
             self.store.close()
             self.store_opened = False
         self.opened = False
+        self._pack_cache.clear()
 
     def info(self):
         """return some infos about the repo (must be opened first)"""
@@ -760,10 +768,14 @@ class Repository:
             raise self.PackLocationUnknown(id, str(self._location))
         pack_id, obj_offset, obj_size = entry.pack_id, entry.obj_offset, entry.obj_size
         id_hex = bin_to_hex(id)
-        key = "packs/" + bin_to_hex(pack_id)
+        # slice from the cached whole pack if get_many (or an earlier get) already loaded it;
+        # otherwise read ranges from the store without loading and caching the whole pack.
+        reader = self._pack_cache.get(pack_id)
+        if reader is None:
+            reader = PackReader(store=self.store, pack_id=pack_id)
         try:
             if read_data:
-                return self.store.load(key, offset=obj_offset, size=obj_size)
+                return reader.read(obj_offset, obj_size)
             else:
                 # RepoObj layout supports separately encrypted metadata and data.
                 # We return enough bytes so the client can decrypt the metadata.
@@ -775,7 +787,7 @@ class Repository:
                 # uses the header's length and ignores trailing bytes -- this is just tidy.) obj_size
                 # comes from the same index we already route with.
                 load_size = min(load_size, obj_size)
-                obj = self.store.load(key, offset=obj_offset, size=load_size)
+                obj = reader.read(obj_offset, load_size)
                 hdr = obj[0:hdr_size]
                 if len(hdr) != hdr_size:
                     raise IntegrityError(f"Object too small [id {id_hex}]: expected {hdr_size}, got {len(hdr)} bytes")
@@ -786,7 +798,7 @@ class Repository:
                     retry_size = hdr_size + meta_size
                     # same boundary as above: normally a no-op, just keeps the retry within this object.
                     retry_size = min(retry_size, obj_size)
-                    obj = self.store.load(key, offset=obj_offset, size=retry_size)
+                    obj = reader.read(obj_offset, retry_size)
                 meta = obj[hdr_size : hdr_size + meta_size]
                 if len(meta) != meta_size:
                     raise IntegrityError(f"Object too small [id {id_hex}]: expected {meta_size}, got {len(meta)} bytes")
@@ -797,9 +809,37 @@ class Repository:
             else:
                 return None
 
+    def _cached_pack_reader(self, pack_id):
+        """Return a PackReader holding the whole pack, loading it into the cache on a miss."""
+        reader = self._pack_cache.get(pack_id)
+        if reader is None:
+            key = "packs/" + bin_to_hex(pack_id)
+            reader = PackReader(pack_id=pack_id, pack_contents=self.store.load(key))
+            self._pack_cache[pack_id] = reader
+        return reader
+
     def get_many(self, ids, read_data=True, raise_missing=True):
+        if not read_data:
+            # read_data=False returns only each object's header+meta, sized per object by get().
+            for id_ in ids:
+                yield self.get(id_, read_data=read_data, raise_missing=raise_missing)
+            return
+
         for id_ in ids:
-            yield self.get(id_, read_data=read_data, raise_missing=raise_missing)
+            self._lock_refresh()
+            entry = self.chunks.get(id_)
+            if entry is None or self.chunks.is_pending(id_):
+                # id unknown or still buffered: get() raises or returns None accordingly
+                yield self.get(id_, read_data=True, raise_missing=raise_missing)
+                continue
+            try:
+                reader = self._cached_pack_reader(entry.pack_id)
+            except StoreObjectNotFound:
+                if raise_missing:
+                    raise self.ObjectNotFound(id_, str(self._location)) from None
+                yield None
+            else:
+                yield reader.read(entry.obj_offset, entry.obj_size)
 
     def put(self, id, data):
         """put a repo object

src/borg/testsuite/repository_test.py

@@ -204,6 +204,116 @@ def test_multi_object_pack_roundtrip(repo_fixtures, request):
         assert repository.get(H(1), read_data=False) == chunk1[: hdr_size + len(meta1)]
 
 
+def test_get_many_one_load_per_pack(repo_fixtures, request):
+    # get_many loads each pack once as a whole and slices every object out of the cached bytes.
+    # stats["load_calls"] counts store.load() calls.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(6)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 3  # three objects per pack -> two packs for the six objects
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+        ids = [chunk_id for chunk_id, _ in objects]
+        assert len({repository.chunks[chunk_id].pack_id for chunk_id in ids}) == 2  # six ids, two packs
+        one_by_one = [repository.get(chunk_id) for chunk_id in ids]  # reference bytes, one store.load each
+
+        loads_before = repository.store.stats["load_calls"]
+        assert list(repository.get_many(ids)) == one_by_one
+        assert repository.store.stats["load_calls"] - loads_before == 2  # one store.load per pack
+
+
+def test_get_many_keeps_request_order(repo_fixtures, request):
+    # Ids requested out of their stored order come back in the requested order with their own bytes.
+    # Interleaving two packs still loads each pack once: the pack cache serves the second visit.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(6)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 3  # two packs: {H0,H1,H2} and {H3,H4,H5}
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+        ids = [H(0), H(3), H(1), H(4), H(2), H(5)]  # interleave the two packs
+        one_by_one = [repository.get(chunk_id) for chunk_id in ids]
+
+        loads_before = repository.store.stats["load_calls"]
+        assert list(repository.get_many(ids)) == one_by_one  # same order, same bytes
+        assert repository.store.stats["load_calls"] - loads_before == 2  # each pack loaded once
+
+
+def test_get_many_missing_id_yields_none(repo_fixtures, request):
+    # With raise_missing=False, an id that was never stored yields None in its position; the ids
+    # before and after it read back unchanged.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(3)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 4  # above the object count, so the pack flushes on flush() only
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+        repository.flush()
+        ids = [H(0), H(99), H(2)]  # H(99) was never put
+        result = list(repository.get_many(ids, raise_missing=False))
+        assert result[0] == repository.get(H(0))
+        assert result[1] is None  # never stored -> None
+        assert result[2] == repository.get(H(2))
+
+
+def test_get_many_repeated_ids(repo_fixtures, request):
+    # A dedup'd item repeats chunk ids, e.g. [A, A, B, C, B]. Each repeat returns the right bytes
+    # and each pack is loaded once: the cached pack serves the later visits.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(4)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 2  # two packs: {H0,H1} and {H2,H3}
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+        assert len({repository.chunks[H(i)].pack_id for i in range(4)}) == 2
+        ids = [H(0), H(0), H(1), H(2), H(1)]  # A A B C B, A/B in one pack, C in the other
+        one_by_one = [repository.get(chunk_id) for chunk_id in ids]
+
+        loads_before = repository.store.stats["load_calls"]
+        assert list(repository.get_many(ids)) == one_by_one  # every position, including repeats
+        assert repository.store.stats["load_calls"] - loads_before == 2  # two packs, one load each
+
+
+def test_get_many_evicts_least_recently_used(repo_fixtures, request):
+    # Visiting more packs than the cache holds evicts the oldest; revisiting an evicted pack reloads it.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(8)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 2  # four packs: {H0,H1} {H2,H3} {H4,H5} {H6,H7}
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+        assert len({repository.chunks[H(i)].pack_id for i in range(8)}) == 4
+        assert repository.PACK_READER_CACHE_SIZE == 3  # cache holds three of the four packs
+        # touch packs 0,1,2 (fills cache), then 3 (evicts pack 0), then pack 0 again (reload).
+        ids = [H(0), H(2), H(4), H(6), H(0)]
+        one_by_one = [repository.get(chunk_id) for chunk_id in ids]
+
+        loads_before = repository.store.stats["load_calls"]
+        assert list(repository.get_many(ids)) == one_by_one
+        assert repository.store.stats["load_calls"] - loads_before == 5  # 4 distinct packs + 1 reload
+
+
+def test_get_reuses_cached_pack(repo_fixtures, request):
+    # get() slices from a pack that get_many already cached, doing no store load; a get() whose pack
+    # is not cached reads only its object's range and does not load the whole pack into the cache.
+    objects = [(H(i), fchunk(b"payload-%02d" % i, chunk_id=H(i))) for i in range(2)]
+    with get_repository_from_fixture(repo_fixtures, request) as repository:
+        repository._pack_writer.max_count = 2  # one pack: {H0, H1}
+        for chunk_id, chunk in objects:
+            repository.put(chunk_id, chunk)
+
+        # cold cache: get() reads one range and leaves the cache empty, so the next get() reads again
+        loads_before = repository.store.stats["load_calls"]
+        reference = repository.get(H(0))
+        assert repository.store.stats["load_calls"] - loads_before == 1  # one ranged read
+        loads_before = repository.store.stats["load_calls"]
+        repository.get(H(1))
+        assert repository.store.stats["load_calls"] - loads_before == 1  # pack was not cached, another range
+
+        list(repository.get_many([H(0), H(1)]))  # loads the whole pack into the cache
+
+        loads_before = repository.store.stats["load_calls"]
+        assert repository.get(H(0)) == reference
+        assert repository.store.stats["load_calls"] - loads_before == 0  # sliced from the cached pack
+        assert repository.get(H(1), read_data=False)  # read_data=False also peeks the cache
+        assert repository.store.stats["load_calls"] - loads_before == 0
+
+
 def build_one_pack(repository, objects):
     with repository:
         repository._pack_writer.max_count = len(objects) + 1  # prevent per-put flush; one pack on flush()