1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

use futures::{future::FusedFuture, FutureExt};
use futures_timer::Delay;
use std::{
	future::Future,
	pin::Pin,
	task::{Context, Poll, Waker},
	time::Duration,
};

enum Inner {
	Infinite {
		/// Waker from the most recent `poll` call. If `None`, either `poll` has not been called
		/// yet, we returned `Poll::Ready` from the last call, or the waker is attached to `delay`.
		waker: Option<Waker>,
		delay: Option<Delay>,
	},
	Finite(Delay),
}

/// Like [`Delay`] but the duration can be infinite (in which case the future will never fire).
/// Unlike [`Delay`], implements [`FusedFuture`], with [`is_terminated`](Self::is_terminated)
/// returning `true` when the delay is infinite. As with [`Delay`], once [`poll`](Self::poll)
/// returns [`Poll::Ready`], it will continue to do so until [`reset`](Self::reset) is called.
pub struct MaybeInfDelay(Inner);

impl MaybeInfDelay {
	/// Create a new `MaybeInfDelay` future. If `duration` is [`Some`], the future will fire after
	/// the given duration has elapsed. If `duration` is [`None`], the future will "never" fire
	/// (although see [`reset`](Self::reset)).
	pub fn new(duration: Option<Duration>) -> Self {
		match duration {
			Some(duration) => Self(Inner::Finite(Delay::new(duration))),
			None => Self(Inner::Infinite { waker: None, delay: None }),
		}
	}

	/// Reset the timer. `duration` is handled just like in [`new`](Self::new). Note that while
	/// this is similar to `std::mem::replace(&mut self, MaybeInfDelay::new(duration))`, with
	/// `replace` you would have to manually ensure [`poll`](Self::poll) was called again; with
	/// `reset` this is not necessary.
	pub fn reset(&mut self, duration: Option<Duration>) {
		match duration {
			Some(duration) => match &mut self.0 {
				Inner::Infinite { waker, delay } => {
					let mut delay = match delay.take() {
						Some(mut delay) => {
							delay.reset(duration);
							delay
						},
						None => Delay::new(duration),
					};
					if let Some(waker) = waker.take() {
						let mut cx = Context::from_waker(&waker);
						match delay.poll_unpin(&mut cx) {
							Poll::Pending => (), // Waker attached to delay
							Poll::Ready(_) => waker.wake(),
						}
					}
					self.0 = Inner::Finite(delay);
				},
				Inner::Finite(delay) => delay.reset(duration),
			},
			None =>
				self.0 = match std::mem::replace(
					&mut self.0,
					Inner::Infinite { waker: None, delay: None },
				) {
					Inner::Finite(delay) => Inner::Infinite { waker: None, delay: Some(delay) },
					infinite => infinite,
				},
		}
	}
}

impl Future for MaybeInfDelay {
	type Output = ();

	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
		match &mut self.0 {
			Inner::Infinite { waker, .. } => {
				*waker = Some(cx.waker().clone());
				Poll::Pending
			},
			Inner::Finite(delay) => delay.poll_unpin(cx),
		}
	}
}

impl FusedFuture for MaybeInfDelay {
	fn is_terminated(&self) -> bool {
		matches!(self.0, Inner::Infinite { .. })
	}
}