gisolate 0.2.15__tar.gz → 0.2.16.dev1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/PKG-INFO +61 -1
  2. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/README.md +60 -0
  3. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/pyproject.toml +1 -1
  4. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/__init__.py +5 -1
  5. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/_internal.py +15 -2
  6. gisolate-0.2.16.dev1/src/gisolate/pubsub.py +437 -0
  7. gisolate-0.2.16.dev1/tests/test_pubsub.py +710 -0
  8. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/uv.lock +1 -1
  9. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/.code-review-graph/.gitignore +0 -0
  10. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/.code-review-graph/graph.db +0 -0
  11. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/.envrc +0 -0
  12. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/.gitignore +0 -0
  13. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/CLAUDE.md +0 -0
  14. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/LICENSE +0 -0
  15. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/lefthook.yml +0 -0
  16. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/_workers.py +0 -0
  17. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/bridge.py +0 -0
  18. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/hub.py +0 -0
  19. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/local.py +0 -0
  20. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/proxy.py +0 -0
  21. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/src/gisolate/subprocess.py +0 -0
  22. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/__init__.py +0 -0
  23. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/__init__.py +0 -0
  24. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/features/process_proxy.feature +0 -0
  25. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/features/serialization.feature +0 -0
  26. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/features/subprocess_run.feature +0 -0
  27. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/features/thread_local.feature +0 -0
  28. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/test_proxy.py +0 -0
  29. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/test_serialization.py +0 -0
  30. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/test_subprocess.py +0 -0
  31. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/bdd/test_thread_local.py +0 -0
  32. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/conftest.py +0 -0
  33. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/helpers.py +0 -0
  34. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_bridge.py +0 -0
  35. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_hub.py +0 -0
  36. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_internal.py +0 -0
  37. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_local.py +0 -0
  38. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_proxy.py +0 -0
  39. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_subprocess.py +0 -0
  40. {gisolate-0.2.15 → gisolate-0.2.16.dev1}/tests/test_workers.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gisolate
3
- Version: 0.2.15
3
+ Version: 0.2.16.dev1
4
4
  Summary: Process isolation for gevent applications — run any object in a clean subprocess, call methods transparently via ZMQ IPC.
5
5
  Project-URL: Repository, https://github.com/wy-z/gisolate
6
6
  License-Expression: MIT
@@ -105,6 +105,47 @@ asyncio.run(main())
105
105
  server.close()
