fix: free memory when block is fully read

Author: Rogdham

CHANGELOG.md

@@ -11,6 +11,7 @@ adheres to [Semantic Versioning](https://semver.org/).
 
 ### :bug: Fixes
 
+- Free memory after a block is fully read
 - Typing: use `BinaryIO` instead of `IO[bytes]`
 
 ## [0.3.1] - 2021-12-26

src/xz/block.py

@@ -148,10 +148,15 @@ def _read(self, size: int) -> bytes:
 
         # read data
         try:
-            return self.operation.decompress(self._pos, size)
+            data = self.operation.decompress(self._pos, size)
         except LZMAError as ex:
             raise XZError(f"block: error while decompressing: {ex}") from ex
 
+        if self._pos + len(data) == self._length:
+            self.operation = None  # free memory
+
+        return data
+
     def writable(self) -> bool:
         return isinstance(self.operation, BlockWrite) or not self._length
 
@@ -174,7 +179,7 @@ def _write_after(self) -> None:
             self.unpadded_size, uncompressed_size = self.operation.finish()
             if uncompressed_size != self.uncompressed_size:
                 raise XZError("block: compressor uncompressed size")
-            self.operation = None
+            self.operation = None  # free memory
 
     def _truncate(self, size: int) -> None:
         # thanks to the writable method, we are sure that length is zero

tests/integration/test_ram_usage.py

@@ -0,0 +1,106 @@
+from io import DEFAULT_BUFFER_SIZE
+from lzma import compress
+from pathlib import Path
+from random import seed
+import sys
+from typing import BinaryIO, Callable, Iterator, Optional, cast
+
+import pytest
+
+from xz import XZFile
+from xz.common import create_xz_index_footer, parse_xz_footer, parse_xz_index
+from xz.io import IOCombiner, IOStatic
+
+if sys.version_info >= (3, 9):
+    from random import randbytes
+else:
+    from random import getrandbits
+
+    def randbytes(length: int) -> bytes:
+        return getrandbits(length * 8).to_bytes(length, "little")
+
+
+@pytest.fixture
+def ram_usage() -> Iterator[Callable[[], int]]:
+    try:
+        import tracemalloc  # pylint: disable=import-outside-toplevel
+    except ImportError:  # e.g. PyPy
+        pytest.skip("tracemalloc module not available")
+
+    try:
+        tracemalloc.start()
+        yield lambda: tracemalloc.get_traced_memory()[1]
+    finally:
+        tracemalloc.stop()
+
+
+BLOCK_SIZE = 1_000_000
+
+
+@pytest.fixture
+def fileobj() -> BinaryIO:
+    # create xz raw data composed of many identical blocks
+    nb_blocks = 50
+
+    seed(0)
+    data = compress(randbytes(BLOCK_SIZE))
+    header = data[:12]
+    footer = data[-12:]
+    check, backward_size = parse_xz_footer(footer)
+    block = data[12 : -12 - backward_size]
+    records = parse_xz_index(data[-12 - backward_size : -12])
+    index_footer = create_xz_index_footer(check, records * nb_blocks)
+
+    return cast(
+        BinaryIO,
+        IOCombiner(
+            IOStatic(header),
+            *[IOStatic(block)] * nb_blocks,
+            IOStatic(index_footer),
+        ),
+    )
+
+
+def test_read_linear(
+    # pylint: disable=redefined-outer-name
+    fileobj: BinaryIO,
+    ram_usage: Callable[[], int],
+) -> None:
+    with XZFile(fileobj) as xz_file:
+        # read almost one block
+        xz_file.read(BLOCK_SIZE - 1)
+        one_block_memory = ram_usage()
+
+        # read all the file
+        while xz_file.read(DEFAULT_BUFFER_SIZE):
+            assert (
+                # should not use much more memory, take 2 as error margin
+                ram_usage()
+                < one_block_memory * 2
+            ), f"Consumes too much RAM (at {100 * xz_file.tell() / len(xz_file):.0f}%)"
+
+
+def test_write(
+    tmp_path: Path,
+    # pylint: disable=redefined-outer-name
+    ram_usage: Callable[[], int],
+) -> None:
+    nb_blocks = 10
+
+    seed(0)
+
+    one_block_memory: Optional[int] = None
+
+    with XZFile(tmp_path / "archive.xz", "w") as xz_file:
+        for i in range(nb_blocks):
+            xz_file.change_block()
+            xz_file.write(randbytes(BLOCK_SIZE))
+
+            if one_block_memory is None:
+                one_block_memory = ram_usage()
+            else:
+                assert (
+                    # should not use much more memory, take 2 as error margin
+                    ram_usage()
+                    < one_block_memory * 2
+                ), f"Consumes too much RAM (at {i / nb_blocks:.0f}%)"

tests/unit/test_block.py

@@ -80,6 +80,10 @@ def test_read_all(
         call.seek(90, SEEK_SET),
         call.read(17),
     ]
+    fileobj.method_calls.clear()
+
+    assert block.read() == b""
+    assert not fileobj.method_calls
 
 
 def test_read_seek_forward(