microsandbox_utils/
wake_pipe.rs

1//! Cross-platform wake notification.
2//!
3//! Works on both Linux and macOS (unlike `eventfd` which is Linux-only).
4//! The write end signals, the read end is pollable via `epoll`/`kqueue`/`poll`.
5
6#[cfg(unix)]
7use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
8#[cfg(windows)]
9use std::os::windows::io::RawHandle;
10use std::time::Duration;
11#[cfg(windows)]
12use windows_sys::Win32::Foundation::{
13    CloseHandle, HANDLE, WAIT_FAILED, WAIT_OBJECT_0, WAIT_TIMEOUT,
14};
15#[cfg(windows)]
16use windows_sys::Win32::System::Threading::{
17    CreateEventW, ResetEvent, SetEvent, WaitForSingleObject,
18};
19
20//--------------------------------------------------------------------------------------------------
21// Types
22//--------------------------------------------------------------------------------------------------
23
24/// Cross-platform wake notification.
25///
26/// On Unix, the write end signals and the read end is pollable via
27/// `epoll`/`kqueue`/`poll`. On Windows, wakeups are coalesced behind a
28/// manual-reset event for code that needs a blocking wait.
29#[cfg(unix)]
30pub struct WakePipe {
31    read_fd: OwnedFd,
32    write_fd: OwnedFd,
33}
34
35/// Cross-platform wake notification.
36///
37/// Windows wakeups are coalesced behind a manual-reset event. Multiple
38/// [`wake`](Self::wake) calls before [`drain`](Self::drain) still represent one
39/// readable/ready state, matching the Unix pipe behavior needed by queue
40/// wakeups.
41#[cfg(windows)]
42pub struct WakePipe {
43    handle: HANDLE,
44    async_wake: tokio::sync::Notify,
45}
46
47//--------------------------------------------------------------------------------------------------
48// Methods
49//--------------------------------------------------------------------------------------------------
50
51impl WakePipe {
52    /// Create a new wake pipe.
53    ///
54    /// Both ends are set to non-blocking and close-on-exec.
55    #[cfg(unix)]
56    pub fn new() -> Self {
57        let mut fds = [0i32; 2];
58
59        // SAFETY: pipe() is a standard POSIX call. We check the return value
60        // and immediately wrap the raw fds in OwnedFd for RAII cleanup.
61        let ret = unsafe { libc::pipe(fds.as_mut_ptr()) };
62        assert!(
63            ret == 0,
64            "pipe() failed: {}",
65            std::io::Error::last_os_error()
66        );
67
68        // Set non-blocking and close-on-exec on both ends.
69        // SAFETY: fds are valid open file descriptors from the pipe() call above.
70        unsafe {
71            set_nonblock_cloexec(fds[0]);
72            set_nonblock_cloexec(fds[1]);
73        }
74
75        Self {
76            // SAFETY: fds are valid and not owned by anything else yet.
77            read_fd: unsafe { OwnedFd::from_raw_fd(fds[0]) },
78            write_fd: unsafe { OwnedFd::from_raw_fd(fds[1]) },
79        }
80    }
81
82    /// Create a new wake primitive.
83    #[cfg(windows)]
84    pub fn new() -> Self {
85        let handle = unsafe { CreateEventW(std::ptr::null(), 1, 0, std::ptr::null()) };
86        if handle.is_null() {
87            panic!("CreateEventW failed: {}", std::io::Error::last_os_error());
88        }
89        Self {
90            handle,
91            async_wake: tokio::sync::Notify::new(),
92        }
93    }
94
95    /// Signal the reader. Safe to call from any thread, multiple times.
96    ///
97    /// Writes a single byte. If the pipe buffer is full the write is silently
98    /// dropped — the reader will still wake because there are unread bytes.
99    #[cfg(unix)]
100    pub fn wake(&self) {
101        // SAFETY: write_fd is a valid, non-blocking file descriptor.
102        // Writing 1 byte to a pipe is atomic on all POSIX systems.
103        unsafe {
104            libc::write(self.write_fd.as_raw_fd(), [1u8].as_ptr().cast(), 1);
105        }
106    }
107
108    /// Signal the reader. Safe to call from any thread, multiple times.
109    #[cfg(windows)]
110    pub fn wake(&self) {
111        unsafe {
112            SetEvent(self.handle);
113        }
114        // Keep an async permit even when the reader has not started waiting yet.
115        // Native waiters still use the manual-reset event independently.
116        self.async_wake.notify_one();
117    }
118
119    /// Wait asynchronously for a wake without occupying a blocking worker.
120    ///
121    /// Intended for one async queue reader. Signals coalesce into one permit;
122    /// dropping this future is cancellation-safe. This does not drain the native
123    /// event: drain it before inspecting the queue, just as with a blocking wait.
124    /// Native drains do not consume async permits, so a reader must tolerate a
125    /// notification whose queued work has already been consumed.
126    #[cfg(windows)]
127    pub async fn notified(&self) {
128        self.async_wake.notified().await;
129    }
130
131    /// Drain all pending wake signals. Call after processing to reset the
132    /// pipe for the next edge-triggered notification.
133    #[cfg(unix)]
134    pub fn drain(&self) {
135        let mut buf = [0u8; 512];
136        loop {
137            // SAFETY: read_fd is a valid, non-blocking file descriptor.
138            let n =
139                unsafe { libc::read(self.read_fd.as_raw_fd(), buf.as_mut_ptr().cast(), buf.len()) };
140            if n <= 0 {
141                break;
142            }
143        }
144    }
145
146    /// Drain all pending wake signals.
147    #[cfg(windows)]
148    pub fn drain(&self) {
149        unsafe {
150            ResetEvent(self.handle);
151        }
152    }
153
154    /// Block until the wake primitive is signaled or the timeout expires.
155    ///
156    /// Returns `true` if the primitive was signaled and `false` on timeout.
157    pub fn wait_timeout(&self, timeout: Duration) -> bool {
158        wait_timeout(self, timeout)
159    }
160
161    /// File descriptor for `epoll`/`kqueue`/`poll(2)` registration.
162    ///
163    /// Becomes readable when [`wake()`](Self::wake) has been called.
164    #[cfg(unix)]
165    pub fn as_raw_fd(&self) -> RawFd {
166        self.read_fd.as_raw_fd()
167    }
168
169    /// Waitable handle for Windows wait APIs.
170    #[cfg(windows)]
171    pub fn as_raw_handle(&self) -> RawHandle {
172        self.handle as RawHandle
173    }
174}
175
176//--------------------------------------------------------------------------------------------------
177// Trait Implementations
178//--------------------------------------------------------------------------------------------------
179
180impl Default for WakePipe {
181    fn default() -> Self {
182        Self::new()
183    }
184}
185
186#[cfg(windows)]
187unsafe impl Send for WakePipe {}
188
189#[cfg(windows)]
190unsafe impl Sync for WakePipe {}
191
192#[cfg(windows)]
193impl Drop for WakePipe {
194    fn drop(&mut self) {
195        unsafe {
196            CloseHandle(self.handle);
197        }
198    }
199}
200
201//--------------------------------------------------------------------------------------------------
202// Functions
203//--------------------------------------------------------------------------------------------------
204
205/// Set `O_NONBLOCK` and `FD_CLOEXEC` on a file descriptor.
206///
207/// # Safety
208///
209/// `fd` must be a valid, open file descriptor.
210#[cfg(unix)]
211unsafe fn set_nonblock_cloexec(fd: RawFd) {
212    unsafe {
213        // Set non-blocking.
214        let flags = libc::fcntl(fd, libc::F_GETFL);
215        assert!(
216            flags >= 0,
217            "fcntl(F_GETFL) failed: {}",
218            std::io::Error::last_os_error()
219        );
220        let ret = libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK);
221        assert!(
222            ret >= 0,
223            "fcntl(F_SETFL) failed: {}",
224            std::io::Error::last_os_error()
225        );
226
227        // Set close-on-exec.
228        let flags = libc::fcntl(fd, libc::F_GETFD);
229        assert!(
230            flags >= 0,
231            "fcntl(F_GETFD) failed: {}",
232            std::io::Error::last_os_error()
233        );
234        let ret = libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC);
235        assert!(
236            ret >= 0,
237            "fcntl(F_SETFD) failed: {}",
238            std::io::Error::last_os_error()
239        );
240    }
241}
242
243#[cfg(unix)]
244fn wait_timeout(pipe: &WakePipe, timeout: Duration) -> bool {
245    let timeout_ms = timeout.as_millis().min(i32::MAX as u128) as i32;
246    poll_fd_readable_timeout(pipe.as_raw_fd(), timeout_ms)
247}
248
249#[cfg(unix)]
250fn poll_fd_readable_timeout(fd: RawFd, timeout_ms: i32) -> bool {
251    loop {
252        let mut pfd = libc::pollfd {
253            fd,
254            events: libc::POLLIN,
255            revents: 0,
256        };
257        // SAFETY: pfd is a valid stack-allocated pollfd.
258        let ret = unsafe { libc::poll(&mut pfd, 1, timeout_ms) };
259        if ret > 0 {
260            return true;
261        }
262        if ret == 0 {
263            return false;
264        }
265
266        let errno = std::io::Error::last_os_error();
267        if errno.raw_os_error() != Some(libc::EINTR) {
268            return false;
269        }
270    }
271}
272
273#[cfg(windows)]
274fn wait_timeout(pipe: &WakePipe, timeout: Duration) -> bool {
275    let timeout_ms = timeout.as_millis().min(u32::MAX as u128) as u32;
276    let result = unsafe { WaitForSingleObject(pipe.handle, timeout_ms) };
277    match result {
278        WAIT_OBJECT_0 => true,
279        WAIT_TIMEOUT | WAIT_FAILED => false,
280        _ => false,
281    }
282}
283
284//--------------------------------------------------------------------------------------------------
285// Tests
286//--------------------------------------------------------------------------------------------------
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291
292    #[cfg(windows)]
293    #[test]
294    fn async_wake_survives_cancellation_and_native_drain() {
295        use std::future::Future;
296        use std::sync::Arc;
297        use std::sync::atomic::{AtomicBool, Ordering};
298        use std::task::{Context, Poll, Wake, Waker};
299
300        struct Flag(AtomicBool);
301        impl Wake for Flag {
302            fn wake(self: Arc<Self>) {
303                self.0.store(true, Ordering::SeqCst);
304            }
305        }
306
307        let flag = Arc::new(Flag(AtomicBool::new(false)));
308        let waker = Waker::from(flag.clone());
309        let mut context = Context::from_waker(&waker);
310        let pipe = WakePipe::new();
311
312        // Pre-registration wakes coalesce, and a native drain cannot steal the
313        // async notification. Neither waiting mechanism consumes the other's signal.
314        pipe.wake();
315        pipe.wake();
316        pipe.drain();
317        assert_eq!(
318            Box::pin(pipe.notified()).as_mut().poll(&mut context),
319            Poll::Ready(())
320        );
321        let mut cancelled = Box::pin(pipe.notified());
322        assert!(cancelled.as_mut().poll(&mut context).is_pending());
323        drop(cancelled);
324
325        let mut waiting = Box::pin(pipe.notified());
326        assert!(waiting.as_mut().poll(&mut context).is_pending());
327        pipe.wake();
328        assert!(flag.0.load(Ordering::SeqCst));
329        // Cancellation after notification must also leave the permit available.
330        drop(waiting);
331        assert_eq!(
332            Box::pin(pipe.notified()).as_mut().poll(&mut context),
333            Poll::Ready(())
334        );
335        assert!(pipe.wait_timeout(Duration::ZERO));
336        pipe.drain();
337        assert!(
338            Box::pin(pipe.notified())
339                .as_mut()
340                .poll(&mut context)
341                .is_pending()
342        );
343
344        // Repeated waits exercise reset/wake ordering without periodic polling.
345        for _ in 0..100 {
346            let mut waiting = Box::pin(pipe.notified());
347            assert!(waiting.as_mut().poll(&mut context).is_pending());
348            pipe.wake();
349            assert!(waiting.as_mut().poll(&mut context).is_ready());
350            pipe.drain();
351        }
352    }
353
354    #[test]
355    fn wake_and_drain() {
356        let pipe = WakePipe::new();
357        // Initially no data — drain is a no-op.
358        pipe.drain();
359
360        // Wake then drain.
361        pipe.wake();
362        pipe.wake();
363        pipe.drain();
364
365        // After drain, another wake should work.
366        pipe.wake();
367        pipe.drain();
368    }
369
370    #[cfg(unix)]
371    #[test]
372    fn fd_is_valid() {
373        let pipe = WakePipe::new();
374        let fd = pipe.as_raw_fd();
375        assert!(fd >= 0);
376    }
377
378    #[test]
379    fn nonblocking_read() {
380        let pipe = WakePipe::new();
381        // Reading from an empty non-blocking pipe should not block.
382        pipe.drain();
383    }
384
385    #[test]
386    fn wait_timeout_observes_wake() {
387        let pipe = WakePipe::new();
388
389        assert!(!pipe.wait_timeout(Duration::from_millis(1)));
390        pipe.wake();
391        assert!(pipe.wait_timeout(Duration::from_secs(1)));
392        pipe.drain();
393        assert!(!pipe.wait_timeout(Duration::from_millis(1)));
394    }
395}
microsandbox_utils/wake_pipe.rs

microsandbox_utils/
wake_pipe.rs