106
106
  ```
107
107
 
108
+ ### ProcessPublisher / ProcessSubscriber — one-way fan-out
109
+
110
+ ZMQ PUB/SUB for one-way data streaming (snapshots, signals, heartbeats). Use this when message loss is acceptable; use `ProcessBridge` when you need request/response with delivery guarantees.
111
+
112
+ ```python
113
+ # Producer (gevent side)
114
+ from gisolate import ProcessPublisher
115
+
116
+ pub = ProcessPublisher("ipc:///tmp/stream.sock").start()
117
+ pub.publish("v1.snapshot.AAPL", {"price": 150.0})
118
+ pub.publish("v1.heartbeat.gevent", {"ts_ns": 1234567890})
119
+ pub.close()
120
+
121
+ # Consumer (asyncio side)
122
+ import asyncio
123
+ from gisolate import ProcessSubscriber
124
+
125
+ async def main():
126
+ sub = ProcessSubscriber("ipc:///tmp/stream.sock")
127
+
128
+ async def on_snapshot(topic, payload):
129
+ print(topic, payload)
130
+
131
+ async def on_heartbeat(topic, payload):
132
+ print("heartbeat", payload)
133
+
134
+ sub.subscribe("v1.snapshot.", on_snapshot)
135
+ sub.subscribe("v1.heartbeat.", on_heartbeat)
136
+ sub.start()
137
+ await asyncio.sleep(10)
138
+ await sub.close()
139
+
140
+ asyncio.run(main())
141
+ ```
142
+
143
+ Notes:
144
+ - **Topic prefix matching** — `sub.subscribe("v1.snapshot.", h)` receives every topic starting with that prefix.
145
+ - **Multiple handlers per prefix** — invoked concurrently with `asyncio.gather`. Exceptions in one handler do not kill the reader task.
146
+ - **Lossy by design** — `publish` is non-blocking; messages are dropped when the send queue is full (slow subscriber). Set `sndhwm=` to tune.
147
+ - **Pluggable serializer** — defaults to `SmartPickle`. Pass any object implementing the `Serializer` protocol (`dumps`/`loads`) to use msgpack, JSON, etc.
148
+
108
149
  ### ThreadLocalProxy — per-thread instances
109
150
 
110
151
  Thread-local proxy using unpatched `threading.local` for true isolation in `gevent.threadpool`:
@@ -153,6 +194,25 @@ Run a function in an isolated subprocess. Blocks with gevent-safe polling.
153
194
  - **`await bridge.call(func, *args, timeout=60, **kwargs)`** — async RPC call (client mode)
154
195
  - **`bridge.close()`** — cleanup resources
155
196
 
197
+ ### `ProcessPublisher(address, *, serializer=SmartPickle, sndhwm=1000)`
198
+
199
+ - **`pub.start()`** — bind the PUB socket (idempotent, returns self)
200
+ - **`pub.publish(topic, payload)`** — non-blocking publish; drops on slow consumers
201
+ - **`pub.close()`** — cleanup (idempotent)
202
+ - Supports context manager (`with` statement)
203
+
204
+ ### `ProcessSubscriber(address, *, serializer=SmartPickle)`
205
+
206
+ - **`sub.subscribe(topic_prefix, handler)`** — register an async handler for a topic prefix
207
+ - **`sub.unsubscribe(topic_prefix, handler=None)`** — remove a handler or all handlers for a prefix
208
+ - **`sub.start()`** — connect and spawn the reader task (idempotent, returns self)
209
+ - **`await sub.close()`** — cancel reader and cleanup (idempotent)
210
+ - Supports async context manager (`async with` statement)
211
+
212
+ ### `Serializer` (Protocol)
213
+
214
+ Anything with `dumps(obj) -> bytes` and `loads(bytes) -> obj` static methods can be used as a serializer for `ProcessPublisher` / `ProcessSubscriber`. Default is `SmartPickle` (pickle, falling back to dill).
215
+
156
216
  ### `ThreadLocalProxy(factory)`
157
217
 
158
218
  Transparent proxy delegating attribute access to a per-thread instance.
@@ -84,6 +84,47 @@ asyncio.run(main())
84
84
  server.close()
85
85
  ```
86
86
 
87
+ ### ProcessPublisher / ProcessSubscriber — one-way fan-out
88
+
89
+ ZMQ PUB/SUB for one-way data streaming (snapshots, signals, heartbeats). Use this when message loss is acceptable; use `ProcessBridge` when you need request/response with delivery guarantees.
90
+
91
+ ```python
92
+ # Producer (gevent side)
93
+ from gisolate import ProcessPublisher
94
+
95
+ pub = ProcessPublisher("ipc:///tmp/stream.sock").start()
96
+ pub.publish("v1.snapshot.AAPL", {"price": 150.0})
97
+ pub.publish("v1.heartbeat.gevent", {"ts_ns": 1234567890})
98
+ pub.close()
99
+
100
+ # Consumer (asyncio side)
101
+ import asyncio
102
+ from gisolate import ProcessSubscriber
103
+
104
+ async def main():
105
+ sub = ProcessSubscriber("ipc:///tmp/stream.sock")
106
+
107
+ async def on_snapshot(topic, payload):
108
+ print(topic, payload)
109
+
110
+ async def on_heartbeat(topic, payload):
111
+ print("heartbeat", payload)
112
+
113
+ sub.subscribe("v1.snapshot.", on_snapshot)
114
+ sub.subscribe("v1.heartbeat.", on_heartbeat)
115
+ sub.start()
116
+ await asyncio.sleep(10)
117
+ await sub.close()
118
+
119
+ asyncio.run(main())
120
+ ```
121
+
122
+ Notes:
123
+ - **Topic prefix matching** — `sub.subscribe("v1.snapshot.", h)` receives every topic starting with that prefix.
124
+ - **Multiple handlers per prefix** — invoked concurrently with `asyncio.gather`. Exceptions in one handler do not kill the reader task.
125
+ - **Lossy by design** — `publish` is non-blocking; messages are dropped when the send queue is full (slow subscriber). Set `sndhwm=` to tune.
126
+ - **Pluggable serializer** — defaults to `SmartPickle`. Pass any object implementing the `Serializer` protocol (`dumps`/`loads`) to use msgpack, JSON, etc.
127
+
87
128
  ### ThreadLocalProxy — per-thread instances
