Merge pull request #9821 from mr-raj12/pack-files-max-size

repository: add max_size byte limit to PackWriter, OR-ed with max_count

src/borg/repository.py

@@ -104,55 +104,51 @@ def build_rest_backend(location):
 
 
 class PackWriter:
-    """Buffers chunks into a pack file and writes to the store when full.
+    """Buffers chunks into a pack file and writes it to the store when full.
 
-    Collects (chunk_id, cdata) pairs in a list and flushes once max_count is
-    reached.  PackWriter maintains the ChunkIndex directly: each add() marks the
-    chunk as pending (pack_id=UNKNOWN_BYTES32); flush() then assigns the real
-    pack_id, offset and size to every pending entry once the pack is on disk.
+    add() buffers a (chunk_id, cdata) pair and marks the chunk pending in the
+    ChunkIndex (pack_id=UNKNOWN_BYTES32); flush() writes the pack and sets each
+    entry's real pack_id and obj_offset.
 
-    The index is not owned here.  Construction requires either a repository or an
-    explicit chunks index; there is no silent default.  With a repository, the writer
-    uses that repository's single, authoritative index (see the chunks property), so
-    there is never a second copy to keep in sync.  Unit tests pass an explicit index.
+    The ChunkIndex comes from the repository, or from an explicit chunks index when
+    there is no repository (see the chunks property).
 
-    max_count bounds how many chunks a pack accumulates before flush() writes it.
-    Raising it produces larger packs without changing this class's interface.
+    max_count bounds how many chunks a pack holds; max_size bounds its byte size.
+    flush() fires when either limit is reached.  Set a limit to None to disable it;
+    at least one must be set, otherwise the pack buffer is unbounded.
     """
 
-    def __init__(self, store, *, max_count=1, chunks=None, repository=None):
+    def __init__(self, store, *, max_count=3, max_size=None, chunks=None, repository=None):
         if repository is None and chunks is None:
             raise ValueError("PackWriter requires either a repository or an explicit chunks index")
+        if max_count is None and max_size is None:
+            raise ValueError("PackWriter needs max_count or max_size, otherwise the pack buffer is unbounded")
         self.store = store
-        self.max_count = max_count
-        self.repository = repository  # when set, the one and only index lives there
-        self._chunks = chunks  # explicit index for repository-less use (tests)
+        self.max_count = max_count  # None = no count limit
+        self.max_size = max_size  # None = no size limit
+        self.repository = repository
+        self._chunks = chunks  # used when there is no repository
         self._pieces = []  # list of (chunk_id, cdata)
+        self._size = 0  # byte size of buffered pieces
 
     @property
     def chunks(self):
-        """The ChunkIndex this writer updates.
-
-        With a repository, this is the repository's single index (shared, not copied).
-        Without one, it is the explicit index passed at construction.
-        """
+        """The ChunkIndex this writer updates: the repository's index, or the
+        explicit index passed at construction when there is no repository."""
         if self.repository is not None:
             return self.repository.chunks
         return self._chunks
 
     def add(self, chunk_id, cdata):
         """Buffer a chunk.  Returns flush results if the pack is now full, else None."""
-        # Mark the chunk as pending (pack_id=UNKNOWN_BYTES32).  flush() assigns the real
-        # pack_id and offset for every piece, so the placeholder offset 0 here is never read:
-        # get() refuses a pending entry (PackLocationUnknown) before any offset would matter.
-        # Precondition: callers add only chunks not already stored (the cache dedups via
-        # seen_chunk() first), so add(chunk_id, 0) never resets a real size on an existing entry.
-        # This is also what keeps ChunkIndex.add's "v.size == 0 or v.size == size" assertion happy:
-        # a fresh id has no entry, so the size=0 we pass here is never compared against a real size.
-        self.chunks.add(chunk_id, 0)  # size filled in by cache layer
+        self.chunks.add(chunk_id, 0)  # size: plaintext chunk size, set by the cache layer
+        # obj_size is final; pack_id and obj_offset get their real values in flush().
         self.chunks.update_pack_info([(chunk_id, UNKNOWN_BYTES32, 0, len(cdata))])
         self._pieces.append((chunk_id, cdata))
-        if len(self._pieces) >= self.max_count:
+        self._size += len(cdata)
+        if (self.max_count is not None and len(self._pieces) >= self.max_count) or (
+            self.max_size is not None and self._size >= self.max_size
+        ):
             return self.flush()
         return None
 
@@ -161,7 +157,7 @@ class PackWriter:
 
         Returns a list of (chunk_id, pack_id, obj_offset, obj_size) tuples --
         one entry per chunk that was written.  Returns None if there was nothing
-        to flush.  Always updates the ChunkIndex with the real pack_id.
+        to flush.  Always updates the ChunkIndex with the real pack_id and obj_offset.
         """
         if not self._pieces:
             return None
