Merge pull request #503 from Razican/random_merge

Randomize Stream::merge to improve the throughput.
new-scheduler
nasa 5 years ago committed by GitHub
commit 3bc4d293dd
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -5,6 +5,7 @@ use pin_project_lite::pin_project;
use crate::prelude::*;
use crate::stream::Fuse;
use crate::utils;
pin_project! {
/// A stream that merges two other streams into a single stream.
@ -27,7 +28,10 @@ pin_project! {
impl<L: Stream, R: Stream> Merge<L, R> {
pub(crate) fn new(left: L, right: R) -> Self {
Self { left: left.fuse(), right: right.fuse() }
Self {
left: left.fuse(),
right: right.fuse(),
}
}
}
@ -40,14 +44,29 @@ where
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.project();
match this.left.poll_next(cx) {
Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
Poll::Ready(None) => this.right.poll_next(cx),
Poll::Pending => match this.right.poll_next(cx) {
Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
Poll::Ready(None) => Poll::Pending,
Poll::Pending => Poll::Pending,
}
if utils::random(1) == 1 {
poll_next_in_order(this.left, this.right, cx)
} else {
poll_next_in_order(this.right, this.left, cx)
}
}
}
fn poll_next_in_order<F, S, I>(
first: Pin<&mut F>,
second: Pin<&mut S>,
cx: &mut Context<'_>,
) -> Poll<Option<I>>
where
F: Stream<Item = I>,
S: Stream<Item = I>,
{
match first.poll_next(cx) {
Poll::Ready(None) => second.poll_next(cx),
Poll::Ready(item) => Poll::Ready(item),
Poll::Pending => match second.poll_next(cx) {
Poll::Ready(None) | Poll::Pending => Poll::Pending,
Poll::Ready(item) => Poll::Ready(item),
},
}
}

@ -1839,18 +1839,17 @@ extension_trait! {
```
# async_std::task::block_on(async {
use async_std::prelude::*;
use async_std::stream;
use async_std::stream::{self, FromStream};
let a = stream::once(1u8);
let b = stream::once(2u8);
let c = stream::once(3u8);
let mut s = a.merge(b).merge(c);
let s = a.merge(b).merge(c);
let mut lst = Vec::from_stream(s).await;
assert_eq!(s.next().await, Some(1u8));
assert_eq!(s.next().await, Some(2u8));
assert_eq!(s.next().await, Some(3u8));
assert_eq!(s.next().await, None);
lst.sort_unstable();
assert_eq!(&lst, &[1u8, 2u8, 3u8]);
# });
```
"#]

Loading…
Cancel
Save