Fix timeout on Windows ARM

Author: lionel-

crates/amalthea/src/fixtures/dummy_frontend.rs

@@ -316,6 +316,7 @@ impl DummyFrontend {
     }
 
     #[track_caller]
+    #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
     pub fn recv(socket: &Socket) -> Message {
         // It's important to wait with a timeout because the kernel thread might have
         // panicked, preventing it from sending the expected message. The tests would then
@@ -335,6 +336,28 @@ impl DummyFrontend {
         panic!("Timeout while expecting message on socket {}", socket.name);
     }
 
+    /// On Windows ARM, ZMQ poll with a non-zero timeout blocks forever
+    /// instead of respecting the timeout. Use non-blocking poll with
+    /// manual timing, matching the workaround in `recv_iopub_with_timeout`.
+    #[track_caller]
+    #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+    pub fn recv(socket: &Socket) -> Message {
+        let start = std::time::Instant::now();
+        let timeout = std::time::Duration::from_secs(30);
+
+        loop {
+            if start.elapsed() >= timeout {
+                panic!("Timeout while expecting message on socket {}", socket.name);
+            }
+
+            match socket.poll_incoming(0) {
+                Ok(true) => return Message::read_from_socket(socket).unwrap(),
+                Ok(false) => std::thread::sleep(std::time::Duration::from_millis(10)),
+                Err(e) => panic!("Error polling socket {}: {e}", socket.name),
+            }
+        }
+    }
+
     /// Receives a Jupyter message from the Control socket
     #[track_caller]
     pub fn recv_control(&self) -> Message {

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

@@ -123,9 +123,25 @@ impl Shell {
                     self.comm_notif_socket.socket.as_poll_item(zmq::POLLIN),
                 ];
 
-                // -1 means block indefinitely
-                match zmq::poll(&mut poll_items, -1) {
-                    Ok(0) => continue,
+                // On Windows ARM, zmq::poll with a non-zero timeout blocks
+                // forever instead of respecting the timeout.
+                #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+                let poll_timeout_ms: i64 = 0;
+                #[cfg(not(all(target_os = "windows", target_arch = "aarch64")))]
+                let poll_timeout_ms: i64 = -1;
+
+                match zmq::poll(&mut poll_items, poll_timeout_ms) {
+                    Ok(0) => {
+                        // On Windows ARM, check for comm events even when
+                        // poll reports nothing ready since inproc
+                        // notifications may not wake the poll.
+                        #[cfg(all(target_os = "windows", target_arch = "aarch64"))]
+                        {
+                            self.process_comm_notification();
+                            std::thread::sleep(std::time::Duration::from_millis(1));
+                        }
+                        continue;
+                    },
                     Ok(_) => (poll_items[0].is_readable(), poll_items[1].is_readable()),
                     Err(err) => {
                         log::warn!("Could not poll shell sockets: {err:?}");