hickory_net/
runtime.rs

1//! Abstractions to deal with different async runtimes.
2
3use core::future::Future;
4use core::marker::Send;
5use core::net::SocketAddr;
6use core::pin::Pin;
7use core::time::Duration;
8#[cfg(feature = "__quic")]
9use std::sync::Arc;
10use std::time::{SystemTime, UNIX_EPOCH};
11use std::{
12    future::poll_fn,
13    io,
14    task::{Context, Poll},
15};
16
17use async_trait::async_trait;
18use futures_io::{AsyncRead, AsyncWrite};
19#[cfg(any(test, feature = "tokio"))]
20use tokio::runtime::Runtime;
21#[cfg(any(test, feature = "tokio"))]
22use tokio::task::JoinHandle;
23
24/// Spawn a background task, if it was present
25#[cfg(any(test, feature = "tokio"))]
26pub fn spawn_bg<F: Future<Output = R> + Send + 'static, R: Send + 'static>(
27    runtime: &Runtime,
28    background: F,
29) -> JoinHandle<R> {
30    runtime.spawn(background)
31}
32
33#[cfg(feature = "tokio")]
34#[doc(hidden)]
35pub mod iocompat {
36    use core::pin::Pin;
37    use core::task::{Context, Poll};
38    use std::io::{self, IoSlice};
39
40    use futures_io::{AsyncRead, AsyncWrite};
41    use tokio::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite, ReadBuf};
42
43    /// Conversion from `tokio::io::{AsyncRead, AsyncWrite}` to `std::io::{AsyncRead, AsyncWrite}`
44    pub struct AsyncIoTokioAsStd<T: TokioAsyncRead + TokioAsyncWrite>(pub T);
45
46    impl<T: TokioAsyncRead + TokioAsyncWrite + Unpin> Unpin for AsyncIoTokioAsStd<T> {}
47    impl<R: TokioAsyncRead + TokioAsyncWrite + Unpin> AsyncRead for AsyncIoTokioAsStd<R> {
48        fn poll_read(
49            mut self: Pin<&mut Self>,
50            cx: &mut Context<'_>,
51            buf: &mut [u8],
52        ) -> Poll<io::Result<usize>> {
53            let mut buf = ReadBuf::new(buf);
54            let polled = Pin::new(&mut self.0).poll_read(cx, &mut buf);
55
56            polled.map_ok(|_| buf.filled().len())
57        }
58    }
59
60    impl<W: TokioAsyncRead + TokioAsyncWrite + Unpin> AsyncWrite for AsyncIoTokioAsStd<W> {
61        fn poll_write(
62            mut self: Pin<&mut Self>,
63            cx: &mut Context<'_>,
64            buf: &[u8],
65        ) -> Poll<io::Result<usize>> {
66            Pin::new(&mut self.0).poll_write(cx, buf)
67        }
68        fn poll_write_vectored(
69            mut self: Pin<&mut Self>,
70            cx: &mut Context<'_>,
71            bufs: &[IoSlice<'_>],
72        ) -> Poll<io::Result<usize>> {
73            Pin::new(&mut self.0).poll_write_vectored(cx, bufs)
74        }
75        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
76            Pin::new(&mut self.0).poll_flush(cx)
77        }
78        fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
79            Pin::new(&mut self.0).poll_shutdown(cx)
80        }
81    }
82
83    /// Conversion from `std::io::{AsyncRead, AsyncWrite}` to `tokio::io::{AsyncRead, AsyncWrite}`
84    pub struct AsyncIoStdAsTokio<T: AsyncRead + AsyncWrite>(pub T);
85
86    impl<T: AsyncRead + AsyncWrite + Unpin> Unpin for AsyncIoStdAsTokio<T> {}
87    impl<R: AsyncRead + AsyncWrite + Unpin> TokioAsyncRead for AsyncIoStdAsTokio<R> {
88        fn poll_read(
89            self: Pin<&mut Self>,
90            cx: &mut Context<'_>,
91            buf: &mut ReadBuf<'_>,
92        ) -> Poll<io::Result<()>> {
93            Pin::new(&mut self.get_mut().0)
94                .poll_read(cx, buf.initialized_mut())
95                .map_ok(|len| buf.advance(len))
96        }
97    }
98
99    impl<W: AsyncRead + AsyncWrite + Unpin> TokioAsyncWrite for AsyncIoStdAsTokio<W> {
100        fn poll_write(
101            self: Pin<&mut Self>,
102            cx: &mut Context<'_>,
103            buf: &[u8],
104        ) -> Poll<Result<usize, io::Error>> {
105            Pin::new(&mut self.get_mut().0).poll_write(cx, buf)
106        }
107
108        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
109            Pin::new(&mut self.get_mut().0).poll_flush(cx)
110        }
111
112        fn poll_shutdown(
113            self: Pin<&mut Self>,
114            cx: &mut Context<'_>,
115        ) -> Poll<Result<(), io::Error>> {
116            Pin::new(&mut self.get_mut().0).poll_close(cx)
117        }
118    }
119}
120
121#[cfg(feature = "tokio")]
122#[allow(unreachable_pub)]
123mod tokio_runtime {
124    use std::sync::Arc;
125    use std::sync::Mutex;
126
127    #[cfg(feature = "__quic")]
128    use quinn::Runtime;
129    use tokio::net::{TcpSocket, TcpStream, UdpSocket as TokioUdpSocket};
130    use tokio::task::JoinSet;
131    use tokio::time::timeout;
132    use tracing::debug;
133
134    use super::iocompat::AsyncIoTokioAsStd;
135    use super::*;
136    use crate::xfer::CONNECT_TIMEOUT;
137
138    /// A handle to the Tokio runtime
139    #[derive(Clone, Default)]
140    pub struct TokioHandle {
141        join_set: Arc<Mutex<JoinSet<()>>>,
142    }
143
144    impl Spawn for TokioHandle {
145        fn spawn_bg(&mut self, future: impl Future<Output = ()> + Send + 'static) {
146            let mut join_set = self.join_set.lock().unwrap();
147            join_set.spawn(future);
148            reap_tasks(&mut join_set);
149        }
150    }
151
152    /// The Tokio Runtime for async execution
153    #[derive(Clone, Default)]
154    pub struct TokioRuntimeProvider(TokioHandle);
155
156    impl TokioRuntimeProvider {
157        /// Create a Tokio runtime
158        pub fn new() -> Self {
159            Self::default()
160        }
161    }
162
163    impl RuntimeProvider for TokioRuntimeProvider {
164        type Handle = TokioHandle;
165        type Timer = TokioTime;
166        type Udp = TokioUdpSocket;
167        type Tcp = AsyncIoTokioAsStd<TcpStream>;
168
169        fn create_handle(&self) -> Self::Handle {
170            self.0.clone()
171        }
172
173        fn connect_tcp(
174            &self,
175            server_addr: SocketAddr,
176            bind_addr: Option<SocketAddr>,
177            wait_for: Option<Duration>,
178        ) -> Pin<Box<dyn Send + Future<Output = Result<Self::Tcp, io::Error>>>> {
179            Box::pin(async move {
180                let socket = match server_addr {
181                    SocketAddr::V4(_) => TcpSocket::new_v4(),
182                    SocketAddr::V6(_) => TcpSocket::new_v6(),
183                }?;
184
185                if let Some(bind_addr) = bind_addr {
186                    socket.bind(bind_addr)?;
187                }
188
189                socket.set_nodelay(true)?;
190                let future = socket.connect(server_addr);
191                let wait_for = wait_for.unwrap_or(CONNECT_TIMEOUT);
192                match timeout(wait_for, future).await {
193                    Ok(Ok(socket)) => Ok(AsyncIoTokioAsStd(socket)),
194                    Ok(Err(e)) => Err(e),
195                    Err(_) => {
196                        debug!(%server_addr, "TCP connect timeout");
197                        Err(io::Error::new(
198                            io::ErrorKind::TimedOut,
199                            "TCP connect timed out",
200                        ))
201                    }
202                }
203            })
204        }
205
206        fn bind_udp(
207            &self,
208            local_addr: SocketAddr,
209            _server_addr: SocketAddr,
210        ) -> Pin<Box<dyn Send + Future<Output = Result<Self::Udp, io::Error>>>> {
211            Box::pin(async move { tokio::net::UdpSocket::bind(local_addr).await })
212        }
213
214        #[cfg(feature = "__quic")]
215        fn quic_binder(&self) -> Option<&dyn QuicSocketBinder> {
216            Some(&TokioQuicSocketBinder)
217        }
218    }
219
220    /// Reap finished tasks from a `JoinSet`, without awaiting or blocking.
221    fn reap_tasks(join_set: &mut JoinSet<()>) {
222        while join_set.try_join_next().is_some() {}
223    }
224
225    #[cfg(feature = "__quic")]
226    struct TokioQuicSocketBinder;
227
228    #[cfg(feature = "__quic")]
229    impl QuicSocketBinder for TokioQuicSocketBinder {
230        fn bind_quic(
231            &self,
232            local_addr: SocketAddr,
233            _server_addr: SocketAddr,
234        ) -> Result<Arc<dyn quinn::AsyncUdpSocket>, io::Error> {
235            let socket = std::net::UdpSocket::bind(local_addr)?;
236            quinn::TokioRuntime.wrap_udp_socket(socket)
237        }
238    }
239}
240
241#[cfg(feature = "tokio")]
242pub use tokio_runtime::{TokioHandle, TokioRuntimeProvider};
243
244/// RuntimeProvider defines which async runtime that handles IO and timers.
245pub trait RuntimeProvider: Clone + Send + Sync + Unpin + 'static {
246    /// Handle to the executor;
247    type Handle: Clone + Send + Spawn + Sync + Unpin;
248
249    /// Timer
250    type Timer: Time;
251
252    /// UdpSocket
253    type Udp: DnsUdpSocket;
254
255    /// TcpStream
256    type Tcp: DnsTcpStream;
257
258    /// Create a runtime handle
259    fn create_handle(&self) -> Self::Handle;
260
261    /// Create a TCP connection with custom configuration.
262    fn connect_tcp(
263        &self,
264        server_addr: SocketAddr,
265        bind_addr: Option<SocketAddr>,
266        timeout: Option<Duration>,
267    ) -> Pin<Box<dyn Send + Future<Output = Result<Self::Tcp, io::Error>>>>;
268
269    /// Create a UDP socket bound to `local_addr`. The returned value should **not** be connected to `server_addr`.
270    /// *Notice: the future should be ready once returned at best effort. Otherwise UDP DNS may need much more retries.*
271    fn bind_udp(
272        &self,
273        local_addr: SocketAddr,
274        server_addr: SocketAddr,
275    ) -> Pin<Box<dyn Send + Future<Output = Result<Self::Udp, io::Error>>>>;
276
277    /// Yields an object that knows how to bind a QUIC socket.
278    //
279    // Use some indirection here to avoid exposing the `quinn` crate in the public API
280    // even for runtimes that might not (want to) provide QUIC support.
281    fn quic_binder(&self) -> Option<&dyn QuicSocketBinder> {
282        None
283    }
284}
285
286/// Trait for DnsUdpSocket
287#[async_trait]
288pub trait DnsUdpSocket
289where
290    Self: Send + Sync + Sized + Unpin,
291{
292    /// Time implementation used for this type
293    type Time: Time;
294
295    /// Poll once Receive data from the socket and returns the number of bytes read and the address from
296    /// where the data came on success.
297    fn poll_recv_from(
298        &self,
299        cx: &mut Context<'_>,
300        buf: &mut [u8],
301    ) -> Poll<io::Result<(usize, SocketAddr)>>;
302
303    /// Receive data from the socket and returns the number of bytes read and the address from
304    /// where the data came on success.
305    async fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
306        poll_fn(|cx| self.poll_recv_from(cx, buf)).await
307    }
308
309    /// Poll once to send data to the given address.
310    fn poll_send_to(
311        &self,
312        cx: &mut Context<'_>,
313        buf: &[u8],
314        target: SocketAddr,
315    ) -> Poll<io::Result<usize>>;
316
317    /// Send data to the given address.
318    async fn send_to(&self, buf: &[u8], target: SocketAddr) -> io::Result<usize> {
319        poll_fn(|cx| self.poll_send_to(cx, buf, target)).await
320    }
321}
322
323/// Noop trait for when the `quinn` dependency is not available.
324#[cfg(not(feature = "__quic"))]
325pub trait QuicSocketBinder {}
326
327/// Create a UDP socket for QUIC usage.
328/// This trait is designed for customization.
329#[cfg(feature = "__quic")]
330pub trait QuicSocketBinder {
331    /// Create a UDP socket for QUIC usage.
332    fn bind_quic(
333        &self,
334        _local_addr: SocketAddr,
335        _server_addr: SocketAddr,
336    ) -> Result<Arc<dyn quinn::AsyncUdpSocket>, io::Error>;
337}
338
339/// Trait for TCP connection
340pub trait DnsTcpStream: AsyncRead + AsyncWrite + Unpin + Send + Sync + Sized + 'static {
341    /// Timer type to use with this TCP stream type
342    type Time: Time;
343}
344
345/// A type defines the Handle which can spawn future.
346pub trait Spawn {
347    /// Spawn a future in the background
348    fn spawn_bg(&mut self, future: impl Future<Output = ()> + Send + 'static);
349}
350
351/// Generic Time for Delay and Timeout.
352// This trait is created to allow to use different types of time systems. It's used in Fuchsia OS, please be mindful when update it.
353#[async_trait]
354pub trait Time: Send + Sync + Unpin {
355    /// Return a type that implements `Future` that will wait until the specified duration has
356    /// elapsed.
357    async fn delay_for(duration: Duration);
358
359    /// Return a type that implement `Future` to complete before the specified duration has elapsed.
360    async fn timeout<F: 'static + Future + Send>(
361        duration: Duration,
362        future: F,
363    ) -> Result<F::Output, io::Error>;
364
365    /// Get the current time as a Unix timestamp.
366    ///
367    /// This returns the number of seconds since the Unix epoch.
368    fn current_time() -> u64 {
369        SystemTime::now()
370            .duration_since(UNIX_EPOCH)
371            .unwrap()
372            .as_secs()
373    }
374}
375
376/// New type which is implemented using tokio::time::{Delay, Timeout}
377#[cfg(any(test, feature = "tokio"))]
378#[derive(Clone, Copy, Debug)]
379pub struct TokioTime;
380
381#[cfg(any(test, feature = "tokio"))]
382#[async_trait]
383impl Time for TokioTime {
384    async fn delay_for(duration: Duration) {
385        tokio::time::sleep(duration).await
386    }
387
388    async fn timeout<F: 'static + Future + Send>(
389        duration: Duration,
390        future: F,
391    ) -> Result<F::Output, io::Error> {
392        tokio::time::timeout(duration, future)
393            .await
394            .map_err(move |_| io::Error::new(io::ErrorKind::TimedOut, "future timed out"))
395    }
396}
hickory_net/runtime.rs

hickory_net/
runtime.rs
