diff --git a/pyproject.toml b/pyproject.toml index d171997..f105810 100644 --- a/pyproject.toml +++ b/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"} diff --git a/src/zmqruntime/__init__.py b/src/zmqruntime/__init__.py index 6fd14f6..6d3c14d 100644 --- a/src/zmqruntime/__init__.py +++ b/src/zmqruntime/__init__.py @@ -58,6 +58,9 @@ "EndpointApplication", "EndpointApplicationCompatibility", "EndpointApplicationCompatibilityError", + "EndpointAudience", + "EndpointLaunch", + "EndpointLaunchMismatchError", "EndpointControlCapability", "EndpointShutdownRequest", "ExecuteRequest", diff --git a/src/zmqruntime/client.py b/src/zmqruntime/client.py index 0395943..a1a8b88 100644 --- a/src/zmqruntime/client.py +++ b/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() diff --git a/src/zmqruntime/config.py b/src/zmqruntime/config.py index 6dbc7a2..3bc22b3 100644 --- a/src/zmqruntime/config.py +++ b/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}" diff --git a/src/zmqruntime/execution/server.py b/src/zmqruntime/execution/server.py index 0b35fdd..d8be036 100644 --- a/src/zmqruntime/execution/server.py +++ b/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: diff --git a/src/zmqruntime/messages.py b/src/zmqruntime/messages.py index 6c0c758..705250f 100644 --- a/src/zmqruntime/messages.py +++ b/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.""" diff --git a/src/zmqruntime/server.py b/src/zmqruntime/server.py index f016bba..db96e98 100644 --- a/src/zmqruntime/server.py +++ b/src/zmqruntime/server.py @@ -7,23 +7,28 @@ import platform import subprocess import threading +import time from abc import ABC, abstractmethod from collections.abc import Callable, Mapping import zmq from metaclass_registry import AutoRegisterMeta +from typing_extensions import Self from zmqruntime.config import TransportMode, ZMQConfig from zmqruntime.messages import ( + AnySessionAudience, ControlErrorResponse, ControlMessageType, ControlResponse, EndpointApplication, + EndpointLaunch, MessageFields, PongResponse, ProcessIdentity, ServerRole, SocketType, + imported_code_root, ) from zmqruntime.transport import ( TransportEndpoint, @@ -51,6 +56,22 @@ def server_type(cls) -> str | None: """Return this server class's registered runtime role.""" return cls._server_type + @classmethod + def code_root(cls) -> str: + """Where this server's code is imported from. + + An application whose import root has a declared owner overrides this, + on its clients too, so both sides state the same fact. + """ + + return imported_code_root(cls) + + def launched_by(self, launch: EndpointLaunch) -> Self: + """Record the launch the starting client handed this server's process.""" + + self.launch = launch + return self + @classmethod def server_role(cls) -> ServerRole: """Return the protocol-level category of this server.""" @@ -79,6 +100,12 @@ def __init__( self.log_file_path = log_file_path self.data_socket_type = data_socket_type if data_socket_type is not None else zmq.PUB self.application = application + # Until a starting client's launch is recorded (``launched_by``), the + # server was started by its own process. + self.launch = EndpointLaunch.by_this_process( + self.code_root(), AnySessionAudience + ) + self._owner_checked_at = time.monotonic() self.zmq_context = None self.data_socket = None self.control_socket = None @@ -176,6 +203,9 @@ def is_running(self): return self._running def process_messages(self): + if not self._running: + return + self._stop_if_orphaned() if not self._running: return @@ -270,6 +300,7 @@ def _create_pong_response(self) -> PongResponse: server_type=self.__class__.server_type(), server_role=self.__class__.server_role(), application=self.application, + launch=self.launch, log_file_path=self.log_file_path, process_identity=ProcessIdentity.current(), ) @@ -277,6 +308,39 @@ def _create_pong_response(self) -> PongResponse: def request_shutdown(self): self._running = False + orphan_check_interval_seconds = 1.0 + + def _stop_if_orphaned(self) -> None: + """Stop once the process that launched this server is gone. + + Nothing else stops a server whose client was killed, or a persistent + server whose session ended: it would hold its code and workers forever. + """ + + now = time.monotonic() + if now - self._owner_checked_at < self.orphan_check_interval_seconds: + return + self._owner_checked_at = now + if not self.launch.owner_is_gone(): + return + if not self.launch.audience_type.orphan_stops(self.has_unfinished_work()): + return + logger.info( + "Process %s, which started this server, is gone; stopping", + self.launch.owner.pid, + ) + self.stop_orphaned() + + def has_unfinished_work(self) -> bool: + """Whether this server still has work an absent owner would lose.""" + + return False + + def stop_orphaned(self) -> None: + """Release this server's work and end its message loop.""" + + self.request_shutdown() + @staticmethod def kill_processes_on_port(port): killed = 0 diff --git a/src/zmqruntime/transport_modes.py b/src/zmqruntime/transport_modes.py index 8a388cc..9b623b5 100644 --- a/src/zmqruntime/transport_modes.py +++ b/src/zmqruntime/transport_modes.py @@ -6,9 +6,10 @@ import json import os import socket +import tempfile from abc import ABC, abstractmethod from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager from ipaddress import ip_address from pathlib import Path from dataclasses import replace @@ -250,6 +251,15 @@ def data_control_pair_is_available( ) -> bool: """Return whether both endpoint identities can be acquired together.""" + @classmethod + @abstractmethod + def private_endpoint(cls, config: ZMQConfig) -> AbstractContextManager[ZMQConfig]: + """Yield ``config`` moved to an address no other caller is given. + + A caller that starts its own server takes its address here instead of + picking a free port in the shared namespace, which reserves nothing. + """ + class TcpTransportDeclaration(TransportDeclaration): """Complete TCP endpoint semantics.""" @@ -372,7 +382,7 @@ def startup_lock_path( port: int, config: ZMQConfig, ) -> Path: - return Path.home() / f".{config.app_name}" / "tcp" / f"{port}.startup.lock" + return config.socket_root() / "tcp" / f"{port}.startup.lock" @classmethod def endpoint_is_stale(cls, port: int, config: ZMQConfig) -> bool: @@ -399,6 +409,19 @@ def data_control_pair_is_available( return False return True + @classmethod + @contextmanager + def private_endpoint(cls, config: ZMQConfig) -> Iterator[ZMQConfig]: + """TCP ports are one namespace per host: take the first unbound pair. + + Another process can still bind the pair before the caller does. The + client's launch check then refuses that server instead of using it. + """ + from .transport import DataControlPortPairAuthority + + pair = DataControlPortPairAuthority.acquire(config, transport_mode=cls.mode) + yield replace(config, default_port=pair.data_port) + class IpcTransportDeclaration(TransportDeclaration): """Complete IPC endpoint semantics.""" @@ -435,7 +458,7 @@ def socket_path( ) -> Path | None: if not cls.is_supported(): return None - ipc_dir = Path.home() / f".{config.app_name}" / config.ipc_socket_dir + ipc_dir = config.socket_root() / config.ipc_socket_dir socket_name = f"{config.ipc_socket_prefix}-{port}{config.ipc_socket_extension}" return ipc_dir / socket_name @@ -561,3 +584,10 @@ def data_control_pair_is_available( return not any( cls.endpoint_in_use(port, host, config) for port in (data_port, control_port) ) + + @classmethod + @contextmanager + def private_endpoint(cls, config: ZMQConfig) -> Iterator[ZMQConfig]: + """A new socket directory, removed when the caller is done with it.""" + with tempfile.TemporaryDirectory(prefix=f"{config.app_name}-") as directory: + yield replace(config, ipc_socket_dir=directory) diff --git a/tests/endpoint_launch_processes.py b/tests/endpoint_launch_processes.py new file mode 100644 index 0000000..bfe42ca --- /dev/null +++ b/tests/endpoint_launch_processes.py @@ -0,0 +1,105 @@ +"""Processes the endpoint-launch tests start: a server, and an owner that starts one.""" + +from __future__ import annotations + +import subprocess +import sys +import time +from dataclasses import replace + +from zmqruntime.config import TransportMode, ZMQConfig +from zmqruntime.execution import ExecutionClient, ExecutionServer +from zmqruntime.messages import EndpointLaunch, ExecuteRequest, ProcessIdentity +from zmqruntime.runner import serve_forever + +TRANSPORT_MODE = TransportMode.default() + + +def address_config(socket_directory: str, port: int) -> ZMQConfig: + """One address in the namespace a test shares between its clients.""" + + return ZMQConfig(ipc_socket_dir=socket_directory, default_port=port) + + +class LaunchedServer(ExecutionServer): + def execute_task(self, execution_id: str, request: ExecuteRequest): + del execution_id, request + return {} + + +class LaunchingClient(ExecutionClient): + """A client whose code root and owner a test may replace to stand for another.""" + + def __init__( + self, + config: ZMQConfig, + *, + persistent: bool = False, + code_root: str | None = None, + owner: ProcessIdentity | None = None, + ) -> None: + super().__init__( + port=config.default_port, + host="127.0.0.1", + persistent=persistent, + transport_mode=TRANSPORT_MODE, + config=config, + ) + self._code_root = code_root + self._owner = owner + + def launch(self) -> EndpointLaunch: + launch = super().launch() + return replace( + launch, + code_root=self._code_root or launch.code_root, + owner=self._owner or launch.owner, + ) + + def _spawn_server_process(self) -> subprocess.Popen: + return subprocess.Popen( + [ + sys.executable, + __file__, + "serve", + self.config.ipc_socket_dir, + str(self.port), + self.launch().to_argument(), + ] + ) + + def send_data(self, data) -> None: + del data + + def serialize_task(self, task, config) -> ExecuteRequest: + del config + return ExecuteRequest(subject_id="subject", pipeline_code=str(task)) + + +def serve(socket_directory: str, port: str, launch: str) -> None: + config = address_config(socket_directory, int(port)) + serve_forever( + LaunchedServer( + port=config.default_port, + host="127.0.0.1", + transport_mode=TRANSPORT_MODE, + config=config, + ).launched_by(EndpointLaunch.from_argument(launch)) + ) + + +def own(socket_directory: str, port: str, persistence: str) -> None: + """Start a server, name its process on stdout, then wait to be killed.""" + + client = LaunchingClient( + address_config(socket_directory, int(port)), + persistent=persistence == "persistent", + ) + if not client.connect(timeout=30): + raise SystemExit("owner could not start its server") + print(client.connected_endpoint.process_identity.pid, flush=True) + time.sleep(600) + + +if __name__ == "__main__": + {"serve": serve, "own": own}[sys.argv[1]](*sys.argv[2:]) diff --git a/tests/test_endpoint_launch.py b/tests/test_endpoint_launch.py new file mode 100644 index 0000000..c0a8633 --- /dev/null +++ b/tests/test_endpoint_launch.py @@ -0,0 +1,224 @@ +"""A client reaches only a server it may use, and an orphaned server exits.""" + +from __future__ import annotations + +import subprocess +import sys +import time +from contextlib import ExitStack +from dataclasses import replace + +import psutil +import pytest +from endpoint_launch_processes import ( + TRANSPORT_MODE, + LaunchedServer, + LaunchingClient, + address_config, +) + +import endpoint_launch_processes +from zmqruntime import DataControlPortPairAuthority, EndpointShutdownMode, ZMQConfig +from zmqruntime.messages import ( + AnySessionAudience, + EndpointAudience, + EndpointLaunch, + EndpointLaunchMismatchError, + OwnerOnlyAudience, + ProcessIdentity, +) + +ANOTHER_PROCESS = ProcessIdentity(pid=1, create_time=0.0) + + +def _stop_server(owner: LaunchingClient) -> None: + """Stop the server a client started, persistent or not.""" + + owner.disconnect() + owner.shutdown_endpoint_on_port( + owner.port, + mode=EndpointShutdownMode.FORCE, + transport_mode=TRANSPORT_MODE, + host=owner.host, + config=owner.config, + ) + + +@pytest.fixture +def two_addresses(): + """Two free addresses in one namespace private to the test.""" + + declaration = TRANSPORT_MODE.declaration + with declaration.private_endpoint(ZMQConfig(default_port=41000)) as first: + second_port = DataControlPortPairAuthority.acquire( + first, + transport_mode=TRANSPORT_MODE, + excluded=frozenset( + (first.default_port, first.default_port + first.control_port_offset) + ), + ).data_port + yield first, address_config(first.ipc_socket_dir, second_port) + + +def test_each_client_reaches_only_the_server_from_its_own_code_root(two_addresses): + first_address, second_address = two_addresses + with ExitStack() as clients: + first = LaunchingClient(first_address, code_root="/checkout-a") + second = LaunchingClient(second_address, code_root="/checkout-b") + for client in (first, second): + clients.callback(client.disconnect) + assert client.connect(timeout=30) + + assert first.connected_endpoint.launch.code_root == "/checkout-a" + assert second.connected_endpoint.launch.code_root == "/checkout-b" + assert ( + first.connected_endpoint.process_identity + != second.connected_endpoint.process_identity + ) + + stray = LaunchingClient(first_address, code_root="/checkout-b") + clients.callback(stray.disconnect) + with pytest.raises(EndpointLaunchMismatchError) as refused: + stray.connect(timeout=5) + + assert "/checkout-a" in str(refused.value) + assert "/checkout-b" in str(refused.value) + assert not stray.is_connected() + assert first.owned_server_process_is_alive() is True + assert first.endpoint.ping(first_address) is not None + + +@pytest.mark.parametrize( + ("persistent", "admitted"), + [(False, False), (True, True)], + ids=["owner_only", "any_session"], +) +def test_only_a_persistent_server_admits_another_process( + two_addresses, persistent, admitted +): + address, _ = two_addresses + owner = LaunchingClient(address, persistent=persistent) + try: + assert owner.connect(timeout=30) + server = owner.connected_endpoint.process_identity + + same_process = LaunchingClient(address, persistent=persistent) + assert same_process.connect_existing(timeout=5) + same_process.disconnect() + + other_process = LaunchingClient( + address, persistent=persistent, owner=ANOTHER_PROCESS + ) + if admitted: + assert other_process.connect_existing(timeout=5) + assert other_process.connected_endpoint.process_identity == server + other_process.disconnect() + else: + with pytest.raises(EndpointLaunchMismatchError, match="for its own use"): + other_process.connect_existing(timeout=5) + finally: + _stop_server(owner) + + +def test_a_server_that_reports_no_launch_is_refused(two_addresses, monkeypatch): + address, _ = two_addresses + owner = LaunchingClient(address, persistent=True) + try: + assert owner.connect(timeout=30) + unreported = LaunchingClient(address, persistent=True) + observed = unreported._try_connect_to_existing + monkeypatch.setattr( + unreported, + "_try_connect_to_existing", + lambda *args, **kwargs: replace(observed(*args, **kwargs), launch=None), + ) + with pytest.raises(EndpointLaunchMismatchError, match="does not report"): + unreported.connect_existing(timeout=5) + finally: + _stop_server(owner) + + +@pytest.mark.parametrize("persistence", ["owner_only", "persistent"]) +def test_a_server_exits_after_its_owner_is_killed(two_addresses, persistence): + address, _ = two_addresses + owner = subprocess.Popen( + [ + sys.executable, + endpoint_launch_processes.__file__, + "own", + address.ipc_socket_dir, + str(address.default_port), + persistence, + ], + stdout=subprocess.PIPE, + text=True, + ) + server = None + try: + server = psutil.Process(int(owner.stdout.readline())) + assert server.is_running() + + owner.kill() + owner.wait(timeout=10) + + server.wait(timeout=15) + finally: + owner.kill() + if server is not None and server.is_running(): + server.kill() + pytest.fail("the server outlived its owner") + + +def _orphaned_server(audience: type[EndpointAudience]) -> LaunchedServer: + server = LaunchedServer(port=41999).launched_by( + EndpointLaunch( + code_root="/checkout-a", + owner=ProcessIdentity(pid=ProcessIdentity.current().pid, create_time=0.0), + audience=audience.wire_name, + ) + ) + server._running = True + return server + + +def _check_owner_now(server: LaunchedServer) -> None: + server._owner_checked_at = time.monotonic() - server.orphan_check_interval_seconds + server._stop_if_orphaned() + + +def test_an_orphaned_owner_only_server_stops_with_work_unfinished(monkeypatch): + server = _orphaned_server(OwnerOnlyAudience) + monkeypatch.setattr(server, "has_unfinished_work", lambda: True) + + _check_owner_now(server) + + assert not server.is_running() + + +def test_an_orphaned_persistent_server_finishes_its_work_first(monkeypatch): + server = _orphaned_server(AnySessionAudience) + monkeypatch.setattr(server, "has_unfinished_work", lambda: True) + _check_owner_now(server) + assert server.is_running() + + monkeypatch.setattr(server, "has_unfinished_work", lambda: False) + _check_owner_now(server) + assert not server.is_running() + + +def test_a_server_whose_owner_is_alive_keeps_running(): + server = LaunchedServer(port=41998) + server._running = True + + _check_owner_now(server) + + assert server.launch.owner == ProcessIdentity.current() + assert server.is_running() + + +def test_every_audience_decodes_from_its_wire_name(): + for wire_name, audience in EndpointAudience.__registry__.items(): + launch = EndpointLaunch.by_this_process(LaunchedServer.code_root(), audience) + assert launch.audience == wire_name + assert EndpointLaunch.from_argument(launch.to_argument()) == launch + assert launch.audience_type is audience diff --git a/tests/test_execution.py b/tests/test_execution.py index 0b2661e..b469abc 100644 --- a/tests/test_execution.py +++ b/tests/test_execution.py @@ -741,6 +741,7 @@ def _try_connect_to_existing(self, port: int, timeout_ms: int = 500): ready=True, server=type(self).__name__, server_role=ServerRole.EXECUTION, + launch=self.launch(), ) def _kill_processes_on_port(self, port: int): @@ -866,6 +867,7 @@ def test_attach_existing_connects_a_ready_endpoint_without_spawning(monkeypatch) ready=True, server="EndpointPolicyExecutionClient", server_role=ServerRole.EXECUTION, + launch=client.launch(), ) monkeypatch.setattr( client, @@ -1200,6 +1202,7 @@ def _try_connect_to_existing(self, port: int, timeout_ms: int = 500): ready=True, server="ConcurrentStartupExecutionClient", server_role=ServerRole.EXECUTION, + launch=self.launch(), ) def _spawn_server_process(self): diff --git a/tests/test_messages.py b/tests/test_messages.py index aeacddc..bf5033a 100644 --- a/tests/test_messages.py +++ b/tests/test_messages.py @@ -220,6 +220,7 @@ def _declared_message_examples(): queued, ProcessResourceUsage(1.0, 2.0), EndpointApplication("app", "1.0"), + messages.EndpointLaunch("/code", identity, "owner_only"), messages.ImageTransferIdentity("image-1", route, identity), messages.ImageAck("image-1", 1, "viewer", return_route=route, producer=identity), messages.TaskProgress("e-1", "running", "running", 10.0, 1.0, 1, 10, subject_id="s"), @@ -246,6 +247,7 @@ def _declared_message_examples(): process_identity=identity, process_usage=ProcessResourceUsage(1.0, 2.0), application=EndpointApplication("app", "1.0"), + launch=messages.EndpointLaunch("/code", identity, "any_session"), ), ControlRequestHeader(ControlMessageType.PING), messages.EndpointShutdownRequest(ControlMessageType.SHUTDOWN, identity), diff --git a/tests/test_startup.py b/tests/test_startup.py index dc3217c..0c37cc1 100644 --- a/tests/test_startup.py +++ b/tests/test_startup.py @@ -106,6 +106,7 @@ def _wait_for_endpoint_ready( ready=True, server=type(self).__name__, server_role=ServerRole.GENERIC, + launch=self.launch(), ) def _setup_client_sockets(self) -> None: diff --git a/tests/test_viewer_protocol.py b/tests/test_viewer_protocol.py index 976f377..dfb5e7c 100644 --- a/tests/test_viewer_protocol.py +++ b/tests/test_viewer_protocol.py @@ -4,7 +4,12 @@ import pytest from zmqruntime import TransportEndpoint, TransportMode, ZMQConfig -from zmqruntime.messages import ProcessResourceUsage, ServerRole +from zmqruntime.messages import ( + AnySessionAudience, + EndpointLaunch, + ProcessResourceUsage, + ServerRole, +) from zmqruntime.streaming import StreamingVisualizerServer from zmqruntime.viewer_protocol import ( ViewerBatchDisplayPayload, @@ -208,6 +213,7 @@ def test_streaming_server_inheritance_owns_viewer_process_usage(monkeypatch): server._ready = True server.log_file_path = None server.application = None + server.launch = EndpointLaunch.by_this_process(server.code_root(), AnySessionAudience) pong = server._create_pong_response()