@@ -188,19 +184,16 @@ class PackWriter:
         try:
             self.store.store(key, pack_data)
         except Exception:
-            # The pack was not durably stored, so every entry add() pre-marked for it now
-            # points at data that does not exist.  Leaving them would make seen_chunk() report
-            # these ids as present, letting a later identical chunk dedup against bytes that were
-            # never written -- silent data loss.  These entries were created this session and never
-            # received a real pack_id, so dropping them restores the index to its pre-add() state
-            # (matching master, where the index only ever reflected successfully stored chunks).
+            # The pack is not stored, so drop the pending entries: keeping them would let
+            # seen_chunk() dedup later chunks against bytes that were never written.
             for chunk_id in pending_ids:
                 entry = self.chunks.get(chunk_id)
                 if entry is not None and entry.pack_id == UNKNOWN_BYTES32:
                     del self.chunks[chunk_id]
             raise
         finally:
-            self._pieces = []  # reset even on failure to prevent re-bundling a failed chunk
+            self._pieces = []  # cleared on success and on failure
+            self._size = 0
         self.chunks.update_pack_info(results)  # replace UNKNOWN_BYTES32 with real pack_id
         return results
 
@@ -507,7 +500,7 @@ class Repository:
         if lock:
             self.lock = Lock(self.store, exclusive, timeout=lock_wait).acquire()
         self._chunks = None
-        self._pack_writer = PackWriter(self.store, max_count=2, repository=self)
+        self._pack_writer = PackWriter(self.store, repository=self)
         self.opened = True
 
     @property

src/borg/testsuite/repository_test.py

@@ -290,6 +290,7 @@ def test_list(repo_fixtures, request):
     with get_repository_from_fixture(repo_fixtures, request) as repository:
         for x in range(100):
             repository.put(H(x), fchunk(b"SOMEDATA", chunk_id=H(x)))  # unique bytes -> unique pack id
+        repository.flush()  # flush the last partial pack so all 100 objects are listable
         repo_list = repository.list()
         assert len(repo_list) == 100
         first_half = repository.list(limit=50)
@@ -387,6 +388,35 @@ def test_pack_writer_n2_flush():
     assert results[1] == (id2, expected_pack_id, len(data1), len(data2))
 
 
+def test_pack_writer_flushes_on_max_size():
+    # max_count is high, so the flush is driven by max_size alone.
+    store = MockStore()
+    pw = PackWriter(store, max_count=100, max_size=10, chunks=ChunkIndex())
+    assert pw.add(b"a" * 32, b"12345") is None
+    results = pw.add(b"b" * 32, b"67890")
+    assert results is not None
+    assert len(results) == 2
+
+
+def test_pack_writer_max_size_none_is_count_only():
+    store = MockStore()
+    pw = PackWriter(store, max_count=2, max_size=None, chunks=ChunkIndex())
+    assert pw.add(b"a" * 32, b"x" * 10_000) is None
+    assert pw.add(b"b" * 32, b"y" * 10_000) is not None
+
+
+def test_pack_writer_max_count_none_is_size_only():
+    store = MockStore()
+    pw = PackWriter(store, max_count=None, max_size=10, chunks=ChunkIndex())
+    assert pw.add(b"a" * 32, b"12345") is None
+    assert pw.add(b"b" * 32, b"67890") is not None
+
+
+def test_pack_writer_requires_a_limit():
+    with pytest.raises(ValueError):
+        PackWriter(MockStore(), max_count=None, max_size=None, chunks=ChunkIndex())
+
+
 def test_pack_writer_rolls_back_index_on_failed_store():
     # If store.store() fails, flush() must drop the entries add() pre-marked, otherwise the index
     # keeps a phantom (indexed but never stored) chunk that seen_chunk() reports as present and a