diff --git a/docs/source/api/viewer_protocol.rst b/docs/source/api/viewer_protocol.rst index 22724a6..68953e5 100644 --- a/docs/source/api/viewer_protocol.rst +++ b/docs/source/api/viewer_protocol.rst @@ -13,6 +13,6 @@ axis values. The protocol classes preserve, validate, and transmit those declared values; they do not inspect image arrays to infer semantic axes. .. automodule:: zmqruntime.viewer_protocol - :members: ViewerProtocolStatus, ViewerControlResponseField, ViewerBatchMessageType, ViewerWireField, ViewerWirePayload, ViewerSourceSpatialWireField, ViewerComponentMode, ViewerComponentModeGroups, ViewerBatchMessageExtraPayload, ViewerSourceSpatialDomainPayload, ViewerBatchMessageWirePayload, viewer_wire_key, ViewerBatchDisplayPayload, viewer_component_mode_value, ViewerBatchItemPayload, ViewerComponentMetadataPayload, ViewerBatchMessagePayload, ViewerTransportEndpoint, ViewerControlReplyHeader, ViewerControlReplyPayload, ViewerAckResponsePayload, ViewerAckPolicy + :members: :member-order: bysource :show-inheritance: diff --git a/pyproject.toml b/pyproject.toml index ffee1f3..d171997 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "zmqruntime" -version = "0.6.1" +version = "0.6.2" 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/transport_modes.py b/src/zmqruntime/transport_modes.py index 21c0daf..8a388cc 100644 --- a/src/zmqruntime/transport_modes.py +++ b/src/zmqruntime/transport_modes.py @@ -4,6 +4,7 @@ import platform import json +import os import socket from abc import ABC, abstractmethod from collections.abc import Iterator @@ -11,7 +12,7 @@ from ipaddress import ip_address from pathlib import Path from dataclasses import replace -from typing import ClassVar +from typing import BinaryIO, ClassVar import portalocker import zmq @@ -21,6 +22,11 @@ from .timeouts import OperationCancellation, OperationDeadline from .messages import ProcessIdentity +# Startup lock handles this process holds, by lock path. Windows locks are +# mandatory, so the owner record inside a held lock file is readable and +# writable only through the handle that holds the lock. +_HELD_STARTUP_LOCKS: dict[Path, BinaryIO] = {} + class TransportDeclaration(ABC, metaclass=AutoRegisterMeta): """Complete behavior owned by one transport implementation.""" @@ -136,7 +142,9 @@ def startup_lock( lock_path = cls.startup_lock_path(port, config) lock_path.parent.mkdir(parents=True, exist_ok=True) - with lock_path.open("a+b") as lock_file: + # Read/write without append: the owner record is rewritten in place. + descriptor = os.open(lock_path, os.O_RDWR | os.O_CREAT | getattr(os, "O_BINARY", 0)) + with os.fdopen(descriptor, "r+b") as lock_file: while True: if cancellation.requested(): yield False @@ -157,19 +165,39 @@ def startup_lock( operation_deadline.remaining_seconds(), ) cancellation.wait(wait_seconds) + _HELD_STARTUP_LOCKS[lock_path] = lock_file try: yield True finally: + del _HELD_STARTUP_LOCKS[lock_path] portalocker.unlock(lock_file) + @classmethod + @contextmanager + def _startup_record(cls, port: int, config: ZMQConfig) -> Iterator[BinaryIO | None]: + """Yield the owner record stream: the held lock handle when this process holds it.""" + path = cls.startup_lock_path(port, config) + held = _HELD_STARTUP_LOCKS.get(path) + if held is not None: + yield held + return + if not path.exists(): + yield None + return + with path.open("r+b") as stream: + yield stream + @classmethod def startup_owner(cls, port: int, config: ZMQConfig) -> ProcessIdentity | None: """Read the exact pre-bind reservation from the existing startup lock. Call under startup_lock. Invalid records fail closed, never as absence. """ - path = cls.startup_lock_path(port, config) - source = path.read_text(encoding="utf-8") if path.exists() else "" + with cls._startup_record(port, config) as stream: + if stream is None: + return None + stream.seek(0) + source = stream.read().decode("utf-8") return ProcessIdentity.from_dict(json.loads(source)) if source else None @classmethod @@ -180,7 +208,9 @@ def record_startup_owner( owner: ProcessIdentity, ) -> None: """Publish an exact invoker or child reservation under the held startup lock.""" - with cls.startup_lock_path(port, config).open("r+b") as stream: + with cls._startup_record(port, config) as stream: + if stream is None: + raise FileNotFoundError(cls.startup_lock_path(port, config)) stream.seek(0) stream.write(json.dumps(owner.to_dict()).encode("utf-8")) stream.truncate() @@ -204,7 +234,7 @@ def release_startup_owner( return False if recorded_owner != owner: return False - with cls.startup_lock_path(port, config).open("r+b") as stream: + with cls._startup_record(port, config) as stream: stream.truncate(0) stream.flush() return True diff --git a/tests/test_owned_startup.py b/tests/test_owned_startup.py index 96359d2..250e0ca 100644 --- a/tests/test_owned_startup.py +++ b/tests/test_owned_startup.py @@ -415,3 +415,30 @@ def test_provisional_release_preserves_unproved_or_changed_owner( if record == "unknown": with pytest.raises(ValueError): declaration.startup_owner(client.port, client.config) + + +@pytest.mark.parametrize("mode", [TransportMode.TCP, TransportMode.IPC]) +def test_owner_record_uses_only_the_held_lock_handle(client, monkeypatch, mode): + """Windows locks are mandatory: a second handle on the held lock file is denied.""" + declaration = mode.declaration + paths = client.transport_mode.declaration.startup_lock_path + monkeypatch.setattr(declaration, "startup_lock_path", paths) + path = declaration.startup_lock_path(client.port, client.config) + owner = ProcessIdentity.current() + original_open = type(path).open + + def denied(self, *args, **kwargs): + if self == path: + raise PermissionError(13, "Permission denied", str(self)) + return original_open(self, *args, **kwargs) + + with declaration.startup_lock(client.port, client.config, None, OperationCancellation()): + with monkeypatch.context() as mandatory: + mandatory.setattr(type(path), "open", denied) + assert declaration.startup_owner(client.port, client.config) is None + declaration.record_startup_owner(client.port, client.config, owner) + assert declaration.startup_owner(client.port, client.config) == owner + assert declaration.release_startup_owner(client.port, client.config, owner) is True + assert declaration.startup_owner(client.port, client.config) is None + declaration.record_startup_owner(client.port, client.config, owner) + assert declaration.startup_owner(client.port, client.config) == owner