Skip backup to worker when low on disk space

Author: CamDavidsonPilon

pioreactor/actions/leader/backup_database.py

@@ -3,6 +3,9 @@
 
 import click
 
+from pathlib import Path
+import subprocess
+
 from pioreactor.cluster_management import get_active_workers_in_inventory
 from pioreactor.config import config
 from pioreactor.exc import RsyncError
@@ -18,6 +21,29 @@
 from pioreactor.whoami import UNIVERSAL_EXPERIMENT
 
 
+def _remote_available_space(address: str, path: str) -> int | None:
+    """Return available bytes on remote machine or ``None`` on failure."""
+    try:
+        result = subprocess.run(
+            ["ssh", "-o", "ConnectTimeout=5", address, "df", "-PB1", path],
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+    except subprocess.CalledProcessError:
+        return None
+
+    try:
+        lines = result.stdout.strip().splitlines()
+        if len(lines) >= 2:
+            return int(lines[1].split()[3])
+    except Exception:
+        return None
+
+    return None
+
+
 def count_writes_occurring() -> int:
     with local_intermittent_storage("mqtt_to_db_streaming") as c:
         return c.get("inserts_in_last_60s", 0)
@@ -89,6 +115,17 @@ def backup_database(output_file: str, force: bool = False, backup_to_workers: in
                 continue
 
             logger.debug(f"Attempting backing up database to {backup_unit}.")
+            available = _remote_available_space(
+                resolve_to_address(backup_unit), str(Path(output_file).parent)
+            )
+            if available is not None and available < Path(output_file).stat().st_size:
+                logger.debug(
+                    f"Skipping backup to {backup_unit}. Not enough disk space."
+                )
+                logger.warning(
+                    f"Unable to backup database to {backup_unit}. Not enough disk space."
+                )
+                continue
             try:
                 rsync(
                     "-hz",

pioreactor/tests/test_backup_database.py

@@ -0,0 +1,48 @@
+import sqlite3
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from pioreactor.actions.leader.backup_database import backup_database
+from pioreactor.config import config
+
+
+@contextmanager
+def dummy_lifecycle(*args, **kwargs):
+    yield SimpleNamespace(job_key="backup_database")
+
+
+def test_skip_backup_when_worker_has_no_space(tmp_path):
+    db_path = tmp_path / "db.sqlite"
+    config["storage"]["database"] = str(db_path)
+
+    conn = sqlite3.connect(db_path)
+    conn.execute("CREATE TABLE t(id INTEGER)")
+    conn.commit()
+    conn.close()
+
+    output = tmp_path / "backup.sqlite"
+
+    with (
+        patch(
+            "pioreactor.actions.leader.backup_database.long_running_managed_lifecycle",
+            dummy_lifecycle,
+        ),
+        patch(
+            "pioreactor.actions.leader.backup_database.create_logger",
+            return_value=MagicMock(),
+        ),
+        patch(
+            "pioreactor.actions.leader.backup_database.get_active_workers_in_inventory",
+            return_value=["worker1"],
+        ),
+        patch(
+            "pioreactor.actions.leader.backup_database._remote_available_space",
+            return_value=0,
+        ),
+        patch(
+            "pioreactor.actions.leader.backup_database.rsync",
+        ) as mock_rsync,
+    ):
+        backup_database(str(output), force=True, backup_to_workers=1)
+        mock_rsync.assert_not_called()