Skip to main content

soil_client/utils/
status_sinks.rs

1// This file is part of Soil.
2
3// Copyright (C) Soil contributors.
4// Copyright (C) Parity Technologies (UK) Ltd.
5// SPDX-License-Identifier: Apache-2.0 OR GPL-3.0-or-later WITH Classpath-exception-2.0
6
7use crate::utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
8use futures::{lock::Mutex, prelude::*};
9use futures_timer::Delay;
10use std::{
11	pin::Pin,
12	task::{Context, Poll},
13	time::Duration,
14};
15
16/// Holds a list of `UnboundedSender`s, each associated with a certain time period. Every time the
17/// period elapses, we push an element on the sender.
18///
19/// Senders are removed only when they are closed.
20pub struct StatusSinks<T> {
21	/// Should only be locked by `next`.
22	inner: Mutex<Inner<T>>,
23	/// Sending side of `Inner::entries_rx`.
24	entries_tx: TracingUnboundedSender<YieldAfter<T>>,
25}
26
27struct Inner<T> {
28	/// The actual entries of the list.
29	entries: stream::FuturesUnordered<YieldAfter<T>>,
30	/// Receives new entries and puts them in `entries`.
31	entries_rx: TracingUnboundedReceiver<YieldAfter<T>>,
32}
33
34struct YieldAfter<T> {
35	delay: Delay,
36	interval: Duration,
37	sender: Option<TracingUnboundedSender<T>>,
38}
39
40impl<T> Default for StatusSinks<T> {
41	fn default() -> Self {
42		Self::new()
43	}
44}
45
46impl<T> StatusSinks<T> {
47	/// Builds a new empty collection.
48	pub fn new() -> StatusSinks<T> {
49		let (entries_tx, entries_rx) = tracing_unbounded("status-sinks-entries", 100_000);
50
51		StatusSinks {
52			inner: Mutex::new(Inner { entries: stream::FuturesUnordered::new(), entries_rx }),
53			entries_tx,
54		}
55	}
56
57	/// Adds a sender to the collection.
58	///
59	/// The `interval` is the time period between two pushes on the sender.
60	pub fn push(&self, interval: Duration, sender: TracingUnboundedSender<T>) {
61		let _ = self.entries_tx.unbounded_send(YieldAfter {
62			delay: Delay::new(interval),
63			interval,
64			sender: Some(sender),
65		});
66	}
67
68	/// Waits until one of the sinks is ready, then returns an object that can be used to send
69	/// an element on said sink.
70	///
71	/// If the object isn't used to send an element, the slot is skipped.
72	pub async fn next(&self) -> ReadySinkEvent<'_, T> {
73		// This is only ever locked by `next`, which means that one `next` at a time can run.
74		let mut inner = self.inner.lock().await;
75		let inner = &mut *inner;
76
77		loop {
78			// Future that produces the next ready entry in `entries`, or doesn't produce anything
79			// if the list is empty.
80			let next_ready_entry = {
81				let entries = &mut inner.entries;
82				async move {
83					if let Some(v) = entries.next().await {
84						v
85					} else {
86						loop {
87							futures::pending!()
88						}
89					}
90				}
91			};
92
93			futures::select! {
94				new_entry = inner.entries_rx.next() => {
95					if let Some(new_entry) = new_entry {
96						inner.entries.push(new_entry);
97					}
98				},
99				(sender, interval) = next_ready_entry.fuse() => {
100					return ReadySinkEvent {
101						sinks: self,
102						sender: Some(sender),
103						interval,
104					}
105				}
106			}
107		}
108	}
109}
110
111/// One of the sinks is ready.
112#[must_use]
113pub struct ReadySinkEvent<'a, T> {
114	sinks: &'a StatusSinks<T>,
115	sender: Option<TracingUnboundedSender<T>>,
116	interval: Duration,
117}
118
119impl<'a, T> ReadySinkEvent<'a, T> {
120	/// Sends an element on the sender.
121	pub fn send(mut self, element: T) {
122		if let Some(sender) = self.sender.take() {
123			if sender.unbounded_send(element).is_ok() {
124				let _ = self.sinks.entries_tx.unbounded_send(YieldAfter {
125					// Note that since there's a small delay between the moment a task is
126					// woken up and the moment it is polled, the period is actually not
127					// `interval` but `interval + <delay>`. We ignore this problem in
128					// practice.
129					delay: Delay::new(self.interval),
130					interval: self.interval,
131					sender: Some(sender),
132				});
133			}
134		}
135	}
136}
137
138impl<'a, T> Drop for ReadySinkEvent<'a, T> {
139	fn drop(&mut self) {
140		if let Some(sender) = self.sender.take() {
141			if sender.is_closed() {
142				return;
143			}
144
145			let _ = self.sinks.entries_tx.unbounded_send(YieldAfter {
146				delay: Delay::new(self.interval),
147				interval: self.interval,
148				sender: Some(sender),
149			});
150		}
151	}
152}
153
154impl<T> futures::Future for YieldAfter<T> {
155	type Output = (TracingUnboundedSender<T>, Duration);
156
157	fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
158		let this = Pin::into_inner(self);
159
160		match Pin::new(&mut this.delay).poll(cx) {
161			Poll::Pending => Poll::Pending,
162			Poll::Ready(()) => {
163				let sender = this
164					.sender
165					.take()
166					.expect("sender is always Some unless the future is finished; qed");
167				Poll::Ready((sender, this.interval))
168			},
169		}
170	}
171}
172
173#[cfg(test)]
174mod tests {
175	use super::StatusSinks;
176	use crate::utils::mpsc::tracing_unbounded;
177	use futures::prelude::*;
178	use std::time::Duration;
179
180	#[test]
181	fn works() {
182		// We're not testing that the `StatusSink` properly enforces an order in the intervals, as
183		// this easily causes test failures on busy CPUs.
184
185		let status_sinks = StatusSinks::new();
186
187		let (tx, rx) = tracing_unbounded("test", 100_000);
188		status_sinks.push(Duration::from_millis(100), tx);
189
190		let mut val_order = 5;
191
192		futures::executor::block_on(futures::future::select(
193			Box::pin(async move {
194				loop {
195					let ev = status_sinks.next().await;
196					val_order += 1;
197					ev.send(val_order);
198				}
199			}),
200			Box::pin(async {
201				let items: Vec<i32> = rx.take(3).collect().await;
202				assert_eq!(items, [6, 7, 8]);
203			}),
204		));
205	}
206}