Release 0.6.2: harden component integration and lifecycle checks

Author: O2eg

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pg-workload"
-version = "0.6.1"
+version = "0.6.2"
 description = "PostgreSQL workload emulator for diagnostic and observability testing"
 readme = "README.md"
 requires-python = ">=3.10"

src/pg_workload/__init__.py

@@ -1,6 +1,6 @@
 """PostgreSQL workload generation and scheduling tools."""
 
-__version__ = "0.6.1"
+__version__ = "0.6.2"
 
 from pg_workload.assets import bundled_profiles_root, initialize_project
 from pg_workload.cli import build_parser

src/pg_workload/assets.py

@@ -50,6 +50,9 @@ def _copy_tree(source_root: Path, destination_root: Path, *, force: bool) -> dic
     destination_root.mkdir(parents=True, exist_ok=True)
     for source in sorted(source_root.rglob("*")):
         relative = source.relative_to(source_root)
+        # Bytecode caches are interpreter by-products of the installation, not assets.
+        if "__pycache__" in relative.parts:
+            continue
         destination = destination_root / relative
         if source.is_symlink():
             raise WorkloadError(f"Bundled asset must not be a symlink: {source}")

src/pg_workload/cli.py

@@ -374,6 +374,11 @@ def build_parser() -> argparse.ArgumentParser:
     stop_cmd.add_argument("--root", default=DEFAULT_ROOT, help="Initialized project directory")
     stop_cmd.add_argument("--daemon-lock-file", default=DEFAULT_DAEMON_LOCK_FILE)
     stop_cmd.add_argument("--scheduler-log-file", default=DEFAULT_SCHEDULER_LOG_FILE)
+    stop_cmd.add_argument(
+        "--expected-identity",
+        type=_scheduler_identity_arg,
+        help="Require the recorded scheduler identity before signalling",
+    )
     stop_cmd.add_argument("--timeout", type=positive_float_arg, default=10.0)
 
     plan_cmd = subparsers.add_parser("plan", help="Build a deterministic execution plan without connecting")
@@ -460,8 +465,18 @@ def _verify_plan_hash(args: argparse.Namespace, operation: str) -> dict[str, obj
     return plan
 
 
+def _scheduler_identity_arg(value: str) -> dict[str, object]:
+    identity = json.loads(value)
+    if not isinstance(identity, dict):
+        raise argparse.ArgumentTypeError("scheduler identity must be a JSON object")
+    return identity
+
+
 def _scheduler_paths(args: argparse.Namespace) -> tuple[Path, Path, Path]:
     root = Path(args.root).resolve()
+    # Without this a mistyped --root reads as "stopped" while the real scheduler keeps running.
+    if not root.is_dir():
+        raise WorkloadError(f"Workload project directory does not exist: {root}")
     lock_path = resolve_relative_path(root, args.daemon_lock_file, "scheduler lock file")
     log_path = resolve_relative_path(root, args.scheduler_log_file, "scheduler log file")
     return root, lock_path, log_path
@@ -593,7 +608,7 @@ def _run(args: argparse.Namespace) -> int:
         return 0
     if args.command == "stop":
         root, lock_path, log_path = _scheduler_paths(args)
-        result = stop_scheduler(root, lock_path, log_path, args.timeout)
+        result = stop_scheduler(root, lock_path, log_path, args.timeout, expected_identity=args.expected_identity)
         if args.machine:
             _emit_machine(args, "succeeded", result=result)
         else:

src/pg_workload/control.py

@@ -46,6 +46,18 @@ def _owned(pid: int, root: Path) -> bool:
     )
 
 
+def scheduler_identity(pid: int) -> dict[str, Any] | None:
+    try:
+        stat = Path(f"/proc/{pid}/stat").read_text().rsplit(")", 1)[1].split()
+        return {
+            "pid": pid,
+            "start_ticks": int(stat[19]),
+            "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip(),
+        }
+    except (OSError, ValueError, IndexError):
+        return None
+
+
 def scheduler_status(root: Path, lock_path: Path, log_path: Path) -> dict[str, Any]:
     pid = _pid(lock_path)
     running = pid is not None and _owned(pid, root)
