Skip to content

Commit 505dd80

Browse files
authored
fix: probe kernel readiness without racing ipykernel initialization (ROADMAP.md #82) (#97)
* fix: probe kernel readiness without racing ipykernel initialization (ROADMAP.md #82) start_session waited for readiness with a single kernel_info_request on the control channel, sent the instant kernel.json appeared. ipykernel registers the control dispatcher (kernelbase.py:551) before it assigns _control_lock (:559), so a probe landing in that window raises TraitError kernel-side and is dropped — and the one request we sent was then never answered, burning the whole remaining budget (the 59.5s CI hang). Readiness now rides the shell channel, whose handler is registered only after _control_lock is assigned, so a shell reply proves initialization finished. No execution can be in flight during a start, so the shell channel's busy-serialization caveat does not apply there. The probe loop the two helpers duplicated is extracted into _kernel_info_reply(kc, channel, timeout); _kernel_alive stays the control wrapper and _kernel_idle becomes _kernel_answers_shell. _kernel_dead, is_session_running, session_status's liveness check and exec_code's silence probe still use control: they must answer for a busy kernel, which only control does. Adds a stand-in kernel that answers kernel_info on shell only and leaves control bound but unserviced, which fails in 30s against the old control probe and passes in ~1s against the new one, plus a real-kernel guard that fires a control probe the instant kernel.json appears. * test: deliver the intruder's control probe in the init-race guard stop_channels() closes the client socket with linger=0, which discards anything ZMQ has not yet handed over — and right after send() on loopback the connect handshake is usually still in flight. Measured against a bound ROUTER with the test's exact code shape: the probe arrived 2/20 times, so test_start_session_survives_control_probe_during_init was mostly firing nothing. Holding the socket open until the start finishes takes it to 20/20. Also from review: - the intruder now retries instead of giving up when it reads a half-written connection file, catches only the (OSError, ValueError) _client documents rather than everything, drops the exists()/stat() TOCTOU, and asserts it actually finished; - start_session's comment names the ipykernel version and line numbers the shell-channel ordering was verified against, since pyproject sets no ipykernel ceiling; - the stand-in's cleanup comment gives the real reason conftest's reaper never sees it (argv is `python -c`, not ipykernel_launcher); - test_start_session_kernel_never_ready_stops_kernel's docstring names _kernel_answers_shell, the helper start_session now probes with.
1 parent dc50953 commit 505dd80

2 files changed

Lines changed: 195 additions & 43 deletions

File tree

‎docx_editor/session.py‎

Lines changed: 53 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -71,16 +71,20 @@ def _client(connection_file: Path) -> "BlockingKernelClient":
7171
return kc
7272

7373

74-
def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
75-
"""True if the kernel answers a kernel_info request on the *control* channel.
74+
def _kernel_info_reply(kc: "BlockingKernelClient", channel: Literal["control", "shell"], timeout: float) -> bool:
75+
"""True if a kernel_info request on ``channel`` is answered within ``timeout``.
7676
77-
ipykernel services the control channel on a dedicated thread, so it replies
78-
even while an execution is in flight. The shell channel (which
79-
wait_for_ready() uses) is serialized behind the running execute_request, so
80-
there a busy kernel is indistinguishable from a dead one.
77+
One request, then poll in 0.5 s slices for a reply whose parent_header names
78+
it — replies to anything else (another client, an earlier probe) are
79+
skipped, not mistaken for ours.
8180
"""
8281
msg = kc.session.msg("kernel_info_request", {})
83-
kc.control_channel.send(msg)
82+
if channel == "control":
83+
kc.control_channel.send(msg)
84+
get_msg = kc.get_control_msg
85+
else:
86+
kc.shell_channel.send(msg)
87+
get_msg = kc.get_shell_msg
8488
msg_id = msg["header"]["msg_id"]
8589

8690
deadline = time.monotonic() + timeout
@@ -89,36 +93,37 @@ def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
8993
if remaining <= 0:
9094
return False
9195
try:
92-
reply = kc.get_control_msg(timeout=min(remaining, 0.5))
96+
reply = get_msg(timeout=min(remaining, 0.5))
9397
except Empty:
9498
continue
9599
if reply.get("parent_header", {}).get("msg_id") == msg_id:
96100
return True
97101

98102

99-
def _kernel_idle(kc: "BlockingKernelClient", timeout: float = 2.0) -> bool:
100-
"""True if the kernel answers a kernel_info request on the *shell* channel.
103+
def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
104+
"""True if the kernel answers a kernel_info request on the *control* channel.
101105
102-
The shell channel is serialized behind any running execute_request, so a
103-
busy kernel leaves this request queued past the window — which is exactly
104-
the idle/busy signal. The queued reply eventually lands on this
105-
disconnected client's socket and is dropped by ZMQ.
106+
ipykernel services the control channel on a dedicated thread, so it replies
107+
even while an execution is in flight. The shell channel (which
108+
wait_for_ready() uses) is serialized behind the running execute_request, so
109+
there a busy kernel is indistinguishable from a dead one.
106110
"""
107-
msg = kc.session.msg("kernel_info_request", {})
108-
kc.shell_channel.send(msg)
109-
msg_id = msg["header"]["msg_id"]
111+
return _kernel_info_reply(kc, "control", timeout)
110112

111-
deadline = time.monotonic() + timeout
112-
while True:
113-
remaining = deadline - time.monotonic()
114-
if remaining <= 0:
115-
return False
116-
try:
117-
reply = kc.get_shell_msg(timeout=min(remaining, 0.5))
118-
except Empty:
119-
continue
120-
if reply.get("parent_header", {}).get("msg_id") == msg_id:
121-
return True
113+
114+
def _kernel_answers_shell(kc: "BlockingKernelClient", timeout: float = 2.0) -> bool:
115+
"""True if the kernel answers a kernel_info request on the *shell* channel.
116+
117+
Two readings, both used:
118+
119+
- Against a running kernel this is the idle/busy signal. The shell channel
120+
is serialized behind any running execute_request, so a busy kernel leaves
121+
this request queued past the window. The queued reply eventually lands on
122+
this disconnected client's socket and is dropped by ZMQ.
123+
- During startup it is the readiness signal, because ipykernel registers the
124+
shell handler only after it finishes initializing (ROADMAP.md #82).
125+
"""
126+
return _kernel_info_reply(kc, "shell", timeout)
122127

123128

124129
def _kernel_dead(kc: "BlockingKernelClient", connection_file: Path) -> bool:
@@ -242,7 +247,7 @@ def start_session(connection_file: Path = DEFAULT_CONNECTION_FILE, timeout: floa
242247
start, the connection file it wrote is unreadable, or
243248
``DOCX_EDITOR_SESSION_START_TIMEOUT`` is not a positive number.
244249
A start that runs out of budget explains itself: which phase
245-
timed out (connection file vs. control-channel reply), whether
250+
timed out (connection file vs. shell-channel reply), whether
246251
the kernel was still alive, the load average, what it wrote to
247252
stderr, and the variable to raise.
248253
ImportError: If the [session] extra is not installed.
@@ -286,8 +291,8 @@ def _abort(summary: str) -> SessionError:
286291
return SessionError(report)
287292

288293
# Two phases share one budget. The connection file is written as soon as
289-
# the sockets are bound — *before* ipykernel imports IPython and wires the
290-
# control channel — so its appearance is an early signal, and the second
294+
# the sockets are bound — *before* ipykernel imports IPython and wires its
295+
# message dispatchers — so its appearance is an early signal, and the second
291296
# phase gets whatever is left (floored at 1 s).
292297
t0 = time.monotonic()
293298
deadline = t0 + timeout
@@ -315,10 +320,24 @@ def _abort(summary: str) -> SessionError:
315320
except (ValueError, OSError) as e:
316321
raise SessionError(f"Kernel wrote an unreadable connection file ({connection_file}): {e}") from e
317322
remaining = max(1.0, deadline - time.monotonic())
318-
if not _kernel_alive(kc, timeout=remaining):
323+
# Readiness is probed on *shell*, not control. No execution can be in
324+
# flight during a start, so the shell channel's busy-serialization
325+
# caveat does not apply here — and ipykernel registers the shell
326+
# handler only after it assigns _control_lock, so a shell reply proves
327+
# initialization finished. A control probe this early can arrive in
328+
# that window and be dropped by the kernel, costing us the whole
329+
# remaining budget for the one request we sent (ROADMAP.md #82).
330+
#
331+
# The ordering is upstream's and pyproject sets no ipykernel ceiling,
332+
# so it is named here: verified against 7.3.0, Kernel.start() in
333+
# kernelbase.py — control on_recv l.551, _control_lock l.559, shell
334+
# on_recv l.564. (l.553's pre-3.10 branch defers the lock to the
335+
# control thread and would break this, but requires-python >= 3.10
336+
# makes it unreachable.)
337+
if not _kernel_answers_shell(kc, timeout=remaining):
319338
raise _abort(
320339
f"Kernel did not become ready within {timeout}s: kernel.json appeared after {t_file:.1f}s, "
321-
f"then no kernel_info reply on the control channel in the remaining {remaining:.1f}s."
340+
f"then no kernel_info reply on the shell channel in the remaining {remaining:.1f}s."
322341
)
323342
except BaseException:
324343
# Never leave a half-started kernel behind with no way to reach it.
@@ -447,7 +466,7 @@ def session_status(connection_file: Path = DEFAULT_CONNECTION_FILE) -> SessionSt
447466
try:
448467
if not _kernel_alive(kc):
449468
return SessionStatus(running=False, pid=pid, state=None, connection_file=connection_file, stale=True)
450-
state: Literal["idle", "busy"] = "idle" if _kernel_idle(kc) else "busy"
469+
state: Literal["idle", "busy"] = "idle" if _kernel_answers_shell(kc) else "busy"
451470
finally:
452471
kc.stop_channels()
453472
return SessionStatus(running=True, pid=pid, state=state, connection_file=connection_file, stale=False)

‎tests/test_session.py‎

Lines changed: 142 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import subprocess
1111
import sys
1212
import textwrap
13+
import threading
1314
import time
1415
import types
1516
from contextlib import contextmanager
@@ -234,9 +235,9 @@ def test_start_session_kernel_never_ready_stops_kernel(tmp_path, monkeypatch):
234235
"""A kernel that writes a valid connection file but never answers is killed, not leaked.
235236
236237
The file names ports nobody listens on: ``_client`` connects (ZeroMQ
237-
connects lazily) and ``_kernel_alive`` times out, so ``start_session``
238-
must report the kernel as not ready, stop the still-running process and
239-
remove both files.
238+
connects lazily) and ``_kernel_answers_shell`` times out, so
239+
``start_session`` must report the kernel as not ready, stop the
240+
still-running process and remove both files.
240241
"""
241242
conn = tmp_path / "kernel.json"
242243
shell, iopub, stdin, control, hb = _free_ports(5)
@@ -265,6 +266,138 @@ def test_start_session_kernel_never_ready_stops_kernel(tmp_path, monkeypatch):
265266
_reap(spawned)
266267

267268

269+
# A kernel stand-in that answers kernel_info on the *shell* channel only. Its
270+
# control channel is bound but never serviced — exactly what ipykernel looks
271+
# like to us when it drops a control request that arrived mid-initialization
272+
# (ROADMAP.md #82). Bind precedes writing the connection file, as in a real
273+
# kernel; a port collision therefore surfaces as a bind traceback on stderr and
274+
# a non-zero exit, not as a mystery timeout.
275+
_SHELL_ONLY_KERNEL = textwrap.dedent(
276+
"""
277+
import json, sys, zmq
278+
from jupyter_client.session import Session
279+
280+
conn_file, info = sys.argv[1], json.loads(sys.argv[2])
281+
ctx = zmq.Context()
282+
shell = ctx.socket(zmq.ROUTER)
283+
shell.bind(f"tcp://127.0.0.1:{info['shell_port']}")
284+
control = ctx.socket(zmq.ROUTER)
285+
control.bind(f"tcp://127.0.0.1:{info['control_port']}")
286+
stdin = ctx.socket(zmq.ROUTER)
287+
stdin.bind(f"tcp://127.0.0.1:{info['stdin_port']}")
288+
iopub = ctx.socket(zmq.PUB)
289+
iopub.bind(f"tcp://127.0.0.1:{info['iopub_port']}")
290+
291+
with open(conn_file, "w") as fh:
292+
json.dump(info, fh)
293+
294+
session = Session(key=b"")
295+
while True:
296+
idents, frames = session.feed_identities(shell.recv_multipart())
297+
msg = session.deserialize(frames, content=False)
298+
if msg["header"]["msg_type"] == "kernel_info_request":
299+
session.send(
300+
shell,
301+
"kernel_info_reply",
302+
content={"status": "ok", "protocol_version": "5.3"},
303+
parent=msg,
304+
ident=idents,
305+
)
306+
"""
307+
)
308+
309+
310+
def test_start_session_ready_when_control_never_answers(tmp_path, monkeypatch):
311+
"""Readiness is decided on the shell channel, so a silent control channel does not hang it.
312+
313+
ipykernel registers the control callback before it assigns
314+
``_control_lock``, so a probe sent the instant kernel.json appears can be
315+
dropped kernel-side; the one request we sent is then never answered and the
316+
whole remaining budget burns (ROADMAP.md #82). The stand-in reproduces that
317+
end state deterministically: shell answers, control never does.
318+
"""
319+
conn = tmp_path / "kernel.json"
320+
shell, iopub, stdin, control, hb = _free_ports(5)
321+
info = {
322+
"transport": "tcp",
323+
"ip": "127.0.0.1",
324+
"key": "",
325+
"signature_scheme": "hmac-sha256",
326+
"kernel_name": "",
327+
"shell_port": shell,
328+
"iopub_port": iopub,
329+
"stdin_port": stdin,
330+
"control_port": control,
331+
"hb_port": hb,
332+
}
333+
spawned = _launch_fake_kernel(monkeypatch, _SHELL_ONLY_KERNEL, str(conn), json.dumps(info))
334+
335+
try:
336+
pid = start_session(conn, timeout=START_BUDGET)
337+
assert pid == spawned[0].pid
338+
assert conn.exists()
339+
assert conn.with_suffix(".pid").exists()
340+
finally:
341+
# conftest's reaper only records spawns whose argv names
342+
# ipykernel_launcher, and _launch_fake_kernel spawns `python -c`, so it
343+
# never sees this stand-in at all. Clean up by hand.
344+
_reap(spawned)
345+
conn.unlink(missing_ok=True)
346+
conn.with_suffix(".pid").unlink(missing_ok=True)
347+
348+
349+
def test_start_session_survives_control_probe_during_init(tmp_path):
350+
"""A real kernel still starts when a control probe lands during its initialization.
351+
352+
A guard, not a reproduction: the window is a single assignment inside
353+
ipykernel, so an intruding probe may well miss it on any given run — the
354+
deterministic half is
355+
``test_start_session_ready_when_control_never_answers``. What this pins is
356+
that when the window *is* hit, start still succeeds and the kernel is
357+
usable afterwards.
358+
"""
359+
conn = tmp_path / "kernel.json"
360+
stop = threading.Event()
361+
362+
def probe_control_at_birth():
363+
while not stop.is_set():
364+
try:
365+
if conn.stat().st_size == 0:
366+
time.sleep(0.0005)
367+
continue
368+
kc = _client(conn)
369+
except (OSError, ValueError):
370+
# No connection file yet, or a half-written one — the two
371+
# states _client documents. Retry: giving up here would leave
372+
# the guard silently probing nothing.
373+
time.sleep(0.0005)
374+
continue
375+
try:
376+
kc.control_channel.send(kc.session.msg("kernel_info_request", {}))
377+
# stop_channels() closes the socket with linger=0, which
378+
# discards anything ZMQ has not yet handed to the kernel — and
379+
# on loopback the connect handshake is usually still in flight
380+
# right after send(). Hold the socket open until the start has
381+
# finished, or the probe this test exists to fire never
382+
# arrives.
383+
stop.wait(timeout=START_BUDGET)
384+
finally:
385+
kc.stop_channels()
386+
return
387+
388+
intruder = threading.Thread(target=probe_control_at_birth, daemon=True)
389+
intruder.start()
390+
try:
391+
assert start_session(conn, timeout=START_BUDGET) > 0
392+
assert is_session_running(conn)
393+
assert eval_code("1 + 1", conn).value == 2
394+
finally:
395+
stop.set()
396+
intruder.join(timeout=5)
397+
assert not intruder.is_alive(), "intruder thread did not finish"
398+
stop_session(conn)
399+
400+
268401
def test_is_session_running_false_without_connection_file(tmp_path):
269402
assert is_session_running(tmp_path / "nope.json") is False
270403

@@ -1043,10 +1176,10 @@ def test_module_entrypoint_runs():
10431176

10441177

10451178
# A kernel that binds its sockets and writes the connection file but never
1046-
# serves the control channel: the exact shape of every start failure in the
1179+
# serves any channel: the exact shape of every start failure in the
10471180
# ROADMAP.md #72 CI logs (ipykernel writes kernel.json before importing IPython
1048-
# and wiring the control dispatcher). Ports 1-5 are unserved on loopback; ZMQ
1049-
# retries the connect silently, so the kernel_info request simply gets no reply.
1181+
# and wiring its dispatchers). Ports 1-5 are unserved on loopback; ZMQ retries
1182+
# the connect silently, so the kernel_info request simply gets no reply.
10501183
_FAKE_KERNEL_SRC = """\
10511184
import json, os, sys, time
10521185
path = sys.argv[sys.argv.index("-f") + 1]
@@ -1075,8 +1208,8 @@ class TestStartFailureDiagnosis:
10751208
tunable for loaded machines (ROADMAP.md #72).
10761209
10771210
The CI failures behind #72 all read "Kernel did not become ready within
1078-
30s" — the connection file had appeared but the control channel never
1079-
answered — with nothing to tell a slow runner from a broken one.
1211+
30s" — the connection file had appeared but the readiness probe never got
1212+
an answer — with nothing to tell a slow runner from a broken one.
10801213
"""
10811214

10821215
@pytest.fixture
@@ -1091,7 +1224,7 @@ def test_start_timeout_explains_itself(self, tmp_path, fake_kernel):
10911224

10921225
assert "did not become ready within 2.0s" in message
10931226
assert "kernel.json appeared after" in message
1094-
assert "no kernel_info reply on the control channel" in message
1227+
assert "no kernel_info reply on the shell channel" in message
10951228
assert "still alive at the timeout" in message
10961229
assert "fake kernel: sockets bound, never served" in message
10971230
assert "tcp 127.0.0.1, shell=1 control=4" in message

0 commit comments

Comments
 (0)