88
129
 
89
130
  Thread-local proxy using unpatched `threading.local` for true isolation in `gevent.threadpool`:
@@ -132,6 +173,25 @@ Run a function in an isolated subprocess. Blocks with gevent-safe polling.
132
173
  - **`await bridge.call(func, *args, timeout=60, **kwargs)`** — async RPC call (client mode)
133
174
  - **`bridge.close()`** — cleanup resources
134
175
 
176
+ ### `ProcessPublisher(address, *, serializer=SmartPickle, sndhwm=1000)`
177
+
178
+ - **`pub.start()`** — bind the PUB socket (idempotent, returns self)
179
+ - **`pub.publish(topic, payload)`** — non-blocking publish; drops on slow consumers
180
+ - **`pub.close()`** — cleanup (idempotent)
181
+ - Supports context manager (`with` statement)
182
+
183
+ ### `ProcessSubscriber(address, *, serializer=SmartPickle)`
184
+
185
+ - **`sub.subscribe(topic_prefix, handler)`** — register an async handler for a topic prefix
186
+ - **`sub.unsubscribe(topic_prefix, handler=None)`** — remove a handler or all handlers for a prefix
187
+ - **`sub.start()`** — connect and spawn the reader task (idempotent, returns self)
188
+ - **`await sub.close()`** — cancel reader and cleanup (idempotent)
189
+ - Supports async context manager (`async with` statement)
190
+
191
+ ### `Serializer` (Protocol)
192
+
193
+ Anything with `dumps(obj) -> bytes` and `loads(bytes) -> obj` static methods can be used as a serializer for `ProcessPublisher` / `ProcessSubscriber`. Default is `SmartPickle` (pickle, falling back to dill).
194
+
135
195
  ### `ThreadLocalProxy(factory)`
136
196
 
137
197
  Transparent proxy delegating attribute access to a per-thread instance.
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "gisolate"
3
- version = "0.2.15"
3
+ version = "0.2.16.dev1"
4
4
  description = "Process isolation for gevent applications — run any object in a clean subprocess, call methods transparently via ZMQ IPC."
5
5
  readme = "README.md"
6
6
  license = "MIT"
@@ -4,11 +4,12 @@ Run any object in a clean subprocess, call methods transparently via ZMQ IPC.
4
4
  Isolates libraries incompatible with gevent monkey-patching.
