109 lines
3.2 KiB
Rust
109 lines
3.2 KiB
Rust
use std::collections::VecDeque;
|
|
use std::fmt;
|
|
use std::future::Future;
|
|
|
|
use futures::FutureExt;
|
|
use tokio::sync::oneshot;
|
|
|
|
/// We define a AsyncStream to replace futures::Stream since we don't want to implement
|
|
/// poll_next nor using async_stream.
|
|
/// Although we use GAT, we don't want the future to capture self's ref. We did like
|
|
/// that before, and this makes it hard to load stream in parallel like Buffered.
|
|
/// Also, our AsyncStream is not like Stream in signature. We return `Option<Future>`
|
|
/// instead of `Future<Output = Option<_>>`.
|
|
pub trait AsyncStream {
|
|
type Item;
|
|
type Future: Future<Output = Self::Item>;
|
|
fn next(&mut self) -> Option<Self::Future>;
|
|
|
|
#[inline]
|
|
fn size_hint(&self) -> (usize, Option<usize>) {
|
|
(0, None)
|
|
}
|
|
}
|
|
|
|
/// Buffered Stream.
|
|
/// By decorating Buffered, the output future of stream will be polled
|
|
/// concurrently.
|
|
/// Here I implement it by spawning tasks. It is indeed not efficient as
|
|
/// `FuturesOrdered` which is used by `futures-util::stream::Buffered`.
|
|
/// As a decorator of an async trait, it is hard to implement it in a poll
|
|
/// way. We can do that, but it breaks the safety boundary which requires
|
|
/// user to make sure that the AsyncStream exists when polling the future
|
|
/// since in our trait definition, the future has no relation with self.
|
|
/// And without poll, we can not drive multiple futures by one future.
|
|
pub struct Buffered<St>
|
|
where
|
|
St: AsyncStream,
|
|
{
|
|
stream: Option<St>,
|
|
queue: VecDeque<oneshot::Receiver<St::Item>>,
|
|
max: usize,
|
|
}
|
|
|
|
impl<St> Buffered<St>
|
|
where
|
|
St: AsyncStream,
|
|
{
|
|
pub fn new(stream: St, buffer_size: usize) -> Self {
|
|
Self {
|
|
stream: Some(stream),
|
|
queue: VecDeque::with_capacity(buffer_size),
|
|
max: buffer_size,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<St> fmt::Debug for Buffered<St>
|
|
where
|
|
St: AsyncStream + fmt::Debug,
|
|
St::Item: fmt::Debug,
|
|
{
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("Buffered")
|
|
.field("stream", &self.stream)
|
|
.field("queue", &self.queue)
|
|
.field("max", &self.max)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl<St> AsyncStream for Buffered<St>
|
|
where
|
|
St: AsyncStream,
|
|
St::Item: Send + 'static,
|
|
St::Future: Send + 'static,
|
|
{
|
|
type Item = St::Item;
|
|
|
|
type Future = impl std::future::Future<Output = Self::Item>;
|
|
|
|
fn next(&mut self) -> Option<Self::Future> {
|
|
while self.queue.len() < self.max {
|
|
let item = match self.stream.as_mut() {
|
|
Some(st) => match st.next() {
|
|
Some(item) => Some(item),
|
|
None => {
|
|
self.stream = None;
|
|
None
|
|
}
|
|
},
|
|
None => None,
|
|
};
|
|
match item {
|
|
Some(f) => {
|
|
let (tx, rx) = oneshot::channel::<Self::Item>();
|
|
tokio::spawn(async move {
|
|
let _ = tx.send(f.await);
|
|
});
|
|
self.queue.push_back(rx);
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
self.queue
|
|
.pop_front()
|
|
.map(|x| x.map(|xx| xx.expect("oneshot tx dropped which is unexpected")))
|
|
}
|
|
}
|