[DCO01A-121] Make snapshot zip smaller by using symlinks to identical entities archives

Author: korbajan

mgmtworker/cloudify_system_workflows/snapshots/snapshot_create.py

@@ -3,6 +3,8 @@
 import pathlib
 import queue
 import shutil
+import hashlib
+import tarfile
 import tempfile
 import zipfile
 from collections import defaultdict
@@ -84,6 +86,11 @@ def __init__(
         self._auditlog_queue = queue.Queue()
         self._auditlog_listener = AuditLogListener(self._client,
                                                    self._auditlog_queue)
+        self._written_archives: dict[str, dict[tuple[str, ...], str]]
+        self._written_archives = {  # track created entities archives
+            'plugins': {},
+            'blueprints': {},  # will do for both blueprints and blueprint_revisions
+        }
 
     def create(self, timeout: float | None = None):
         """Dumps manager's data and some metadata into a single zip file"""
@@ -263,6 +270,7 @@ def _write_files(
                 if _should_append_entity(dump_type, entity):
                     self._auditlog_listener.append_entity(
                         tenant_name, dump_type, entity)
+            self._update_written_archives(entity_id, dump_type, output_dir)
         # Dump the data as JSON files
         filenum = _get_max_filenum_in_dir(output_dir) or 0
         for (source, source_id), items in data_buckets.items():
@@ -308,17 +316,24 @@ def _create_archive(self):
         ) as zf:
             base_dir = os.path.join(root_dir, os.curdir)
             base_dir = os.path.normpath(base_dir)
-            for dirpath, dirnames, filenames in os.walk(base_dir):
+            for dirpath, dirnames, filenames in os.walk(base_dir, followlinks=False):
+                root_path = Path(dirpath)
                 arcdirpath = os.path.relpath(dirpath, root_dir)
                 for name in sorted(dirnames):
                     path = os.path.join(dirpath, name)
                     arcname = os.path.join(arcdirpath, name)
                     zf.write(path, arcname)
                 for name in filenames:
-                    path = os.path.join(dirpath, name)
-                    path = os.path.normpath(path)
-                    if os.path.isfile(path):
-                        arcname = os.path.join(arcdirpath, name)
+                    path = root_path / name
+                    arcname = path.relative_to(root_dir)
+                    if path.is_symlink():
+                        zip_info = zipfile.ZipInfo(str(arcname))
+                        zip_info.create_system = 3  # Unix
+                        st = os.lstat(path)
+                        zip_info.external_attr = st.st_mode << 16
+                        link_target = os.readlink(path)
+                        zf.writestr(zip_info, link_target)
+                    elif os.path.isfile(path):
                         zf.write(path, arcname)
 
     def _upload_archive(self):
@@ -392,6 +407,29 @@ def _update_snapshot_status(self, status, error=None):
             error=error
         )
 
+    def _update_written_archives(self, entity_id, dump_type, output_dir):
+        dest_dir = (output_dir / f'{dump_type}').resolve()
+        suffix = {
+            'plugins': '.zip',
+            'blueprints': '.tar.gz',
+        }.get(dump_type)
+        if not suffix:
+            return
+        entity_archive = dest_dir / f'{entity_id}{suffix}'
+        if dump_type == 'plugins':
+            content_hashes = get_zip_content_hashes(entity_archive)
+        else:
+            content_hashes = get_tar_gz_content_hashes(entity_archive)
+        if existing_path := self._written_archives[dump_type].get(content_hashes):
+            entity_archive.unlink(missing_ok=False)
+            os.symlink(
+                os.path.relpath(existing_path, entity_archive).split("/", 1)[-1],
+                entity_archive,
+            )
+            ctx.logger.debug("Created symlink: %s to %s", entity_archive, existing_path)
+            return
+        self._written_archives[dump_type][content_hashes] = entity_archive
+
 
 def _prepare_temp_dir() -> Path:
     """Prepare temporary (working) directory structure"""
@@ -516,3 +554,54 @@ def get_all(method, kwargs=None):
         kwargs['_offset'] = len(data)
 
     return data
