ipykernel/subshell_manager.py

"""Manager of subshells in a kernel."""

from __future__ import annotations

import asyncio
import json
import typing as t
import uuid
from functools import partial
from threading import Lock, current_thread

import zmq
from tornado.ioloop import IOLoop

from .socket_pair import SocketPair
from .subshell import SubshellThread
from .thread import SHELL_CHANNEL_THREAD_NAME
from .utils import _async_in_context


class SubshellManager:
    """A manager of subshells.

    Controls the lifetimes of subshell threads and their associated ZMQ sockets and
    streams. Runs mostly in the shell channel thread.

    Care needed with threadsafe access here.  All write access to the cache occurs in
    the shell channel thread so there is only ever one write access at any one time.
    Reading of cache information can be performed by other threads, so all reads are
    protected by a lock so that they are atomic.

    Sending reply messages via the shell_socket is wrapped by another lock to protect
    against multiple subshells attempting to send at the same time.

    .. versionadded:: 7
    """

    def __init__(
        self,
        context: zmq.Context[t.Any],
        shell_channel_io_loop: IOLoop,
        shell_socket: zmq.Socket[t.Any],
    ):
        """Initialize the subshell manager."""
        self._parent_thread = current_thread()

        self._context: zmq.Context[t.Any] = context
        self._shell_channel_io_loop = shell_channel_io_loop
        self._shell_socket = shell_socket
        self._cache: dict[str, SubshellThread] = {}
        self._lock_cache = Lock()  # Sync lock across threads when accessing cache.

        # Inproc socket pair for communication from control thread to shell channel thread,
        # such as for create_subshell_request messages. Reply messages are returned straight away.
        self.control_to_shell_channel = SocketPair(self._context, "control")
        self.control_to_shell_channel.on_recv(
            self._shell_channel_io_loop, self._process_control_request, copy=True
        )

        # Inproc socket pair for communication from shell channel thread to main thread,
        # such as for execute_request messages.
        self._shell_channel_to_main = SocketPair(self._context, "main")

        # Inproc socket pair for communication from main thread to shell channel thread.
        # such as for execute_reply messages.
        self._main_to_shell_channel = SocketPair(self._context, "main-reverse")
        self._main_to_shell_channel.on_recv(
            self._shell_channel_io_loop, self._send_on_shell_channel
        )

    def close(self) -> None:
        """Stop all subshells and close all resources."""
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
        with self._lock_cache:
            while True:
                try:
                    _, subshell_thread = self._cache.popitem()
                except KeyError:
                    break
                self._stop_subshell(subshell_thread)

        self.control_to_shell_channel.close()
        self._main_to_shell_channel.close()
        self._shell_channel_to_main.close()

    def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketPair:
        """Return the inproc socket pair used to send messages from the shell channel
        to a particular subshell or main shell."""
        if subshell_id is None:
            return self._shell_channel_to_main
        with self._lock_cache:
            return self._cache[subshell_id].shell_channel_to_subshell

    def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
        """Return the socket used by a particular subshell or main shell to send
        messages to the shell channel.
        """
        if subshell_id is None:
            return self._main_to_shell_channel.from_socket
        with self._lock_cache:
            return self._cache[subshell_id].subshell_to_shell_channel.from_socket

    def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
        """Return the socket used by the shell channel to send messages to a particular
        subshell or main shell.
        """
        return self.get_shell_channel_to_subshell_pair(subshell_id).from_socket

    def get_subshell_aborting(self, subshell_id: str) -> bool:
        """Get the boolean aborting flag of the specified subshell."""
        with self._lock_cache:
            return self._cache[subshell_id].aborting

    def get_subshell_asyncio_lock(self, subshell_id: str) -> asyncio.Lock:
        """Return the asyncio lock belonging to the specified subshell."""
        with self._lock_cache:
            return self._cache[subshell_id].asyncio_lock

    def list_subshell(self) -> list[str]:
        """Return list of current subshell ids.

        Can be called by any subshell using %subshell magic.
        """
        with self._lock_cache:
            return list(self._cache)

    def set_on_recv_callback(self, on_recv_callback):
        """Set the callback used by the main shell and all subshells to receive
        messages sent from the shell channel thread.
        """
        assert current_thread() == self._parent_thread
        self._on_recv_callback = on_recv_callback
        self._shell_channel_to_main.on_recv(
            IOLoop.current(), _async_in_context(partial(on_recv_callback, None))
        )

    def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None:
        """Set the aborting flag of the specified subshell."""
        with self._lock_cache:
            self._cache[subshell_id].aborting = aborting

    def subshell_id_from_thread_id(self, thread_id: int) -> str | None:
        """Return subshell_id of the specified thread_id.

        Raises RuntimeError if thread_id is not the main shell or a subshell.

        Only used by %subshell magic so does not have to be fast/cached.
        """
        with self._lock_cache:
            if thread_id == self._parent_thread.ident:
                return None
            for id, subshell in self._cache.items():
                if subshell.ident == thread_id:
                    return id
            msg = f"Thread id {thread_id!r} does not correspond to a subshell of this kernel"
            raise RuntimeError(msg)

    def _create_subshell(self) -> str:
        """Create and start a new subshell thread."""
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME

        subshell_id = str(uuid.uuid4())
        subshell_thread = SubshellThread(subshell_id, self._context)

        with self._lock_cache:
            assert subshell_id not in self._cache
            self._cache[subshell_id] = subshell_thread

        subshell_thread.shell_channel_to_subshell.on_recv(
            subshell_thread.io_loop,
            _async_in_context(partial(self._on_recv_callback, subshell_id)),
        )

        subshell_thread.subshell_to_shell_channel.on_recv(
            self._shell_channel_io_loop, self._send_on_shell_channel
        )

        subshell_thread.start()
        return subshell_id

    def _delete_subshell(self, subshell_id: str) -> None:
        """Delete subshell identified by subshell_id.

        Raises key error if subshell_id not in cache.
        """
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME

        with self._lock_cache:
            subshell_threwad = self._cache.pop(subshell_id)

        self._stop_subshell(subshell_threwad)

    def _process_control_request(
        self,
        request: list[t.Any],
    ) -> None:
        """Process a control request message received on the control inproc
        socket and return the reply.  Runs in the shell channel thread.
        """
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME

        try:
            decoded = json.loads(request[0])
            type = decoded["type"]
            reply: dict[str, t.Any] = {"status": "ok"}

            if type == "create":
                reply["subshell_id"] = self._create_subshell()
            elif type == "delete":
                subshell_id = decoded["subshell_id"]
                self._delete_subshell(subshell_id)
            elif type == "list":
                reply["subshell_id"] = self.list_subshell()
            else:
                msg = f"Unrecognised message type {type!r}"
                raise RuntimeError(msg)
        except BaseException as err:
            reply = {
                "status": "error",
                "evalue": str(err),
            }

        # Return the reply to the control thread.
        self.control_to_shell_channel.to_socket.send_json(reply)

    def _send_on_shell_channel(self, msg) -> None:
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
        self._shell_socket.send_multipart(msg)

    def _stop_subshell(self, subshell_thread: SubshellThread) -> None:
        """Stop a subshell thread and close all of its resources."""
        assert current_thread().name == SHELL_CHANNEL_THREAD_NAME

        if subshell_thread.is_alive():
            subshell_thread.stop()
            subshell_thread.join()