5
5
  """
6
6
 
7
- from ._internal import ProcessError, RemoteError
7
+ from ._internal import ProcessError, RemoteError, Serializer
8
8
  from .bridge import ProcessBridge
9
9
  from .hub import ensure_hub_started, shutdown as shutdown_hub, spawn_on_main_hub
10
10
  from .local import ThreadLocalProxy
11
11
  from .proxy import ProcessProxy, get_default_mp_context, set_default_mp_context
12
+ from .pubsub import ProcessPublisher, ProcessSubscriber
12
13
  from .subprocess import run_in_subprocess
13
14
 
14
15
  # Pre-initialize threadpoolctl on main thread to cache library info.
@@ -24,7 +25,10 @@ __all__ = [
24
25
  "ProcessBridge",
25
26
  "ProcessError",
26
27
  "ProcessProxy",
28
+ "ProcessPublisher",
29
+ "ProcessSubscriber",
27
30
  "RemoteError",
31
+ "Serializer",
28
32
  "ThreadLocalProxy",
29
33
  "ensure_hub_started",
30
34
  "get_default_mp_context",
@@ -4,7 +4,7 @@ import contextlib
4
4
  import io
5
5
  import logging
6
6
  import pickle
7
- from typing import Any
7
+ from typing import Any, Protocol
8
8
 
9
9
  import dill
10
10
  import gevent.monkey
@@ -44,8 +44,21 @@ class RemoteError(RuntimeError):
44
44
  # ---------------------------------------------------------------------------
45
45
 
46
46
 
47
+ class Serializer(Protocol):
48
+ """Pluggable serializer protocol: ``dumps(obj) -> bytes`` / ``loads(bytes) -> obj``."""
49
+
50
+ @staticmethod
51
+ def dumps(obj: Any) -> bytes: ...
52
+
53
+ @staticmethod
54
+ def loads(data: bytes) -> Any: ...
55
+
56
+
47
57
  class SmartPickle:
48
- """Serializer preferring pickle, falling back to dill. Learns from failures."""
58
+ """Serializer preferring pickle, falling back to dill. Learns from failures.
59
+
60
+ Implements the :class:`Serializer` protocol.
61
+ """
49
62
 
50
63
  _PICKLE = b"P"
51
64
  _DILL = b"D"
@@ -0,0 +1,437 @@
1
+ """ProcessPublisher / ProcessSubscriber: ZMQ PUB/SUB one-way fan-out.
2
+
3
+ Use ``ProcessBridge`` when you need request/response RPC.
4
+ Use these when you need one-way fan-out (snapshots, signals, heartbeats) where
5
+ message loss is acceptable — PUB drops messages for slow subscribers.
6
+ """
7
+
8
+ import asyncio
9
+ import contextlib
10
+ import logging
11
+ import os
12
+ from typing import Any, Awaitable, Callable
13
+
14
+ from ._internal import Serializer, SmartPickle
15
+
16
+ log = logging.getLogger(__name__)
17
+
18
+ Handler = Callable[[str, Any], Awaitable[None]]
19
+
20
+
21
+ def _safe_close(sock: Any, ctx: Any) -> None:
22
+ """Best-effort tear down of a ZMQ socket + context. Swallows errors."""
23
+ with contextlib.suppress(Exception):
24
+ if sock is not None:
25
+ sock.close(linger=0)
26
+ with contextlib.suppress(Exception):
27
+ if ctx is not None:
28
+ ctx.term()
29
+
30
+
31
+ # ---------------------------------------------------------------------------
32
+ # Publisher (gevent side)
33
+ # ---------------------------------------------------------------------------
34
+
35
+
36
+ class ProcessPublisher:
37
+ """ZMQ PUB socket for one-way fan-out from a gevent producer.
38
+
39
+ Topic-based dispatch with a pluggable serializer (default
40
+ :class:`SmartPickle`). ``publish`` is non-blocking — slow subscribers
41
+ cause messages to be dropped once the high-water mark is hit, matching
42
+ standard PUB semantics.
43
+
44
+ Args:
45
+ address: IPC/TCP address (e.g., ``"ipc:///tmp/stream.sock"``).
46
+ serializer: Optional serializer; defaults to :class:`SmartPickle`.
47
+ sndhwm: Send high-water mark. Beyond this, messages are dropped.
48
+
49
+ Example::
50
+
51
+ pub = ProcessPublisher("ipc:///tmp/stream.sock").start()
52
+ pub.publish("v1.snapshot.AAPL", {"price": 150.0})
53
+ pub.close()
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ address: str,
59
+ *,
60
+ serializer: Serializer = SmartPickle,
61
+ sndhwm: int = 1000,
62
+ ):
63
+ self._addr = address
64
+ self._serializer = serializer
65
+ self._sndhwm = sndhwm
66
+ self._started = False
67
+ self._sock: Any = None
68
+ self._ctx: Any = None
69
+ self._send_lock: Any = None
70
+
71
+ def __del__(self):
72
+ # Best-effort sync cleanup. Avoid full close() — at GC time the
73
+ # gevent hub may be torn down, so acquiring the send-lock can fail.
74
+ # The ipc file is left on disk (process is exiting anyway).
75
+ _safe_close(getattr(self, "_sock", None), getattr(self, "_ctx", None))
76
+
77
+ def __enter__(self) -> "ProcessPublisher":
78
+ return self.start()
79
+
80
+ def __exit__(self, *_) -> None:
81
+ self.close()
82
+
83
+ @property
84
+ def address(self) -> str:
85
+ """IPC/TCP address."""
86
+ return self._addr
87
+
88
+ def start(self) -> "ProcessPublisher":
89
+ """Bind the PUB socket. Idempotent. Returns self for chaining."""
90
+ if self._started:
91
+ return self
92
+
93
+ import gevent.lock
94
+ import zmq.green as zmq_mod
95
+
96
+ ctx = zmq_mod.Context()
97
+ sock = None
98
+ try:
99
+ sock = ctx.socket(zmq_mod.PUB)
100
+ sock.setsockopt(zmq_mod.LINGER, 0)
101
+ sock.setsockopt(zmq_mod.SNDHWM, self._sndhwm)
102
+ sock.bind(self._addr)
103
+ except Exception:
104
+ _safe_close(sock, ctx)
105
+ raise
106
+
107
+ self._ctx = ctx
108
+ self._sock = sock
109
+ self._send_lock = gevent.lock.Semaphore()
110
+ self._started = True
111
+ return self
112
+
113
+ def publish(self, topic: str, payload: Any) -> None:
114
+ """Publish ``payload`` under ``topic``. Non-blocking.
115
+
116
+ Drops the message silently if the send queue is full (slow subscribers).
117
+ Safe to call concurrently from multiple greenlets.
118
+ """
119
+ if not self._started:
120
+ raise RuntimeError("ProcessPublisher.publish() called before start()")
121
+
122
+ import zmq.green as zmq_mod
123
+
124
+ data = self._serializer.dumps(payload)
125
+ with self._send_lock:
126
+ # Concurrent close() may have torn the socket down between our
127
+ # _started check above and acquiring the lock; re-check.
128
+ if not self._started:
129
+ return
130
+ try:
131
+ self._sock.send_multipart(
132
+ [topic.encode("utf-8"), data], flags=zmq_mod.NOBLOCK
133
+ )
134
+ except zmq_mod.Again:
135
+ # SNDHWM hit — slow subscribers. Drop, matching PUB semantics.
136
+ log.debug("publisher dropped message for topic %s (HWM hit)", topic)
137
+ except zmq_mod.ZMQError as exc:
138
+ log.warning("publisher send failed for topic %s: %s", topic, exc)
139
+
140
+ def close(self) -> None:
141
+ """Tear down the socket. Idempotent."""
142
+ if not self._started:
143
+ return
144
+ # Serialize with publish(): closing while a greenlet is mid-send is
145
+ # undefined. publish() re-checks _started inside the same lock.
146
+ with self._send_lock:
147
+ self._started = False
148
+ _safe_close(self._sock, self._ctx)
149
+ if self._addr.startswith("ipc://"):
150
+ with contextlib.suppress(OSError):
151
+ os.unlink(self._addr[6:])
152
+
153
+
154
+ # ---------------------------------------------------------------------------
155
+ # Subscriber (asyncio side)
156
+ # ---------------------------------------------------------------------------
157
+
158
+
159
+ class ProcessSubscriber:
160
+ """ZMQ SUB socket for asyncio consumers.
161
+
162
+ Register topic-prefix handlers; a single reader task dispatches incoming
163
+ messages. Multiple handlers may share a prefix and are invoked
164
+ concurrently. An exception in one handler is logged but does not kill
165
+ the reader.
166
+
167
+ Args:
168
+ address: IPC/TCP address (e.g., ``"ipc:///tmp/stream.sock"``).
169
+ serializer: Optional serializer; defaults to :class:`SmartPickle`.
170
+
171
+ Example::
172
+
173
+ sub = ProcessSubscriber("ipc:///tmp/stream.sock")
174
+
175
+ async def on_snapshot(topic, payload): ...
176
+
177
+ sub.subscribe("v1.snapshot.", on_snapshot)
178
+ sub.start()
179
+ ...
180
+ await sub.close()
181
+ """
182
+
183
+ def __init__(
184
+ self,
185
+ address: str,
186
+ *,
187
+ serializer: Serializer = SmartPickle,
188
+ ):
189
+ self._addr = address
190
+ self._serializer = serializer
191
+ self._started = False
192
+ self._sock: Any = None
193
+ self._ctx: Any = None
194
+ self._reader_task: Any = None
195
+ # Tasks currently running ProcessSubscriber._invoke. close() uses
196
+ # this to detect "called from a handler my reader is awaiting" and
197
+ # skip the reader-join (which would self-deadlock).
198
+ self._handler_tasks: set[Any] = set()
199
+ self._handlers: dict[str, list[Handler]] = {}
200
+
201
+ def __del__(self):
202
+ # Best-effort sync cleanup from finalizer. Reader task is leaked here;
203
+ # users should ``await close()`` explicitly for deterministic cleanup.
204
+ _safe_close(getattr(self, "_sock", None), getattr(self, "_ctx", None))
205
+
206
+ async def __aenter__(self) -> "ProcessSubscriber":
207
+ return self.start()
208
+
209
+ async def __aexit__(self, *_) -> None:
210
+ await self.close()
211
+
212
+ @property
213
+ def address(self) -> str:
214
+ """IPC/TCP address."""
215
+ return self._addr
216
+
217
+ def start(self) -> "ProcessSubscriber":
218
+ """Connect the SUB socket and spawn the reader task. Idempotent.
219
+
220
+ Must be called with a running asyncio loop. The reader task is bound
221
+ to that loop; subsequent :meth:`subscribe`, :meth:`unsubscribe`, and
222
+ :meth:`close` calls must run on the same loop/thread (ZMQ sockets
223
+ are not thread-safe).
224
+ """
225
+ if self._started:
226
+ return self
227
+
228
+ import zmq.asyncio
229
+
230
+ # Require a running loop *before* allocating ZMQ resources, so a
231
+ # caller misusing the API doesn't leave the subscriber half-built
232
+ # (sock/ctx allocated, _started=True, no reader task).
233
+ loop = asyncio.get_running_loop()
234
+
235
+ # Fresh context per subscriber (mirrors ProcessBridge), so ``close()``
236
+ # fully releases libzmq resources and restart is clean.
237
+ ctx = zmq.asyncio.Context()
238
+ sock = None
239
+ try:
240
+ sock = ctx.socket(zmq.SUB)
241
+ sock.setsockopt(zmq.LINGER, 0)
242
+ sock.connect(self._addr)
243
+ # Re-subscribe to any prefixes registered before start().
244
+ for prefix in self._handlers:
245
+ sock.setsockopt(zmq.SUBSCRIBE, prefix.encode("utf-8"))
246
+ except Exception:
247
+ _safe_close(sock, ctx)
248
+ raise
249
+
250
+ self._ctx = ctx
251
+ self._sock = sock
252
+ self._started = True
253
+ # Bind the reader to *this* socket. A close()+start() restart from
254
+ # inside a handler swaps ``self._sock``; the stale reader must not
255
+ # resume against the new socket (would race the new reader's recv).
256
+ self._reader_task = loop.create_task(self._read_loop(sock))
257
+ return self
258
+
259
+ def subscribe(self, topic_prefix: str, handler: Handler) -> None:
260
+ """Register ``handler`` for messages whose topic starts with ``topic_prefix``.
261
+
262
+ Multiple handlers may share a prefix; they are invoked concurrently.
263
+ Safe to call before or after :meth:`start`.
264
+
265
+ After :meth:`start`, call only from the same thread that owns the
266
+ asyncio loop (ZMQ sockets are not thread-safe).
267
+ """
268
+ new_prefix = topic_prefix not in self._handlers
269
+ self._handlers.setdefault(topic_prefix, []).append(handler)
270
+ if new_prefix and self._started:
271
+ import zmq
272
+
273
+ self._sock.setsockopt(zmq.SUBSCRIBE, topic_prefix.encode("utf-8"))
274
+
275
+ def unsubscribe(
276
+ self, topic_prefix: str, handler: Handler | None = None
277
+ ) -> None:
278
+ """Remove ``handler`` (or all handlers) for ``topic_prefix``.
279
+
280
+ When the last handler for a prefix is removed, the ZMQ-level
281
+ subscription is also dropped. After :meth:`start`, call only from
282
+ the same thread that owns the asyncio loop.
283
+ """
284
+ handlers = self._handlers.get(topic_prefix)
285
+ if not handlers:
286
+ return
287
+ if handler is None:
288
+ handlers.clear()
289
+ else:
290
+ with contextlib.suppress(ValueError):
291
+ handlers.remove(handler)
292
+ if not handlers:
293
+ self._handlers.pop(topic_prefix, None)
294
+ if self._started:
295
+ import zmq
296
+
297
+ with contextlib.suppress(Exception):
298
+ self._sock.setsockopt(
299
+ zmq.UNSUBSCRIBE, topic_prefix.encode("utf-8")
300
+ )
301
+
302
+ def _match(self, topic: str) -> list[Handler]:
303
+ return [
304
+ h
305
+ for prefix, handlers in self._handlers.items()
306
+ if topic.startswith(prefix)
307
+ for h in handlers
308
+ ]
309
+
310
+ async def _read_loop(self, sock: Any) -> None:
311
+ """Single reader task: dispatch messages to matched handlers.
312
+
313
+ ``sock`` is captured at task creation time so a close+restart cycle
314
+ leaves stale readers pointing at the already-closed socket — their
315
+ recv fails immediately and they exit without touching the new socket.
316
+ """
317
+ try:
318
+ while True:
319
+ # Check before each recv: close() or a close+start restart
320
+ # may have swapped ``self._sock`` while we were suspended in
321
+ # the previous gather() (pyzmq does not necessarily wake a
322
+ # pending recv future when the socket is closed, so we can't
323
+ # rely solely on recv raising).
324
+ if not self._started or sock is not self._sock:
325
+ return
326
+ try:
327
+ parts = await sock.recv_multipart()
328
+ except Exception:
329
+ if not self._started or sock is not self._sock:
330
+ return
331
+ raise
332
+ if len(parts) < 2:
333
+ continue
334
+ topic_bytes, data, *_ = parts
335
+ topic = topic_bytes.decode("utf-8", errors="replace")
336
+ handlers = self._match(topic)
337
+ if not handlers:
338
+ continue
339
+ try:
340
+ payload = self._serializer.loads(data)
341
+ except Exception:
342
+ log.exception(
343
+ "subscriber failed to deserialize topic %s", topic
344
+ )
345
+ continue
346
+ # return_exceptions=True isolates the reader from a
347
+ # handler's CancelledError or any BaseException leaking
348
+ # past _invoke.
349
+ results = await asyncio.gather(
350
+ *(self._invoke(h, topic, payload) for h in handlers),
351
+ return_exceptions=True,
352
+ )
353
+ for r in results:
354
+ if isinstance(r, (SystemExit, KeyboardInterrupt)):
355
+ # Honor process-exit intent from a handler.
356
+ raise r
357
+ # CancelledError is BaseException (not Exception) since
358
+ # 3.8, so ``isinstance(r, Exception)`` naturally skips it.
359
+ if isinstance(r, Exception):
360
+ log.error(
361
+ "subscriber handler raised %s for topic %s",
362
+ type(r).__name__,
363
+ topic,
364
+ exc_info=r,
365
+ )
366
+ except asyncio.CancelledError:
367
+ pass
368
+ except Exception:
369
+ log.exception("subscriber reader task crashed")
370
+ # Reader is dying for an unrelated reason (malformed frame,
371
+ # libzmq error, etc.). Tear our resources down so the
372
+ # subscriber doesn't look alive while silently dropping
373
+ # messages, and so ``start()`` can rebuild from scratch.
374
+ if sock is self._sock and self._started:
375
+ self._started = False
376
+ ctx = self._ctx
377
+ self._sock = None
378
+ self._ctx = None
379
+ self._reader_task = None
380
+ _safe_close(sock, ctx)
381
+
382
+ async def _invoke(self, handler: Handler, topic: str, payload: Any) -> None:
383
+ task = asyncio.current_task()
384
+ if task is not None:
385
+ self._handler_tasks.add(task)
386
+ try:
387
+ await handler(topic, payload)
388
+ except Exception:
389
+ log.exception("subscriber handler failed for topic %s", topic)
390
+ finally:
391
+ if task is not None:
392
+ self._handler_tasks.discard(task)
393
+
394
+ async def close(self) -> None:
395
+ """Tear down the socket and join the reader task. Idempotent.
396
+
397
+ Must be awaited from the asyncio loop that owns this subscriber
398
+ (the one that called :meth:`start`). Calling concurrently from
399
+ multiple coroutines is safe; the second caller returns immediately.
400
+
401
+ Safe to call from inside a handler: the reader task is not joined
402
+ in that case (joining yourself would deadlock), and sibling handlers
403
+ in the current dispatch are allowed to finish — we never
404
+ ``task.cancel()`` the reader, so ``asyncio.gather`` is not torn down.
405
+ """
406
+ if not self._started:
407
+ return
408
+ # Snapshot owned resources into locals, then null on self so a
409
+ # concurrent start() cannot have its fresh ctx/sock closed by us.
410
+ self._started = False
411
+ sock, ctx, task = self._sock, self._ctx, self._reader_task
412
+ self._sock = None
413
+ self._ctx = None
414
+ self._reader_task = None
415
+
416
+ # Close socket first. Any in-flight recv fails; the reader sees
417
+ # _started=False and exits cleanly. Avoid task.cancel() — it would
418
+ # propagate through asyncio.gather and cancel sibling handlers
419
+ # mid-execution when close() is called from inside a handler.
420
+ _safe_close(sock, ctx)
421
+
422
+ # Join the reader for deterministic cleanup. Skip when the caller
423
+ # IS the reader, or is a handler the reader is currently awaiting —
424
+ # both deadlock. In those cases the reader exits on its own once
425
+ # gather completes (closed socket + _started=False).
426
+ current = asyncio.current_task()
427
+ if (
428
+ task is not None
429
+ and task is not current
430
+ and current not in self._handler_tasks
431
+ and not task.done()
432
+ ):
433
+ with contextlib.suppress(BaseException):
434
+ await asyncio.wait({task}, timeout=2.0)
435
+
436
+
437
+ __all__ = ["ProcessPublisher", "ProcessSubscriber"]