soil_client/utils/
status_sinks.rs1use 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
16pub struct StatusSinks<T> {
21 inner: Mutex<Inner<T>>,
23 entries_tx: TracingUnboundedSender<YieldAfter<T>>,
25}
26
27struct Inner<T> {
28 entries: stream::FuturesUnordered<YieldAfter<T>>,
30 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 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 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 pub async fn next(&self) -> ReadySinkEvent<'_, T> {
73 let mut inner = self.inner.lock().await;
75 let inner = &mut *inner;
76
77 loop {
78 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#[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 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 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 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}