compio_runtime/fd/poll_fd/
mod.rs

1cfg_select! {
2    windows => {
3        #[path = "windows.rs"]
4        mod sys;
5    }
6    unix => {
7        #[path = "unix.rs"]
8        mod sys;
9    }
10    _ => {}
11}
12
13#[cfg(windows)]
14use std::os::windows::io::{AsRawSocket, RawSocket};
15use std::{
16    future::poll_fn,
17    io,
18    net::Shutdown,
19    ops::Deref,
20    pin::Pin,
21    task::{Context, Poll},
22};
23
24use compio_buf::IntoInner;
25use compio_driver::{AsFd, AsRawFd, BorrowedFd, RawFd, SharedFd, ToSharedFd};
26use socket2::{SockAddr, Socket};
27
28/// Providing functionalities to wait for readiness.
29///
30/// ## Platform specific
31/// * Windows: only supports sockets.
32#[derive(Debug)]
33pub struct PollFd<T: AsFd>(sys::PollFd<T>);
34
35impl<T: AsFd> PollFd<T> {
36    fn run_socket_op<R>(&self, f: impl FnOnce(&Socket) -> io::Result<R>) -> io::Result<R> {
37        sys::run_socket_op(self.0.as_fd(), f)
38    }
39
40    /// Create [`PollFd`] without attaching the source.
41    ///
42    /// Ready-based sources does not need to be attached.
43    pub fn new(source: T) -> io::Result<Self> {
44        Self::from_shared_fd(SharedFd::new(source))
45    }
46
47    /// Create [`PollFd`] from a shared file descriptor.
48    pub fn from_shared_fd(inner: SharedFd<T>) -> io::Result<Self> {
49        sys::run_socket_op(inner.as_fd(), |socket| socket.set_nonblocking(true))?;
50        Ok(Self(sys::PollFd::new(inner)?))
51    }
52}
53
54impl<T: AsFd + 'static> PollFd<T> {
55    /// Accept a connection from this socket.
56    pub async fn accept(&self) -> io::Result<(PollFd<Socket>, SockAddr)> {
57        poll_fn(|cx| self.poll_accept(cx)).await
58    }
59
60    /// Poll to accept a connection from this socket.
61    pub fn poll_accept(
62        &self,
63        cx: &mut Context<'_>,
64    ) -> Poll<io::Result<(PollFd<Socket>, SockAddr)>> {
65        self.poll_accept_with(cx, |source| {
66            let (socket, addr) = sys::run_socket_op(source.as_fd(), Socket::accept)?;
67            Ok((PollFd::new(socket)?, addr))
68        })
69    }
70
71    /// Connect this socket to the specified address.
72    pub async fn connect(&self, addr: &SockAddr) -> io::Result<()> {
73        match self.run_socket_op(|socket| socket.connect(addr)) {
74            Ok(()) => return Ok(()),
75            Err(e) if is_connect_pending(&e) => {}
76            Err(e) => return Err(e),
77        }
78
79        self.connect_ready().await?;
80        self.run_socket_op(|socket| match socket.take_error()? {
81            Some(e) => Err(e),
82            None => Ok(()),
83        })
84    }
85
86    /// Wait for accept readiness, before calling `accept`, or after `accept`
87    /// returns `WouldBlock`.
88    pub async fn accept_ready(&self) -> io::Result<()> {
89        poll_fn(|cx| self.poll_accept_ready(cx)).await
90    }
91
92    /// Wait for connect readiness.
93    pub async fn connect_ready(&self) -> io::Result<()> {
94        poll_fn(|cx| self.poll_connect_ready(cx)).await
95    }
96
97    /// Wait for read readiness.
98    pub async fn read_ready(&self) -> io::Result<()> {
99        poll_fn(|cx| self.poll_read_ready(cx)).await
100    }
101
102    /// Wait for write readiness.
103    pub async fn write_ready(&self) -> io::Result<()> {
104        poll_fn(|cx| self.poll_write_ready(cx)).await
105    }
106
107    /// Poll for accept readiness.
108    pub fn poll_accept_ready(&self, cx: &mut Context) -> Poll<io::Result<()>> {
109        self.0.poll_accept_ready(cx)
110    }
111
112    /// Poll for connect readiness.
113    pub fn poll_connect_ready(&self, cx: &mut Context) -> Poll<io::Result<()>> {
114        self.0.poll_connect_ready(cx)
115    }
116
117    /// Poll for read readiness.
118    pub fn poll_read_ready(&self, cx: &mut Context) -> Poll<io::Result<()>> {
119        self.0.poll_read_ready(cx)
120    }
121
122    /// Poll for write readiness.
123    pub fn poll_write_ready(&self, cx: &mut Context) -> Poll<io::Result<()>> {
124        self.0.poll_write_ready(cx)
125    }
126
127    /// Poll for accept readiness and call the provided function.
128    pub fn poll_accept_with<R>(
129        &self,
130        cx: &mut Context,
131        mut f: impl FnMut(&T) -> io::Result<R>,
132    ) -> Poll<io::Result<R>> {
133        loop {
134            match f(&self.0) {
135                Ok(result) => break Poll::Ready(Ok(result)),
136                Err(e) if is_would_block(&e) => {
137                    std::task::ready!(self.poll_accept_ready(cx))?;
138                }
139                Err(e) => break Poll::Ready(Err(e)),
140            }
141        }
142    }
143
144    /// Poll for read readiness and call the provided function.
145    pub fn poll_read_with<R>(
146        &self,
147        cx: &mut Context,
148        mut f: impl FnMut(&T) -> io::Result<R>,
149    ) -> Poll<io::Result<R>> {
150        loop {
151            match f(&self.0) {
152                Ok(result) => break Poll::Ready(Ok(result)),
153                Err(e) if is_would_block(&e) => {
154                    std::task::ready!(self.poll_read_ready(cx))?;
155                }
156                Err(e) => break Poll::Ready(Err(e)),
157            }
158        }
159    }
160
161    /// Poll for write readiness and call the provided function.
162    pub fn poll_write_with<R>(
163        &self,
164        cx: &mut Context,
165        mut f: impl FnMut(&T) -> io::Result<R>,
166    ) -> Poll<io::Result<R>> {
167        loop {
168            match f(&self.0) {
169                Ok(result) => break Poll::Ready(Ok(result)),
170                Err(e) if is_would_block(&e) => {
171                    std::task::ready!(self.poll_write_ready(cx))?;
172                }
173                Err(e) => break Poll::Ready(Err(e)),
174            }
175        }
176    }
177}
178
179impl<T: AsFd + 'static> PollFd<T> {
180    /// Poll for read readiness and read data.
181    pub fn poll_read(&self, cx: &mut Context, buf: &mut [u8]) -> Poll<io::Result<usize>> {
182        self.poll_read_with(cx, |fd| sys::read(fd.as_fd(), buf))
183    }
184
185    /// Poll for read readiness and read data into a slice of buffers.
186    pub fn poll_read_vectored(
187        &self,
188        cx: &mut Context<'_>,
189        bufs: &mut [io::IoSliceMut<'_>],
190    ) -> Poll<io::Result<usize>> {
191        self.poll_read_with(cx, |fd| sys::readv(fd.as_fd(), bufs))
192    }
193
194    /// Poll for read readiness and read data into an uninitialized buffer.
195    #[cfg(feature = "read_buf")]
196    pub fn poll_read_buf(
197        &self,
198        cx: &mut Context,
199        mut buf: std::io::BorrowedCursor<'_, u8>,
200    ) -> Poll<io::Result<()>> {
201        self.poll_read_with(cx, |fd| {
202            // SAFETY: platform reads only initialize the bytes they report.
203            let read = sys::read_uninit(fd.as_fd(), unsafe { buf.as_mut() })?;
204            unsafe { buf.advance(read) };
205            Ok(())
206        })
207    }
208}
209
210impl<T: AsFd + 'static> PollFd<T> {
211    /// Poll for write readiness and write data.
212    pub fn poll_write(&self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
213        self.poll_write_with(cx, |fd| sys::write(fd.as_fd(), buf))
214    }
215
216    /// Poll for write readiness and write data from a slice of buffers.
217    ///
218    /// Whether this is more efficient than [`poll_write`] depends on the
219    /// source: it is a single `writev` for sockets and files, while other
220    /// sources may fall back to writing the first non-empty buffer.
221    ///
222    /// [`poll_write`]: Self::poll_write
223    pub fn poll_write_vectored(
224        &self,
225        cx: &mut Context<'_>,
226        bufs: &[io::IoSlice<'_>],
227    ) -> Poll<io::Result<usize>> {
228        self.poll_write_with(cx, |fd| sys::writev(fd.as_fd(), bufs))
229    }
230
231    /// Flush pending writes.
232    ///
233    /// [`PollFd`] does not buffer writes, so this is a no-op.
234    pub fn poll_flush(&self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
235        Poll::Ready(Ok(()))
236    }
237}
238
239impl<T: AsFd> PollFd<T> {
240    /// Shut down the write half, so the peer of a connected socket observes the
241    /// end of the stream while this side can still read.
242    ///
243    /// Sources that cannot be half-closed report success and are left as they
244    /// are, since there is no write half to shut down. Shutting down twice is
245    /// successful as well.
246    ///
247    /// Like the `shutdown` methods in `std`, this does not flush the source.
248    fn shutdown_write(&self) -> io::Result<()> {
249        match self.run_socket_op(|socket| socket.shutdown(Shutdown::Write)) {
250            Err(e) if is_not_a_connected_socket(&e) => Ok(()),
251            result => result,
252        }
253    }
254}
255
256impl<T: AsFd> IntoInner for PollFd<T> {
257    type Inner = SharedFd<T>;
258
259    fn into_inner(self) -> Self::Inner {
260        self.0.into_inner()
261    }
262}
263
264impl<T: AsFd> ToSharedFd<T> for PollFd<T> {
265    fn to_shared_fd(&self) -> SharedFd<T> {
266        self.0.to_shared_fd()
267    }
268}
269
270impl<T: AsFd> AsFd for PollFd<T> {
271    fn as_fd(&self) -> BorrowedFd<'_> {
272        self.0.as_fd()
273    }
274}
275
276impl<T: AsFd> AsRawFd for PollFd<T> {
277    fn as_raw_fd(&self) -> RawFd {
278        self.0.as_raw_fd()
279    }
280}
281
282#[cfg(windows)]
283impl<T: AsFd + AsRawSocket> AsRawSocket for PollFd<T> {
284    fn as_raw_socket(&self) -> RawSocket {
285        self.0.as_raw_socket()
286    }
287}
288
289impl<T: AsFd> Deref for PollFd<T> {
290    type Target = T;
291
292    fn deref(&self) -> &Self::Target {
293        &self.0
294    }
295}
296
297fn is_would_block(e: &io::Error) -> bool {
298    cfg_select! {
299        unix => {
300            matches!(
301                e.kind(),
302                io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
303            ) || e.raw_os_error() == Some(libc::EINPROGRESS)
304        }
305        windows => {
306            use windows_sys::Win32::Networking::WinSock::WSAEINPROGRESS;
307            matches!(
308                e.kind(),
309                io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
310            ) || e.raw_os_error() == Some(WSAEINPROGRESS)
311        }
312        _ => {
313            matches!(
314                e.kind(),
315                io::ErrorKind::WouldBlock | io::ErrorKind::Interrupted
316            )
317        }
318    }
319}
320
321fn is_connect_pending(e: &io::Error) -> bool {
322    if is_would_block(e) {
323        return true;
324    }
325
326    cfg_select! {
327        unix => e.raw_os_error() == Some(libc::EALREADY),
328        windows => {
329            use windows_sys::Win32::Networking::WinSock::WSAEALREADY;
330            e.raw_os_error() == Some(WSAEALREADY)
331        }
332        _ => false,
333    }
334}
335
336fn is_not_a_connected_socket(e: &io::Error) -> bool {
337    cfg_select! {
338        unix => {
339            matches!(
340                e.raw_os_error(),
341                Some(libc::ENOTSOCK) | Some(libc::ENOTCONN)
342            )
343        }
344        windows => {
345            use windows_sys::Win32::Networking::WinSock::{WSAENOTCONN, WSAENOTSOCK};
346
347            matches!(e.raw_os_error(), Some(WSAENOTSOCK) | Some(WSAENOTCONN))
348        }
349    }
350}
351
352impl<T: AsFd + 'static> futures_util::AsyncRead for &PollFd<T> {
353    fn poll_read(
354        self: Pin<&mut Self>,
355        cx: &mut Context<'_>,
356        buf: &mut [u8],
357    ) -> Poll<io::Result<usize>> {
358        (*self).poll_read(cx, buf)
359    }
360
361    fn poll_read_vectored(
362        self: Pin<&mut Self>,
363        cx: &mut Context<'_>,
364        bufs: &mut [io::IoSliceMut<'_>],
365    ) -> Poll<io::Result<usize>> {
366        (*self).poll_read_vectored(cx, bufs)
367    }
368}
369
370impl<T: AsFd + 'static> futures_util::AsyncRead for PollFd<T> {
371    fn poll_read(
372        self: Pin<&mut Self>,
373        cx: &mut Context<'_>,
374        buf: &mut [u8],
375    ) -> Poll<io::Result<usize>> {
376        (*self).poll_read(cx, buf)
377    }
378
379    fn poll_read_vectored(
380        self: Pin<&mut Self>,
381        cx: &mut Context<'_>,
382        bufs: &mut [io::IoSliceMut<'_>],
383    ) -> Poll<io::Result<usize>> {
384        (*self).poll_read_vectored(cx, bufs)
385    }
386}
387
388impl<T: AsFd + 'static> futures_util::AsyncWrite for &PollFd<T> {
389    fn poll_write(
390        self: Pin<&mut Self>,
391        cx: &mut Context<'_>,
392        buf: &[u8],
393    ) -> Poll<io::Result<usize>> {
394        (*self).poll_write(cx, buf)
395    }
396
397    fn poll_write_vectored(
398        self: Pin<&mut Self>,
399        cx: &mut Context<'_>,
400        bufs: &[io::IoSlice<'_>],
401    ) -> Poll<io::Result<usize>> {
402        (*self).poll_write_vectored(cx, bufs)
403    }
404
405    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
406        (*self).poll_flush(cx)
407    }
408
409    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
410        Poll::Ready(self.shutdown_write())
411    }
412}
413
414impl<T: AsFd + 'static> futures_util::AsyncWrite for PollFd<T> {
415    fn poll_write(
416        self: Pin<&mut Self>,
417        cx: &mut Context<'_>,
418        buf: &[u8],
419    ) -> Poll<io::Result<usize>> {
420        (*self).poll_write(cx, buf)
421    }
422
423    fn poll_write_vectored(
424        self: Pin<&mut Self>,
425        cx: &mut Context<'_>,
426        bufs: &[io::IoSlice<'_>],
427    ) -> Poll<io::Result<usize>> {
428        (*self).poll_write_vectored(cx, bufs)
429    }
430
431    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
432        (*self).poll_flush(cx)
433    }
434
435    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
436        Poll::Ready(self.shutdown_write())
437    }
438}
compio_runtime/fd/poll_fd/mod.rs

compio_runtime/fd/poll_fd/
mod.rs
