Fix polling on Windows ARM

Author: lionel-

crates/amalthea/src/kernel.rs

@@ -531,37 +531,69 @@ fn socket_bridge_thread(
     };
 
     loop {
-        let n = unwrap!(
-            zmq::poll(&mut poll_items, -1),
-            Err(err) => {
-                debug_panic!("While polling 0MQ items: {err:?}");
-                0
-            }
-        );
-
-        for _ in 0..n {
-            if consume_outbound_notification() {
-                forward_outbound();
-                continue;
-            }
-
+        // On Windows ARM, zmq::poll with a non-zero timeout blocks forever
+        // and inproc notification sockets may not wake the poll at all.
+        // Use a fully separate polling path that doesn't rely on ZMQ
+        // readability reporting: non-blocking poll + unconditional drain
+        // of all sources + short sleep when idle.
+        #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+        {
+            // Drain outbound messages (IOPub, StdIn) unconditionally
+            consume_outbound_notification();
+            forward_outbound();
+
+            // Check inbound sockets with non-blocking poll
             if has_inbound(&stdin_socket) {
                 unwrap!(
                     forward_inbound(&stdin_socket, &stdin_inbound_tx),
                     Err(err) => debug_panic!("While forwarding inbound message: {err:?}")
                 );
-                continue;
             }
-
             if has_inbound(&iopub_socket) {
                 unwrap!(
                     forward_inbound_subscription(&iopub_socket, &iopub_inbound_tx),
                     Err(err) => debug_panic!("While forwarding inbound message: {err:?}")
                 );
-                continue;
             }
 
-            debug_panic!("Could not find readable message");
+            std::thread::sleep(std::time::Duration::from_millis(1));
+            continue;
+        }
+
+        #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
+        {
+            let n = unwrap!(
+                zmq::poll(&mut poll_items, -1),
+                Err(err) => {
+                    debug_panic!("While polling 0MQ items: {err:?}");
+                    0
+                }
+            );
+
+            for _ in 0..n {
+                if consume_outbound_notification() {
+                    forward_outbound();
+                    continue;
+                }
+
+                if has_inbound(&stdin_socket) {
+                    unwrap!(
+                        forward_inbound(&stdin_socket, &stdin_inbound_tx),
+                        Err(err) => debug_panic!("While forwarding inbound message: {err:?}")
+                    );
+                    continue;
+                }
+
+                if has_inbound(&iopub_socket) {
+                    unwrap!(
+                        forward_inbound_subscription(&iopub_socket, &iopub_inbound_tx),
+                        Err(err) => debug_panic!("While forwarding inbound message: {err:?}")
+                    );
+                    continue;
+                }
+
+                debug_panic!("Could not find readable message");
+            }
         }
     }
 }

crates/amalthea/src/socket.rs

@@ -214,10 +214,38 @@ impl Socket {
         }
     }
 
+    #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
     pub fn poll_incoming(&self, timeout_ms: i64) -> zmq::Result<bool> {
         Ok(self.socket.poll(zmq::PollEvents::POLLIN, timeout_ms)? != 0)
     }
 
+    /// On Windows ARM, ZMQ poll with a non-zero timeout blocks forever
+    /// instead of respecting the timeout. Use non-blocking poll with
+    /// manual timing.
+    #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+    pub fn poll_incoming(&self, timeout_ms: i64) -> zmq::Result<bool> {
+        if timeout_ms == 0 {
+            return Ok(self.socket.poll(zmq::PollEvents::POLLIN, 0)? != 0);
+        }
+
+        let start = std::time::Instant::now();
+        let timeout = if timeout_ms < 0 {
+            std::time::Duration::from_secs(u64::MAX / 2)
+        } else {
+            std::time::Duration::from_millis(timeout_ms as u64)
+        };
+
+        loop {
+            if self.socket.poll(zmq::PollEvents::POLLIN, 0)? != 0 {
+                return Ok(true);
+            }
+            if start.elapsed() >= timeout {
+                return Ok(false);
+            }
+            std::thread::sleep(std::time::Duration::from_millis(1));
+        }
+    }
+
     pub fn has_incoming_data(&self) -> zmq::Result<bool> {
         self.poll_incoming(0)
     }

