AnyImage/src/snapform/blendjob/service.py

import socket
import subprocess
import time
import uuid
from dataclasses import dataclass

from .client import JobClient


@dataclass
class ServiceConnection:
    port: int
    pid: int
    instance_id: str
    client: JobClient


class ServiceController:
    """Own a dedicated local JobServer process and its HTTP connection."""

    def __init__(
        self,
        command_factory,
        *,
        cwd_factory,
        log_path_factory,
        host="127.0.0.1",
        creationflags=0,
    ):
        self.command_factory = command_factory
        self.cwd_factory = cwd_factory
        self.log_path_factory = log_path_factory
        self.host = host
        self.creationflags = creationflags
        self.connection = None
        self.process = None
        self.log_file = None
        self.auto_start = False
        self.snapshot = {"state": "STOPPED", "message": "Stopped"}

    def start(self):
        if self.process is not None and self.process.poll() is None:
            return self.connection
        self._close_process()
        port = self._available_port()
        instance_id = uuid.uuid4().hex
        log_path = self.log_path_factory()
        log_path.parent.mkdir(parents=True, exist_ok=True)
        self.log_file = log_path.open("a", encoding="utf-8")
        command = self.command_factory(port, instance_id)
        try:
            self.process = subprocess.Popen(
                command,
                stdout=self.log_file,
                stderr=subprocess.STDOUT,
                cwd=self.cwd_factory(),
                creationflags=self.creationflags,
            )
        except OSError:
            self.log_file.close()
            self.log_file = None
            raise
        client = JobClient(self.host, port)
        self.connection = ServiceConnection(
            port,
            self.process.pid,
            instance_id,
            client,
        )
        self.snapshot = {"state": "STARTING", "message": "Starting", "port": port}
        return self.connection

    def health(self, timeout=0.3):
        if self.connection is None:
            raise RuntimeError("Service is not running")
        health = self.connection.client.request("GET", "/health", timeout=timeout)
        if health.get("instance_id") != self.connection.instance_id:
            raise RuntimeError("Service instance does not match")
        state = "BUSY" if health.get("busy") else "READY"
        self.snapshot = {
            **health,
            "state": state,
            "message": "Busy" if state == "BUSY" else "Ready",
            "port": self.connection.port,
        }
        return health

    def ensure(self, timeout=60.0):
        self.auto_start = True
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                self.health(timeout=0.5)
                return self.connection
            except RuntimeError:
                self._close_finished_process()
            if self.process is None:
                self.start()
            time.sleep(0.05)
        raise RuntimeError(f"Service did not become ready; log: {self.log_path_factory()}")

    def enable(self):
        self.auto_start = True
        try:
            self.health(timeout=0.2)
        except RuntimeError:
            self.start()
        return self.connection

    def stop(self, timeout=10.0):
        self.auto_start = False
        if self.connection is not None:
            try:
                self.connection.client.shutdown()
            except RuntimeError:
                pass
        deadline = time.monotonic() + timeout
        while self.process is not None and self.process.poll() is None:
            if time.monotonic() >= deadline:
                self._close_process(terminate=True)
                break
            time.sleep(0.05)
        self._close_process()
        self.connection = None
        self.snapshot = {"state": "STOPPED", "message": "Stopped"}

    def restart(self):
        self.stop()
        self.auto_start = True
        return self.start()

    def submit(self, request):
        return self.ensure().client.submit_request(request)

    def status(self, job_id):
        return self.ensure().client.status(job_id)

    def cancel(self, job_id):
        return self.ensure().client.cancel(job_id)

    def resource(self, name):
        return self.ensure().client.resource(name)

    def clear_resource(self, name):
        return self.ensure().client.clear_resource(name)

    def detach(self):
        self.auto_start = False
        self.connection = None
        self._close_process(terminate=True)
        self.snapshot = {"state": "STOPPED", "message": "Stopped"}

    def poll(self, *, available=True):
        if not available:
            self.snapshot = {"state": "UNAVAILABLE", "message": "Environment is not installed"}
            return 2.0
        try:
            self.health(timeout=0.5)
            return 5.0
        except RuntimeError as error:
            self._close_finished_process()
            if not self.auto_start:
                self.snapshot = {"state": "STOPPED", "message": "Stopped"}
                return 2.0
            try:
                self.start()
            except OSError as start_error:
                self.snapshot = {"state": "ERROR", "message": str(start_error)}
                return 2.0
            self.snapshot = {"state": "STARTING", "message": "Starting"}
            return 1.0

    def mark_job_complete(self, job_id):
        active = self.snapshot.get("active_job")
        if active and active.get("job_id") != job_id:
            return
        if self.snapshot.get("state") == "BUSY":
            self.snapshot = {
                **self.snapshot,
                "busy": False,
                "active_job": None,
                "state": "READY",
                "message": "Ready",
            }

    def _available_port(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
            listener.bind((self.host, 0))
            return listener.getsockname()[1]

    def _close_finished_process(self):
        if self.process is not None and self.process.poll() is not None:
            self._close_process()
            self.connection = None

    def _close_process(self, terminate=False):
        process = self.process
        self.process = None
        if process is not None and terminate and process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=3.0)
        if self.log_file is not None:
            self.log_file.close()
            self.log_file = None