-
Notifications
You must be signed in to change notification settings - Fork 625
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat(futures-util/stream): implement stream.unzip adapter #2256
Closed
Closed
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,177 @@ | ||
use crate::task::AtomicWaker; | ||
use alloc::sync::{Arc, Weak}; | ||
use core::pin::Pin; | ||
use futures_core::stream::{FusedStream, Stream}; | ||
use futures_core::task::{Context, Poll}; | ||
use pin_project::{pin_project, pinned_drop}; | ||
use std::sync::mpsc; | ||
|
||
/// SAFETY: safe because only one of two unzipped streams is guaranteed | ||
/// to be accessing underlying stream. This is guaranteed by mpsc. Right | ||
/// stream will access underlying stream only if Sender (or left stream) | ||
/// is dropped in which case try_recv returns disconnected error. | ||
unsafe fn poll_unzipped<S, T1, T2>( | ||
stream: Pin<&mut Arc<S>>, | ||
cx: &mut Context<'_>, | ||
) -> Poll<Option<S::Item>> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
stream | ||
.map_unchecked_mut(|x| &mut *(Arc::as_ptr(x) as *mut S)) | ||
.poll_next(cx) | ||
} | ||
|
||
#[pin_project(PinnedDrop)] | ||
#[derive(Debug)] | ||
#[must_use = "streams do nothing unless polled"] | ||
pub struct UnzipLeft<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
#[pin] | ||
stream: Arc<S>, | ||
right_waker: Weak<AtomicWaker>, | ||
right_queue: mpsc::Sender<Option<T2>>, | ||
} | ||
|
||
impl<S, T1, T2> UnzipLeft<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
fn send_to_right(&self, value: Option<T2>) { | ||
if let Some(right_waker) = self.right_waker.upgrade() { | ||
// if right_waker.upgrade() succeeds, then right is not | ||
// dropped so send won't fail. | ||
let _ = self.right_queue.send(value); | ||
right_waker.wake(); | ||
} | ||
} | ||
} | ||
|
||
impl<S, T1, T2> FusedStream for UnzipLeft<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)> + FusedStream, | ||
{ | ||
fn is_terminated(&self) -> bool { | ||
self.stream.as_ref().is_terminated() | ||
} | ||
} | ||
|
||
impl<S, T1, T2> Stream for UnzipLeft<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
type Item = T1; | ||
|
||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { | ||
let this = self.as_mut().project(); | ||
|
||
// SAFETY: for safety details see comment for function: poll_unzipped | ||
if let Some(value) = ready!(unsafe { poll_unzipped(this.stream, cx) }) { | ||
self.send_to_right(Some(value.1)); | ||
return Poll::Ready(Some(value.0)); | ||
} | ||
self.send_to_right(None); | ||
Poll::Ready(None) | ||
} | ||
} | ||
|
||
#[pinned_drop] | ||
impl<S, T1, T2> PinnedDrop for UnzipLeft<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
fn drop(self: Pin<&mut Self>) { | ||
let this = self.project(); | ||
// wake right stream if it isn't dropped | ||
if let Some(right_waker) = this.right_waker.upgrade() { | ||
// drop right_queue sender to cause rx.try_recv to return | ||
// TryRecvError::Disconnected, so that right knows left is | ||
// dropped and now it should take over polling the base stream. | ||
drop(this.right_queue); | ||
right_waker.wake(); | ||
} | ||
} | ||
} | ||
|
||
#[pin_project] | ||
#[derive(Debug)] | ||
#[must_use = "streams do nothing unless polled"] | ||
pub struct UnzipRight<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
#[pin] | ||
stream: Arc<S>, | ||
waker: Arc<AtomicWaker>, | ||
queue: mpsc::Receiver<Option<T2>>, | ||
is_done: bool, | ||
} | ||
|
||
impl<S, T1, T2> FusedStream for UnzipRight<S, T1, T2> | ||
where | ||
S: FusedStream<Item = (T1, T2)>, | ||
{ | ||
fn is_terminated(&self) -> bool { | ||
self.is_done | ||
} | ||
} | ||
|
||
impl<S, T1, T2> Stream for UnzipRight<S, T1, T2> | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
type Item = T2; | ||
|
||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { | ||
let this = self.project(); | ||
this.waker.register(&cx.waker().clone()); | ||
|
||
match this.queue.try_recv() { | ||
Ok(value) => { | ||
// can't know if more items are in the queue so wake the task | ||
// again while there are items. Will cause extra wake though. | ||
cx.waker().clone().wake(); | ||
if value.is_none() { | ||
*this.is_done = true; | ||
} | ||
Poll::Ready(value) | ||
} | ||
Err(mpsc::TryRecvError::Disconnected) => { | ||
// if left is dropped, it is no longer polling the base stream | ||
// so right should poll it instead. | ||
// SAFETY: for safety details see comment for function: poll_unzipped | ||
if let Some(value) = ready!(unsafe { poll_unzipped(this.stream, cx) }) { | ||
return Poll::Ready(Some(value.1)); | ||
} | ||
*this.is_done = true; | ||
Poll::Ready(None) | ||
} | ||
_ => Poll::Pending, | ||
} | ||
} | ||
} | ||
|
||
pub fn unzip<S, T1, T2>(stream: S) -> (UnzipLeft<S, T1, T2>, UnzipRight<S, T1, T2>) | ||
where | ||
S: Stream<Item = (T1, T2)>, | ||
{ | ||
let base_stream = Arc::new(stream); | ||
let waker = Arc::new(AtomicWaker::new()); | ||
let (tx, rx) = mpsc::channel::<Option<T2>>(); | ||
|
||
( | ||
UnzipLeft { | ||
stream: base_stream.clone(), | ||
right_waker: Arc::downgrade(&waker), | ||
right_queue: tx, | ||
}, | ||
UnzipRight { | ||
stream: base_stream.clone(), | ||
waker: waker, | ||
queue: rx, | ||
is_done: false, | ||
}, | ||
) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,78 @@ | ||
use futures::stream::{self, StreamExt}; | ||
use futures::executor::{block_on_stream}; | ||
|
||
fn fail_on_thread_panic() { | ||
std::panic::set_hook(Box::new(move |panic_info: &std::panic::PanicInfo| { | ||
println!("{}", panic_info.to_string()); | ||
std::process::exit(1); | ||
})); | ||
} | ||
|
||
fn sample_stream(start: usize, end: usize) -> futures_util::stream::Iter<std::vec::IntoIter<(usize, usize)>> { | ||
let list_iter = (start..end) | ||
.filter(|&x| x % 2 == 1) | ||
.map(|x| (x, x + 1)); | ||
|
||
return stream::iter(list_iter.collect::<Vec<_>>()); | ||
} | ||
|
||
#[test] | ||
fn left_dropped_before_first_poll() { | ||
let (_, mut s2) = { | ||
let (s1, s2) = sample_stream(1, 2).unzip(); | ||
(block_on_stream(s1), block_on_stream(s2)) | ||
}; | ||
|
||
assert_eq!(s2.next(), Some(2)); | ||
assert_eq!(s2.next(), None); | ||
} | ||
|
||
#[test] | ||
fn left_dropped_after_polled() { | ||
fail_on_thread_panic(); | ||
|
||
let (mut s1, mut s2) = { | ||
let (s1, s2) = sample_stream(1, 4).unzip(); | ||
(block_on_stream(s1), block_on_stream(s2)) | ||
}; | ||
|
||
|
||
let t1 = std::thread::spawn(move || { | ||
assert_eq!(s1.next(), Some(1)); | ||
drop(s1); | ||
}); | ||
|
||
let t2 = std::thread::spawn(move || { | ||
assert_eq!(s2.next(), Some(2)); | ||
assert_eq!(s2.next(), Some(4)); | ||
assert_eq!(s2.next(), None); | ||
}); | ||
|
||
let _ = t1.join(); | ||
let _ = t2.join(); | ||
} | ||
|
||
#[test] | ||
fn right_dropped_before_first_poll() { | ||
let (mut s1, _) = { | ||
let (s1, s2) = sample_stream(1, 2).unzip(); | ||
(block_on_stream(s1), block_on_stream(s2)) | ||
}; | ||
|
||
assert_eq!(s1.next(), Some(1)); | ||
assert_eq!(s1.next(), None); | ||
} | ||
|
||
#[test] | ||
fn right_dropped_after_polled() { | ||
let (mut s1, mut s2) = { | ||
let (s1, s2) = sample_stream(1, 4).unzip(); | ||
(block_on_stream(s1), block_on_stream(s2)) | ||
}; | ||
|
||
assert_eq!(s1.next(), Some(1)); | ||
assert_eq!(s2.next(), Some(2)); | ||
drop(s2); | ||
assert_eq!(s1.next(), Some(3)); | ||
assert_eq!(s1.next(), None); | ||
} |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't think this is correct. You can call
UnzipLeft(Right)::poll_next
andUnzipRight(Left)::is_terminated
at the same time.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Oh right! But I can simply maintain
is_done
boolean inUnzipRight
and set it totrue
if I receivePoll::Ready(None)
right? For exampleFuse<S>
seems to do the same.Or is there a case when:
is_terminated()
returns true for underline stream but it still hasn't emittedPoll::Ready(None)
and won't ever emit it?If thats the case, would
UnzipLeft
be dropped if itis_terminated
? So would it be sufficient ifis_terminated()
returned:is_done
if left is not dropped.stream.is_terminated()
if left is dropped, since it's now safe to call that method.