Add AssertMinFreeDiskSpace

pleskdistup/actions/common_checks.py

@@ -1,6 +1,8 @@
 # Copyright 2023-2024. Plesk International GmbH. All rights reserved.
+import json
 import os
 import subprocess
+import typing
 
 from pleskdistup.common import action, log, packages, version
 
@@ -142,3 +144,109 @@ def __init__(self):
 
     def _do_check(self) -> bool:
         return subprocess.run(["/bin/fuser", "/var/lib/apt/lists/lock"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode != 0
+
+
+class MinFreeDiskSpaceViolation(typing.NamedTuple):
+    """Information about a filesystem with insufficient free disk space."""
+    dev: str
+    """Device name."""
+    req_bytes: int
+    """Required space on the device (in bytes)."""
+    avail_bytes: int
+    """Available space on the device (in bytes)."""
+    paths: typing.Set[str]
+    """Paths belonging to this device."""
+
+
+class AssertMinFreeDiskSpace(action.CheckAction):
+    """Check if there's enough free disk space.
+
+    Args:
+        requirements: A dictionary mapping paths to minimum free disk
+            space (in bytes) on the devices containing them.
+        name: Name of the check.
+        description: Description to show if the assertion is violated. If ends
+            with a colon, list of violations will be added automatically.
+    """
+    violations: typing.List[MinFreeDiskSpaceViolation]
+    """List of filesystems with insiffucient free disk space."""
+    _description: str
+
+    def __init__(
+        self,
+        requirements: typing.Dict[str, int],
+        name: str = "check if there's enough free disk space",
+        description: str = "There's not enough free disk space:",
+    ):
+        self.requirements = requirements
+        self.name = name
+        self._description = description
+        self.violations = []
+
+    @property
+    def description(self) -> str:
+        """Description of violations or empty string.
+
+        Can be set to change the message template.
+        """
+        if not self.violations:
+            return ""
+        res = self._description
+        if res.endswith(":"):
+            res += " " + ", ".join(
+                f"on filesystem {v.dev!r} for "
+                f"{', '.join(repr(p) for p in sorted(v.paths))} "
+                f"(need {v.req_bytes / 1024**2} MiB, "
+                f"got {v.avail_bytes / 1024**2} MiB)" for v in self.violations
+            )
+        return res
+
+    @description.setter
+    def description(self, val: str) -> None:
+        self._description = val
+
+    def _do_check(self) -> bool:
+        """Perform the check."""
+        log.debug("Checking minimum free disk space")
+        cmd = [
+            "findmnt", "--output", "source,target,avail",
+            "--bytes", "--json", "-T",
+        ]
+        self.violations = []
+        filesystems: typing.Dict[str, dict] = {}
+        for path, req in self.requirements.items():
+            log.debug(f"Checking {path!r} minimum free disk space requirement of {req}")
+            proc = subprocess.run(
+                cmd + [path],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                check=True,
+                universal_newlines=True,
+            )
+            log.debug(
+                f"Command {cmd + [path]} returned {proc.returncode}, "
+                f"stdout: '{proc.stdout}', stderr: '{proc.stderr}'"
+            )
+            fs_data = json.loads(proc.stdout)["filesystems"][0]
+            fs_data["req"] = 0
+            fs_data["paths"] = []
+            if fs_data["source"] not in filesystems:
+                log.debug(f"Discovered new filesystem {fs_data}")
+                filesystems[fs_data["source"]] = fs_data
+            log.debug(
+                f"Adding space requirement of {req} to "
+                f"{filesystems[fs_data['source']]}"
+            )
+            filesystems[fs_data["source"]]["req"] += req
+            filesystems[fs_data["source"]]["paths"].append(path)
+        for dev, fs_data in filesystems.items():
+            if fs_data["req"] > fs_data["avail"]:
+                self.violations.append(
+                    MinFreeDiskSpaceViolation(
+                        dev,
+                        fs_data["req"],
+                        fs_data["avail"],
+                        set(fs_data["paths"]),
+                    )
+                )
+        return len(self.violations) == 0