use crate::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
use futures::{lock::Mutex, prelude::*};
use futures_timer::Delay;
use std::{
pin::Pin,
task::{Context, Poll},
time::Duration,
};
pub struct StatusSinks<T> {
inner: Mutex<Inner<T>>,
entries_tx: TracingUnboundedSender<YieldAfter<T>>,
}
struct Inner<T> {
entries: stream::FuturesUnordered<YieldAfter<T>>,
entries_rx: TracingUnboundedReceiver<YieldAfter<T>>,
}
struct YieldAfter<T> {
delay: Delay,
interval: Duration,
sender: Option<TracingUnboundedSender<T>>,
}
impl<T> Default for StatusSinks<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> StatusSinks<T> {
pub fn new() -> StatusSinks<T> {
let (entries_tx, entries_rx) = tracing_unbounded("status-sinks-entries", 100_000);
StatusSinks {
inner: Mutex::new(Inner { entries: stream::FuturesUnordered::new(), entries_rx }),
entries_tx,
}
}
pub fn push(&self, interval: Duration, sender: TracingUnboundedSender<T>) {
let _ = self.entries_tx.unbounded_send(YieldAfter {
delay: Delay::new(interval),
interval,
sender: Some(sender),
});
}
pub async fn next(&self) -> ReadySinkEvent<'_, T> {
let mut inner = self.inner.lock().await;
let inner = &mut *inner;
loop {
let next_ready_entry = {
let entries = &mut inner.entries;
async move {
if let Some(v) = entries.next().await {
v
} else {
loop {
futures::pending!()
}
}
}
};
futures::select! {
new_entry = inner.entries_rx.next() => {
if let Some(new_entry) = new_entry {
inner.entries.push(new_entry);
}
},
(sender, interval) = next_ready_entry.fuse() => {
return ReadySinkEvent {
sinks: self,
sender: Some(sender),
interval,
}
}
}
}
}
}
#[must_use]
pub struct ReadySinkEvent<'a, T> {
sinks: &'a StatusSinks<T>,
sender: Option<TracingUnboundedSender<T>>,
interval: Duration,
}
impl<'a, T> ReadySinkEvent<'a, T> {
pub fn send(mut self, element: T) {
if let Some(sender) = self.sender.take() {
if sender.unbounded_send(element).is_ok() {
let _ = self.sinks.entries_tx.unbounded_send(YieldAfter {
delay: Delay::new(self.interval),
interval: self.interval,
sender: Some(sender),
});
}
}
}
}
impl<'a, T> Drop for ReadySinkEvent<'a, T> {
fn drop(&mut self) {
if let Some(sender) = self.sender.take() {
if sender.is_closed() {
return
}
let _ = self.sinks.entries_tx.unbounded_send(YieldAfter {
delay: Delay::new(self.interval),
interval: self.interval,
sender: Some(sender),
});
}
}
}
impl<T> futures::Future for YieldAfter<T> {
type Output = (TracingUnboundedSender<T>, Duration);
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let this = Pin::into_inner(self);
match Pin::new(&mut this.delay).poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(()) => {
let sender = this
.sender
.take()
.expect("sender is always Some unless the future is finished; qed");
Poll::Ready((sender, this.interval))
},
}
}
}
#[cfg(test)]
mod tests {
use super::StatusSinks;
use crate::mpsc::tracing_unbounded;
use futures::prelude::*;
use std::time::Duration;
#[test]
fn works() {
let status_sinks = StatusSinks::new();
let (tx, rx) = tracing_unbounded("test", 100_000);
status_sinks.push(Duration::from_millis(100), tx);
let mut val_order = 5;
futures::executor::block_on(futures::future::select(
Box::pin(async move {
loop {
let ev = status_sinks.next().await;
val_order += 1;
ev.send(val_order);
}
}),
Box::pin(async {
let items: Vec<i32> = rx.take(3).collect().await;
assert_eq!(items, [6, 7, 8]);
}),
));
}
}