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    /// Clone the pollable read end for an independently owned registration.
170    #[cfg(unix)]
171    pub fn try_clone_read_fd(&self) -> std::io::Result<OwnedFd> {
172        self.read_fd.try_clone()
173    }
174
175    /// Waitable handle for Windows wait APIs.
176    #[cfg(windows)]
177    pub fn as_raw_handle(&self) -> RawHandle {
178        self.handle as RawHandle
179    }
180}
181
182//--------------------------------------------------------------------------------------------------
183// Trait Implementations
184//--------------------------------------------------------------------------------------------------
185
186impl Default for WakePipe {
187    fn default() -> Self {
188        Self::new()
189    }
190}
191
192#[cfg(windows)]
193unsafe impl Send for WakePipe {}
194
195#[cfg(windows)]
196unsafe impl Sync for WakePipe {}
197
198#[cfg(windows)]
199impl Drop for WakePipe {
200    fn drop(&mut self) {
201        unsafe {
202            CloseHandle(self.handle);
203        }
204    }
205}
206
207//--------------------------------------------------------------------------------------------------
208// Functions
209//--------------------------------------------------------------------------------------------------
210
211/// Set `O_NONBLOCK` and `FD_CLOEXEC` on a file descriptor.
212///
213/// # Safety
214///
215/// `fd` must be a valid, open file descriptor.
216#[cfg(unix)]
217unsafe fn set_nonblock_cloexec(fd: RawFd) {
218    unsafe {
219        // Set non-blocking.
220        let flags = libc::fcntl(fd, libc::F_GETFL);
221        assert!(
222            flags >= 0,
223            "fcntl(F_GETFL) failed: {}",
224            std::io::Error::last_os_error()
225        );
226        let ret = libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK);
227        assert!(
228            ret >= 0,
229            "fcntl(F_SETFL) failed: {}",
230            std::io::Error::last_os_error()
231        );
232
233        // Set close-on-exec.
234        let flags = libc::fcntl(fd, libc::F_GETFD);
235        assert!(
236            flags >= 0,
237            "fcntl(F_GETFD) failed: {}",
238            std::io::Error::last_os_error()
239        );
240        let ret = libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC);
241        assert!(
242            ret >= 0,
243            "fcntl(F_SETFD) failed: {}",
244            std::io::Error::last_os_error()
245        );
246    }
247}
248
249#[cfg(unix)]
250fn wait_timeout(pipe: &WakePipe, timeout: Duration) -> bool {
251    let timeout_ms = timeout.as_millis().min(i32::MAX as u128) as i32;
252    poll_fd_readable_timeout(pipe.as_raw_fd(), timeout_ms)
253}
254
255#[cfg(unix)]
256fn poll_fd_readable_timeout(fd: RawFd, timeout_ms: i32) -> bool {
257    loop {
258        let mut pfd = libc::pollfd {
259            fd,
260            events: libc::POLLIN,
261            revents: 0,
262        };
263        // SAFETY: pfd is a valid stack-allocated pollfd.
264        let ret = unsafe { libc::poll(&mut pfd, 1, timeout_ms) };
265        if ret > 0 {
266            return true;
267        }
268        if ret == 0 {
269            return false;
270        }
271
272        let errno = std::io::Error::last_os_error();
273        if errno.raw_os_error() != Some(libc::EINTR) {
274            return false;
275        }
276    }
277}
278
279#[cfg(windows)]
280fn wait_timeout(pipe: &WakePipe, timeout: Duration) -> bool {
281    let timeout_ms = timeout.as_millis().min(u32::MAX as u128) as u32;
282    let result = unsafe { WaitForSingleObject(pipe.handle, timeout_ms) };
283    match result {
284        WAIT_OBJECT_0 => true,
285        WAIT_TIMEOUT | WAIT_FAILED => false,
286        _ => false,
287    }
288}
289
290//--------------------------------------------------------------------------------------------------
291// Tests
292//--------------------------------------------------------------------------------------------------
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[cfg(windows)]
299    #[test]
300    fn async_wake_survives_cancellation_and_native_drain() {
301        use std::future::Future;
302        use std::sync::Arc;
303        use std::sync::atomic::{AtomicBool, Ordering};
304        use std::task::{Context, Poll, Wake, Waker};
305
306        struct Flag(AtomicBool);
307        impl Wake for Flag {
308            fn wake(self: Arc<Self>) {
309                self.0.store(true, Ordering::SeqCst);
310            }
311        }
312
313        let flag = Arc::new(Flag(AtomicBool::new(false)));
314        let waker = Waker::from(flag.clone());
315        let mut context = Context::from_waker(&waker);
316        let pipe = WakePipe::new();
317
318        // Pre-registration wakes coalesce, and a native drain cannot steal the
319        // async notification. Neither waiting mechanism consumes the other's signal.
320        pipe.wake();
321        pipe.wake();
322        pipe.drain();
323        assert_eq!(
324            Box::pin(pipe.notified()).as_mut().poll(&mut context),
325            Poll::Ready(())
326        );
327        let mut cancelled = Box::pin(pipe.notified());
328        assert!(cancelled.as_mut().poll(&mut context).is_pending());
329        drop(cancelled);
330
331        let mut waiting = Box::pin(pipe.notified());
332        assert!(waiting.as_mut().poll(&mut context).is_pending());
333        pipe.wake();
334        assert!(flag.0.load(Ordering::SeqCst));
335        // Cancellation after notification must also leave the permit available.
336        drop(waiting);
337        assert_eq!(
338            Box::pin(pipe.notified()).as_mut().poll(&mut context),
339            Poll::Ready(())
340        );
341        assert!(pipe.wait_timeout(Duration::ZERO));
342        pipe.drain();
343        assert!(
344            Box::pin(pipe.notified())
345                .as_mut()
346                .poll(&mut context)
347                .is_pending()
348        );
349
350        // Repeated waits exercise reset/wake ordering without periodic polling.
351        for _ in 0..100 {
352            let mut waiting = Box::pin(pipe.notified());
353            assert!(waiting.as_mut().poll(&mut context).is_pending());
354            pipe.wake();
355            assert!(waiting.as_mut().poll(&mut context).is_ready());
356            pipe.drain();
357        }
358    }
359
360    #[test]
361    fn wake_and_drain() {
362        let pipe = WakePipe::new();
363        // Initially no data — drain is a no-op.
364        pipe.drain();
365
366        // Wake then drain.
367        pipe.wake();
368        pipe.wake();
369        pipe.drain();
370
371        // After drain, another wake should work.
372        pipe.wake();
373        pipe.drain();
374    }
375
376    #[cfg(unix)]
377    #[test]
378    fn fd_is_valid() {
379        let pipe = WakePipe::new();
380        let fd = pipe.as_raw_fd();
381        assert!(fd >= 0);
382    }
383
384    #[test]
385    fn nonblocking_read() {
386        let pipe = WakePipe::new();
387        // Reading from an empty non-blocking pipe should not block.
388        pipe.drain();
389    }
390
391    #[test]
392    fn wait_timeout_observes_wake() {
393        let pipe = WakePipe::new();
394
395        assert!(!pipe.wait_timeout(Duration::from_millis(1)));
396        pipe.wake();
397        assert!(pipe.wait_timeout(Duration::from_secs(1)));
398        pipe.drain();
399        assert!(!pipe.wait_timeout(Duration::from_millis(1)));
400    }
401}
microsandbox_utils/wake_pipe.rs

microsandbox_utils/
wake_pipe.rs
