swansong/implementation/
guarded.rs

1use super::Inner;
2use crate::Guard;
3use futures_core::Stream;
4use std::{
5    future::Future,
6    ops::{Deref, DerefMut},
7    pin::Pin,
8    sync::Arc,
9    task::{Context, Poll},
10};
11
12pin_project_lite::pin_project! {
13    /// Guarded is a convenient way to attach a [`Guard`] to another type.
14    ///
15    /// Guarded does not stop the wrapped type on shutdown, but will delay shutdown until it is
16    /// dropped. To stop the wrapped type, use
17    /// [`Swansong::interrupt`][crate::Swansong::interrupt]. To both stop the wrapped type and
18    /// also act as a guard, use [`Interrupt::guarded`][crate::Interrupt::guarded].
19    ///
20    /// Guarded implements Future, Stream, Clone, Debug, AsyncRead, and AsyncWrite when the wrapped
21    /// type also does.
22    ///
23    /// Guarded implements [`Deref`] and [`DerefMut`] to the wrapped type.
24    #[derive(Clone, Debug, PartialEq, Eq)]
25    pub struct Guarded<T> {
26        guard: Guard,
27        #[pin]
28        wrapped_type: T
29    }
30}
31
32impl<T> Guarded<T> {
33    #[track_caller]
34    pub(crate) fn new(inner: &Arc<Inner>, wrapped_type: T) -> Self {
35        Self {
36            guard: Guard::new(inner),
37            wrapped_type,
38        }
39    }
40
41    /// Transform this `Guarded<T>` into the inner `T`, dropping the [`Guard`] in the process.
42    ///
43    /// Doing this allows shutdown to proceed if no other guards exist and shutdown is initiated.
44    pub fn into_inner(self) -> T {
45        self.wrapped_type
46    }
47}
48
49impl<T: Future> Future for Guarded<T> {
50    type Output = T::Output;
51
52    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
53        self.project().wrapped_type.poll(cx)
54    }
55}
56
57impl<T: Stream> Stream for Guarded<T> {
58    type Item = T::Item;
59
60    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
61        self.project().wrapped_type.poll_next(cx)
62    }
63
64    fn size_hint(&self) -> (usize, Option<usize>) {
65        self.wrapped_type.size_hint()
66    }
67}
68
69impl<T> Deref for Guarded<T> {
70    type Target = T;
71
72    fn deref(&self) -> &Self::Target {
73        &self.wrapped_type
74    }
75}
76
77impl<T> DerefMut for Guarded<T> {
78    fn deref_mut(&mut self) -> &mut Self::Target {
79        &mut self.wrapped_type
80    }
81}
82
83#[cfg(feature = "futures-io")]
84impl<T: futures_io::AsyncRead> futures_io::AsyncRead for Guarded<T> {
85    fn poll_read(
86        self: Pin<&mut Self>,
87        cx: &mut Context<'_>,
88        buf: &mut [u8],
89    ) -> Poll<std::io::Result<usize>> {
90        self.project().wrapped_type.poll_read(cx, buf)
91    }
92
93    fn poll_read_vectored(
94        self: Pin<&mut Self>,
95        cx: &mut Context<'_>,
96        bufs: &mut [std::io::IoSliceMut<'_>],
97    ) -> Poll<std::io::Result<usize>> {
98        self.project().wrapped_type.poll_read_vectored(cx, bufs)
99    }
100}
101
102#[cfg(feature = "futures-io")]
103impl<T: futures_io::AsyncWrite> futures_io::AsyncWrite for Guarded<T> {
104    fn poll_write(
105        self: Pin<&mut Self>,
106        cx: &mut Context<'_>,
107        buf: &[u8],
108    ) -> Poll<std::io::Result<usize>> {
109        self.project().wrapped_type.poll_write(cx, buf)
110    }
111
112    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
113        self.project().wrapped_type.poll_flush(cx)
114    }
115
116    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
117        self.project().wrapped_type.poll_close(cx)
118    }
119
120    fn poll_write_vectored(
121        self: Pin<&mut Self>,
122        cx: &mut Context<'_>,
123        bufs: &[std::io::IoSlice<'_>],
124    ) -> Poll<std::io::Result<usize>> {
125        self.project().wrapped_type.poll_write_vectored(cx, bufs)
126    }
127}
128
129#[cfg(feature = "futures-io")]
130impl<T: futures_io::AsyncBufRead> futures_io::AsyncBufRead for Guarded<T> {
131    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
132        self.project().wrapped_type.poll_fill_buf(cx)
133    }
134
135    fn consume(self: Pin<&mut Self>, amt: usize) {
136        self.project().wrapped_type.consume(amt);
137    }
138}
139
140#[cfg(feature = "tokio")]
141impl<T: tokio::io::AsyncRead> tokio::io::AsyncRead for Guarded<T> {
142    fn poll_read(
143        self: Pin<&mut Self>,
144        cx: &mut Context<'_>,
145        buf: &mut tokio::io::ReadBuf<'_>,
146    ) -> Poll<std::io::Result<()>> {
147        self.project().wrapped_type.poll_read(cx, buf)
148    }
149}
150#[cfg(feature = "tokio")]
151impl<T: tokio::io::AsyncWrite> tokio::io::AsyncWrite for Guarded<T> {
152    fn poll_write(
153        self: Pin<&mut Self>,
154        cx: &mut Context<'_>,
155        buf: &[u8],
156    ) -> Poll<Result<usize, std::io::Error>> {
157        self.project().wrapped_type.poll_write(cx, buf)
158    }
159
160    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
161        self.project().wrapped_type.poll_flush(cx)
162    }
163
164    fn poll_shutdown(
165        self: Pin<&mut Self>,
166        cx: &mut Context<'_>,
167    ) -> Poll<Result<(), std::io::Error>> {
168        self.project().wrapped_type.poll_shutdown(cx)
169    }
170
171    fn poll_write_vectored(
172        self: Pin<&mut Self>,
173        cx: &mut Context<'_>,
174        bufs: &[std::io::IoSlice<'_>],
175    ) -> Poll<Result<usize, std::io::Error>> {
176        self.project().wrapped_type.poll_write_vectored(cx, bufs)
177    }
178
179    fn is_write_vectored(&self) -> bool {
180        self.wrapped_type.is_write_vectored()
181    }
182}
183
184#[cfg(feature = "tokio")]
185impl<T: tokio::io::AsyncBufRead> tokio::io::AsyncBufRead for Guarded<T> {
186    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
187        self.project().wrapped_type.poll_fill_buf(cx)
188    }
189
190    fn consume(self: Pin<&mut Self>, amt: usize) {
191        self.project().wrapped_type.consume(amt);
192    }
193}
swansong/implementation/guarded.rs

swansong/implementation/
guarded.rs
