fix: probe kernel readiness without racing ipykernel initialization (ROADMAP.md #82)

docx_editor/session.py

@@ -71,16 +71,20 @@ def _client(connection_file: Path) -> "BlockingKernelClient":
     return kc
 
 
-def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
-    """True if the kernel answers a kernel_info request on the *control* channel.
+def _kernel_info_reply(kc: "BlockingKernelClient", channel: Literal["control", "shell"], timeout: float) -> bool:
+    """True if a kernel_info request on ``channel`` is answered within ``timeout``.
 
-    ipykernel services the control channel on a dedicated thread, so it replies
-    even while an execution is in flight. The shell channel (which
-    wait_for_ready() uses) is serialized behind the running execute_request, so
-    there a busy kernel is indistinguishable from a dead one.
+    One request, then poll in 0.5 s slices for a reply whose parent_header names
+    it — replies to anything else (another client, an earlier probe) are
+    skipped, not mistaken for ours.
     """
     msg = kc.session.msg("kernel_info_request", {})
-    kc.control_channel.send(msg)
+    if channel == "control":
+        kc.control_channel.send(msg)
+        get_msg = kc.get_control_msg
+    else:
+        kc.shell_channel.send(msg)
+        get_msg = kc.get_shell_msg
     msg_id = msg["header"]["msg_id"]
 
     deadline = time.monotonic() + timeout
@@ -89,36 +93,37 @@ def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
         if remaining <= 0:
             return False
         try:
-            reply = kc.get_control_msg(timeout=min(remaining, 0.5))
+            reply = get_msg(timeout=min(remaining, 0.5))
         except Empty:
             continue
         if reply.get("parent_header", {}).get("msg_id") == msg_id:
             return True
 
 
-def _kernel_idle(kc: "BlockingKernelClient", timeout: float = 2.0) -> bool:
-    """True if the kernel answers a kernel_info request on the *shell* channel.
+def _kernel_alive(kc: "BlockingKernelClient", timeout: float = 5.0) -> bool:
+    """True if the kernel answers a kernel_info request on the *control* channel.
 
-    The shell channel is serialized behind any running execute_request, so a
-    busy kernel leaves this request queued past the window — which is exactly
-    the idle/busy signal. The queued reply eventually lands on this
-    disconnected client's socket and is dropped by ZMQ.
+    ipykernel services the control channel on a dedicated thread, so it replies
+    even while an execution is in flight. The shell channel (which
+    wait_for_ready() uses) is serialized behind the running execute_request, so
+    there a busy kernel is indistinguishable from a dead one.
     """
-    msg = kc.session.msg("kernel_info_request", {})
-    kc.shell_channel.send(msg)
-    msg_id = msg["header"]["msg_id"]
+    return _kernel_info_reply(kc, "control", timeout)
 
-    deadline = time.monotonic() + timeout
-    while True:
-        remaining = deadline - time.monotonic()
-        if remaining <= 0:
-            return False
-        try:
-            reply = kc.get_shell_msg(timeout=min(remaining, 0.5))
-        except Empty:
-            continue
-        if reply.get("parent_header", {}).get("msg_id") == msg_id:
-            return True
+
+def _kernel_answers_shell(kc: "BlockingKernelClient", timeout: float = 2.0) -> bool:
+    """True if the kernel answers a kernel_info request on the *shell* channel.
+
+    Two readings, both used:
+
+    - Against a running kernel this is the idle/busy signal. The shell channel
+      is serialized behind any running execute_request, so a busy kernel leaves
+      this request queued past the window. The queued reply eventually lands on
+      this disconnected client's socket and is dropped by ZMQ.
+    - During startup it is the readiness signal, because ipykernel registers the
+      shell handler only after it finishes initializing (ROADMAP.md #82).
+    """
+    return _kernel_info_reply(kc, "shell", timeout)
 
 
 def _kernel_dead(kc: "BlockingKernelClient", connection_file: Path) -> bool:
@@ -242,7 +247,7 @@ def start_session(connection_file: Path = DEFAULT_CONNECTION_FILE, timeout: floa
             start, the connection file it wrote is unreadable, or
             ``DOCX_EDITOR_SESSION_START_TIMEOUT`` is not a positive number.
             A start that runs out of budget explains itself: which phase
-            timed out (connection file vs. control-channel reply), whether
+            timed out (connection file vs. shell-channel reply), whether
             the kernel was still alive, the load average, what it wrote to
             stderr, and the variable to raise.
         ImportError: If the [session] extra is not installed.
@@ -286,8 +291,8 @@ def _abort(summary: str) -> SessionError:
         return SessionError(report)
 
     # Two phases share one budget. The connection file is written as soon as
-    # the sockets are bound — *before* ipykernel imports IPython and wires the
-    # control channel — so its appearance is an early signal, and the second
+    # the sockets are bound — *before* ipykernel imports IPython and wires its
+    # message dispatchers — so its appearance is an early signal, and the second
     # phase gets whatever is left (floored at 1 s).
     t0 = time.monotonic()
     deadline = t0 + timeout
@@ -315,10 +320,24 @@ def _abort(summary: str) -> SessionError:
         except (ValueError, OSError) as e:
             raise SessionError(f"Kernel wrote an unreadable connection file ({connection_file}): {e}") from e
         remaining = max(1.0, deadline - time.monotonic())
-        if not _kernel_alive(kc, timeout=remaining):
+        # Readiness is probed on *shell*, not control. No execution can be in
+        # flight during a start, so the shell channel's busy-serialization
+        # caveat does not apply here — and ipykernel registers the shell
+        # handler only after it assigns _control_lock, so a shell reply proves
+        # initialization finished. A control probe this early can arrive in
+        # that window and be dropped by the kernel, costing us the whole
+        # remaining budget for the one request we sent (ROADMAP.md #82).
+        #
+        # The ordering is upstream's and pyproject sets no ipykernel ceiling,
+        # so it is named here: verified against 7.3.0, Kernel.start() in
+        # kernelbase.py — control on_recv l.551, _control_lock l.559, shell
+        # on_recv l.564. (l.553's pre-3.10 branch defers the lock to the
+        # control thread and would break this, but requires-python >= 3.10
+        # makes it unreachable.)
+        if not _kernel_answers_shell(kc, timeout=remaining):
             raise _abort(
                 f"Kernel did not become ready within {timeout}s: kernel.json appeared after {t_file:.1f}s, "
-                f"then no kernel_info reply on the control channel in the remaining {remaining:.1f}s."
+                f"then no kernel_info reply on the shell channel in the remaining {remaining:.1f}s."
             )
     except BaseException:
         # 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
     try:
         if not _kernel_alive(kc):
             return SessionStatus(running=False, pid=pid, state=None, connection_file=connection_file, stale=True)
-        state: Literal["idle", "busy"] = "idle" if _kernel_idle(kc) else "busy"
+        state: Literal["idle", "busy"] = "idle" if _kernel_answers_shell(kc) else "busy"
     finally:
         kc.stop_channels()
     return SessionStatus(running=True, pid=pid, state=state, connection_file=connection_file, stale=False)

tests/test_session.py

@@ -10,6 +10,7 @@
 import subprocess
 import sys
 import textwrap
+import threading
 import time
 import types
 from contextlib import contextmanager
@@ -234,9 +235,9 @@ def test_start_session_kernel_never_ready_stops_kernel(tmp_path, monkeypatch):
     """A kernel that writes a valid connection file but never answers is killed, not leaked.
 
     The file names ports nobody listens on: ``_client`` connects (ZeroMQ
-    connects lazily) and ``_kernel_alive`` times out, so ``start_session``
-    must report the kernel as not ready, stop the still-running process and
-    remove both files.
+    connects lazily) and ``_kernel_answers_shell`` times out, so
+    ``start_session`` must report the kernel as not ready, stop the
+    still-running process and remove both files.
     """
     conn = tmp_path / "kernel.json"
     shell, iopub, stdin, control, hb = _free_ports(5)
@@ -265,6 +266,138 @@ def test_start_session_kernel_never_ready_stops_kernel(tmp_path, monkeypatch):
         _reap(spawned)
 
 
+# A kernel stand-in that answers kernel_info on the *shell* channel only. Its
+# control channel is bound but never serviced — exactly what ipykernel looks
+# like to us when it drops a control request that arrived mid-initialization
+# (ROADMAP.md #82). Bind precedes writing the connection file, as in a real
+# kernel; a port collision therefore surfaces as a bind traceback on stderr and
+# a non-zero exit, not as a mystery timeout.
+_SHELL_ONLY_KERNEL = textwrap.dedent(
+    """
+    import json, sys, zmq
+    from jupyter_client.session import Session
+
+    conn_file, info = sys.argv[1], json.loads(sys.argv[2])
+    ctx = zmq.Context()
+    shell = ctx.socket(zmq.ROUTER)
+    shell.bind(f"tcp://127.0.0.1:{info['shell_port']}")
+    control = ctx.socket(zmq.ROUTER)
+    control.bind(f"tcp://127.0.0.1:{info['control_port']}")
+    stdin = ctx.socket(zmq.ROUTER)
+    stdin.bind(f"tcp://127.0.0.1:{info['stdin_port']}")
+    iopub = ctx.socket(zmq.PUB)
+    iopub.bind(f"tcp://127.0.0.1:{info['iopub_port']}")
+
+    with open(conn_file, "w") as fh:
+        json.dump(info, fh)
+
+    session = Session(key=b"")
+    while True:
+        idents, frames = session.feed_identities(shell.recv_multipart())
+        msg = session.deserialize(frames, content=False)
+        if msg["header"]["msg_type"] == "kernel_info_request":
+            session.send(
+                shell,
+                "kernel_info_reply",
+                content={"status": "ok", "protocol_version": "5.3"},
+                parent=msg,
+                ident=idents,
+            )
+    """
+)
+
+
+def test_start_session_ready_when_control_never_answers(tmp_path, monkeypatch):
+    """Readiness is decided on the shell channel, so a silent control channel does not hang it.
+
+    ipykernel registers the control callback before it assigns
+    ``_control_lock``, so a probe sent the instant kernel.json appears can be
+    dropped kernel-side; the one request we sent is then never answered and the
+    whole remaining budget burns (ROADMAP.md #82). The stand-in reproduces that
+    end state deterministically: shell answers, control never does.
+    """
+    conn = tmp_path / "kernel.json"
+    shell, iopub, stdin, control, hb = _free_ports(5)
+    info = {
+        "transport": "tcp",
+        "ip": "127.0.0.1",
+        "key": "",
+        "signature_scheme": "hmac-sha256",
+        "kernel_name": "",
+        "shell_port": shell,
+        "iopub_port": iopub,
+        "stdin_port": stdin,
+        "control_port": control,
+        "hb_port": hb,
+    }
+    spawned = _launch_fake_kernel(monkeypatch, _SHELL_ONLY_KERNEL, str(conn), json.dumps(info))
+
+    try:
+        pid = start_session(conn, timeout=START_BUDGET)
+        assert pid == spawned[0].pid
+        assert conn.exists()
+        assert conn.with_suffix(".pid").exists()
+    finally:
+        # conftest's reaper only records spawns whose argv names
+        # ipykernel_launcher, and _launch_fake_kernel spawns `python -c`, so it
+        # never sees this stand-in at all. Clean up by hand.
+        _reap(spawned)
+        conn.unlink(missing_ok=True)
+        conn.with_suffix(".pid").unlink(missing_ok=True)
+
+
+def test_start_session_survives_control_probe_during_init(tmp_path):
+    """A real kernel still starts when a control probe lands during its initialization.
+
+    A guard, not a reproduction: the window is a single assignment inside
+    ipykernel, so an intruding probe may well miss it on any given run — the
+    deterministic half is
+    ``test_start_session_ready_when_control_never_answers``. What this pins is
+    that when the window *is* hit, start still succeeds and the kernel is
+    usable afterwards.
+    """
+    conn = tmp_path / "kernel.json"
+    stop = threading.Event()
+
+    def probe_control_at_birth():
+        while not stop.is_set():
+            try:
+                if conn.stat().st_size == 0:
+                    time.sleep(0.0005)
+                    continue
+                kc = _client(conn)
+            except (OSError, ValueError):
+                # No connection file yet, or a half-written one — the two
+                # states _client documents. Retry: giving up here would leave
+                # the guard silently probing nothing.
+                time.sleep(0.0005)
+                continue
+            try:
+                kc.control_channel.send(kc.session.msg("kernel_info_request", {}))
+                # stop_channels() closes the socket with linger=0, which
+                # discards anything ZMQ has not yet handed to the kernel — and
+                # on loopback the connect handshake is usually still in flight
+                # right after send(). Hold the socket open until the start has
+                # finished, or the probe this test exists to fire never
+                # arrives.
+                stop.wait(timeout=START_BUDGET)
+            finally:
+                kc.stop_channels()
+            return
+
+    intruder = threading.Thread(target=probe_control_at_birth, daemon=True)
+    intruder.start()
+    try:
+        assert start_session(conn, timeout=START_BUDGET) > 0
+        assert is_session_running(conn)
+        assert eval_code("1 + 1", conn).value == 2
+    finally:
+        stop.set()
+        intruder.join(timeout=5)
+        assert not intruder.is_alive(), "intruder thread did not finish"
+        stop_session(conn)
+
+
 def test_is_session_running_false_without_connection_file(tmp_path):
     assert is_session_running(tmp_path / "nope.json") is False
 
@@ -1043,10 +1176,10 @@ def test_module_entrypoint_runs():
 
 
 # A kernel that binds its sockets and writes the connection file but never
-# serves the control channel: the exact shape of every start failure in the
+# serves any channel: the exact shape of every start failure in the
 # ROADMAP.md #72 CI logs (ipykernel writes kernel.json before importing IPython
-# and wiring the control dispatcher). Ports 1-5 are unserved on loopback; ZMQ
-# retries the connect silently, so the kernel_info request simply gets no reply.
+# and wiring its dispatchers). Ports 1-5 are unserved on loopback; ZMQ retries
+# the connect silently, so the kernel_info request simply gets no reply.
 _FAKE_KERNEL_SRC = """\
 import json, os, sys, time
 path = sys.argv[sys.argv.index("-f") + 1]
@@ -1075,8 +1208,8 @@ class TestStartFailureDiagnosis:
     tunable for loaded machines (ROADMAP.md #72).
 
     The CI failures behind #72 all read "Kernel did not become ready within
-    30s" — the connection file had appeared but the control channel never
-    answered — with nothing to tell a slow runner from a broken one.
+    30s" — the connection file had appeared but the readiness probe never got
+    an answer — with nothing to tell a slow runner from a broken one.
     """
 
     @pytest.fixture
@@ -1091,7 +1224,7 @@ def test_start_timeout_explains_itself(self, tmp_path, fake_kernel):
 
         assert "did not become ready within 2.0s" in message
         assert "kernel.json appeared after" in message
-        assert "no kernel_info reply on the control channel" in message
+        assert "no kernel_info reply on the shell channel" in message
         assert "still alive at the timeout" in message
         assert "fake kernel: sockets bound, never served" in message
         assert "tcp 127.0.0.1, shell=1 control=4" in message