feat: add strategies to free LZMA decompressors

Author: Rogdham

CHANGELOG.md

@@ -9,9 +9,19 @@ adheres to [Semantic Versioning](https://semver.org/).
 
 [unreleased]: https://github.com/rogdham/bigxml/compare/v0.3.1...HEAD
 
+### :rocket: Added
+
+- Advanced users may use the new `block_read_strategy` argument of `XZFile`/`xz.open` to
+  customize the strategy for freeing block readers, and implement a different tradeoff
+  between memory consumption and read speed when alternating reads between several
+  blocks; the following strategies are provided: `RollingBlockReadStrategy` and
+  `KeepBlockReadStrategy`
+
 ### :bug: Fixes
 
 - Free memory after a block is fully read
+- Free memory of LZMA decompressors when many blocks are partially read; this is a
+  tradeoff defaulting to keeping the last 8 LZMA decompressors used
 - Typing: use `BinaryIO` instead of `IO[bytes]`
 
 ## [0.3.1] - 2021-12-26

src/xz/__init__.py

@@ -37,6 +37,7 @@
 from xz.common import XZError
 from xz.file import XZFile
 from xz.open import xz_open
+from xz.strategy import KeepBlockReadStrategy, RollingBlockReadStrategy
 
 # pylint: disable=redefined-builtin
 open = xz_open
@@ -45,6 +46,8 @@
 
 __all__ = (
     "__version__",
+    "KeepBlockReadStrategy",
+    "RollingBlockReadStrategy",
     "XZError",
     "XZFile",
     "open",

src/xz/block.py

@@ -10,7 +10,13 @@
     parse_xz_index,
 )
 from xz.io import IOAbstract, IOCombiner, IOStatic
-from xz.typing import _LZMAFiltersType, _LZMAPresetType
+from xz.strategy import KeepBlockReadStrategy
+from xz.typing import (
+    Optional,
+    _BlockReadStrategyType,
+    _LZMAFiltersType,
+    _LZMAPresetType,
+)
 
 
 class BlockRead:
@@ -122,12 +128,14 @@ def __init__(
         uncompressed_size: int,
         preset: _LZMAPresetType = None,
         filters: _LZMAFiltersType = None,
+        block_read_strategy: Optional[_BlockReadStrategyType] = None,
     ):
         super().__init__(uncompressed_size)
         self.fileobj = fileobj
         self.check = check
         self.preset = preset
         self.filters = filters
+        self.block_read_strategy = block_read_strategy or KeepBlockReadStrategy()
         self.unpadded_size = unpadded_size
         self.operation: Union[BlockRead, BlockWrite, None] = None
 
@@ -139,6 +147,8 @@ def _read(self, size: int) -> bytes:
         # enforce read mode
         if not isinstance(self.operation, BlockRead):
             self._write_end()
+            self.clear()
+            self.block_read_strategy.on_create(self)
             self.operation = BlockRead(
                 self.fileobj,
                 self.check,
@@ -147,13 +157,14 @@ def _read(self, size: int) -> bytes:
             )
 
         # read data
+        self.block_read_strategy.on_read(self)
         try:
             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
+            self.clear()
 
         return data
 
@@ -163,6 +174,7 @@ def writable(self) -> bool:
     def _write(self, data: bytes) -> int:
         # enforce write mode
         if not isinstance(self.operation, BlockWrite):
+            self.clear()
             self.operation = BlockWrite(
                 self.fileobj,
                 self.check,
@@ -179,10 +191,15 @@ 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  # free memory
+            self.clear()
 
     def _truncate(self, size: int) -> None:
         # thanks to the writable method, we are sure that length is zero
         # so we don't need to handle the case of truncating in middle of the block
         self.seek(size)
         self.write(b"")
+
+    def clear(self) -> None:
+        if isinstance(self.operation, BlockRead):
+            self.block_read_strategy.on_delete(self)
+        self.operation = None  # free memory

src/xz/file.py

@@ -5,8 +5,14 @@
 
 from xz.common import DEFAULT_CHECK, XZError
 from xz.io import IOCombiner, IOProxy
+from xz.strategy import RollingBlockReadStrategy
 from xz.stream import XZStream
-from xz.typing import _LZMAFilenameType, _LZMAFiltersType, _LZMAPresetType
+from xz.typing import (
+    _BlockReadStrategyType,
+    _LZMAFilenameType,
+    _LZMAFiltersType,
+    _LZMAPresetType,
+)
 from xz.utils import parse_mode, proxy_property
 
 
@@ -28,7 +34,8 @@ def __init__(
         *,
         check: int = -1,
         preset: _LZMAPresetType = None,
-        filters: _LZMAFiltersType = None
+        filters: _LZMAFiltersType = None,
+        block_read_strategy: Optional[_BlockReadStrategyType] = None,
     ) -> None:
         """Open an XZ file in binary mode.
 
@@ -52,6 +59,11 @@ def __init__(
 
         For more information about the check/preset/filters arguments,
         refer to the documentation of the lzma module.
+
+        The block_read_strategy argument allows to specify a strategy
+        for freeing block readers, and implement a different tradeoff
+        between memory consumption and read speed when alternating reads
+        between several blocks.
         """
         self._close_fileobj = False
         self._close_check_empty = False
@@ -60,6 +72,14 @@ def __init__(
 
         self.mode, self._readable, self._writable = parse_mode(mode)
 
+        # create strategy
+        if block_read_strategy is None:
+            self.block_read_strategy: _BlockReadStrategyType = (
+                RollingBlockReadStrategy()
+            )
+        else:
+            self.block_read_strategy = block_read_strategy
+
         # get fileobj
         if isinstance(filename, (str, bytes, os.PathLike)):
             # pylint: disable=consider-using-with, unspecified-encoding
@@ -143,7 +163,7 @@ def _init_parse(self) -> None:
                 raise XZError("file: invalid size")
             self.fileobj.seek(-4, SEEK_CUR)
             if any(self.fileobj.read(4)):
-                streams.append(XZStream.parse(self.fileobj))
+                streams.append(XZStream.parse(self.fileobj, self.block_read_strategy))
             else:
                 self.fileobj.seek(-4, SEEK_CUR)  # stream padding
 
@@ -161,6 +181,7 @@ def _create_fileobj(self) -> XZStream:
             self.check,
             self.preset,
             self.filters,
+            self.block_read_strategy,
         )
 
     def change_stream(self) -> None:

src/xz/open.py

@@ -4,6 +4,7 @@
 
 from xz.file import XZFile
 from xz.typing import (
+    _BlockReadStrategyType,
     _LZMAFilenameType,
     _LZMAFiltersType,
     _LZMAPresetType,
@@ -22,6 +23,7 @@ def __init__(
         check: int = -1,
         preset: _LZMAPresetType = None,
         filters: _LZMAFiltersType = None,
+        block_read_strategy: Optional[_BlockReadStrategyType] = None,
         encoding: Optional[str] = None,
         errors: Optional[str] = None,
         newline: Optional[str] = None,
@@ -32,6 +34,7 @@ def __init__(
             check=check,
             preset=preset,
             filters=filters,
+            block_read_strategy=block_read_strategy,
         )
         super().__init__(
             cast(BinaryIO, self.xz_file),
@@ -45,6 +48,9 @@ def __init__(
     filters: _LZMAFiltersType = proxy_property("filters", "xz_file")
     stream_boundaries: List[int] = proxy_property("stream_boundaries", "xz_file")
     block_boundaries: List[int] = proxy_property("block_boundaries", "xz_file")
+    block_read_strategy: _BlockReadStrategyType = proxy_property(
+        "block_read_strategy", "xz_file"
+    )
 
     @wraps(XZFile.change_stream)
     def change_stream(self) -> None:
@@ -66,6 +72,7 @@ def xz_open(
     check: int = -1,
     preset: _LZMAPresetType = None,
     filters: _LZMAFiltersType = None,
+    block_read_strategy: Optional[_BlockReadStrategyType] = None,
     # text-mode kwargs
     encoding: Optional[str] = None,
     errors: Optional[str] = None,
@@ -83,6 +90,7 @@ def xz_open(
     check: int = -1,
     preset: _LZMAPresetType = None,
     filters: _LZMAFiltersType = None,
+    block_read_strategy: Optional[_BlockReadStrategyType] = None,
     # text-mode kwargs
     encoding: Optional[str] = None,
     errors: Optional[str] = None,
@@ -100,6 +108,7 @@ def xz_open(
     check: int = -1,
     preset: _LZMAPresetType = None,
     filters: _LZMAFiltersType = None,
+    block_read_strategy: Optional[_BlockReadStrategyType] = None,
     # text-mode kwargs
     encoding: Optional[str] = None,
     errors: Optional[str] = None,
@@ -116,6 +125,7 @@ def xz_open(
     check: int = -1,
     preset: _LZMAPresetType = None,
     filters: _LZMAFiltersType = None,
+    block_read_strategy: Optional[_BlockReadStrategyType] = None,
     # text-mode kwargs
     encoding: Optional[str] = None,
     errors: Optional[str] = None,
@@ -145,6 +155,7 @@ def xz_open(
             check=check,
             preset=preset,
             filters=filters,
+            block_read_strategy=block_read_strategy,
             encoding=encoding,
             errors=errors,
             newline=newline,
@@ -163,4 +174,5 @@ def xz_open(
         check=check,
         preset=preset,
         filters=filters,
+        block_read_strategy=block_read_strategy,
     )

src/xz/strategy.py

@@ -0,0 +1,41 @@
+import time
+from typing import TYPE_CHECKING, Dict
+
+if TYPE_CHECKING:  # pragma: no cover
+    # avoid circular dependency
+    from xz.block import XZBlock
+
+
+class KeepBlockReadStrategy:
+    def on_create(self, block: "XZBlock") -> None:
+        pass  # do nothing
+
+    def on_delete(self, block: "XZBlock") -> None:
+        pass  # do nothing
+
+    def on_read(self, block: "XZBlock") -> None:
+        pass  # do nothing
+
+
+class RollingBlockReadStrategy:
+    def __init__(self, max_block_read_nb: int = 8) -> None:
+        self.block_reads: Dict["XZBlock", float] = {}
+        self.max_block_read_nb = max_block_read_nb
+
+    def _freshly_used(self, block: "XZBlock") -> None:
+        self.block_reads[block] = time.monotonic()
+
+    def on_create(self, block: "XZBlock") -> None:
+        self._freshly_used(block)
+        if len(self.block_reads) > self.max_block_read_nb:
+            to_clear = min(
+                self.block_reads.items(),
+                key=lambda item: item[1],
+            )[0]
+            to_clear.clear()  # will call on_delete
+
+    def on_delete(self, block: "XZBlock") -> None:
+        del self.block_reads[block]
+
+    def on_read(self, block: "XZBlock") -> None:
+        self._freshly_used(block)

src/xz/stream.py

@@ -12,7 +12,12 @@
     round_up,
 )
 from xz.io import IOCombiner, IOProxy
-from xz.typing import _LZMAFiltersType, _LZMAPresetType
+from xz.typing import (
+    Optional,
+    _BlockReadStrategyType,
+    _LZMAFiltersType,
+    _LZMAPresetType,
+)
 
 
 class XZStream(IOCombiner[XZBlock]):
@@ -22,12 +27,14 @@ def __init__(
         check: int,
         preset: _LZMAPresetType = None,
         filters: _LZMAFiltersType = None,
+        block_read_strategy: Optional[_BlockReadStrategyType] = None,
     ) -> None:
         super().__init__()
         self.fileobj = fileobj
         self._check = check
         self.preset = preset
         self.filters = filters
+        self.block_read_strategy = block_read_strategy
 
     @property
     def check(self) -> int:
@@ -44,7 +51,11 @@ def _fileobj_blocks_end_pos(self) -> int:
         )
 
     @classmethod
-    def parse(cls, fileobj: BinaryIO) -> "XZStream":
+    def parse(
+        cls,
+        fileobj: BinaryIO,
+        block_read_strategy: Optional[_BlockReadStrategyType] = None,
+    ) -> "XZStream":
         """Parse one XZ stream from a fileobj.
 
         fileobj position should be right at the end of the stream when calling
@@ -70,6 +81,7 @@ def parse(cls, fileobj: BinaryIO) -> "XZStream":
                     check,
                     unpadded_size,
                     uncompressed_size,
+                    block_read_strategy=block_read_strategy,
                 )
             )
             block_start = block_end
@@ -84,7 +96,7 @@ def parse(cls, fileobj: BinaryIO) -> "XZStream":
         header_start_pos = fileobj.seek(-12, SEEK_CUR)
 
         stream_fileobj = IOProxy(fileobj, header_start_pos, footer_end_pos)
-        stream = cls(stream_fileobj, check)
+        stream = cls(stream_fileobj, check, block_read_strategy=block_read_strategy)
         for block in blocks:
             stream._append(block)
         return stream
@@ -102,6 +114,7 @@ def _create_fileobj(self) -> XZBlock:
             0,
             self.preset,
             self.filters,
+            self.block_read_strategy,
         )
 
     def _write_before(self) -> None: