sc_mixnet/
maybe_inf_delay.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19use futures::{future::FusedFuture, FutureExt};
20use futures_timer::Delay;
21use std::{
22	future::Future,
23	pin::Pin,
24	task::{Context, Poll, Waker},
25	time::Duration,
26};
27
28enum Inner {
29	Infinite {
30		/// Waker from the most recent `poll` call. If `None`, either `poll` has not been called
31		/// yet, we returned `Poll::Ready` from the last call, or the waker is attached to `delay`.
32		waker: Option<Waker>,
33		delay: Option<Delay>,
34	},
35	Finite(Delay),
36}
37
38/// Like [`Delay`] but the duration can be infinite (in which case the future will never fire).
39/// Unlike [`Delay`], implements [`FusedFuture`], with [`is_terminated`](Self::is_terminated)
40/// returning `true` when the delay is infinite. As with [`Delay`], once [`poll`](Self::poll)
41/// returns [`Poll::Ready`], it will continue to do so until [`reset`](Self::reset) is called.
42pub struct MaybeInfDelay(Inner);
43
44impl MaybeInfDelay {
45	/// Create a new `MaybeInfDelay` future. If `duration` is [`Some`], the future will fire after
46	/// the given duration has elapsed. If `duration` is [`None`], the future will "never" fire
47	/// (although see [`reset`](Self::reset)).
48	pub fn new(duration: Option<Duration>) -> Self {
49		match duration {
50			Some(duration) => Self(Inner::Finite(Delay::new(duration))),
51			None => Self(Inner::Infinite { waker: None, delay: None }),
52		}
53	}
54
55	/// Reset the timer. `duration` is handled just like in [`new`](Self::new). Note that while
56	/// this is similar to `std::mem::replace(&mut self, MaybeInfDelay::new(duration))`, with
57	/// `replace` you would have to manually ensure [`poll`](Self::poll) was called again; with
58	/// `reset` this is not necessary.
59	pub fn reset(&mut self, duration: Option<Duration>) {
60		match duration {
61			Some(duration) => match &mut self.0 {
62				Inner::Infinite { waker, delay } => {
63					let mut delay = match delay.take() {
64						Some(mut delay) => {
65							delay.reset(duration);
66							delay
67						},
68						None => Delay::new(duration),
69					};
70					if let Some(waker) = waker.take() {
71						let mut cx = Context::from_waker(&waker);
72						match delay.poll_unpin(&mut cx) {
73							Poll::Pending => (), // Waker attached to delay
74							Poll::Ready(_) => waker.wake(),
75						}
76					}
77					self.0 = Inner::Finite(delay);
78				},
79				Inner::Finite(delay) => delay.reset(duration),
80			},
81			None =>
82				self.0 = match std::mem::replace(
83					&mut self.0,
84					Inner::Infinite { waker: None, delay: None },
85				) {
86					Inner::Finite(delay) => Inner::Infinite { waker: None, delay: Some(delay) },
87					infinite => infinite,
88				},
89		}
90	}
91}
92
93impl Future for MaybeInfDelay {
94	type Output = ();
95
96	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
97		match &mut self.0 {
98			Inner::Infinite { waker, .. } => {
99				*waker = Some(cx.waker().clone());
100				Poll::Pending
101			},
102			Inner::Finite(delay) => delay.poll_unpin(cx),
103		}
104	}
105}
106
107impl FusedFuture for MaybeInfDelay {
108	fn is_terminated(&self) -> bool {
109		matches!(self.0, Inner::Infinite { .. })
110	}
111}
sc_mixnet/maybe_inf_delay.rs

sc_mixnet/
maybe_inf_delay.rs
