forked from mirror/async-std
Merge pull request #503 from Razican/random_merge
Randomize Stream::merge to improve the throughput.
This commit is contained in:
commit
3bc4d293dd
2 changed files with 33 additions and 15 deletions
|
@ -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…
Reference in a new issue