A client reaches only a server it may use; an orphaned server exits (0.6.2)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "zmqruntime"
-version = "0.6.2"
+version = "0.6.3"
 description = "Generic ZMQ-based distributed execution framework"
 authors = [{name = "Tristan Simas", email = "tristan.simas@mail.mcgill.ca"}]
 license = {text = "MIT"}

src/zmqruntime/__init__.py

@@ -58,6 +58,9 @@
                 "EndpointApplication",
                 "EndpointApplicationCompatibility",
                 "EndpointApplicationCompatibilityError",
+                "EndpointAudience",
+                "EndpointLaunch",
+                "EndpointLaunchMismatchError",
                 "EndpointControlCapability",
                 "EndpointShutdownRequest",
                 "ExecuteRequest",

src/zmqruntime/client.py

@@ -25,13 +25,17 @@
 from zmqruntime.messages import (
     ControlMessageType,
     EndpointApplicationCompatibility,
+    EndpointAudience,
     EndpointControlCapability,
+    EndpointLaunch,
+    EndpointLaunchMismatchError,
     EndpointShutdownRequest,
     MessageFields,
     PongResponse,
     ProcessExit,
     ProcessIdentity,
     ResponseType,
+    imported_code_root,
 )
 from zmqruntime.startup import (
     IDLE_ENDPOINT_STARTUP_OBSERVER,
@@ -920,6 +924,52 @@ def _connection_cancelled(self) -> bool:
         cancellation = self._connection_cancellation.get()
         return cancellation is not None and cancellation.requested()
 
+    def launch(self) -> EndpointLaunch:
+        """The launch of a server this client starts.
+
+        ``_spawn_server_process`` hands it to the server, which reports it in
+        every heartbeat. It is also this client's side of the comparison when
+        it attaches to a server it did not start.
+        """
+
+        return EndpointLaunch.by_this_process(
+            self.code_root(), EndpointAudience.of_persistent(self.persistent)
+        )
+
+    @classmethod
+    def code_root(cls) -> str:
+        """Where this client's code is imported from (see ``ZMQServer.code_root``)."""
+
+        return imported_code_root(cls)
+
+    def _require_started_endpoint(self, endpoint: PongResponse) -> None:
+        """A client that starts a server connects to that server and no other."""
+
+        if endpoint.launch != self.launch():
+            raise EndpointLaunchMismatchError(
+                self.endpoint.control_url(self.config),
+                "it is not the server this client just started",
+            )
+
+    def _require_usable_endpoint(self, endpoint: PongResponse) -> None:
+        """Attach only to a local server whose code and owner admit this client.
+
+        A remote host is an address the user named; its paths and process
+        ids cannot be compared with this machine's.
+        """
+
+        if not self.transport_mode.declaration.endpoint_is_local(self.host, self.port):
+            return
+        reason = (
+            "it does not report where its code is or which process started it"
+            if endpoint.launch is None
+            else endpoint.launch.refusal(self.launch())
+        )
+        if reason is not None:
+            raise EndpointLaunchMismatchError(
+                self.endpoint.control_url(self.config), reason
+            )
+
     def connect(
         self,
         timeout: float = 10.0,
@@ -1137,6 +1187,13 @@ def _connect_locked(
                 process.stop()
                 self.endpoint.cleanup(self.config)
                 return self._cancelled_connection_result()
+            try:
+                self._require_started_endpoint(endpoint)
+            except EndpointLaunchMismatchError:
+                # Another server holds the address: stop only the child, and
+                # leave that server's socket files alone.
+                process.stop()
+                raise
             owned_connection = OwnedEndpointConnection(
                 process=process,
                 target=self.endpoint,
@@ -1159,6 +1216,7 @@ def _attach_existing_endpoint(self, timeout: float) -> bool:
         )
         if endpoint is None or self._connection_cancelled():
             return False
+        self._require_usable_endpoint(endpoint)
         self._setup_client_sockets()
         self._connection = AttachedEndpointConnection(endpoint)
         self._emit_connected_status()

src/zmqruntime/config.py

@@ -4,6 +4,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from typing import TYPE_CHECKING, Annotated
 
 from annotated_types import Ge, Gt, Le, MinLen, Predicate
@@ -94,8 +95,8 @@ class ZMQConfig:
     ipc_socket_dir: NonBlankString = "ipc"
     """Directory containing IPC socket files.
 
-    Relative values are resolved beneath the runtime socket root; absolute
-    values select an explicit directory. TCP transport ignores this field.
+    Relative values are resolved beneath ``socket_root()``; absolute values
+    select an explicit directory. TCP transport ignores this field.
     """
 
     ipc_socket_prefix: NonBlankString = "zmq"
@@ -123,3 +124,13 @@ def __post_init__(self) -> None:
         validate_annotated_dataclass(self)
         if self.default_port + self.control_port_offset > 65535:
             raise ValueError("default_port plus control_port_offset must not exceed 65535.")
+
+    def socket_root(self) -> Path:
+        """Directory under which this application's socket and lock files live.
+
+        Every process that resolves the same root shares one address
+        namespace: port ``n`` names the same socket file for all of them. An
+        application overrides this to keep installations or sessions apart.
+        """
+
+        return Path.home() / f".{self.app_name}"

src/zmqruntime/execution/server.py

@@ -383,6 +383,14 @@ def _interrupt_execution(self, execution_id: str) -> int:
 
         return self._kill_worker_processes()
 
+    def has_unfinished_work(self) -> bool:
+        return self._lifecycle.snapshot(uptime=0.0).active_executions > 0
+
+    def stop_orphaned(self) -> None:
+        self._cancel_all_executions()
+        self._kill_worker_processes()
+        super().stop_orphaned()
+
     def _shutdown_workers(self, msg: dict[str, Any], force: bool = False) -> dict[str, Any]:
         _request, error = self._validate_and_parse(msg, EndpointShutdownRequest)
         if error:

src/zmqruntime/messages.py

@@ -7,18 +7,22 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import os
 import pickle
 import signal
 import sys
+from abc import ABC, abstractmethod
 from collections.abc import Iterable, Mapping
 from dataclasses import dataclass, field, fields, replace
 from enum import Enum
+from pathlib import Path
 from typing import Any, ClassVar, Dict, Optional, Tuple
 from uuid import UUID
 
 import psutil
+from metaclass_registry import AutoRegisterMeta
 from python_introspect import dataclass_from_mapping
 from python_introspect.jsonable import JsonValue, to_jsonable
 from typing_extensions import Self
@@ -890,6 +894,153 @@ def __str__(self) -> str:
         )
 
 
+class EndpointAudience(ABC, metaclass=AutoRegisterMeta):
+    """Which processes may use a server, and what it does once its owner is gone.
+
+    The wire carries ``wire_name``; the registry is the only list of audiences.
+    """
+
+    __registry_key__ = "wire_name"
+    __skip_if_no_key__ = True
+
+    wire_name: ClassVar[Optional[str]] = None
+
+    @classmethod
+    def of_persistent(cls, persistent: bool) -> "type[EndpointAudience]":
+        """A persistent server outlives its client, so other sessions may use it."""
+
+        return AnySessionAudience if persistent else OwnerOnlyAudience
+
+    @classmethod
+    @abstractmethod
+    def admits(cls, owner: "ProcessIdentity", client: "ProcessIdentity") -> bool:
+        """Whether ``client`` may use a server that ``owner`` started."""
+
+    @classmethod
+    @abstractmethod
+    def orphan_stops(cls, unfinished_work: bool) -> bool:
+        """Whether a server whose owner is gone stops now."""
+
+
+class OwnerOnlyAudience(EndpointAudience):
+    """A server for the process that started it: its owner stops it on disconnect."""
+
+    wire_name = "owner_only"
+
+    @classmethod
+    def admits(cls, owner: "ProcessIdentity", client: "ProcessIdentity") -> bool:
+        return client == owner
+
+    @classmethod
+    def orphan_stops(cls, unfinished_work: bool) -> bool:
+        return True
+
+
+class AnySessionAudience(EndpointAudience):
+    """A server other sessions running the same code may attach to."""
+
+    wire_name = "any_session"
+
+    @classmethod
+    def admits(cls, owner: "ProcessIdentity", client: "ProcessIdentity") -> bool:
+        return True
+
+    @classmethod
+    def orphan_stops(cls, unfinished_work: bool) -> bool:
+        return not unfinished_work
+
+
+def imported_code_root(declared: type) -> str:
+    """Directory the top-level package that defines ``declared`` is imported from.
+
+    Two checkouts of one application import the same class from different
+    directories; this is what tells their servers apart.
+    """
+
+    top_level = sys.modules[declared.__module__.partition(".")[0]]
+    package_paths = tuple(getattr(top_level, "__path__", ()))
+    if package_paths:
+        origin = Path(package_paths[0])
+    else:
+        module_file = getattr(top_level, "__file__", None)
+        origin = Path(module_file) if module_file else Path.cwd() / top_level.__name__
+    return str(origin.resolve().parent)
+
+
+@dataclass(frozen=True)
+class EndpointLaunch(WireMessage):
+    """Where a server's code is, which process started it, and who may use it."""
+
+    code_root: str
+    owner: ProcessIdentity
+    audience: str
+
+    def __post_init__(self) -> None:
+        if not self.code_root:
+            raise ValueError("Endpoint launch requires a code root")
+        if self.audience not in EndpointAudience.__registry__:
+            raise ValueError(f"Unknown endpoint audience {self.audience!r}")
+
+    @classmethod
+    def by_this_process(
+        cls,
+        code_root: str,
+        audience: "type[EndpointAudience]",
+    ) -> Self:
+        """The launch of an endpoint started by the calling process."""
+
+        return cls(
+            code_root=code_root,
+            owner=ProcessIdentity.current(),
+            audience=audience.wire_name,
+        )
+
+    @property
+    def audience_type(self) -> "type[EndpointAudience]":
+        return EndpointAudience.__registry__[self.audience]
+
+    def refusal(self, client: "EndpointLaunch") -> Optional[str]:
+        """Why a client launched as ``client`` may not use this server, if so."""
+
+        if self.code_root != client.code_root:
+            return (
+                f"it runs code from {self.code_root}, "
+                f"and this client runs code from {client.code_root}"
+            )
+        if not self.audience_type.admits(self.owner, client.owner):
+            return (
+                f"process {self.owner.pid} started it for its own use, "
+                f"and this client is process {client.owner.pid}"
+            )
+        return None
+
+    def owner_is_gone(self) -> bool:
+        """Whether the starting process has exited (unknown counts as present)."""
+
+        return self.owner.is_alive() is False
+
+    def to_argument(self) -> str:
+        """One command-line argument a spawned server decodes with ``from_argument``."""
+
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_argument(cls, argument: str) -> Self:
+        return cls.from_dict(json.loads(argument))
+
+
+class EndpointLaunchMismatchError(RuntimeError):
+    """A client refused a server it did not start and may not use."""
+
+    def __init__(self, address: str, reason: str) -> None:
+        self.address = address
+        self.reason = reason
+        super().__init__(
+            f"Refusing the server at {address}: {reason}. "
+            "Stop that server, or give this client an address of its own."
+        )
+
+
 @dataclass(frozen=True)
 class PongResponse(TypedWireMessage, WireView, ControlResponse):
     """Complete typed server-heartbeat response.
@@ -921,6 +1072,7 @@ class PongResponse(TypedWireMessage, WireView, ControlResponse):
     compile_message: Optional[str] = None
     process_usage: Optional[ProcessResourceUsage] = None
     application: Optional[EndpointApplication] = None
+    launch: Optional[EndpointLaunch] = None
 
     def __post_init__(self) -> None:
         """Hold advertised capabilities as one ordered set."""