@@ -55,6 +67,7 @@ def scheduler_status(root: Path, lock_path: Path, log_path: Path) -> dict[str, A
         "pid": pid if running else None,
         "stale_pid": pid if pid is not None and not running else None,
         "owned": running,
+        "identity": scheduler_identity(pid) if running else None,
         "root": str(root.resolve()),
         "lock_file": str(lock_path),
         "log_file": str(log_path),
@@ -139,6 +152,7 @@ def stop_scheduler(
     lock_path: Path,
     log_path: Path,
     timeout: float,
+    expected_identity: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     pid = _pid(lock_path)
     if pid is None:
@@ -148,7 +162,20 @@ def stop_scheduler(
             raise WorkloadOwnershipError(f"Refusing to signal pid {pid}: it is not a pg-workload scheduler for {root}")
         _clear_stale_pid(lock_path, pid)
         return scheduler_status(root, lock_path, log_path)
-    os.kill(pid, signal.SIGTERM)
+    if expected_identity is not None:
+        # Pin the process before checking its identity: PID reuse must never redirect the signal.
+        try:
+            fd = os.pidfd_open(pid)
+        except ProcessLookupError:
+            return scheduler_status(root, lock_path, log_path)
+        try:
+            if scheduler_identity(pid) != expected_identity:
+                raise WorkloadOwnershipError("Scheduler identity differs from the recorded run; refusing to stop it")
+            signal.pidfd_send_signal(fd, signal.SIGTERM)
+        finally:
+            os.close(fd)
+    else:
+        os.kill(pid, signal.SIGTERM)
     deadline = time.monotonic() + timeout
     while time.monotonic() < deadline:
         if not _cmdline(pid):

src/pg_workload/orchestration.py

@@ -54,7 +54,9 @@ def _file_hash(path: Path) -> str:
 def profile_descriptor(profile: Profile) -> dict[str, Any]:
     files = []
     for path in sorted(profile.path.rglob("*")):
-        if not path.is_file() or "log" in path.relative_to(profile.path).parts:
+        parts = path.relative_to(profile.path).parts
+        # Job logs and interpreter bytecode caches are run by-products, not profile content.
+        if not path.is_file() or "log" in parts or "__pycache__" in parts:
             continue
         files.append(
             {

tests/test_scheduler_identity.py

@@ -0,0 +1,57 @@
+import os
+import subprocess
+import sys
+import time
+
+import pytest
+
+from pg_workload.control import WorkloadOwnershipError, scheduler_status, stop_scheduler
+
+
+def test_guarded_stop_leaves_replacement_scheduler_alive(tmp_path):
+    (tmp_path / "data").mkdir()
+    lock = tmp_path / "state/scheduler.lock"
+    log = tmp_path / "scheduler.log"
+    with log.open("w") as output:
+        process = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "pg_workload",
+                "scheduler",
+                "--root",
+                str(tmp_path),
+                "--target=external",
+                "--host=127.0.0.1",
+                "--dry-run",
+            ],
+            stdout=output,
+            stderr=output,
+        )
+    try:
+        deadline = time.monotonic() + 10
+        while time.monotonic() < deadline:
+            status = scheduler_status(tmp_path, lock, log)
+            if status["running"]:
+                break
+            assert process.poll() is None, log.read_text()
+            time.sleep(0.05)
+        assert status["running"], log.read_text()
+        identity = status["identity"]
+        assert identity["pid"] == process.pid
+        for old in [
+            {},
+            dict(identity, start_ticks=identity["start_ticks"] - 1),
+            dict(identity, boot_id="old-boot"),
+            dict(identity, pid=os.getpid()),
+        ]:
+            with pytest.raises(WorkloadOwnershipError, match="identity differs"):
+                stop_scheduler(tmp_path, lock, log, 3, expected_identity=old)
+            assert process.poll() is None
+        stopped = stop_scheduler(tmp_path, lock, log, 5, expected_identity=identity)
+        assert not stopped["running"]
+        process.wait(timeout=5)
+    finally:
+        if process.poll() is None:
+            process.terminate()
+            process.wait(timeout=5)

tests/test_workload.py

@@ -4,10 +4,12 @@
 import io
 import json
 import os
+import queue
 import re
 import subprocess
 import sys
 import tempfile
+import threading
 import time
 import types
 import unittest
@@ -157,13 +159,30 @@ def plan_hash(*, disk_max_used_pct: int = 90) -> str:
             first = plan_hash()
             self.assertEqual(first, plan_hash())
             self.assertNotEqual(first, plan_hash(disk_max_used_pct=95))
+            cache = root / "data" / "pagila" / "__pycache__"
+            cache.mkdir(exist_ok=True)
+            (cache / "generator.cpython-399.pyc").write_bytes(b"interpreter by-product")
+            self.assertEqual(first, plan_hash())
             profile = root / "data" / "pagila" / "profile.yml"
             profile.write_text(
                 profile.read_text(encoding="utf-8") + "\n# local experiment\n",
                 encoding="utf-8",
             )
             self.assertNotEqual(first, plan_hash())
 
+    def test_scheduler_commands_reject_a_missing_project_root(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            missing = Path(tmpdir) / "no-such-project"
+            for command in ("status", "stop"):
+                stdout = io.StringIO()
+                with contextlib.redirect_stdout(stdout):
+                    result = main(["--machine", command, "--root", str(missing)])
+                envelope = json.loads(stdout.getvalue())
+                self.assertNotEqual(0, result)
+                self.assertEqual("failed", envelope["status"])
+                self.assertIn("does not exist", envelope["error"]["message"])
+            self.assertFalse(missing.exists())
+
     def test_detached_scheduler_has_owned_start_status_stop_lifecycle(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             root = Path(tmpdir)
@@ -764,6 +783,7 @@ def test_init_is_idempotent_and_protects_modified_assets(self):
             second = self.workload.initialize_project(root)
             self.assertGreater(first["assets"]["created"], 0)
             self.assertGreater(second["assets"]["unchanged"], 0)
+            self.assertEqual([], sorted(root.rglob("__pycache__")))
             self.assertTrue((root / "schema" / "pg_workload-v1.schema.json").is_file())
             manifest = root / "data" / "simple_stock" / "profile.yml"
             manifest.write_text(manifest.read_text(encoding="utf-8") + "# local edit\n", encoding="utf-8")
@@ -1379,8 +1399,30 @@ def test_running_scheduler_picks_up_new_profile_from_state_without_restart(self)
                 stderr=subprocess.PIPE,
                 text=True,
             )
+            lines = []
+            events = queue.Queue()
+
+            def read_stdout():
+                for line in proc.stdout:
+                    lines.append(line)
+                    events.put(line)
+
+            reader = threading.Thread(target=read_stdout, daemon=True)
+            reader.start()
+
+            def wait_for(text):
+                deadline = time.monotonic() + 15
+                while time.monotonic() < deadline:
+                    try:
+                        line = events.get(timeout=max(0.01, deadline - time.monotonic()))
+                    except queue.Empty:
+                        break
+                    if text in line:
+                        return
+                self.fail(f"Timed out waiting for {text!r}: {''.join(lines)}")
+
             try:
-                time.sleep(0.2)
+                wait_for("Scheduler started.")
                 self.create_minimal_profile(root, "dynamic_profile")
                 self.workload.update_job_state(
                     root / "state" / "workloads.yml",
@@ -1389,10 +1431,14 @@ def test_running_scheduler_picks_up_new_profile_from_state_without_restart(self)
                     enabled=True,
                     interval=60,
                 )
-                time.sleep(0.5)
+                wait_for("Running dynamic_profile:main")
+                wait_for("pgbench")
             finally:
                 proc.terminate()
-            stdout, stderr = proc.communicate(timeout=10)
+                proc.wait(timeout=10)
+                reader.join(timeout=5)
+            stdout = "".join(lines)
+            stderr = proc.stderr.read()
 
             self.assertIn("Scheduler started.", stdout)
             self.assertIn("Running dynamic_profile:main", stdout)