use crate::Result;
use futures::stream::Stream;
use serde_json::Value;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::task::{Context, Poll};
pub type Predicate = Box<dyn Fn(&Value) -> bool + Send>;
pub struct FilteredStream<S> {
inner: S,
predicate: Predicate,
eval_count: AtomicU64,
}
impl<S> FilteredStream<S> {
pub fn new(inner: S, predicate: Predicate) -> Self {
Self {
inner,
predicate,
eval_count: AtomicU64::new(0),
}
}
}
impl<S> Stream for FilteredStream<S>
where
S: Stream<Item = Result<Value>> + Unpin,
{
type Item = Result<Value>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match Pin::new(&mut self.inner).poll_next(cx) {
Poll::Ready(Some(Ok(value))) => {
let eval_idx = self.eval_count.fetch_add(1, Ordering::Relaxed);
let passed = if eval_idx.is_multiple_of(1000) {
let filter_start = std::time::Instant::now();
let result = (self.predicate)(&value);
let filter_duration = filter_start.elapsed().as_millis() as u64;
crate::metrics::histograms::filter_duration("unknown", filter_duration);
result
} else {
(self.predicate)(&value)
};
if passed {
return Poll::Ready(Some(Ok(value)));
}
crate::metrics::counters::rows_filtered("unknown", 1);
continue;
}
Poll::Ready(Some(Err(e))) => {
return Poll::Ready(Some(Err(e)));
}
Poll::Ready(None) => {
return Poll::Ready(None);
}
Poll::Pending => {
return Poll::Pending;
}
}
}
}
}
#[cfg(test)]
mod tests;