mirror/async-std
2
0
Fork 1
mirror of https://github.com/async-rs/async-std.git synced 2025-04-23 16:56:46 +00:00
Code Issues Releases Wiki Activity

Merge pull request #503 from Razican/random_merge

Browse source 
Randomize Stream::merge to improve the throughput.
This commit is contained in:
nasa 2019-11-20 08:42:50 +09:00 committed by GitHub
commit 3bc4d293dd
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 33 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
src/stream/stream/merge.rs
View file

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

11
src/stream/stream/mod.rs
View file

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