Fix synchronisation between comm opening and plot updates

Author: lionel-

crates/amalthea/src/comm/event.rs

@@ -5,6 +5,7 @@
  *
  */
 
+use crossbeam::channel::Sender;
 use serde_json::Value;
 
 use crate::comm::comm_channel::CommMsg;
@@ -21,4 +22,10 @@ pub enum CommEvent {
 
     /// A Comm was closed
     Closed(String),
+
+    /// Synchronisation barrier. Shell signals the sender after processing all
+    /// preceding events. The caller blocks on the paired receiver to guarantee
+    /// that earlier events (e.g. `Opened`) have been fully handled before
+    /// continuing.
+    Barrier(Sender<()>),
 }

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;
@@ -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.
     ///

crates/amalthea/src/socket/shell.rs

@@ -37,6 +37,8 @@ use crate::wire::comm_info_request::CommInfoRequest;
 use crate::wire::comm_msg::CommWireMsg;
 use crate::wire::comm_open::CommOpen;
 use crate::wire::exception::Exception;
+use crate::wire::execute_reply::ExecuteReply;
+use crate::wire::execute_request::ExecuteRequest;
 use crate::wire::header::JupyterHeader;
 use crate::wire::jupyter_message::JupyterMessage;
 use crate::wire::jupyter_message::Message;
@@ -206,6 +208,10 @@ impl Shell {
                 );
             },
 
