Avoids memory issues with big files while extracting archives (#2192)

* no longer requires to create folder before extracting

* adding chunked extractor to avoid huge mermory usage

* this is till required

* no need to create destination

* moved archiving to servicelib

* fixing tests

* pylint dose not undestand

* replaced archiving and unarchivng
- fixes an error with with high memory consumption

* disabled compression

* adding extra test information to error

* adding extra tests

* remove extra line

* removing error from destination_folder

* properly fixed error

* only extract files

* handle directories differently, just create them

* fixing very badly written test

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

packages/service-library/src/servicelib/archiving_utils.py

@@ -0,0 +1,130 @@
+import asyncio
+import logging
+import zipfile
+from pathlib import Path
+from concurrent.futures import ProcessPoolExecutor
+from typing import Iterator
+
+
+log = logging.getLogger(__name__)
+
+
+def _full_file_path_from_dir_and_subdirs(dir_path: Path) -> Iterator[Path]:
+    for path in dir_path.rglob("*"):
+        if path.is_file():
+            yield path
+
+
+def _strip_directory_from_path(input_path: Path, to_strip: Path) -> Path:
+    to_strip = f"{str(to_strip)}/"
+    return Path(str(input_path).replace(to_strip, ""))
+
+
+def _read_in_chunks(file_object, chunk_size=1024 * 8):
+    """Lazy function (generator) to read a file piece by piece.
+    Default chunk size: 8k."""
+    while True:
+        data = file_object.read(chunk_size)
+        if not data:
+            break
+        yield data
+
+
+def _zipfile_single_file_extract_worker(
+    zip_file_path: Path, file_in_archive: str, destination_folder: Path, is_dir: bool
+) -> None:
+    """Extracing in chunks to avoid memory pressure on zip/unzip"""
+    with zipfile.ZipFile(zip_file_path) as zf:
+        # assemble destination and ensure it exits
+        destination_path = destination_folder / file_in_archive
+
+        if is_dir:
+            destination_path.mkdir(parents=True, exist_ok=True)
+            return
+
+        with zf.open(name=file_in_archive) as zip_fp:
+            with open(destination_path, "wb") as destination_fp:
+                for chunk in _read_in_chunks(zip_fp):
+                    destination_fp.write(chunk)
+
+
+def ensure_destination_subdirectories_exist(
+    zip_file_handler: zipfile.ZipFile, destination_folder: Path
+) -> None:
+    # assemble full destination paths
+    full_destination_paths = {
+        destination_folder / entry.filename for entry in zip_file_handler.infolist()
+    }
+    # extract all possible subdirectories
+    subdirectories = {x.parent for x in full_destination_paths}
+    # create all subdirectories before extracting
+    for subdirectory in subdirectories:
+        Path(subdirectory).mkdir(parents=True, exist_ok=True)
+
+
+async def unarchive_dir(archive_to_extract: Path, destination_folder: Path) -> None:
+    with zipfile.ZipFile(archive_to_extract, mode="r") as zip_file_handler:
+        with ProcessPoolExecutor() as pool:
+            loop = asyncio.get_event_loop()
+
+            # running in process poll is not ideal for concurrency issues
+            # to avoid race conditions all subdirectories where files will be extracted need to exist
+            # creating them before the extraction is under way avoids the issue
+            # the following avoids race conditions while unzippin in parallel
+            ensure_destination_subdirectories_exist(
+                zip_file_handler=zip_file_handler,
+                destination_folder=destination_folder,
+            )
+
+            tasks = [
+                loop.run_in_executor(
+                    pool,
+                    _zipfile_single_file_extract_worker,
+                    archive_to_extract,
+                    zip_entry.filename,
+                    destination_folder,
+                    zip_entry.is_dir(),
+                )
+                for zip_entry in zip_file_handler.infolist()
+            ]
+
+            await asyncio.gather(*tasks)
+
+
+def _serial_add_to_archive(
+    dir_to_compress: Path, destination: Path, compress: bool, store_relative_path: bool
+) -> None:
+    compression = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED
+    with zipfile.ZipFile(destination, "w", compression=compression) as zip_file_handler:
+        files_to_compress_generator = _full_file_path_from_dir_and_subdirs(
+            dir_to_compress
+        )
+        for file_to_add in files_to_compress_generator:
+            try:
+                file_name_in_archive = (
+                    _strip_directory_from_path(file_to_add, dir_to_compress)
+                    if store_relative_path
+                    else file_to_add
+                )
+                zip_file_handler.write(file_to_add, file_name_in_archive)
+            except ValueError:
+                log.exception("Could write files to archive, please check logs")
+                return False
+    return True
+
+
+async def archive_dir(
+    dir_to_compress: Path, destination: Path, compress: bool, store_relative_path: bool
+) -> bool:
+    with ProcessPoolExecutor(max_workers=1) as pool:
+        return await asyncio.get_event_loop().run_in_executor(
+            pool,
+            _serial_add_to_archive,
+            dir_to_compress,
+            destination,
+            compress,
+            store_relative_path,
+        )
+
+
+__all__ = ["archive_dir", "unarchive_dir"]

packages/service-library/tests/test_archiving_utils.py

@@ -0,0 +1,237 @@
+# pylint:disable=redefined-outer-name,unused-argument
+
+import os
+import tempfile
+import hashlib
+import random
+from pathlib import Path
+import asyncio
+from typing import Set, List, Dict, Iterator, Tuple
+from concurrent.futures import ProcessPoolExecutor
+import string
+import secrets
+
+
+import pytest
+
+from servicelib.archiving_utils import archive_dir, unarchive_dir
+
+
+@pytest.fixture
+def temp_dir_one() -> Path:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        yield Path(temp_dir)
+
+
+@pytest.fixture
+def temp_dir_two(tmpdir) -> Path:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        yield Path(temp_dir)
+
+
+@pytest.fixture
+def dir_with_random_content() -> Path:
+    def random_string(length: int) -> str:
+        return "".join(secrets.choice(string.ascii_letters) for i in range(length))
+
+    def make_files_in_dir(dir_path: Path, file_count: int) -> None:
+        for _ in range(file_count):
+            (dir_path / f"{random_string(8)}.bin").write_bytes(
+                os.urandom(random.randint(1, 10))
+            )
+
+    def ensure_dir(path_to_ensure: Path) -> Path:
+        path_to_ensure.mkdir(parents=True, exist_ok=True)
+        return path_to_ensure
+
+    def make_subdirectory_with_content(subdir_name: Path, max_file_count: int) -> None:
+        subdir_name = ensure_dir(subdir_name)
+        make_files_in_dir(
+            dir_path=subdir_name,
+            file_count=random.randint(1, max_file_count),
+        )
+
+    def make_subdirectories_with_content(
+        subdir_name: Path, max_subdirectories_count: int, max_file_count: int
+    ) -> None:
+        subdirectories_count = random.randint(1, max_subdirectories_count)
+        for _ in range(subdirectories_count):
+            make_subdirectory_with_content(
+                subdir_name=subdir_name / f"{random_string(4)}",
+                max_file_count=max_file_count,
+            )
+
+    def get_dirs_and_subdris_in_path(path_to_scan: Path) -> Iterator[Path]:
+        return [path for path in path_to_scan.rglob("*") if path.is_dir()]
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir_path = Path(temp_dir)
+        data_container = ensure_dir(temp_dir_path / "study_data")
+
+        make_subdirectories_with_content(
+            subdir_name=data_container, max_subdirectories_count=5, max_file_count=5
+        )
+        make_files_in_dir(dir_path=data_container, file_count=5)
+
+        # creates a good amount of files
+        for _ in range(4):
+            for subdirectory_path in get_dirs_and_subdris_in_path(data_container):
+                make_subdirectories_with_content(
+                    subdir_name=subdirectory_path,
+                    max_subdirectories_count=3,
+                    max_file_count=3,
+                )
+
+        yield temp_dir_path
+
+
+def strip_directory_from_path(input_path: Path, to_strip: Path) -> Path:
+    to_strip = f"{str(to_strip)}/"
+    return Path(str(input_path).replace(to_strip, ""))
+
+
+def get_all_files_in_dir(dir_path: Path) -> Set[Path]:
+    return {
+        strip_directory_from_path(x, dir_path)
+        for x in dir_path.rglob("*")
+        if x.is_file()
+    }
+
+
+def _compute_hash(file_path: Path) -> Tuple[Path, str]:
+    with open(file_path, "rb") as file_to_hash:
+        file_hash = hashlib.md5()
+        chunk = file_to_hash.read(8192)
+        while chunk:
+            file_hash.update(chunk)
+            chunk = file_to_hash.read(8192)
+
+    return file_path, file_hash.hexdigest()
+
+
+async def compute_hashes(file_paths: List[Path]) -> Dict[Path, str]:
+    """given a list of files computes hashes for the files on a process pool"""
+
+    loop = asyncio.get_event_loop()
+
+    with ProcessPoolExecutor() as prcess_pool_executor:
+        tasks = [
+            loop.run_in_executor(prcess_pool_executor, _compute_hash, file_path)
+            for file_path in file_paths
+        ]
+        # pylint: disable=unnecessary-comprehension
+        # see return value of _compute_hash it is a tuple, mapping list[Tuple[Path,str]] to Dict[Path, str] here
+        return {k: v for k, v in await asyncio.gather(*tasks)}
+
+
+def full_file_path_from_dir_and_subdirs(dir_path: Path) -> List[Path]:
+    return [x for x in dir_path.rglob("*") if x.is_file()]
+
+
+async def assert_same_directory_content(
+    dir_to_compress: Path, output_dir: Path, inject_relative_path: Path = None
+) -> None:
+    def _relative_path(input_path: Path) -> Path:
+        return Path(str(inject_relative_path / str(input_path))[1:])
+
+    input_set = get_all_files_in_dir(dir_to_compress)
+    output_set = get_all_files_in_dir(output_dir)
+
+    if inject_relative_path is not None:
+        input_set = {_relative_path(x) for x in input_set}
+
+    assert (
+        input_set == output_set
+    ), f"There following files are missing {input_set - output_set}"
+
+    # computing the hashes for dir_to_compress and map in a dict
+    # with the name starting from the root of the directory and md5sum
+    dir_to_compress_hashes = {
+        strip_directory_from_path(k, dir_to_compress): v
+        for k, v in (
+            await compute_hashes(full_file_path_from_dir_and_subdirs(dir_to_compress))
+        ).items()
+    }
+
+    # computing the hashes for output_dir and map in a dict
+    # with the name starting from the root of the directory and md5sum
+    output_dir_hashes = {
+        strip_directory_from_path(k, output_dir): v
+        for k, v in (
+            await compute_hashes(full_file_path_from_dir_and_subdirs(output_dir))
+        ).items()
+    }
+
+    # finally check if hashes are mapped 1 to 1 in order to verify
+    # that the compress/decompress worked correctly
+    for key in dir_to_compress_hashes:
+        assert (
+            dir_to_compress_hashes[key]
+            == output_dir_hashes[_relative_path(key) if inject_relative_path else key]
+        )
+
+
+# end utils
+
+
+@pytest.mark.parametrize(
+    "compress,store_relative_path",
+    [[True, True], [True, False], [False, True], [False, False]],
+)
+async def test_archive_unarchive_same_structure_dir(
+    dir_with_random_content: Path,
+    temp_dir_one: Path,
+    temp_dir_two: Path,
+    compress: bool,
+    store_relative_path: bool,
+):
+    archive_file = temp_dir_one / "archive.zip"
+
+    archive_result = await archive_dir(
+        dir_to_compress=dir_with_random_content,
+        destination=archive_file,
+        store_relative_path=store_relative_path,
+        compress=compress,
+    )
+    assert archive_result is True
+
+    await unarchive_dir(
+        archive_to_extract=archive_file, destination_folder=temp_dir_two
+    )
+
+    await assert_same_directory_content(
+        dir_with_random_content,
+        temp_dir_two,
+        None if store_relative_path else dir_with_random_content,
+    )
+
+
+@pytest.mark.parametrize(
+    "compress,store_relative_path",
+    [[True, True], [True, False], [False, True], [False, False]],
+)
+async def test_unarchive_in_same_dir_as_archive(
+    dir_with_random_content: Path,
+    temp_dir_one: Path,
+    compress: bool,
+    store_relative_path: bool,
+):
+    archive_file = temp_dir_one / "archive.zip"
+
+    archive_result = await archive_dir(
+        dir_to_compress=dir_with_random_content,
+        destination=archive_file,
+        store_relative_path=store_relative_path,
+        compress=compress,
+    )
+    assert archive_result is True
+
+    await unarchive_dir(
+        archive_to_extract=archive_file, destination_folder=temp_dir_one
+    )
+    archive_file.unlink()
+    await assert_same_directory_content(
+        dir_with_random_content,
+        temp_dir_one,
+        None if store_relative_path else dir_with_random_content,
+    )