Run plot comms on the R thread

crates/amalthea/src/comm/event.rs

@@ -5,15 +5,20 @@
  *
  */
 
+use crossbeam::channel::Sender;
 use serde_json::Value;
 
 use crate::comm::comm_channel::CommMsg;
 use crate::socket::comm::CommSocket;
 
 /// Comm events sent to the frontend via Shell.
 pub enum CommEvent {
-    /// A new Comm was opened
-    Opened(CommSocket, Value),
+    /// A new Comm was opened. The optional `Sender` is a synchronisation
+    /// barrier: if provided, Shell signals it after processing the open
+    /// (sending `comm_open` on IOPub). The caller blocks on the paired
+    /// receiver to guarantee that the `comm_open` message has been sent
+    /// before any subsequent messages.
+    Opened(CommSocket, Value, Option<Sender<()>>),
 
     /// A message was received on a Comm; the first value is the comm ID, and the
     /// second value is the message.

crates/amalthea/src/language/shell_handler.rs

@@ -6,6 +6,7 @@
  */
 
 use async_trait::async_trait;
+use crossbeam::channel::Receiver;
 
 use crate::comm::comm_channel::Comm;
 use crate::comm::comm_channel::CommMsg;
@@ -25,10 +26,10 @@ use crate::wire::kernel_info_request::KernelInfoRequest;
 use crate::wire::originator::Originator;
 
 /// Result of a `handle_comm_msg` or `handle_comm_close` call on the
-/// `ShellHandler`. `Handled` means the kernel processed the message
-/// synchronously (blocking Shell until done). `NotHandled` means amalthea
-/// should fall back to the historical `incoming_tx` path. This fallback is
-/// temporary until all comms are migrated to the blocking path.
+/// `ShellHandler`. `Handled` means the kernel dispatched the message
+/// (possibly asynchronously via a completion receiver). `NotHandled` means
+/// amalthea should fall back to the historical `incoming_tx` path. This
+/// fallback is temporary until all comms are migrated to the new path.
 pub enum CommHandled {
     Handled,
     NotHandled,
@@ -53,17 +54,17 @@ pub trait ShellHandler: Send {
         req: &IsCompleteRequest,
     ) -> crate::Result<IsCompleteReply>;
 
-    /// Handles a request to execute code.
-    ///
-    /// The `originator` is an opaque byte array identifying the peer that sent
-    /// the request; it is needed to perform an input request during execution.
+    /// Kicks off execution of the given request and returns a channel that
+    /// will receive the reply once execution completes. Shell select-loops
+    /// on this receiver together with `comm_event_rx` so it can process
+    /// comm events (e.g. barrier handshakes) while execution is in progress.
     ///
     /// Docs: https://jupyter-client.readthedocs.io/en/stable/messaging.html#execute
-    async fn handle_execute_request(
+    fn start_execute_request(
         &mut self,
         originator: Originator,
         req: &ExecuteRequest,
-    ) -> crate::Result<ExecuteReply>;
+    ) -> Receiver<crate::Result<ExecuteReply>>;
 
     /// Handles a request to provide completions for the given code fragment.
     ///
@@ -80,26 +81,30 @@ pub trait ShellHandler: Send {
     /// Docs: https://jupyter-client.readthedocs.io/en/stable/messaging.html#history
     async fn handle_history_request(&self, req: &HistoryRequest) -> crate::Result<HistoryReply>;
 
-    /// Handles a request to open a comm.
-    ///
-    /// https://jupyter-client.readthedocs.io/en/stable/messaging.html#opening-a-comm
+    /// Handle a request to open a comm.
     ///
-    /// Returns true if the handler handled the request (and opened the comm), false if it did not.
+    /// Returns `(true, Some(receiver))` if the comm was opened and the handler
+    /// was dispatched asynchronously. Shell will select-loop on the receiver
+    /// and `comm_event_rx` to drain comm events while the handler runs.
     ///
-    /// * `target` - The target name of the comm, such as `positron.variables`
-    /// * `comm` - The comm channel to use to communicate with the frontend
-    /// * `data` - The `data` payload from the `comm_open` message
-    async fn handle_comm_open(
-        &self,
+    /// Returns `(true, None)` if the comm was opened synchronously.
+    /// Returns `(false, None)` if the comm was not handled.
+    fn handle_comm_open(
+        &mut self,
         target: Comm,
         comm: CommSocket,
         data: serde_json::Value,
-    ) -> crate::Result<bool>;
+    ) -> crate::Result<(bool, Option<Receiver<()>>)>;
 
-    /// Handle an incoming comm message (RPC or data). Return
-    /// `CommHandled::Handled` if the message was processed, or
-    /// `CommHandled::NotHandled` to fall back to the existing
-    /// `incoming_tx` path.
+    /// Handle an incoming comm message (RPC or data).
+    ///
+    /// Returns `(CommHandled::Handled, Some(receiver))` if the message was
+    /// dispatched to the R thread. Shell will select-loop on the receiver
+    /// and `comm_event_rx` to drain comm events (e.g. barriers from
+    /// `comm_open_backend`) while the handler runs.
+    ///
+    /// Returns `(CommHandled::NotHandled, None)` to fall back to the
+    /// existing `incoming_tx` path.
     ///
     /// * `comm_id` - The comm's unique identifier
     /// * `comm_name` - The comm's target name (e.g. `"positron.dataExplorer"`)
@@ -112,21 +117,22 @@ pub trait ShellHandler: Send {
         _comm_name: &str,
         _msg: CommMsg,
         _originator: Originator,
-    ) -> crate::Result<CommHandled> {
-        Ok(CommHandled::NotHandled)
+    ) -> crate::Result<(CommHandled, Option<Receiver<()>>)> {
+        Ok((CommHandled::NotHandled, None))
     }
 
-    /// Handle a comm close. Return `CommHandled::Handled` if the close
-    /// was processed, or `CommHandled::NotHandled` to fall back to the
-    /// existing `incoming_tx` path.
+    /// Handle a comm close.
+    ///
+    /// Same pattern as `handle_comm_msg`: returns a completion receiver
+    /// so Shell can drain comm events while the handler runs.
     ///
     /// * `comm_id` - The comm's unique identifier
     /// * `comm_name` - The comm's target name
     fn handle_comm_close(
         &mut self,
         _comm_id: &str,
         _comm_name: &str,
-    ) -> crate::Result<CommHandled> {
-        Ok(CommHandled::NotHandled)
+    ) -> crate::Result<(CommHandled, Option<Receiver<()>>)> {
+        Ok((CommHandled::NotHandled, None))
     }
 }