+            CommEvent::Barrier(done_tx) => {
+                done_tx.send(()).log_err();
+            },
+
             CommEvent::Message(comm_id, msg) => {
                 let Some(comm) = self.open_comms.iter().find(|c| c.comm_id == comm_id) else {
                     log::warn!("Received message for unknown comm channel {comm_id}: {msg:?}");
@@ -241,6 +247,14 @@ impl Shell {
     /// Process a message received from the front-end, optionally dispatching
     /// messages to the IOPub or execution threads
     fn process_message(&mut self, msg: Message) -> crate::Result<()> {
+        // Execute requests get special handling: Shell select-loops on both
+        // the execute response and comm events, allowing it to process
+        // backend-initiated comm opens (with barrier handshakes) without
+        // deadlocking.
+        if let Message::ExecuteRequest(req) = msg {
+            return self.handle_execute_request(req);
+        }
+
         // Extract references to the components we need to pass to handlers.
         // This allows us to borrow different fields of self independently.
         let iopub_tx = &self.iopub_tx;
@@ -257,13 +271,6 @@ impl Shell {
             Message::IsCompleteRequest(req) => Self::handle_request(iopub_tx, socket, req, |msg| {
                 block_on(shell_handler.handle_is_complete_request(msg))
             }),
-            Message::ExecuteRequest(req) => {
-                // FIXME: We should ideally not pass the originator to the language kernel
-                let originator = Originator::from(&req);
-                Self::handle_request(iopub_tx, socket, req, |msg| {
-                    block_on(shell_handler.handle_execute_request(originator, msg))
-                })
-            },
             Message::CompleteRequest(req) => Self::handle_request(iopub_tx, socket, req, |msg| {
                 block_on(shell_handler.handle_complete_request(msg))
             }),
@@ -364,6 +371,68 @@ impl Shell {
         result.and(Ok(()))
     }
 
+    /// Handle an execute request. Unlike other requests that use the generic
+    /// `handle_request`, this method select-loops on both the execute response
+    /// and `comm_event_rx`. This allows Shell to process comm events (e.g.
+    /// `CommEvent::Barrier` from `comm_open_backend`) while the R thread is
+    /// still executing, preventing a deadlock where the R thread waits for
+    /// Shell to drain comm events while Shell waits for the execute response.
+    fn handle_execute_request(&mut self, req: JupyterMessage<ExecuteRequest>) -> crate::Result<()> {
+        use crossbeam::channel::Select;
+
+        self.iopub_tx
+            .send(status(req.clone(), ExecutionState::Busy))
+            .unwrap();
+
+        log::info!("Received shell request: {req:?}");
+
+        // FIXME: We should ideally not pass the originator to the language kernel
+        let originator = Originator::from(&req);
+        let response_rx = self
+            .shell_handler
+            .start_execute_request(originator, &req.content);
+
+        // Select-loop: drain comm events while waiting for the execute reply.
+        let result = loop {
+            let mut sel = Select::new();
+            let resp_idx = sel.recv(&response_rx);
+            let comm_idx = sel.recv(&self.comm_event_rx);
+
+            let ready = sel.ready();
+            if ready == resp_idx {
+                // Drain any comm events that arrived before or alongside the response
+                while let Ok(event) = self.comm_event_rx.try_recv() {
+                    self.process_comm_event(event);
+                }
+                break response_rx.recv().unwrap();
+            } else if ready == comm_idx {
+                if let Ok(event) = self.comm_event_rx.try_recv() {
+                    self.process_comm_event(event);
+                }
+            }
+        };
+
+        let result = match result {
+            Ok(reply) => req.send_reply(reply, &self.socket),
+            Err(crate::Error::ShellErrorReply(error)) => {
+                req.send_error::<ExecuteReply>(error, &self.socket)
+            },
+            Err(crate::Error::ShellErrorExecuteReply(error, exec_count)) => {
+                req.send_execute_error(error, exec_count, &self.socket)
+            },
+            Err(err) => {
+                let error = Exception::internal_error(format!("{err:?}"));
+                req.send_error::<ExecuteReply>(error, &self.socket)
+            },
+        };
+
+        self.iopub_tx
+            .send(status(req.clone(), ExecutionState::Idle))
+            .unwrap();
+
+        result.and(Ok(()))
+    }
+
     fn handle_notification<Not, Handler>(
         iopub_tx: &Sender<IOPubMessage>,
         not: JupyterMessage<Not>,

crates/amalthea/tests/client/shell.rs

@@ -82,63 +82,8 @@ impl Shell {
             warn!("Could not prompt for input: {}", err);
         }
     }
-}
-
-#[async_trait]
-impl ShellHandler for Shell {
-    async fn handle_info_request(
-        &mut self,
-        _req: &KernelInfoRequest,
-    ) -> amalthea::Result<KernelInfoReply> {
-        let info = LanguageInfo {
-            name: String::from("Test"),
-            version: String::from("1.0"),
-            file_extension: String::from(".ech"),
-            mimetype: String::from("text/echo"),
-            pygments_lexer: None,
-            codemirror_mode: None,
-            nbconvert_exporter: None,
-            positron: None,
-        };
-        Ok(KernelInfoReply {
-            status: Status::Ok,
-            banner: format!("Amalthea Echo {}", env!("CARGO_PKG_VERSION")),
-            implementation: String::from("echo"),
-            implementation_version: String::from(env!("CARGO_PKG_VERSION")),
-            debugger: false,
-            help_links: Vec::new(),
-            language_info: info,
-        })
-    }
-
-    async fn handle_complete_request(
-        &self,
-        _req: &CompleteRequest,
-    ) -> amalthea::Result<CompleteReply> {
-        // No matches in this toy implementation.
-        Ok(CompleteReply {
-            matches: Vec::new(),
-            status: Status::Ok,
-            cursor_start: 0,
-            cursor_end: 0,
-            metadata: json!({}),
-        })
-    }
 
-    /// Handle a request to test code for completion.
-    async fn handle_is_complete_request(
-        &self,
-        _req: &IsCompleteRequest,
-    ) -> amalthea::Result<IsCompleteReply> {
-        // In this echo example, the code is always complete!
-        Ok(IsCompleteReply {
-            status: IsComplete::Complete,
-            indent: String::from(""),
-        })
-    }
-
-    /// Handles an ExecuteRequest; "executes" the code by echoing it.
-    async fn handle_execute_request(
+    fn execute(
         &mut self,
         originator: Originator,
         req: &ExecuteRequest,
@@ -227,6 +172,72 @@ impl ShellHandler for Shell {
             user_expressions: serde_json::Value::Null,
         })
     }
+}
+
+#[async_trait]
+impl ShellHandler for Shell {
+    async fn handle_info_request(
+        &mut self,
+        _req: &KernelInfoRequest,
+    ) -> amalthea::Result<KernelInfoReply> {
+        let info = LanguageInfo {
+            name: String::from("Test"),
+            version: String::from("1.0"),
+            file_extension: String::from(".ech"),
+            mimetype: String::from("text/echo"),
+            pygments_lexer: None,
+            codemirror_mode: None,
+            nbconvert_exporter: None,
+            positron: None,
+        };
+        Ok(KernelInfoReply {
+            status: Status::Ok,
+            banner: format!("Amalthea Echo {}", env!("CARGO_PKG_VERSION")),
+            implementation: String::from("echo"),
+            implementation_version: String::from(env!("CARGO_PKG_VERSION")),
+            debugger: false,
+            help_links: Vec::new(),
+            language_info: info,
+        })
+    }
+
+    async fn handle_complete_request(
+        &self,
+        _req: &CompleteRequest,
+    ) -> amalthea::Result<CompleteReply> {
+        // No matches in this toy implementation.
+        Ok(CompleteReply {
+            matches: Vec::new(),
+            status: Status::Ok,
+            cursor_start: 0,
+            cursor_end: 0,
+            metadata: json!({}),
+        })
+    }
+
+    /// Handle a request to test code for completion.
+    async fn handle_is_complete_request(
+        &self,
+        _req: &IsCompleteRequest,
+    ) -> amalthea::Result<IsCompleteReply> {
+        // In this echo example, the code is always complete!
+        Ok(IsCompleteReply {
+            status: IsComplete::Complete,
+            indent: String::from(""),
+        })
+    }
+
+    /// Handles an ExecuteRequest; "executes" the code by echoing it.
+    fn start_execute_request(
+        &mut self,
+        originator: Originator,
+        req: &ExecuteRequest,
+    ) -> crossbeam::channel::Receiver<amalthea::Result<ExecuteReply>> {
+        let (tx, rx) = crossbeam::channel::bounded(1);
+        let result = self.execute(originator, req);
+        tx.send(result).unwrap();
+        rx
+    }
 
     /// Handles an introspection request
     async fn handle_inspect_request(&self, req: &InspectRequest) -> amalthea::Result<InspectReply> {

crates/ark/src/console/console_comm.rs

@@ -40,7 +40,12 @@ impl Console {
     /// Register a backend-initiated comm on the R thread.
     ///
     /// Creates the `CommSocket` and `CommHandlerContext`, calls `handle_open`,
-    /// sends `CommEvent::Opened` to amalthea, and returns the comm ID.
+    /// sends `CommEvent::Opened` to Shell, and returns the comm ID.
+    ///
+    /// Blocks until Shell has fully processed the open (sent `comm_open` on
+    /// IOPub and registered the comm for routing). This guarantees that any
+    /// `comm_msg` updates the caller sends afterwards are ordered after the
+    /// `comm_open` on IOPub.
     pub(crate) fn comm_open_backend(
         &mut self,
         comm_name: &str,
@@ -65,6 +70,13 @@ impl Console {
         self.comm_event_tx
             .send(CommEvent::Opened(comm, open_metadata))?;
 
+        // Block until Shell has processed the Opened event, ensuring the
+        // `comm_open` message is on IOPub before we return. Any updates
+        // the caller sends after this point are guaranteed to follow it.
+        let (done_tx, done_rx) = crossbeam::channel::bounded(0);
+        self.comm_event_tx.send(CommEvent::Barrier(done_tx))?;
+        done_rx.recv()?;
+
         Ok(comm_id)
     }
 

crates/ark/src/modules/positron/graphics.R

@@ -80,6 +80,10 @@ setHook("before.grid.newpage", action = "replace", function(...) {
     grDevices::deviceIsInteractive(ARK_GRAPHICS_DEVICE_NAME)
 }
 
+current_plot_id <- function() {
+    .ps.Call("ps_graphics_current_plot_id")
+}
+
 # Create a recording of the current plot.
 #
 # This saves the plot's display list, so it can be used to re-render plots as

crates/ark/src/plots/graphics_device.rs

@@ -1358,6 +1358,14 @@ unsafe extern "C-unwind" fn ps_graphics_get_metadata(id: SEXP) -> anyhow::Result
     }
 }
 
+/// Return the current plot ID. Used by tests to verify that layout panels
+/// share the same page (same ID) and that overflow creates a new page.
+#[harp::register]
+unsafe extern "C-unwind" fn ps_graphics_current_plot_id() -> anyhow::Result<SEXP> {
+    let id = Console::get().device_context().id();
+    Ok(RObject::from(&id).sexp)
+}
+
 /// Push a source file URI onto the source context stack.
 /// Called from the `source()` hook when entering a sourced file.
 #[harp::register]

crates/ark/src/shell.rs

@@ -175,11 +175,11 @@ impl ShellHandler for Shell {
 
     /// Handles an ExecuteRequest by sending the code to the R execution thread
     /// for processing.
-    async fn handle_execute_request(
+    fn start_execute_request(
         &mut self,
         originator: Originator,
         req: &ExecuteRequest,
-    ) -> amalthea::Result<ExecuteReply> {
+    ) -> crossbeam::channel::Receiver<amalthea::Result<ExecuteReply>> {
         let (response_tx, response_rx) = unbounded::<amalthea::Result<ExecuteReply>>();
         let mut req_clone = req.clone();
         req_clone.code = convert_line_endings(&req_clone.code, LineEnding::Posix);
@@ -196,7 +196,7 @@ impl ShellHandler for Shell {
 
         trace!("Code sent to R: {}", req_clone.code);
 
-        response_rx.recv().unwrap()
+        response_rx
     }
 
     /// Handles an introspection request