+
+
+def get_zip_content_hashes(zip_path) -> tuple[str, ...]:
+    hashes: set[str] = set()
+    all_dirs: set[str] = set()
+    not_empty_dirs: set[str] = set()
+    filenames: set[str] = set()
+    with zipfile.ZipFile(zip_path, 'r') as archive:
+        for info in archive.infolist():
+            if not info.is_dir():
+                filenames.add(info.filename)
+                parts = info.filename.split('/')
+                for i in range(1, len(parts)):
+                    not_empty_dirs.add('/'.join(parts[:i]))
+                with archive.open(info) as fileobj:
+                    content = fileobj.read()
+                    sha256 = hashlib.sha256(content).hexdigest()
+                    hashes.add(sha256)
+            else:
+                all_dirs.add(info.filename.rstrip("/"))
+    if filenames:
+        hashes.add(hashlib.sha256(":".join(filenames).encode("utf-8")).hexdigest())
+    if empty_dirs := all_dirs - not_empty_dirs:
+        hashes.add(hashlib.sha256(":".join(empty_dirs).encode("utf-8")).hexdigest())
+    return tuple(hashes)
+
+
+def get_tar_gz_content_hashes(tar_gz_path) -> tuple[str, ...]:
+    hashes: set[str] = set()
+    all_dirs: set[str] = set()
+    not_empty_dirs: set[str] = set()
+    filenames: set[str] = set()
+    with tarfile.open(tar_gz_path, 'r:gz') as archive:
+        for member in archive.getmembers():
+            if member.isfile():
+                filenames.add(member.name)
+                parts = member.name.split('/')
+                for i in range(1, len(parts)):
+                    not_empty_dirs.add('/'.join(parts[:i]))
+                fileobj = archive.extractfile(member)
+                if fileobj:
+                    content = fileobj.read()
+                    sha256 = hashlib.sha256(content).hexdigest()
+                    hashes.add(sha256)
+            else:
+                all_dirs.add(member.name.rstrip("/"))
+    if filenames:
+        hashes.add(hashlib.sha256(":".join(filenames).encode("utf-8")).hexdigest())
+    if empty_dirs := all_dirs - not_empty_dirs:
+        hashes.add(hashlib.sha256(":".join(empty_dirs).encode("utf-8")).hexdigest())
+    return tuple(hashes)

mgmtworker/cloudify_system_workflows/snapshots/snapshot_restore.py

@@ -6,12 +6,14 @@
 import uuid
 import base64
 import shutil
+import stat
 import zipfile
 import tempfile
 import threading
 import subprocess
 from contextlib import contextmanager
 from functools import partial
+from pathlib import Path
 from typing import Any
 
 from cloudify.workflows import ctx
@@ -76,6 +78,11 @@
 
 # Reproduced/modified from patch for https://bugs.python.org/issue15795
 class ZipFile(zipfile.ZipFile):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._all_entries = {info.filename.rstrip('/'): info for info in self.infolist()}
+    
     def _extract_member(self, member, targetpath, pwd):
         """Extract the ZipInfo object 'member' to a physical
            file on the path targetpath.
@@ -112,11 +119,18 @@ def _extract_member(self, member, targetpath, pwd):
                 os.mkdir(targetpath)
             return targetpath
 
-        with self.open(member, pwd=pwd) as source, \
-                open(targetpath, "wb") as target:
+        _mode = member.external_attr >> 16
+        if stat.S_ISLNK(_mode):
+            link = self.read(member.filename).decode('utf-8')
+            source_path = Path(member.filename).parent / link
+            member_to_extract = self._all_entries[os.path.normpath(source_path)]
+        else:
+            member_to_extract = member
+
+        with self.open(member_to_extract, pwd=pwd) as source, open(targetpath, "wb") as target:
             shutil.copyfileobj(source, target)
 
-        mode = member.external_attr >> 16 & 0xFFF
+        mode = member_to_extract.external_attr >> 16 & 0xFFF
         os.chmod(targetpath, mode)
         return targetpath