p2panda_stream/processors/
stream.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use std::pin::Pin;
4use std::task::{Context, Poll};
5
6use futures_core::Stream;
7use pin_project::pin_project;
8
9use crate::processors::Processor;
10use crate::processors::buffered::{Buffer, BufferReceiver, BufferSender};
11
12/// Extension trait for `Processor` that provides convenient methods for `Stream` processing.
13pub trait ProcessorExt<T>: Processor<T>
14where
15    Self: 'static,
16    T: 'static,
17{
18    /// Converts this processor into a `Stream` that takes an input stream and yields the
19    /// processor's output items.
20    fn into_stream<S: Stream<Item = T>>(self, input_stream: S) -> ProcessorStream<Self, S, T>
21    where
22        Self: Sized,
23    {
24        ProcessorStream::new(self, input_stream)
25    }
26}
27
28impl<P, T: 'static> ProcessorExt<T> for P where P: Processor<T> + 'static {}
29
30#[pin_project]
31#[must_use = "streams do nothing unless polled"]
32pub struct ProcessorStream<P, S, T>
33where
34    P: Processor<T>,
35    S: Stream<Item = T>,
36{
37    #[pin]
38    input_stream: S,
39    buffer: Buffer,
40    tx: BufferSender<T>,
41    rx: BufferReceiver<P, T>,
42}
43
44impl<P, S, T> ProcessorStream<P, S, T>
45where
46    P: Processor<T> + 'static,
47    S: Stream<Item = T>,
48    T: 'static,
49{
50    pub(crate) fn new(processor: P, input_stream: S) -> ProcessorStream<P, S, T> {
51        // Wrap processor around an unbounded buffer making sure we're correctly awaiting the
52        // futures returned by "process" and "next", correctly handling processors which might have
53        // "expensive" async behavior.
54        let (buffer, tx, rx) = Buffer::new(processor);
55        ProcessorStream {
56            input_stream,
57            buffer,
58            tx,
59            rx,
60        }
61    }
62}
63
64impl<P, S, T> Stream for ProcessorStream<P, S, T>
65where
66    P: Processor<T>,
67    S: Stream<Item = T>,
68{
69    type Item = Result<P::Output, P::Error>;
70
71    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
72        let mut this = self.project();
73        let mut made_progress = false;
74
75        // Check input stream.
76        //
77        // If there's an item in the previous `Stream`, we take it and push it into our internal
78        // buffer queue. If there's nothing in it (Poll::Pending), we continue.
79        //
80        // If the input stream terminated, we ignore this and continue as well since we don't know
81        // if the processor has more work to do. Processors never cease operation and are only
82        // stopped via higher level logic. This is why we also continue here in the `Stream` and
83        // never terminate.
84        match this.input_stream.as_mut().poll_next(cx) {
85            Poll::Ready(Some(input)) => {
86                let _ = this.tx.send(input);
87                made_progress = true;
88            }
89            Poll::Ready(None) | Poll::Pending => (),
90        };
91
92        // Check output of processor.
93        //
94        // If an output item was given, we yield it in this Stream. If an error was returned we do
95        // the same.
96        //
97        // If the next method returns a pending state we assume this is because there are no items
98        // given.
99        match this.rx.poll_recv(cx) {
100            Poll::Ready(item) => Poll::Ready(item),
101            Poll::Pending => {
102                if made_progress {
103                    // We sent new input, try again immediately.
104                    cx.waker().wake_by_ref();
105                }
106                Poll::Pending
107            }
108        }
109    }
110}
111
112/// Extension trait for `Stream` that provides convenient methods to layer up processors.
113pub trait StreamLayerExt<T>: Stream<Item = T> + Sized
114where
115    T: 'static,
116{
117    /// Processes this stream through the provided layer.
118    fn layer<P>(self, processor: P) -> ProcessorStream<P, Self, T>
119    where
120        P: Processor<T> + 'static,
121    {
122        ProcessorStream::new(processor, self)
123    }
124}
125
126impl<S, T: 'static> StreamLayerExt<T> for S where S: Stream<Item = T> {}
p2panda_stream/processors/stream.rs

p2panda_stream/processors/
stream.rs
