debug_env/server/services/manager.py

"""
ServiceManager — manages the lifecycle of a buggy server subprocess.

Used by incident-type tasks to spin up a realistic service environment
that the agent can query, diagnose, fix, and restart.
"""

import logging
import os
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.request

logger = logging.getLogger(__name__)


class ServiceManager:
    """Manages a single buggy-server subprocess for SRE incident scenarios."""

    def __init__(self) -> None:
        self._process: subprocess.Popen | None = None
        self._log_file = None
        self._port: int | None = None
        self._workdir: str | None = None
        self._entrypoint: str | None = None
        self._traffic: "TrafficGenerator | None" = None

    # ── Public API ─────────────────────────────────────────────────────────────

    def find_free_port(self) -> int:
        """Return a free TCP port by binding to port 0 and releasing it."""
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            return s.getsockname()[1]

    def start(self, workdir: str, entrypoint: str, port: int) -> None:
        """
        Spawn the buggy server subprocess.

        Opens ``{workdir}/app.log`` for stdout/stderr, injects
        ``BUGGY_SERVER_PORT`` into the environment, then polls ``GET /health``
        with exponential backoff (100 ms → 1 600 ms) until the server is
        healthy or a 5-second timeout elapses.

        Args:
            workdir:    Absolute path to the task working directory.
            entrypoint: Filename of the server script (e.g. ``"buggy_server.py"``).
            port:       TCP port to bind to.
        """
        self._workdir = workdir
        self._entrypoint = entrypoint
        self._port = port

        log_path = os.path.join(workdir, "app.log")
        self._log_file = open(log_path, "w", encoding="utf-8")

        env = os.environ.copy()
        env["BUGGY_SERVER_PORT"] = str(port)

        self._process = subprocess.Popen(
            [sys.executable, entrypoint],
            cwd=workdir,
            env=env,
            stdout=self._log_file,
            stderr=self._log_file,
        )
        logger.info("Started buggy server PID=%s on port %s", self._process.pid, port)

        self._wait_for_health(port, timeout=5.0)

        # Start background traffic so the log fills with realistic entries
        from debug_env.server.services.traffic import TrafficGenerator
        self._traffic = TrafficGenerator(port)
        self._traffic.start()

    def stop(self) -> None:
        """Terminate the subprocess and release resources."""
        if self._traffic is not None:
            self._traffic.stop()
            self._traffic = None

        if self._process is not None:
            try:
                self._process.terminate()
                self._process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                logger.warning("Process did not terminate gracefully — killing")
                self._process.kill()
                self._process.wait()
            except OSError:
                pass
            logger.info("Stopped buggy server PID=%s", self._process.pid)
            self._process = None

        if self._log_file is not None:
            try:
                self._log_file.close()
            except OSError:
                pass
            self._log_file = None

    def restart(self, workdir: str, entrypoint: str) -> int:
        """
        Stop the current service and start a fresh one on a new free port.

        Returns:
            The new port the service is listening on.
        """
        self.stop()
        port = self.find_free_port()
        self.start(workdir, entrypoint, port)
        return port

    def is_running(self) -> bool:
        """Return True if the subprocess is alive."""
        return self._process is not None and self._process.poll() is None

    @property
    def port(self) -> int | None:
        """The TCP port the service is (or was last) listening on."""
        return self._port

    # ── Internal helpers ───────────────────────────────────────────────────────

    def _wait_for_health(self, port: int, timeout: float = 5.0) -> None:
        """Poll GET /health with exponential backoff until healthy or timeout."""
        deadline = time.monotonic() + timeout
        delay = 0.1
        url = f"http://localhost:{port}/health"

        while time.monotonic() < deadline:
            try:
                with urllib.request.urlopen(url, timeout=1) as resp:
                    if resp.status == 200:
                        logger.info("Service is healthy on port %s", port)
                        return
            except Exception:
                pass
            time.sleep(delay)
            delay = min(delay * 2, 1.6)

        raise TimeoutError(
            f"Buggy server did not become healthy on port {port} within {timeout}s. "
            "Check app.log for startup errors."
        )