Skip to main content

compio_io/compat/
tokio_io.rs

1//! [`tokio::io`] traits for the [`futures_util::io`] adapters, so that crates
2//! written against them (`h2`, `tokio-rustls`, `tonic`...) run on compio
3//! streams without a tokio runtime.
4
5use std::{
6    io,
7    pin::Pin,
8    task::{Context, Poll, ready},
9};
10
11use tokio::io::ReadBuf;
12
13use crate::{
14    AsyncRead, AsyncWrite,
15    compat::{AsyncReadStream, AsyncStream, AsyncWriteStream},
16    util::Splittable,
17};
18
19/// Reads into the unfilled part of `buf` without initializing it first.
20///
21/// # Safety
22///
23/// `read` must not de-initialize any byte of the slice it is given, and when
24/// it returns `Ok(n)`, it must have initialized the first `n` bytes of it.
25unsafe fn read_uninit(
26    buf: &mut ReadBuf<'_>,
27    read: impl FnOnce(&mut [std::mem::MaybeUninit<u8>]) -> Poll<io::Result<usize>>,
28) -> Poll<io::Result<()>> {
29    // SAFETY: the caller guarantees `read` does not de-initialize the slice.
30    let n = ready!(read(unsafe { buf.unfilled_mut() }))?;
31    // SAFETY: the caller guarantees `read` initialized the first `n` bytes.
32    unsafe { buf.assume_init(n) };
33    buf.advance(n);
34    Poll::Ready(Ok(()))
35}
36
37impl<S: AsyncRead + Unpin + 'static> tokio::io::AsyncRead for AsyncReadStream<S> {
38    fn poll_read(
39        self: Pin<&mut Self>,
40        cx: &mut Context<'_>,
41        buf: &mut ReadBuf<'_>,
42    ) -> Poll<io::Result<()>> {
43        // SAFETY: `poll_read_uninit` only copies initialized bytes into the
44        // slice, and returns how many it copied to its start.
45        unsafe { read_uninit(buf, |unfilled| self.poll_read_uninit(cx, unfilled)) }
46    }
47}
48
49impl<S: AsyncRead + Unpin + 'static> tokio::io::AsyncBufRead for AsyncReadStream<S> {
50    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
51        futures_util::AsyncBufRead::poll_fill_buf(self, cx)
52    }
53
54    fn consume(self: Pin<&mut Self>, amt: usize) {
55        futures_util::AsyncBufRead::consume(self, amt)
56    }
57}
58
59impl<S: AsyncWrite + Unpin + 'static> tokio::io::AsyncWrite for AsyncWriteStream<S> {
60    fn poll_write(
61        self: Pin<&mut Self>,
62        cx: &mut Context<'_>,
63        buf: &[u8],
64    ) -> Poll<io::Result<usize>> {
65        futures_util::AsyncWrite::poll_write(self, cx, buf)
66    }
67
68    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
69        futures_util::AsyncWrite::poll_flush(self, cx)
70    }
71
72    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
73        futures_util::AsyncWrite::poll_close(self, cx)
74    }
75}
76
77impl<S: Splittable + 'static> tokio::io::AsyncRead for AsyncStream<S>
78where
79    S::ReadHalf: AsyncRead + Unpin,
80{
81    fn poll_read(
82        self: Pin<&mut Self>,
83        cx: &mut Context<'_>,
84        buf: &mut ReadBuf<'_>,
85    ) -> Poll<io::Result<()>> {
86        // SAFETY: `poll_read_uninit` only copies initialized bytes into the
87        // slice, and returns how many it copied to its start.
88        unsafe { read_uninit(buf, |unfilled| self.poll_read_uninit(cx, unfilled)) }
89    }
90}
91
92impl<S: Splittable + 'static> tokio::io::AsyncBufRead for AsyncStream<S>
93where
94    S::ReadHalf: AsyncRead + Unpin,
95{
96    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
97        futures_util::AsyncBufRead::poll_fill_buf(self, cx)
98    }
99
100    fn consume(self: Pin<&mut Self>, amt: usize) {
101        futures_util::AsyncBufRead::consume(self, amt)
102    }
103}
104
105impl<S: Splittable + 'static> tokio::io::AsyncWrite for AsyncStream<S>
106where
107    S::WriteHalf: AsyncWrite + Unpin,
108{
109    fn poll_write(
110        self: Pin<&mut Self>,
111        cx: &mut Context<'_>,
112        buf: &[u8],
113    ) -> Poll<io::Result<usize>> {
114        futures_util::AsyncWrite::poll_write(self, cx, buf)
115    }
116
117    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
118        futures_util::AsyncWrite::poll_flush(self, cx)
119    }
120
121    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
122        futures_util::AsyncWrite::poll_close(self, cx)
123    }
124}