discarding_stream.rs

use std::{
    pin::Pin,
    task::{Context, Poll},
};

use chrono::{DateTime, Duration, Utc};
use futures_core::{Stream, stream::LocalBoxStream};

pub struct DiscardingStream<'a, T> {
    last_message: Option<DateTime<Utc>>,
    discard_interval: Duration,
    stream: LocalBoxStream<'a, T>,
}

impl<'a, T> DiscardingStream<'a, T> {
    pub fn new(
        stream: LocalBoxStream<'a, T>,
        discard_interval: Duration,
    ) -> DiscardingStream<'a, T> {
        DiscardingStream {
            last_message: None,
            discard_interval,
            stream,
        }
    }
}

impl<'a, T> Stream for DiscardingStream<'a, T> {
    type Item = T;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.get_mut();
        match this.stream.as_mut().poll_next(cx) {
            Poll::Ready(Some(item)) => {
                let now = Utc::now();
                if let Some(last) = this.last_message {
                    if (now - last) < this.discard_interval {
                        // Discard the message
                        return Poll::Pending;
                    }
                }
                this.last_message = Some(now);
                Poll::Ready(Some(item))
            }
            other => other,
        }
    }
}