Add anyio Event for _main_shell_ready to aid with start and test reliability.

Author: fleming79

ipykernel/kernelapp.py

@@ -20,7 +20,7 @@
 
 import zmq
 import zmq_anyio
-from anyio import create_task_group, run, to_thread
+from anyio import create_task_group, run
 from IPython.core.application import (  # type:ignore[attr-defined]
     BaseIPythonApplication,
     base_aliases,
@@ -764,15 +764,9 @@ def start(self) -> None:
         backend = "trio" if self.trio_loop else "asyncio"
         run(partial(self._start, backend), backend=backend)
 
-    async def _wait_to_enter_eventloop(self) -> None:
-        await to_thread.run_sync(self.kernel._eventloop_set.wait)
-        await self.kernel.enter_eventloop()
-
     async def main(self) -> None:
         async with create_task_group() as tg:
-            tg.start_soon(self._wait_to_enter_eventloop)
             tg.start_soon(self.kernel.start)
-
         if self.kernel.eventloop:
             self.kernel._eventloop_set.set()
 

ipykernel/kernelbase.py

@@ -40,6 +40,7 @@
 import zmq_anyio
 from anyio import (
     TASK_STATUS_IGNORED,
+    Event,
     create_memory_object_stream,
     create_task_group,
     sleep,
@@ -126,6 +127,7 @@ class Kernel(SingletonConfigurable):
     stdin_socket = Any()
 
     _send_exec_request: Dict[dict[zmq_anyio.Socket, MemoryObjectSendStream]] = Dict()
+    _main_shell_ready = Instance(Event, ())
 
     log: logging.Logger = Instance(logging.Logger, allow_none=True)  # type:ignore[assignment]
 
@@ -436,13 +438,15 @@ async def shell_main(self, subshell_id: str | None):
             async with create_task_group() as tg:
                 if not socket.started.is_set():
                     await tg.start(socket.start)
-                tg.start_soon(self.process_shell, socket)
+                tg.start_soon(self._process_shell, socket)
                 tg.start_soon(self._execute_request_handler, receive_stream)
                 if subshell_id is None:
                     # Main subshell.
+                    self._main_shell_ready.set()
                     await to_thread.run_sync(self.shell_stop.wait)
                     tg.cancel_scope.cancel()
             self._send_exec_request.pop(socket, None)
+            await send_stream.aclose()
 
     async def _execute_request_handler(self, receive_stream: MemoryObjectReceiveStream):
         async with receive_stream:
@@ -461,8 +465,9 @@ async def _execute_request_handler(self, receive_stream: MemoryObjectReceiveStre
                 except BaseException as e:
                     self.log.exception("Execute request", exc_info=e)
 
-    async def process_shell(self, socket):
+    async def _process_shell(self, socket):
         # socket=None is valid if kernel subshells are not supported.
+        await self._main_shell_ready.wait()
         try:
             while True:
                 await self.process_shell_message(socket=socket)
@@ -476,15 +481,13 @@ async def process_shell_message(self, msg=None, socket=None):
         # If msg is set, process that message.
         # If msg is None, await the next message to arrive on the socket.
         assert self.session is not None
+        socket = socket or self.shell_socket
         if self._supports_kernel_subshells:
             assert threading.current_thread() not in (
                 self.control_thread,
                 self.shell_channel_thread,
             )
             assert socket is not None
-        else:
-            assert threading.current_thread() == threading.main_thread()
-            socket = self.shell_socket
 
         msg = msg or await socket.arecv_multipart(copy=False).wait()
 
@@ -532,8 +535,8 @@ async def process_shell_message(self, msg=None, socket=None):
                     result = handler(socket, idents, msg)
                     if inspect.isawaitable(result):
                         await result
-            except Exception:
-                self.log.error("Exception in message handler:", exc_info=True)  # noqa: G201
+            except Exception as e:
+                self.log.error("Exception in message handler:", exc_info=e)
             except KeyboardInterrupt:
                 # Ctrl-c shouldn't crash the kernel here.
                 self.log.error("KeyboardInterrupt caught in kernel.")
@@ -583,6 +586,7 @@ async def start(self, *, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None:
             self.shell_stop = threading.Event()
 
             tg.start_soon(self.shell_main, None)
+            await self._main_shell_ready.wait()
             if self.shell_channel_thread:
                 # Assign tasks to and start shell channel thread.
                 manager = self.shell_channel_thread.manager

tests/conftest.py

@@ -207,6 +207,7 @@ async def kernel(anyio_backend):
     async with create_task_group() as tg:
         kernel = MockKernel()
         tg.start_soon(kernel.start)
+        await kernel._main_shell_ready.wait()
         try:
             yield kernel
         finally:
@@ -218,6 +219,7 @@ async def ipkernel(anyio_backend):
     async with create_task_group() as tg:
         kernel = MockIPyKernel()
         tg.start_soon(kernel.start)
+        await kernel._main_shell_ready.wait()
         try:
             yield kernel
         finally:

tests/test_ipkernel_direct.py

@@ -34,14 +34,19 @@ async def test_direct_kernel_info_request(ipkernel):
 
 
 async def test_direct_execute_request(ipkernel: MockIPyKernel) -> None:
-    reply = await ipkernel.test_shell_message("execute_request", dict(code="hello", silent=False))
+    reply = await ipkernel.test_shell_message(
+        "execute_request", dict(code="invalid_call()", silent=False)
+    )
     assert reply["header"]["msg_type"] == "execute_reply"
+    ipkernel._aborted_time += 10
     reply = await ipkernel.test_shell_message(
         "execute_request", dict(code="trigger_error", silent=False)
     )
     assert reply["content"]["status"] == "aborted"
-
-    reply = await ipkernel.test_shell_message("execute_request", dict(code="hello", silent=False))
+    ipkernel._aborted_time = time.monotonic()
+    reply = await ipkernel.test_shell_message(
+        "execute_request", dict(code="okay=True", silent=False)
+    )
     assert reply["header"]["msg_type"] == "execute_reply"