crates/amalthea/src/socket/shell.rs

@@ -111,6 +111,27 @@ impl Shell {
 
     /// Main loop for the Shell thread; to be invoked by the kernel.
     pub fn listen(&mut self) {
+        #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+        self.listen_polling();
+
+        #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
+        self.listen_blocking();
+    }
+
+    /// On Windows ARM, zmq::poll with a non-zero timeout blocks forever
+    /// and inproc notification sockets may not wake the poll. Use
+    /// non-blocking poll with unconditional comm checks and a short sleep.
+    #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+    fn listen_polling(&mut self) {
+        loop {
+            self.process_comm_notification();
+            self.process_shell_socket();
+            std::thread::sleep(std::time::Duration::from_millis(1));
+        }
+    }
+
+    #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
+    fn listen_blocking(&mut self) {
         loop {
             log::trace!("Waiting for shell messages or comm events");
 
@@ -139,24 +160,32 @@ impl Shell {
             }
 
             if shell_readable {
-                let message = match Message::read_from_socket(&self.socket) {
-                    Ok(m) => m,
-                    Err(err) => {
-                        log::warn!("Could not read message from shell socket: {err:?}");
-                        continue;
-                    },
-                };
-
-                // Handle the message; any failures while handling the messages are
-                // delivered to the client instead of reported up the stack, so the
-                // only errors likely here are "can't deliver to client"
-                if let Err(err) = self.process_message(message) {
-                    log::error!("Could not handle shell message: {err:?}");
-                }
+                self.process_shell_socket();
             }
         }
     }
 
+    fn process_shell_socket(&mut self) {
+        if !self.socket.has_incoming_data().unwrap_or(false) {
+            return;
+        }
+
+        let message = match Message::read_from_socket(&self.socket) {
+            Ok(m) => m,
+            Err(err) => {
+                log::warn!("Could not read message from shell socket: {err:?}");
+                return;
+            },
+        };
+
+        // Handle the message; any failures while handling the messages are
+        // delivered to the client instead of reported up the stack, so the
+        // only errors likely here are "can't deliver to client"
+        if let Err(err) = self.process_message(message) {
+            log::error!("Could not handle shell message: {err:?}");
+        }
+    }
+
     /// Process comm event notifications from the notifier thread.
     /// Drains all pending notifications and all pending events.
     fn process_comm_notification(&mut self) {

crates/ark_test/src/dummy_frontend.rs

@@ -248,7 +248,6 @@ impl DummyArkFrontend {
 
     /// Receive from IOPub with a timeout.
     /// Returns `None` if the timeout expires before a message arrives.
-    #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
     pub fn recv_iopub_with_timeout(&self, timeout: Duration) -> Option<Message> {
         let timeout_ms = timeout.as_millis() as i64;
         if self.guard.iopub_socket.poll_incoming(timeout_ms).unwrap() {
@@ -258,34 +257,6 @@ impl DummyArkFrontend {
         }
     }
 
-    /// Receive from IOPub with a timeout.
-    /// Returns `None` if the timeout expires before a message arrives.
-    ///
-    /// On Windows ARM, ZMQ poll with timeout blocks forever instead of
-    /// respecting the timeout. Use non-blocking poll with manual timing.
-    #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
-    pub fn recv_iopub_with_timeout(&self, timeout: Duration) -> Option<Message> {
-        let start = std::time::Instant::now();
-
-        loop {
-            if start.elapsed() >= timeout {
-                return None;
-            }
-
-            // Use non-blocking poll (timeout=0) to avoid ZMQ blocking forever
-            match self.guard.iopub_socket.poll_incoming(0) {
-                Ok(true) => {
-                    return Some(Message::read_from_socket(&self.guard.iopub_socket).unwrap());
-                },
-                Ok(false) => {
-                    // No message available, sleep briefly and try again
-                    std::thread::sleep(Duration::from_millis(10));
-                },
-                Err(_) => return None,
-            }
-        }
-    }
-
     /// Core primitive: receive the next non-stream, non-variables-comm message
     /// from IOPub.
     ///