+220

libs/kasync/src/io.rs

reviewed

··· 1 1 + // Copyright 2025. Jonas Kruckenberg 2 2 + // 3 3 + // Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or 4 4 + // http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or 5 5 + // http://opensource.org/licenses/MIT>, at your option. This file may not be 6 6 + // copied, modified, or distributed except according to those terms. 7 7 + 8 8 + use core::cmp; 9 9 + use core::convert::Infallible; 10 10 + use core::ops::{Deref, DerefMut}; 11 11 + use core::pin::Pin; 12 12 + use core::task::{ready, Context, Poll}; 13 13 + 14 14 + pub trait Read { 15 15 + type Err: core::error::Error; 16 16 + 17 17 + // Pull some bytes from this source into the specified buffer. 18 18 + fn poll_read( 19 19 + self: Pin<&mut Self>, 20 20 + cx: &mut Context<'_>, 21 21 + buf: &mut [u8], 22 22 + ) -> Poll<Result<usize, Self::Err>>; 23 23 + 24 24 + // Attempt to read from the AsyncRead into bufs using vectored IO operations. 25 25 + // 26 26 + // This method is similar to poll_read, but allows data to be read into multiple buffers using a single operation. 27 27 + // 28 28 + // On success, returns Poll::Ready(Ok(num_bytes_read)). 29 29 + fn poll_read_vectored( 30 30 + mut self: Pin<&mut Self>, 31 31 + cx: &mut Context<'_>, 32 32 + bufs: &mut [&mut [u8]], 33 33 + ) -> Poll<Result<usize, Self::Err>> { 34 34 + let mut nread = 0; 35 35 + for b in bufs { 36 36 + if !b.is_empty() { 37 37 + nread += ready!(self.as_mut().poll_read(cx, b)?); 38 38 + } 39 39 + } 40 40 + Poll::Ready(Ok(nread)) 41 41 + } 42 42 + } 43 43 + 44 44 + pub trait Write { 45 45 + type Err: core::error::Error; 46 46 + 47 47 + // Writes a buffer into this writer. 48 48 + // 49 49 + // returning how many bytes were written. 50 50 + fn poll_write( 51 51 + self: Pin<&mut Self>, 52 52 + cx: &mut Context<'_>, 53 53 + buf: &[u8], 54 54 + ) -> Poll<Result<usize, Self::Err>>; 55 55 + 56 56 + // Flushes this output stream, ensuring that all intermediately buffered contents reach their destination. 57 57 + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>>; 58 58 + 59 59 + // Attempt to close the object. 60 60 + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>>; 61 61 + 62 62 + // Attempt to write bytes from bufs into the object using vectored IO operations. 63 63 + // 64 64 + // This method is similar to poll_write, but allows data from multiple buffers to be written using a single operation. 65 65 + // 66 66 + // On success, returns Poll::Ready(Ok(num_bytes_written)). 67 67 + fn poll_write_vectored( 68 68 + mut self: Pin<&mut Self>, 69 69 + cx: &mut Context<'_>, 70 70 + bufs: &[&[u8]], 71 71 + ) -> Poll<Result<usize, Self::Err>> { 72 72 + let mut nwritten = 0; 73 73 + for b in bufs { 74 74 + if !b.is_empty() { 75 75 + nwritten += ready!(self.as_mut().poll_write(cx, b)?); 76 76 + } 77 77 + } 78 78 + Poll::Ready(Ok(nwritten)) 79 79 + } 80 80 + } 81 81 + 82 82 + // ===== impl Read ===== 83 83 + 84 84 + impl Read for &[u8] { 85 85 + type Err = Infallible; 86 86 + 87 87 + fn poll_read( 88 88 + self: Pin<&mut Self>, 89 89 + _cx: &mut Context<'_>, 90 90 + buf: &mut [u8], 91 91 + ) -> Poll<Result<usize, Self::Err>> { 92 92 + let amt = cmp::min(buf.len(), self.len()); 93 93 + let (a, b) = self.split_at(amt); 94 94 + 95 95 + // First check if the amount of bytes we want to read is small: 96 96 + // `copy_from_slice` will generally expand to a call to `memcpy`, and 97 97 + // for a single byte the overhead is significant. 98 98 + if amt == 1 { 99 99 + buf[0] = a[0]; 100 100 + } else { 101 101 + buf[..amt].copy_from_slice(a); 102 102 + } 103 103 + 104 104 + *self.get_mut() = b; 105 105 + 106 106 + Poll::Ready(Ok(amt)) 107 107 + } 108 108 + } 109 109 + 110 110 + impl<P> Read for Pin<P> 111 111 + where 112 112 + P: DerefMut + Unpin, 113 113 + <P as Deref>::Target: Read, 114 114 + { 115 115 + type Err = <P::Target as Read>::Err; 116 116 + 117 117 + fn poll_read( 118 118 + self: Pin<&mut Self>, 119 119 + cx: &mut Context<'_>, 120 120 + buf: &mut [u8], 121 121 + ) -> Poll<Result<usize, Self::Err>> { 122 122 + self.poll_read(cx, buf) 123 123 + } 124 124 + 125 125 + fn poll_read_vectored( 126 126 + self: Pin<&mut Self>, 127 127 + cx: &mut Context<'_>, 128 128 + bufs: &mut [&mut [u8]], 129 129 + ) -> Poll<Result<usize, Self::Err>> { 130 130 + self.poll_read_vectored(cx, bufs) 131 131 + } 132 132 + } 133 133 + 134 134 + impl<T> Read for &mut T 135 135 + where 136 136 + T: Read + Unpin + ?Sized, 137 137 + { 138 138 + type Err = T::Err; 139 139 + 140 140 + fn poll_read( 141 141 + self: Pin<&mut Self>, 142 142 + cx: &mut Context<'_>, 143 143 + buf: &mut [u8], 144 144 + ) -> Poll<Result<usize, Self::Err>> { 145 145 + self.poll_read(cx, buf) 146 146 + } 147 147 + 148 148 + fn poll_read_vectored( 149 149 + self: Pin<&mut Self>, 150 150 + cx: &mut Context<'_>, 151 151 + bufs: &mut [&mut [u8]], 152 152 + ) -> Poll<Result<usize, Self::Err>> { 153 153 + self.poll_read_vectored(cx, bufs) 154 154 + } 155 155 + } 156 156 + 157 157 + // ===== impl Write ===== 158 158 + 159 159 + impl<P> Write for Pin<P> 160 160 + where 161 161 + P: DerefMut + Unpin, 162 162 + <P as Deref>::Target: Write, 163 163 + { 164 164 + type Err = <P::Target as Write>::Err; 165 165 + 166 166 + fn poll_write( 167 167 + self: Pin<&mut Self>, 168 168 + cx: &mut Context<'_>, 169 169 + buf: &[u8], 170 170 + ) -> Poll<Result<usize, Self::Err>> { 171 171 + self.poll_write(cx, buf) 172 172 + } 173 173 + 174 174 + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>> { 175 175 + self.poll_flush(cx) 176 176 + } 177 177 + 178 178 + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>> { 179 179 + self.poll_close(cx) 180 180 + } 181 181 + 182 182 + fn poll_write_vectored( 183 183 + self: Pin<&mut Self>, 184 184 + cx: &mut Context<'_>, 185 185 + bufs: &[&[u8]], 186 186 + ) -> Poll<Result<usize, Self::Err>> { 187 187 + self.poll_write_vectored(cx, bufs) 188 188 + } 189 189 + } 190 190 + 191 191 + impl<T> Write for &mut T 192 192 + where 193 193 + T: Write + Unpin + ?Sized, 194 194 + { 195 195 + type Err = T::Err; 196 196 + 197 197 + fn poll_write( 198 198 + self: Pin<&mut Self>, 199 199 + cx: &mut Context<'_>, 200 200 + buf: &[u8], 201 201 + ) -> Poll<Result<usize, Self::Err>> { 202 202 + self.poll_write(cx, buf) 203 203 + } 204 204 + 205 205 + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>> { 206 206 + self.poll_flush(cx) 207 207 + } 208 208 + 209 209 + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Err>> { 210 210 + self.poll_close(cx) 211 211 + } 212 212 + 213 213 + fn poll_write_vectored( 214 214 + self: Pin<&mut Self>, 215 215 + cx: &mut Context<'_>, 216 216 + bufs: &[&[u8]], 217 217 + ) -> Poll<Result<usize, Self::Err>> { 218 218 + self.poll_write_vectored(cx, bufs) 219 219 + } 220 220 + }

+1

libs/kasync/src/lib.rs

reviewed

··· 22 22 pub mod sync; 23 23 pub mod task; 24 24 pub mod time; 25 25 + pub mod io; 25 26 26 27 pub use error::{Closed, SpawnError}; 27 28 pub use futures::future;