1 | use crate::io::AsyncWrite; |
2 | |
3 | use bytes::Buf; |
4 | use pin_project_lite::pin_project; |
5 | use std::future::Future; |
6 | use std::io; |
7 | use std::marker::PhantomPinned; |
8 | use std::pin::Pin; |
9 | use std::task::{Context, Poll}; |
10 | |
11 | pin_project! { |
12 | /// A future to write some of the buffer to an `AsyncWrite`. |
13 | #[derive(Debug)] |
14 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
15 | pub struct WriteBuf<'a, W, B> { |
16 | writer: &'a mut W, |
17 | buf: &'a mut B, |
18 | #[pin] |
19 | _pin: PhantomPinned, |
20 | } |
21 | } |
22 | |
23 | /// Tries to write some bytes from the given `buf` to the writer in an |
24 | /// asynchronous manner, returning a future. |
25 | pub(crate) fn write_buf<'a, W, B>(writer: &'a mut W, buf: &'a mut B) -> WriteBuf<'a, W, B> |
26 | where |
27 | W: AsyncWrite + Unpin, |
28 | B: Buf, |
29 | { |
30 | WriteBuf { |
31 | writer, |
32 | buf, |
33 | _pin: PhantomPinned, |
34 | } |
35 | } |
36 | |
37 | impl<W, B> Future for WriteBuf<'_, W, B> |
38 | where |
39 | W: AsyncWrite + Unpin, |
40 | B: Buf, |
41 | { |
42 | type Output = io::Result<usize>; |
43 | |
44 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> { |
45 | let me = self.project(); |
46 | |
47 | if !me.buf.has_remaining() { |
48 | return Poll::Ready(Ok(0)); |
49 | } |
50 | |
51 | let n = ready!(Pin::new(me.writer).poll_write(cx, me.buf.chunk()))?; |
52 | me.buf.advance(n); |
53 | Poll::Ready(Ok(n)) |
54 | } |
55 | } |
56 | |