write_chunkindex_to_repo: reduce memory needs

pyproject.toml

@@ -30,7 +30,7 @@ classifiers = [
 license = "BSD-3-Clause"
 license-files = ["LICENSE", "AUTHORS"]
 dependencies = [
-  "borghash ~= 0.1.0",
+  "borghash ~= 0.2.0",
   "borgstore[rest,blake3] ~= 0.6.0",
   "msgpack >=1.0.3, <=1.2.1",
   "packaging",

src/borg/cache.py

@@ -662,16 +662,45 @@ def write_chunkindex_to_repo(
     new_hashes = set()  # content hashes of the fragments that make up the index we are writing now
     fragments_written = 0
 
-    # sort the selected keys, so that an identical set of entries always produces identical
-    # fragments (identical content hashes), no matter in which order the entries were inserted
-    # into the hash table. this makes writing/repacking idempotent and convergent across clients:
-    # a fragment that already exists in the repo is not stored again (see _store_chunkindex_fragment)
-    # and no differently-partitioned duplicates of the same entries can pile up.
-    keys = sorted(key for key, _ in chunks.iteritems(only_new=incremental))
-    total = len(keys)
-    if keys:
-        # partition the selected entries into batches of at most max_entries entries:
-        batches = chunkit(keys, max_entries)
+    total = chunks.new_count if incremental else len(chunks)
+    if total > max_entries:
+        # to keep memory usage low, we don't build one huge, sorted list of all selected keys
+        # (~90 bytes per key!), see #9886: as the keys are uniformly distributed hash digests, we can
+        # partition them into 2 ** prefix_bits similarly sized, disjoint sets by their leading key bits
+        # and select / sort / write one partition's keys at a time. selecting a partition is a cheap,
+        # C-level filtering scan of the in-memory hash table, see borghash.HashTable.items.
+        # aim a bit below max_entries: without that headroom, a total of exactly
+        # 2 ** prefix_bits * max_entries would put the expected partition size right AT max_entries,
+        # so about half the partitions would overshoot it by a little.
+        target_entries = max_entries - max_entries // 20  # 5% headroom
+        partitions = -(-total // target_entries)  # == ceil(total / target_entries)
+        prefix_bits = (partitions - 1).bit_length()  # smallest prefix_bits with 2 ** prefix_bits >= partitions
+    else:
+        prefix_bits = 0  # all selected keys are one partition (and no filtering scans are needed)
+
+    def gen_batches():
+        # sort the selected keys per partition, so that an identical set of entries always produces
+        # identical fragments (identical content hashes), no matter in which order the entries were
+        # inserted into the hash table: partition membership and prefix_bits (chosen by entry count)
+        # only depend on the selected entries. this makes writing/repacking idempotent and convergent
+        # across clients: a fragment that already exists in the repo is not stored again (see
+        # _store_chunkindex_fragment) and no differently-partitioned duplicates of the same entries
+        # can pile up. as the prefix compares the keys' leading bits, ascending prefixes yield the
+        # same globally sorted key sequence a single all-keys sort would have produced.
+        for prefix in range(2**prefix_bits):
+            keys = sorted(
+                key for key, _ in chunks.iteritems(only_new=incremental, prefix_bits=prefix_bits, prefix=prefix)
+            )
+            # Usually a no-op splitwise: with hash digests as keys, prefix_bits was chosen so that a
+            # partition comes out well below max_entries (>100 sigma of margin), so this yields the
+            # partition as a single batch. It is what makes the "no fragment has more than
+            # max_entries entries" invariant a hard guarantee rather than an assumption about the
+            # key distribution - keys that are not uniformly distributed (e.g. in tests) can put
+            # everything into one partition.
+            yield from chunkit(keys, max_entries)
+
+    if total:
+        batches = gen_batches()
     elif force_write:
         # write a single empty fragment (e.g. at repo creation or after delete_chunkindex_from_repo()):
         batches = [[]]

src/borg/hashindex.pyi

@@ -32,11 +32,14 @@ class ChunkIndex:
     M_SYSTEM: int
     def add(self, key: bytes, size: int) -> None: ...
     def update_pack_info(self, pack_results: list | None) -> None: ...
-    def iteritems(self, *, only_new: bool = ...) -> Iterator: ...
+    def iteritems(self, *, only_new: bool = ..., prefix_bits: int = ..., prefix: int = ...) -> Iterator: ...
+    @property
+    def new_count(self) -> int: ...
     def clear_new(self) -> None: ...
     def __contains__(self, key: bytes) -> bool: ...
     def __getitem__(self, key: bytes) -> type[ChunkIndexEntry]: ...
     def __setitem__(self, key: bytes, value: CIE) -> None: ...
+    def __delitem__(self, key: bytes) -> None: ...
 
 class FuseVersionsIndexEntry(NamedTuple):
     version: int

src/borg/hashindex.pyx

@@ -59,19 +59,34 @@ class ChunkIndex(HTProxyMixin, MutableMapping):
     def __init__(self, capacity=1000, path=None, usable=None):
         if path:
             self.ht = HashTableNT.read(path)
+            self._new_count = None  # unknown, computed lazily (see new_count)
         else:
             if usable is not None:
                 capacity = usable * 2  # load factor 0.5
             self.ht = HashTableNT(key_size=32, value_type=ChunkIndexEntry, value_format=ChunkIndexEntryFormat,
                                   capacity=capacity)
+            self._new_count = 0
 
     def hide_system_flags(self, value):
         user_flags = value.flags & self.M_USER
         return value._replace(flags=user_flags)
 
-    def iteritems(self, *, only_new=False):
-        """Iterates items (optionally only new items); hides system flags."""
-        for key, value in self.ht.items():
+    @property
+    def new_count(self):
+        """Count of new entries (entries that have the F_NEW flag set)."""
+        if self._new_count is None:
+            # loaded from a file: compute once, then maintain incrementally.
+            self._new_count = sum(1 for key, value in self.ht.items() if value.flags & self.F_NEW)
+        return self._new_count
+
+    def iteritems(self, *, only_new=False, prefix_bits=0, prefix=0):
+        """
+        Iterates items (optionally only new items and/or only items of one key-prefix partition);
+        hides system flags.
+
+        See borghash.HashTable.items for the prefix_bits / prefix semantics.
+        """
+        for key, value in self.ht.items(prefix_bits=prefix_bits, prefix=prefix):
             if not only_new or (value.flags & self.F_NEW):
                 yield key, self.hide_system_flags(value)
 
@@ -101,16 +116,29 @@ class ChunkIndex(HTProxyMixin, MutableMapping):
         except KeyError:
             prev_flags = self.F_NONE
             is_new = True
+            inserting = True
         else:
             prev_flags = prev.flags
             is_new = bool(prev_flags & self.F_NEW)  # was new? stays new!
+            inserting = False
         system_flags = prev_flags & self.M_SYSTEM
         if is_new:
             system_flags |= self.F_NEW
         else:
             system_flags &= ~self.F_NEW
         user_flags = value.flags & self.M_USER
         self.ht[key] = value._replace(flags=system_flags | user_flags)
+        if inserting and self._new_count is not None:
+            self._new_count += 1  # inserted a new entry (overwriting keeps the F_NEW state)
+
+    def __delitem__(self, key):
+        value = self.ht.pop(key)  # raises KeyError if not present, like del would
+        if self._new_count is not None and (value.flags & self.F_NEW):
+            self._new_count -= 1
+
+    def clear(self):
+        self.ht.clear()
+        self._new_count = 0
 
     def update_pack_info(self, pack_results):
         """Set pack_id, obj_offset and obj_size from a list of (chunk_id, pack_id, obj_offset, obj_size)
@@ -133,6 +161,7 @@ class ChunkIndex(HTProxyMixin, MutableMapping):
             if value.flags & self.F_NEW:
                 flags = value.flags & ~self.F_NEW
                 self.ht[key] = value._replace(flags=flags)
+        self._new_count = 0
 
     @classmethod
     def read(cls, path):

src/borg/testsuite/cache_test.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 import time
 from datetime import UTC, datetime
@@ -379,6 +380,63 @@ def test_write_chunkindex_deterministic_fragments(tmp_path, monkeypatch):
     assert hashes[0] == hashes[1]  # identical fragment sets from differently-ordered inputs
 
 
+def _u_key(i):
+    """A pseudo-random (uniformly distributed) 32-byte chunk id for entry number i."""
+    return hashlib.sha256(i.to_bytes(4, "big")).digest()
+
+
+def test_write_chunkindex_partitioned_write(tmp_path, monkeypatch):
+    """A big write with uniformly distributed keys is partitioned by leading key bits.
+
+    When more than MAX entries are selected, the write path processes the keys one key-prefix
+    partition at a time instead of building one huge all-keys list (#9886). The resulting
+    fragments must still be bounded by MAX and together hold every entry exactly once; as the
+    partitions compare the keys' leading bits, the fragments are contiguous ranges of the
+    globally sorted key space.
+    """
+    monkeypatch.setattr(cache_mod, "CHUNKINDEX_FRAGMENT_ENTRIES_MAX", 100)
+    monkeypatch.setattr(cache_mod, "CHUNKINDEX_FRAGMENT_ENTRIES_MIN", 50)
+
+    keys = [_u_key(i) for i in range(1000)]
+    for incremental in (False, True):
+        repository_location = os.fspath(tmp_path / f"repository{incremental}")
+        with Repository(repository_location, exclusive=True, create=True) as repository:
+            delete_chunkindex_from_repo(repository)  # start from a known-empty fragment set
+            write_chunkindex_to_repo(
+                repository, _make_chunkindex(keys), incremental=incremental, force_write=not incremental
+            )
+            fragment_keys = [
+                sorted(read_chunkindex_from_repo(repository, name)) for name, _ in list_chunkindex_fragments(repository)
+            ]
+            # 1000 uniformly distributed entries with MAX=100 must have been really partitioned:
+            assert len(fragment_keys) >= 10
+            assert all(len(fk) <= 100 for fk in fragment_keys)
+            # ordering the fragments by their first key must yield the globally sorted key sequence,
+            # so the fragments are disjoint, complete and contiguous:
+            fragment_keys.sort(key=lambda fk: fk[0])
+            assert [k for fk in fragment_keys for k in fk] == sorted(keys)
+            # the index rebuilt from the fragments is complete:
+            chunks = build_chunkindex_from_repo(repository)
+            assert len(chunks) == 1000
+            assert set(chunks) == set(keys)
+
+
+def test_write_chunkindex_partitioned_deterministic(tmp_path, monkeypatch):
+    """Partitioned writes are convergent, too: identical entries -> identical fragments."""
+    monkeypatch.setattr(cache_mod, "CHUNKINDEX_FRAGMENT_ENTRIES_MAX", 100)
+
+    keys = [_u_key(i) for i in range(1000)]
+    shuffled = keys[311:] + keys[:311]  # same set, different insertion order
+    hashes = []
+    for key_list in (keys, shuffled):
+        repository_location = os.fspath(tmp_path / f"repository{len(hashes)}")
+        with Repository(repository_location, exclusive=True, create=True) as repository:
+            delete_chunkindex_from_repo(repository)
+            write_chunkindex_to_repo(repository, _make_chunkindex(key_list), incremental=False, force_write=True)
+            hashes.append({name for name, _ in list_chunkindex_fragments(repository)})
+    assert hashes[0] == hashes[1]  # identical fragment sets from differently-ordered inputs
+
+
 def test_close_consolidates_fragments_across_sessions(tmp_path, monkeypatch):
     """End-to-end: repeated create-like sessions leave a bounded, consolidated set of fragments."""
     monkeypatch.setenv("BORG_PASSPHRASE", "test")

src/borg/testsuite/hashindex_test.py

@@ -92,3 +92,73 @@ def new_chunks():
     assert new_chunks() == [(key2, value2a)]
     chunks.clear_new()
     assert new_chunks() == []
+
+
+def _cie(key):
+    return ChunkIndexEntry(flags=ChunkIndex.F_USED, size=1, pack_id=key, obj_offset=0, obj_size=0)
+
+
+def test_new_count():
+    chunks = ChunkIndex()
+    assert chunks.new_count == 0
+    keys = [H2(x) for x in range(10)]
+    for i, key in enumerate(keys):
+        chunks[key] = _cie(key)
+        assert chunks.new_count == i + 1
+    chunks[keys[0]] = _cie(keys[0])  # overwriting a new entry: it stays new, not counted twice
+    assert chunks.new_count == 10
+    del chunks[keys[9]]  # deleting a new entry decrements
+    assert chunks.new_count == 9
+    chunks.clear_new()
+    assert chunks.new_count == 0
+    chunks[keys[9]] = _cie(keys[9])  # re-inserting: new again
+    assert chunks.new_count == 1
+    chunks[keys[0]] = _cie(keys[0])  # overwriting a not-new entry: it stays not-new
+    assert chunks.new_count == 1
+    del chunks[keys[0]]  # deleting a not-new entry: no change
+    assert chunks.new_count == 1
+    with pytest.raises(struct.error):  # a failing insert must not bump the count
+        chunks[H2(11)] = ChunkIndexEntry(flags=ChunkIndex.F_NONE, size=2**33, pack_id=H2(11), obj_offset=0, obj_size=0)
+    assert chunks.new_count == 1
+    chunks.clear()
+    assert chunks.new_count == 0
+
+
+def test_new_count_after_read(tmp_path):
+    # .write() persists the raw entries including the F_NEW flag, so a freshly loaded
+    # ChunkIndex must compute new_count (lazily) from the loaded entries.
+    chunks = ChunkIndex()
+    keys = [H2(x) for x in range(5)]
+    for key in keys:
+        chunks[key] = _cie(key)
+    chunks.clear_new()
+    new_key = H2(1000)
+    chunks[new_key] = _cie(new_key)
+    path = str(tmp_path / "chunks")
+    chunks.write(path)
+    loaded = ChunkIndex.read(path)
+    assert loaded.new_count == 1
+    loaded.clear_new()
+    assert loaded.new_count == 0
+
+
+def test_iteritems_prefix():
+    prefix_bits = 3
+    chunks = ChunkIndex()
+    keys = [H2(x) for x in range(100)]
+    for key in keys:
+        chunks[key] = _cie(key)
+    collected = []
+    for prefix in range(2**prefix_bits):
+        part = [key for key, _ in chunks.iteritems(prefix_bits=prefix_bits, prefix=prefix)]
+        assert all(key[0] >> (8 - prefix_bits) == prefix for key in part)
+        collected += part
+    assert sorted(collected) == sorted(keys)  # complete and disjoint
+    # combining a prefix filter with only_new:
+    chunks.clear_new()
+    new_key = H2(1000)
+    chunks[new_key] = _cie(new_key)
+    new_prefix = new_key[0] >> (8 - prefix_bits)
+    for prefix in range(2**prefix_bits):
+        part = [key for key, _ in chunks.iteritems(only_new=True, prefix_bits=prefix_bits, prefix=prefix)]
+        assert part == ([new_key] if prefix == new_prefix else [])