futures_test/io/
limited.rs

1use futures_io::{self as io, AsyncBufRead, AsyncRead, AsyncWrite};
2use pin_utils::{unsafe_pinned, unsafe_unpinned};
3use std::{
4    cmp,
5    pin::Pin,
6    task::{Context, Poll},
7};
8
9/// I/O wrapper that limits the number of bytes written or read on each call.
10///
11/// See the [`limited`] and [`limited_write`] methods.
12///
13/// [`limited`]: super::AsyncReadTestExt::limited
14/// [`limited_write`]: super::AsyncWriteTestExt::limited_write
15#[derive(Debug)]
16pub struct Limited<Io> {
17    io: Io,
18    limit: usize,
19}
20
21impl<Io: Unpin> Unpin for Limited<Io> {}
22
23impl<Io> Limited<Io> {
24    unsafe_pinned!(io: Io);
25    unsafe_unpinned!(limit: usize);
26
27    pub(crate) fn new(io: Io, limit: usize) -> Limited<Io> {
28        Limited { io, limit }
29    }
30
31    /// Acquires a reference to the underlying I/O object that this adaptor is
32    /// wrapping.
33    pub fn get_ref(&self) -> &Io {
34        &self.io
35    }
36
37    /// Acquires a mutable reference to the underlying I/O object that this
38    /// adaptor is wrapping.
39    pub fn get_mut(&mut self) -> &mut Io {
40        &mut self.io
41    }
42
43    /// Acquires a pinned mutable reference to the underlying I/O object that
44    /// this adaptor is wrapping.
45    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Io> {
46        self.io()
47    }
48
49    /// Consumes this adaptor returning the underlying I/O object.
50    pub fn into_inner(self) -> Io {
51        self.io
52    }
53}
54
55impl<W: AsyncWrite> AsyncWrite for Limited<W> {
56    fn poll_write(
57        mut self: Pin<&mut Self>,
58        cx: &mut Context<'_>,
59        buf: &[u8],
60    ) -> Poll<io::Result<usize>> {
61        let limit = *self.as_mut().limit();
62        self.io().poll_write(cx, &buf[..cmp::min(limit, buf.len())])
63    }
64
65    fn poll_flush(
66        self: Pin<&mut Self>,
67        cx: &mut Context<'_>,
68    ) -> Poll<io::Result<()>> {
69        self.io().poll_flush(cx)
70    }
71
72    fn poll_close(
73        self: Pin<&mut Self>,
74        cx: &mut Context<'_>,
75    ) -> Poll<io::Result<()>> {
76        self.io().poll_close(cx)
77    }
78}
79
80impl<R: AsyncRead> AsyncRead for Limited<R> {
81    fn poll_read(
82        mut self: Pin<&mut Self>,
83        cx: &mut Context<'_>,
84        buf: &mut [u8],
85    ) -> Poll<io::Result<usize>> {
86        let limit = cmp::min(*self.as_mut().limit(), buf.len());
87        self.io().poll_read(cx, &mut buf[..limit])
88    }
89}
90
91impl<R: AsyncBufRead> AsyncBufRead for Limited<R> {
92    fn poll_fill_buf(
93        self: Pin<&mut Self>,
94        cx: &mut Context<'_>,
95    ) -> Poll<io::Result<&[u8]>> {
96        self.io().poll_fill_buf(cx)
97    }
98
99    fn consume(self: Pin<&mut Self>, amount: usize) {
100        self.io().consume(amount)
101    }
102}
futures_test/io/limited.rs

futures_test/io/
limited.rs
