summaryrefslogtreecommitdiffstats
path: root/tokio-util/src/lib.rs
diff options
context:
space:
mode:
Diffstat (limited to 'tokio-util/src/lib.rs')
-rw-r--r--tokio-util/src/lib.rs34
1 files changed, 34 insertions, 0 deletions
diff --git a/tokio-util/src/lib.rs b/tokio-util/src/lib.rs
index 1e4b9d40..10b828ef 100644
--- a/tokio-util/src/lib.rs
+++ b/tokio-util/src/lib.rs
@@ -57,3 +57,37 @@ pub mod either;
#[cfg(feature = "time")]
pub mod time;
+
+#[cfg(any(feature = "io", feature = "codec"))]
+mod util {
+ use tokio::io::{AsyncRead, ReadBuf};
+
+ use bytes::BufMut;
+ use futures_core::ready;
+ use std::io;
+ use std::pin::Pin;
+ use std::task::{Context, Poll};
+
+ pub(crate) fn poll_read_buf<T: AsyncRead>(
+ cx: &mut Context<'_>,
+ io: Pin<&mut T>,
+ buf: &mut impl BufMut,
+ ) -> Poll<io::Result<usize>> {
+ if !buf.has_remaining_mut() {
+ return Poll::Ready(Ok(0));
+ }
+
+ let orig = buf.bytes_mut().as_ptr() as *const u8;
+ let mut b = ReadBuf::uninit(buf.bytes_mut());
+
+ ready!(io.poll_read(cx, &mut b))?;
+ let n = b.filled().len();
+
+ // Safety: we can assume `n` bytes were read, since they are in`filled`.
+ assert_eq!(orig, b.filled().as_ptr());
+ unsafe {
+ buf.advance_mut(n);
+ }
+ Poll::Ready(Ok(n))
+ }
+}