tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-09-06 00:00:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

Move stream items into tokio-stream (#3277)

Browse Source 
This change removes all references to `Stream` from
within the `tokio` crate and moves them into a new
`tokio-stream` crate. Most types have had their
`impl Stream` removed as well in-favor of their
inherent methods.

Closes #2870
This commit is contained in:
Lucio Franco
2020-12-15 20:24:38 -08:00
committed by GitHub
parent fcce78b33a
commit 8efa62013b
101 changed files with 710 additions and 577 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio-stream/src/all.rs
+55
View File
@@ -0,0 +1,55 @@
use crate::Stream;
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Future for the [`all`](super::StreamExt::all) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct AllFuture<'a, St: ?Sized, F> {
stream: &'a mut St,
f: F,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
impl<'a, St: ?Sized, F> AllFuture<'a, St, F> {
pub(super) fn new(stream: &'a mut St, f: F) -> Self {
Self {
stream,
f,
_pin: PhantomPinned,
}
}
}
impl<St, F> Future for AllFuture<'_, St, F>
where
St: ?Sized + Stream + Unpin,
F: FnMut(St::Item) -> bool,
{
type Output = bool;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.project();
let next = futures_core::ready!(Pin::new(me.stream).poll_next(cx));
match next {
Some(v) => {
if !(me.f)(v) {
Poll::Ready(false)
} else {
cx.waker().wake_by_ref();
Poll::Pending
}
}
None => Poll::Ready(true),
}
}
}
tokio-stream/src/any.rs
+55
View File
@@ -0,0 +1,55 @@
use crate::Stream;
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Future for the [`any`](super::StreamExt::any) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct AnyFuture<'a, St: ?Sized, F> {
stream: &'a mut St,
f: F,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
impl<'a, St: ?Sized, F> AnyFuture<'a, St, F> {
pub(super) fn new(stream: &'a mut St, f: F) -> Self {
Self {
stream,
f,
_pin: PhantomPinned,
}
}
}
impl<St, F> Future for AnyFuture<'_, St, F>
where
St: ?Sized + Stream + Unpin,
F: FnMut(St::Item) -> bool,
{
type Output = bool;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.project();
let next = futures_core::ready!(Pin::new(me.stream).poll_next(cx));
match next {
Some(v) => {
if (me.f)(v) {
Poll::Ready(true)
} else {
cx.waker().wake_by_ref();
Poll::Pending
}
}
None => Poll::Ready(false),
}
}
}
tokio-stream/src/chain.rs
+49
View File
@@ -0,0 +1,49 @@
use crate::{Fuse, Stream};
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream returned by the [`chain`](super::StreamExt::chain) method.
pub struct Chain<T, U> {
#[pin]
a: Fuse<T>,
#[pin]
b: U,
}
}
impl<T, U> Chain<T, U> {
pub(super) fn new(a: T, b: U) -> Chain<T, U>
where
T: Stream,
U: Stream,
{
Chain { a: Fuse::new(a), b }
}
}
impl<T, U> Stream for Chain<T, U>
where
T: Stream,
U: Stream<Item = T::Item>,
{
type Item = T::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
use Poll::Ready;
let me = self.project();
if let Some(v) = ready!(me.a.poll_next(cx)) {
return Ready(Some(v));
}
me.b.poll_next(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
super::merge_size_hints(self.a.size_hint(), self.b.size_hint())
}
}
tokio-stream/src/collect.rs
+233
View File
@@ -0,0 +1,233 @@
use crate::Stream;
use core::future::Future;
use core::marker::PhantomPinned;
use core::mem;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
// Do not export this struct until `FromStream` can be unsealed.
pin_project! {
/// Future returned by the [`collect`](super::StreamExt::collect) method.
#[must_use = "futures do nothing unless you `.await` or poll them"]
#[derive(Debug)]
pub struct Collect<T, U>
where
T: Stream,
U: FromStream<T::Item>,
{
#[pin]
stream: T,
collection: U::InternalCollection,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
/// Convert from a [`Stream`](crate::Stream).
///
/// This trait is not intended to be used directly. Instead, call
/// [`StreamExt::collect()`](super::StreamExt::collect).
///
/// # Implementing
///
/// Currently, this trait may not be implemented by third parties. The trait is
/// sealed in order to make changes in the future. Stabilization is pending
/// enhancements to the Rust language.
pub trait FromStream<T>: sealed::FromStreamPriv<T> {}
impl<T, U> Collect<T, U>
where
T: Stream,
U: FromStream<T::Item>,
{
pub(super) fn new(stream: T) -> Collect<T, U> {
let (lower, upper) = stream.size_hint();
let collection = U::initialize(sealed::Internal, lower, upper);
Collect {
stream,
collection,
_pin: PhantomPinned,
}
}
}
impl<T, U> Future for Collect<T, U>
where
T: Stream,
U: FromStream<T::Item>,
{
type Output = U;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<U> {
use Poll::Ready;
loop {
let mut me = self.as_mut().project();
let item = match ready!(me.stream.poll_next(cx)) {
Some(item) => item,
None => {
return Ready(U::finalize(sealed::Internal, &mut me.collection));
}
};
if !U::extend(sealed::Internal, &mut me.collection, item) {
return Ready(U::finalize(sealed::Internal, &mut me.collection));
}
}
}
}
// ===== FromStream implementations
impl FromStream<()> for () {}
impl sealed::FromStreamPriv<()> for () {
type InternalCollection = ();
fn initialize(_: sealed::Internal, _lower: usize, _upper: Option<usize>) {}
fn extend(_: sealed::Internal, _collection: &mut (), _item: ()) -> bool {
true
}
fn finalize(_: sealed::Internal, _collection: &mut ()) {}
}
impl<T: AsRef<str>> FromStream<T> for String {}
impl<T: AsRef<str>> sealed::FromStreamPriv<T> for String {
type InternalCollection = String;
fn initialize(_: sealed::Internal, _lower: usize, _upper: Option<usize>) -> String {
String::new()
}
fn extend(_: sealed::Internal, collection: &mut String, item: T) -> bool {
collection.push_str(item.as_ref());
true
}
fn finalize(_: sealed::Internal, collection: &mut String) -> String {
mem::replace(collection, String::new())
}
}
impl<T> FromStream<T> for Vec<T> {}
impl<T> sealed::FromStreamPriv<T> for Vec<T> {
type InternalCollection = Vec<T>;
fn initialize(_: sealed::Internal, lower: usize, _upper: Option<usize>) -> Vec<T> {
Vec::with_capacity(lower)
}
fn extend(_: sealed::Internal, collection: &mut Vec<T>, item: T) -> bool {
collection.push(item);
true
}
fn finalize(_: sealed::Internal, collection: &mut Vec<T>) -> Vec<T> {
mem::replace(collection, vec![])
}
}
impl<T> FromStream<T> for Box<[T]> {}
impl<T> sealed::FromStreamPriv<T> for Box<[T]> {
type InternalCollection = Vec<T>;
fn initialize(_: sealed::Internal, lower: usize, upper: Option<usize>) -> Vec<T> {
<Vec<T> as sealed::FromStreamPriv<T>>::initialize(sealed::Internal, lower, upper)
}
fn extend(_: sealed::Internal, collection: &mut Vec<T>, item: T) -> bool {
<Vec<T> as sealed::FromStreamPriv<T>>::extend(sealed::Internal, collection, item)
}
fn finalize(_: sealed::Internal, collection: &mut Vec<T>) -> Box<[T]> {
<Vec<T> as sealed::FromStreamPriv<T>>::finalize(sealed::Internal, collection)
.into_boxed_slice()
}
}
impl<T, U, E> FromStream<Result<T, E>> for Result<U, E> where U: FromStream<T> {}
impl<T, U, E> sealed::FromStreamPriv<Result<T, E>> for Result<U, E>
where
U: FromStream<T>,
{
type InternalCollection = Result<U::InternalCollection, E>;
fn initialize(
_: sealed::Internal,
lower: usize,
upper: Option<usize>,
) -> Result<U::InternalCollection, E> {
Ok(U::initialize(sealed::Internal, lower, upper))
}
fn extend(
_: sealed::Internal,
collection: &mut Self::InternalCollection,
item: Result<T, E>,
) -> bool {
assert!(collection.is_ok());
match item {
Ok(item) => {
let collection = collection.as_mut().ok().expect("invalid state");
U::extend(sealed::Internal, collection, item)
}
Err(err) => {
*collection = Err(err);
false
}
}
}
fn finalize(_: sealed::Internal, collection: &mut Self::InternalCollection) -> Result<U, E> {
if let Ok(collection) = collection.as_mut() {
Ok(U::finalize(sealed::Internal, collection))
} else {
let res = mem::replace(collection, Ok(U::initialize(sealed::Internal, 0, Some(0))));
if let Err(err) = res {
Err(err)
} else {
unreachable!();
}
}
}
}
pub(crate) mod sealed {
#[doc(hidden)]
pub trait FromStreamPriv<T> {
/// Intermediate type used during collection process
///
/// The name of this type is internal and cannot be relied upon.
type InternalCollection;
/// Initialize the collection
fn initialize(
internal: Internal,
lower: usize,
upper: Option<usize>,
) -> Self::InternalCollection;
/// Extend the collection with the received item
///
/// Return `true` to continue streaming, `false` complete collection.
fn extend(internal: Internal, collection: &mut Self::InternalCollection, item: T) -> bool;
/// Finalize collection into target type.
fn finalize(internal: Internal, collection: &mut Self::InternalCollection) -> Self;
}
#[allow(missing_debug_implementations)]
pub struct Internal;
}
tokio-stream/src/empty.rs
+50
View File
@@ -0,0 +1,50 @@
use crate::Stream;
use core::marker::PhantomData;
use core::pin::Pin;
use core::task::{Context, Poll};
/// Stream for the [`empty`](fn@empty) function.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Empty<T>(PhantomData<T>);
impl<T> Unpin for Empty<T> {}
unsafe impl<T> Send for Empty<T> {}
unsafe impl<T> Sync for Empty<T> {}
/// Creates a stream that yields nothing.
///
/// The returned stream is immediately ready and returns `None`. Use
/// [`stream::pending()`](super::pending()) to obtain a stream that is never
/// ready.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// use tokio_stream::{self as stream, StreamExt};
///
/// #[tokio::main]
/// async fn main() {
/// let mut none = stream::empty::<i32>();
///
/// assert_eq!(None, none.next().await);
/// }
/// ```
pub const fn empty<T>() -> Empty<T> {
Empty(PhantomData)
}
impl<T> Stream for Empty<T> {
type Item = T;
fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<T>> {
Poll::Ready(None)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, Some(0))
}
}
tokio-stream/src/filter.rs
+58
View File
@@ -0,0 +1,58 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream returned by the [`filter`](super::StreamExt::filter) method.
#[must_use = "streams do nothing unless polled"]
pub struct Filter<St, F> {
#[pin]
stream: St,
f: F,
}
}
impl<St, F> fmt::Debug for Filter<St, F>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Filter")
.field("stream", &self.stream)
.finish()
}
}
impl<St, F> Filter<St, F> {
pub(super) fn new(stream: St, f: F) -> Self {
Self { stream, f }
}
}
impl<St, F> Stream for Filter<St, F>
where
St: Stream,
F: FnMut(&St::Item) -> bool,
{
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<St::Item>> {
loop {
match ready!(self.as_mut().project().stream.poll_next(cx)) {
Some(e) => {
if (self.as_mut().project().f)(&e) {
return Poll::Ready(Some(e));
}
}
None => return Poll::Ready(None),
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, self.stream.size_hint().1) // can't know a lower bound, due to the predicate
}
}
tokio-stream/src/filter_map.rs
+58
View File
@@ -0,0 +1,58 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream returned by the [`filter_map`](super::StreamExt::filter_map) method.
#[must_use = "streams do nothing unless polled"]
pub struct FilterMap<St, F> {
#[pin]
stream: St,
f: F,
}
}
impl<St, F> fmt::Debug for FilterMap<St, F>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FilterMap")
.field("stream", &self.stream)
.finish()
}
}
impl<St, F> FilterMap<St, F> {
pub(super) fn new(stream: St, f: F) -> Self {
Self { stream, f }
}
}
impl<St, F, T> Stream for FilterMap<St, F>
where
St: Stream,
F: FnMut(St::Item) -> Option<T>,
{
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
loop {
match ready!(self.as_mut().project().stream.poll_next(cx)) {
Some(e) => {
if let Some(e) = (self.as_mut().project().f)(e) {
return Poll::Ready(Some(e));
}
}
None => return Poll::Ready(None),
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, self.stream.size_hint().1) // can't know a lower bound, due to the predicate
}
}
tokio-stream/src/fold.rs
+57
View File
@@ -0,0 +1,57 @@
use crate::Stream;
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Future returned by the [`fold`](super::StreamExt::fold) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct FoldFuture<St, B, F> {
#[pin]
stream: St,
acc: Option<B>,
f: F,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
impl<St, B, F> FoldFuture<St, B, F> {
pub(super) fn new(stream: St, init: B, f: F) -> Self {
Self {
stream,
acc: Some(init),
f,
_pin: PhantomPinned,
}
}
}
impl<St, B, F> Future for FoldFuture<St, B, F>
where
St: Stream,
F: FnMut(B, St::Item) -> B,
{
type Output = B;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut me = self.project();
loop {
let next = ready!(me.stream.as_mut().poll_next(cx));
match next {
Some(v) => {
let old = me.acc.take().unwrap();
let new = (me.f)(old, v);
*me.acc = Some(new);
}
None => return Poll::Ready(me.acc.take().unwrap()),
}
}
}
}
tokio-stream/src/fuse.rs
+53
View File
@@ -0,0 +1,53 @@
use crate::Stream;
use pin_project_lite::pin_project;
use std::pin::Pin;
use std::task::{Context, Poll};
pin_project! {
/// Stream returned by [`fuse()`][super::StreamExt::fuse].
#[derive(Debug)]
pub struct Fuse<T> {
#[pin]
stream: Option<T>,
}
}
impl<T> Fuse<T>
where
T: Stream,
{
pub(crate) fn new(stream: T) -> Fuse<T> {
Fuse {
stream: Some(stream),
}
}
}
impl<T> Stream for Fuse<T>
where
T: Stream,
{
type Item = T::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
let res = match Option::as_pin_mut(self.as_mut().project().stream) {
Some(stream) => ready!(stream.poll_next(cx)),
None => return Poll::Ready(None),
};
if res.is_none() {
// Do not poll the stream anymore
self.as_mut().project().stream.set(None);
}
Poll::Ready(res)
}
fn size_hint(&self) -> (usize, Option<usize>) {
match self.stream {
Some(ref stream) => stream.size_hint(),
None => (0, Some(0)),
}
}
}
tokio-stream/src/iter.rs
+67
View File
@@ -0,0 +1,67 @@
use crate::Stream;
use core::pin::Pin;
use core::task::{Context, Poll};
/// Stream for the [`iter`](fn@iter) function.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Iter<I> {
iter: I,
yield_amt: usize,
}
impl<I> Unpin for Iter<I> {}
/// Converts an `Iterator` into a `Stream` which is always ready
/// to yield the next value.
///
/// Iterators in Rust don't express the ability to block, so this adapter
/// simply always calls `iter.next()` and returns that.
///
/// ```
/// # async fn dox() {
/// use tokio_stream::{self as stream, StreamExt};
///
/// let mut stream = stream::iter(vec![17, 19]);
///
/// assert_eq!(stream.next().await, Some(17));
/// assert_eq!(stream.next().await, Some(19));
/// assert_eq!(stream.next().await, None);
/// # }
/// ```
pub fn iter<I>(i: I) -> Iter<I::IntoIter>
where
I: IntoIterator,
{
Iter {
iter: i.into_iter(),
yield_amt: 0,
}
}
impl<I> Stream for Iter<I>
where
I: Iterator,
{
type Item = I::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<I::Item>> {
// TODO: add coop back
if self.yield_amt >= 32 {
self.yield_amt = 0;
cx.waker().wake_by_ref();
Poll::Pending
} else {
self.yield_amt += 1;
Poll::Ready(self.iter.next())
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}
tokio-stream/src/lib.rs
+1015
View File
File diff suppressed because it is too large Load Diff
tokio-stream/src/macros.rs
+18
View File
@@ -0,0 +1,18 @@
macro_rules! cfg_time {
($($item:item)*) => {
$(
#[cfg(feature = "time")]
#[cfg_attr(docsrs, doc(cfg(feature = "time")))]
$item
)*
}
}
macro_rules! ready {
($e:expr $(,)?) => {
match $e {
std::task::Poll::Ready(t) => t,
std::task::Poll::Pending => return std::task::Poll::Pending,
}
};
}
tokio-stream/src/map.rs
+51
View File
@@ -0,0 +1,51 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream for the [`map`](super::StreamExt::map) method.
#[must_use = "streams do nothing unless polled"]
pub struct Map<St, F> {
#[pin]
stream: St,
f: F,
}
}
impl<St, F> fmt::Debug for Map<St, F>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Map").field("stream", &self.stream).finish()
}
}
impl<St, F> Map<St, F> {
pub(super) fn new(stream: St, f: F) -> Self {
Map { stream, f }
}
}
impl<St, F, T> Stream for Map<St, F>
where
St: Stream,
F: FnMut(St::Item) -> T,
{
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
self.as_mut()
.project()
.stream
.poll_next(cx)
.map(|opt| opt.map(|x| (self.as_mut().project().f)(x)))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
tokio-stream/src/merge.rs
+89
View File
@@ -0,0 +1,89 @@
use crate::{Fuse, Stream};
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream returned by the [`merge`](super::StreamExt::merge) method.
pub struct Merge<T, U> {
#[pin]
a: Fuse<T>,
#[pin]
b: Fuse<U>,
// When `true`, poll `a` first, otherwise, `poll` b`.
a_first: bool,
}
}
impl<T, U> Merge<T, U> {
pub(super) fn new(a: T, b: U) -> Merge<T, U>
where
T: Stream,
U: Stream,
{
Merge {
a: Fuse::new(a),
b: Fuse::new(b),
a_first: true,
}
}
}
impl<T, U> Stream for Merge<T, U>
where
T: Stream,
U: Stream<Item = T::Item>,
{
type Item = T::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T::Item>> {
let me = self.project();
let a_first = *me.a_first;
// Toggle the flag
*me.a_first = !a_first;
if a_first {
poll_next(me.a, me.b, cx)
} else {
poll_next(me.b, me.a, cx)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
super::merge_size_hints(self.a.size_hint(), self.b.size_hint())
}
}
fn poll_next<T, U>(
first: Pin<&mut T>,
second: Pin<&mut U>,
cx: &mut Context<'_>,
) -> Poll<Option<T::Item>>
where
T: Stream,
U: Stream<Item = T::Item>,
{
use Poll::*;
let mut done = true;
match first.poll_next(cx) {
Ready(Some(val)) => return Ready(Some(val)),
Ready(None) => {}
Pending => done = false,
}
match second.poll_next(cx) {
Ready(Some(val)) => return Ready(Some(val)),
Ready(None) => {}
Pending => done = false,
}
if done {
Ready(None)
} else {
Pending
}
}
tokio-stream/src/next.rs
+37
View File
@@ -0,0 +1,37 @@
use crate::Stream;
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Future for the [`next`](super::StreamExt::next) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Next<'a, St: ?Sized> {
stream: &'a mut St,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
impl<'a, St: ?Sized> Next<'a, St> {
pub(super) fn new(stream: &'a mut St) -> Self {
Next {
stream,
_pin: PhantomPinned,
}
}
}
impl<St: ?Sized + Stream + Unpin> Future for Next<'_, St> {
type Output = Option<St::Item>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.project();
Pin::new(me.stream).poll_next(cx)
}
}
tokio-stream/src/once.rs
+52
View File
@@ -0,0 +1,52 @@
use crate::{Iter, Stream};
use core::option;
use core::pin::Pin;
use core::task::{Context, Poll};
/// Stream for the [`once`](fn@once) function.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Once<T> {
iter: Iter<option::IntoIter<T>>,
}
impl<I> Unpin for Once<I> {}
/// Creates a stream that emits an element exactly once.
///
/// The returned stream is immediately ready and emits the provided value once.
///
/// # Examples
///
/// ```
/// use tokio_stream::{self as stream, StreamExt};
///
/// #[tokio::main]
/// async fn main() {
/// // one is the loneliest number
/// let mut one = stream::once(1);
///
/// assert_eq!(Some(1), one.next().await);
///
/// // just one, that's all we get
/// assert_eq!(None, one.next().await);
/// }
/// ```
pub fn once<T>(value: T) -> Once<T> {
Once {
iter: crate::iter(Some(value).into_iter()),
}
}
impl<T> Stream for Once<T> {
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
Pin::new(&mut self.iter).poll_next(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}
tokio-stream/src/pending.rs
+54
View File
@@ -0,0 +1,54 @@
use crate::Stream;
use core::marker::PhantomData;
use core::pin::Pin;
use core::task::{Context, Poll};
/// Stream for the [`pending`](fn@pending) function.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Pending<T>(PhantomData<T>);
impl<T> Unpin for Pending<T> {}
unsafe impl<T> Send for Pending<T> {}
unsafe impl<T> Sync for Pending<T> {}
/// Creates a stream that is never ready
///
/// The returned stream is never ready. Attempting to call
/// [`next()`](crate::StreamExt::next) will never complete. Use
/// [`stream::empty()`](super::empty()) to obtain a stream that is is
/// immediately empty but returns no values.
///
/// # Examples
///
/// Basic usage:
///
/// ```no_run
/// use tokio_stream::{self as stream, StreamExt};
///
/// #[tokio::main]
/// async fn main() {
/// let mut never = stream::pending::<i32>();
///
/// // This will never complete
/// never.next().await;
///
/// unreachable!();
/// }
/// ```
pub const fn pending<T>() -> Pending<T> {
Pending(PhantomData)
}
impl<T> Stream for Pending<T> {
type Item = T;
fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<T>> {
Poll::Pending
}
fn size_hint(&self) -> (usize, Option<usize>) {
(0, None)
}
}
tokio-stream/src/skip.rs
+63
View File
@@ -0,0 +1,63 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream for the [`skip`](super::StreamExt::skip) method.
#[must_use = "streams do nothing unless polled"]
pub struct Skip<St> {
#[pin]
stream: St,
remaining: usize,
}
}
impl<St> fmt::Debug for Skip<St>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Skip")
.field("stream", &self.stream)
.finish()
}
}
impl<St> Skip<St> {
pub(super) fn new(stream: St, remaining: usize) -> Self {
Self { stream, remaining }
}
}
impl<St> Stream for Skip<St>
where
St: Stream,
{
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match ready!(self.as_mut().project().stream.poll_next(cx)) {
Some(e) => {
if self.remaining == 0 {
return Poll::Ready(Some(e));
}
*self.as_mut().project().remaining -= 1;
}
None => return Poll::Ready(None),
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let (lower, upper) = self.stream.size_hint();
let lower = lower.saturating_sub(self.remaining);
let upper = upper.map(|x| x.saturating_sub(self.remaining));
(lower, upper)
}
}
tokio-stream/src/skip_while.rs
+73
View File
@@ -0,0 +1,73 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream for the [`skip_while`](super::StreamExt::skip_while) method.
#[must_use = "streams do nothing unless polled"]
pub struct SkipWhile<St, F> {
#[pin]
stream: St,
predicate: Option<F>,
}
}
impl<St, F> fmt::Debug for SkipWhile<St, F>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SkipWhile")
.field("stream", &self.stream)
.finish()
}
}
impl<St, F> SkipWhile<St, F> {
pub(super) fn new(stream: St, predicate: F) -> Self {
Self {
stream,
predicate: Some(predicate),
}
}
}
impl<St, F> Stream for SkipWhile<St, F>
where
St: Stream,
F: FnMut(&St::Item) -> bool,
{
type Item = St::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.project();
if let Some(predicate) = this.predicate {
loop {
match ready!(this.stream.as_mut().poll_next(cx)) {
Some(item) => {
if !(predicate)(&item) {
*this.predicate = None;
return Poll::Ready(Some(item));
}
}
None => return Poll::Ready(None),
}
}
} else {
this.stream.poll_next(cx)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let (lower, upper) = self.stream.size_hint();
if self.predicate.is_some() {
return (0, upper);
}
(lower, upper)
}
}
tokio-stream/src/stream_map.rs
+664
View File
@@ -0,0 +1,664 @@
use crate::Stream;
use std::borrow::Borrow;
use std::hash::Hash;
use std::pin::Pin;
use std::task::{Context, Poll};
/// Combine many streams into one, indexing each source stream with a unique
/// key.
///
/// `StreamMap` is similar to [`StreamExt::merge`] in that it combines source
/// streams into a single merged stream that yields values in the order that
/// they arrive from the source streams. However, `StreamMap` has a lot more
/// flexibility in usage patterns.
///
/// `StreamMap` can:
///
/// * Merge an arbitrary number of streams.
/// * Track which source stream the value was received from.
/// * Handle inserting and removing streams from the set of managed streams at
/// any point during iteration.
///
/// All source streams held by `StreamMap` are indexed using a key. This key is
/// included with the value when a source stream yields a value. The key is also
/// used to remove the stream from the `StreamMap` before the stream has
/// completed streaming.
///
/// # `Unpin`
///
/// Because the `StreamMap` API moves streams during runtime, both streams and
/// keys must be `Unpin`. In order to insert a `!Unpin` stream into a
/// `StreamMap`, use [`pin!`] to pin the stream to the stack or [`Box::pin`] to
/// pin the stream in the heap.
///
/// # Implementation
///
/// `StreamMap` is backed by a `Vec<(K, V)>`. There is no guarantee that this
/// internal implementation detail will persist in future versions, but it is
/// important to know the runtime implications. In general, `StreamMap` works
/// best with a "smallish" number of streams as all entries are scanned on
/// insert, remove, and polling. In cases where a large number of streams need
/// to be merged, it may be advisable to use tasks sending values on a shared
/// [`mpsc`] channel.
///
/// [`StreamExt::merge`]: crate::StreamExt::merge
/// [`mpsc`]: https://docs.rs/tokio/1.0/tokio/sync/mpsc/index.html
/// [`pin!`]: https://docs.rs/tokio/1.0/tokio/macro.pin.html
/// [`Box::pin`]: std::boxed::Box::pin
///
/// # Examples
///
/// Merging two streams, then remove them after receiving the first value
///
/// ```
/// use tokio_stream::{StreamExt, StreamMap, Stream};
/// use tokio::sync::mpsc;
/// use std::pin::Pin;
///
/// #[tokio::main]
/// async fn main() {
/// let (tx1, mut rx1) = mpsc::channel::<usize>(10);
/// let (tx2, mut rx2) = mpsc::channel::<usize>(10);
///
/// // Convert the channels to a `Stream`.
/// let rx1 = Box::pin(async_stream::stream! {
/// while let Some(item) = rx1.recv().await {
/// yield item;
/// }
/// }) as Pin<Box<dyn Stream<Item = usize> + Send>>;
///
/// let rx2 = Box::pin(async_stream::stream! {
/// while let Some(item) = rx2.recv().await {
/// yield item;
/// }
/// }) as Pin<Box<dyn Stream<Item = usize> + Send>>;
///
/// tokio::spawn(async move {
/// tx1.send(1).await.unwrap();
///
/// // This value will never be received. The send may or may not return
/// // `Err` depending on if the remote end closed first or not.
/// let _ = tx1.send(2).await;
/// });
///
/// tokio::spawn(async move {
/// tx2.send(3).await.unwrap();
/// let _ = tx2.send(4).await;
/// });
///
/// let mut map = StreamMap::new();
///
/// // Insert both streams
/// map.insert("one", rx1);
/// map.insert("two", rx2);
///
/// // Read twice
/// for _ in 0..2 {
/// let (key, val) = map.next().await.unwrap();
///
/// if key == "one" {
/// assert_eq!(val, 1);
/// } else {
/// assert_eq!(val, 3);
/// }
///
/// // Remove the stream to prevent reading the next value
/// map.remove(key);
/// }
/// }
/// ```
///
/// This example models a read-only client to a chat system with channels. The
/// client sends commands to join and leave channels. `StreamMap` is used to
/// manage active channel subscriptions.
///
/// For simplicity, messages are displayed with `println!`, but they could be
/// sent to the client over a socket.
///
/// ```no_run
/// use tokio_stream::{Stream, StreamExt, StreamMap};
///
/// enum Command {
/// Join(String),
/// Leave(String),
/// }
///
/// fn commands() -> impl Stream<Item = Command> {
/// // Streams in user commands by parsing `stdin`.
/// # tokio_stream::pending()
/// }
///
/// // Join a channel, returns a stream of messages received on the channel.
/// fn join(channel: &str) -> impl Stream<Item = String> + Unpin {
/// // left as an exercise to the reader
/// # tokio_stream::pending()
/// }
///
/// #[tokio::main]
/// async fn main() {
/// let mut channels = StreamMap::new();
///
/// // Input commands (join / leave channels).
/// let cmds = commands();
/// tokio::pin!(cmds);
///
/// loop {
/// tokio::select! {
/// Some(cmd) = cmds.next() => {
/// match cmd {
/// Command::Join(chan) => {
/// // Join the channel and add it to the `channels`
/// // stream map
/// let msgs = join(&chan);
/// channels.insert(chan, msgs);
/// }
/// Command::Leave(chan) => {
/// channels.remove(&chan);
/// }
/// }
/// }
/// Some((chan, msg)) = channels.next() => {
/// // Received a message, display it on stdout with the channel
/// // it originated from.
/// println!("{}: {}", chan, msg);
/// }
/// // Both the `commands` stream and the `channels` stream are
/// // complete. There is no more work to do, so leave the loop.
/// else => break,
/// }
/// }
/// }
/// ```
#[derive(Debug)]
pub struct StreamMap<K, V> {
/// Streams stored in the map
entries: Vec<(K, V)>,
}
impl<K, V> StreamMap<K, V> {
/// An iterator visiting all key-value pairs in arbitrary order.
///
/// The iterator element type is &'a (K, V).
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// map.insert("a", pending::<i32>());
/// map.insert("b", pending());
/// map.insert("c", pending());
///
/// for (key, stream) in map.iter() {
/// println!("({}, {:?})", key, stream);
/// }
/// ```
pub fn iter(&self) -> impl Iterator<Item = &(K, V)> {
self.entries.iter()
}
/// An iterator visiting all key-value pairs mutably in arbitrary order.
///
/// The iterator element type is &'a mut (K, V).
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// map.insert("a", pending::<i32>());
/// map.insert("b", pending());
/// map.insert("c", pending());
///
/// for (key, stream) in map.iter_mut() {
/// println!("({}, {:?})", key, stream);
/// }
/// ```
pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut (K, V)> {
self.entries.iter_mut()
}
/// Creates an empty `StreamMap`.
///
/// The stream map is initially created with a capacity of `0`, so it will
/// not allocate until it is first inserted into.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, Pending};
///
/// let map: StreamMap<&str, Pending<()>> = StreamMap::new();
/// ```
pub fn new() -> StreamMap<K, V> {
StreamMap { entries: vec![] }
}
/// Creates an empty `StreamMap` with the specified capacity.
///
/// The stream map will be able to hold at least `capacity` elements without
/// reallocating. If `capacity` is 0, the stream map will not allocate.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, Pending};
///
/// let map: StreamMap<&str, Pending<()>> = StreamMap::with_capacity(10);
/// ```
pub fn with_capacity(capacity: usize) -> StreamMap<K, V> {
StreamMap {
entries: Vec::with_capacity(capacity),
}
}
/// Returns an iterator visiting all keys in arbitrary order.
///
/// The iterator element type is &'a K.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// map.insert("a", pending::<i32>());
/// map.insert("b", pending());
/// map.insert("c", pending());
///
/// for key in map.keys() {
/// println!("{}", key);
/// }
/// ```
pub fn keys(&self) -> impl Iterator<Item = &K> {
self.iter().map(|(k, _)| k)
}
/// An iterator visiting all values in arbitrary order.
///
/// The iterator element type is &'a V.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// map.insert("a", pending::<i32>());
/// map.insert("b", pending());
/// map.insert("c", pending());
///
/// for stream in map.values() {
/// println!("{:?}", stream);
/// }
/// ```
pub fn values(&self) -> impl Iterator<Item = &V> {
self.iter().map(|(_, v)| v)
}
/// An iterator visiting all values mutably in arbitrary order.
///
/// The iterator element type is &'a mut V.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// map.insert("a", pending::<i32>());
/// map.insert("b", pending());
/// map.insert("c", pending());
///
/// for stream in map.values_mut() {
/// println!("{:?}", stream);
/// }
/// ```
pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
self.iter_mut().map(|(_, v)| v)
}
/// Returns the number of streams the map can hold without reallocating.
///
/// This number is a lower bound; the `StreamMap` might be able to hold
/// more, but is guaranteed to be able to hold at least this many.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, Pending};
///
/// let map: StreamMap<i32, Pending<()>> = StreamMap::with_capacity(100);
/// assert!(map.capacity() >= 100);
/// ```
pub fn capacity(&self) -> usize {
self.entries.capacity()
}
/// Returns the number of streams in the map.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut a = StreamMap::new();
/// assert_eq!(a.len(), 0);
/// a.insert(1, pending::<i32>());
/// assert_eq!(a.len(), 1);
/// ```
pub fn len(&self) -> usize {
self.entries.len()
}
/// Returns `true` if the map contains no elements.
///
/// # Examples
///
/// ```
/// use std::collections::HashMap;
///
/// let mut a = HashMap::new();
/// assert!(a.is_empty());
/// a.insert(1, "a");
/// assert!(!a.is_empty());
/// ```
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
/// Clears the map, removing all key-stream pairs. Keeps the allocated
/// memory for reuse.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut a = StreamMap::new();
/// a.insert(1, pending::<i32>());
/// a.clear();
/// assert!(a.is_empty());
/// ```
pub fn clear(&mut self) {
self.entries.clear();
}
/// Insert a key-stream pair into the map.
///
/// If the map did not have this key present, `None` is returned.
///
/// If the map did have this key present, the new `stream` replaces the old
/// one and the old stream is returned.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
///
/// assert!(map.insert(37, pending::<i32>()).is_none());
/// assert!(!map.is_empty());
///
/// map.insert(37, pending());
/// assert!(map.insert(37, pending()).is_some());
/// ```
pub fn insert(&mut self, k: K, stream: V) -> Option<V>
where
K: Hash + Eq,
{
let ret = self.remove(&k);
self.entries.push((k, stream));
ret
}
/// Removes a key from the map, returning the stream at the key if the key was previously in the map.
///
/// The key may be any borrowed form of the map's key type, but `Hash` and
/// `Eq` on the borrowed form must match those for the key type.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
/// map.insert(1, pending::<i32>());
/// assert!(map.remove(&1).is_some());
/// assert!(map.remove(&1).is_none());
/// ```
pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: Hash + Eq,
{
for i in 0..self.entries.len() {
if self.entries[i].0.borrow() == k {
return Some(self.entries.swap_remove(i).1);
}
}
None
}
/// Returns `true` if the map contains a stream for the specified key.
///
/// The key may be any borrowed form of the map's key type, but `Hash` and
/// `Eq` on the borrowed form must match those for the key type.
///
/// # Examples
///
/// ```
/// use tokio_stream::{StreamMap, pending};
///
/// let mut map = StreamMap::new();
/// map.insert(1, pending::<i32>());
/// assert_eq!(map.contains_key(&1), true);
/// assert_eq!(map.contains_key(&2), false);
/// ```
pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
where
K: Borrow<Q>,
Q: Hash + Eq,
{
for i in 0..self.entries.len() {
if self.entries[i].0.borrow() == k {
return true;
}
}
false
}
}
impl<K, V> StreamMap<K, V>
where
K: Unpin,
V: Stream + Unpin,
{
/// Polls the next value, includes the vec entry index
fn poll_next_entry(&mut self, cx: &mut Context<'_>) -> Poll<Option<(usize, V::Item)>> {
use Poll::*;
let start = self::rand::thread_rng_n(self.entries.len() as u32) as usize;
let mut idx = start;
for _ in 0..self.entries.len() {
let (_, stream) = &mut self.entries[idx];
match Pin::new(stream).poll_next(cx) {
Ready(Some(val)) => return Ready(Some((idx, val))),
Ready(None) => {
// Remove the entry
self.entries.swap_remove(idx);
// Check if this was the last entry, if so the cursor needs
// to wrap
if idx == self.entries.len() {
idx = 0;
} else if idx < start && start <= self.entries.len() {
// The stream being swapped into the current index has
// already been polled, so skip it.
idx = idx.wrapping_add(1) % self.entries.len();
}
}
Pending => {
idx = idx.wrapping_add(1) % self.entries.len();
}
}
}
// If the map is empty, then the stream is complete.
if self.entries.is_empty() {
Ready(None)
} else {
Pending
}
}
}
impl<K, V> Default for StreamMap<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K, V> Stream for StreamMap<K, V>
where
K: Clone + Unpin,
V: Stream + Unpin,
{
type Item = (K, V::Item);
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if let Some((idx, val)) = ready!(self.poll_next_entry(cx)) {
let key = self.entries[idx].0.clone();
Poll::Ready(Some((key, val)))
} else {
Poll::Ready(None)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let mut ret = (0, Some(0));
for (_, stream) in &self.entries {
let hint = stream.size_hint();
ret.0 += hint.0;
match (ret.1, hint.1) {
(Some(a), Some(b)) => ret.1 = Some(a + b),
(Some(_), None) => ret.1 = None,
_ => {}
}
}
ret
}
}
mod rand {
use std::cell::Cell;
mod loom {
#[cfg(not(loom))]
pub(crate) mod rand {
use std::collections::hash_map::RandomState;
use std::hash::{BuildHasher, Hash, Hasher};
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::Relaxed;
static COUNTER: AtomicU32 = AtomicU32::new(1);
pub(crate) fn seed() -> u64 {
let rand_state = RandomState::new();
let mut hasher = rand_state.build_hasher();
// Hash some unique-ish data to generate some new state
COUNTER.fetch_add(1, Relaxed).hash(&mut hasher);
// Get the seed
hasher.finish()
}
}
#[cfg(loom)]
pub(crate) mod rand {
pub(crate) fn seed() -> u64 {
1
}
}
}
/// Fast random number generate
///
/// Implement xorshift64+: 2 32-bit xorshift sequences added together.
/// Shift triplet [17,7,16] was calculated as indicated in Marsaglia's
/// Xorshift paper: https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf
/// This generator passes the SmallCrush suite, part of TestU01 framework:
/// http://simul.iro.umontreal.ca/testu01/tu01.html
#[derive(Debug)]
pub(crate) struct FastRand {
one: Cell<u32>,
two: Cell<u32>,
}
impl FastRand {
/// Initialize a new, thread-local, fast random number generator.
pub(crate) fn new(seed: u64) -> FastRand {
let one = (seed >> 32) as u32;
let mut two = seed as u32;
if two == 0 {
// This value cannot be zero
two = 1;
}
FastRand {
one: Cell::new(one),
two: Cell::new(two),
}
}
pub(crate) fn fastrand_n(&self, n: u32) -> u32 {
// This is similar to fastrand() % n, but faster.
// See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
let mul = (self.fastrand() as u64).wrapping_mul(n as u64);
(mul >> 32) as u32
}
fn fastrand(&self) -> u32 {
let mut s1 = self.one.get();
let s0 = self.two.get();
s1 ^= s1 << 17;
s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16;
self.one.set(s0);
self.two.set(s1);
s0.wrapping_add(s1)
}
}
// Used by `StreamMap`
pub(crate) fn thread_rng_n(n: u32) -> u32 {
thread_local! {
static THREAD_RNG: FastRand = FastRand::new(loom::rand::seed());
}
THREAD_RNG.with(|rng| rng.fastrand_n(n))
}
}
tokio-stream/src/take.rs
+76
View File
@@ -0,0 +1,76 @@
use crate::Stream;
use core::cmp;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream for the [`take`](super::StreamExt::take) method.
#[must_use = "streams do nothing unless polled"]
pub struct Take<St> {
#[pin]
stream: St,
remaining: usize,
}
}
impl<St> fmt::Debug for Take<St>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Take")
.field("stream", &self.stream)
.finish()
}
}
impl<St> Take<St> {
pub(super) fn new(stream: St, remaining: usize) -> Self {
Self { stream, remaining }
}
}
impl<St> Stream for Take<St>
where
St: Stream,
{
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if *self.as_mut().project().remaining > 0 {
self.as_mut().project().stream.poll_next(cx).map(|ready| {
match &ready {
Some(_) => {
*self.as_mut().project().remaining -= 1;
}
None => {
*self.as_mut().project().remaining = 0;
}
}
ready
})
} else {
Poll::Ready(None)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
if self.remaining == 0 {
return (0, Some(0));
}
let (lower, upper) = self.stream.size_hint();
let lower = cmp::min(lower, self.remaining as usize);
let upper = match upper {
Some(x) if x < self.remaining as usize => Some(x),
_ => Some(self.remaining as usize),
};
(lower, upper)
}
}
tokio-stream/src/take_while.rs
+79
View File
@@ -0,0 +1,79 @@
use crate::Stream;
use core::fmt;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Stream for the [`take_while`](super::StreamExt::take_while) method.
#[must_use = "streams do nothing unless polled"]
pub struct TakeWhile<St, F> {
#[pin]
stream: St,
predicate: F,
done: bool,
}
}
impl<St, F> fmt::Debug for TakeWhile<St, F>
where
St: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("TakeWhile")
.field("stream", &self.stream)
.field("done", &self.done)
.finish()
}
}
impl<St, F> TakeWhile<St, F> {
pub(super) fn new(stream: St, predicate: F) -> Self {
Self {
stream,
predicate,
done: false,
}
}
}
impl<St, F> Stream for TakeWhile<St, F>
where
St: Stream,
F: FnMut(&St::Item) -> bool,
{
type Item = St::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if !*self.as_mut().project().done {
self.as_mut().project().stream.poll_next(cx).map(|ready| {
let ready = ready.and_then(|item| {
if !(self.as_mut().project().predicate)(&item) {
None
} else {
Some(item)
}
});
if ready.is_none() {
*self.as_mut().project().done = true;
}
ready
})
} else {
Poll::Ready(None)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
if self.done {
return (0, Some(0));
}
let (_, upper) = self.stream.size_hint();
(0, upper)
}
}
tokio-stream/src/throttle.rs
+97
View File
@@ -0,0 +1,97 @@
//! Slow down a stream by enforcing a delay between items.
use crate::Stream;
use tokio::time::{Duration, Instant, Sleep};
use std::future::Future;
use std::marker::Unpin;
use std::pin::Pin;
use std::task::{self, Poll};
use pin_project_lite::pin_project;
pub(super) fn throttle<T>(duration: Duration, stream: T) -> Throttle<T>
where
T: Stream,
{
let delay = if duration == Duration::from_millis(0) {
None
} else {
Some(tokio::time::sleep_until(Instant::now() + duration))
};
Throttle {
delay,
duration,
has_delayed: true,
stream,
}
}
pin_project! {
/// Stream for the [`throttle`](throttle) function.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Throttle<T> {
// `None` when duration is zero.
delay: Option<Sleep>,
duration: Duration,
// Set to true when `delay` has returned ready, but `stream` hasn't.
has_delayed: bool,
// The stream to throttle
#[pin]
stream: T,
}
}
// XXX: are these safe if `T: !Unpin`?
impl<T: Unpin> Throttle<T> {
/// Acquires a reference to the underlying stream that this combinator is
/// pulling from.
pub fn get_ref(&self) -> &T {
&self.stream
}
/// Acquires a mutable reference to the underlying stream that this combinator
/// is pulling from.
///
/// Note that care must be taken to avoid tampering with the state of the stream
/// which may otherwise confuse this combinator.
pub fn get_mut(&mut self) -> &mut T {
&mut self.stream
}
/// Consumes this combinator, returning the underlying stream.
///
/// Note that this may discard intermediate state of this combinator, so care
/// should be taken to avoid losing resources when this is called.
pub fn into_inner(self) -> T {
self.stream
}
}
impl<T: Stream> Stream for Throttle<T> {
type Item = T::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
if !self.has_delayed && self.delay.is_some() {
ready!(Pin::new(self.as_mut().project().delay.as_mut().unwrap()).poll(cx));
*self.as_mut().project().has_delayed = true;
}
let value = ready!(self.as_mut().project().stream.poll_next(cx));
if value.is_some() {
let dur = self.duration;
if let Some(ref mut delay) = self.as_mut().project().delay {
delay.reset(Instant::now() + dur);
}
*self.as_mut().project().has_delayed = false;
}
Poll::Ready(value)
}
}
tokio-stream/src/timeout.rs
+92
View File
@@ -0,0 +1,92 @@
use crate::{Fuse, Stream};
use tokio::time::{Instant, Sleep};
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
use std::fmt;
use std::time::Duration;
pin_project! {
/// Stream returned by the [`timeout`](super::StreamExt::timeout) method.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct Timeout<S> {
#[pin]
stream: Fuse<S>,
deadline: Sleep,
duration: Duration,
poll_deadline: bool,
}
}
/// Error returned by `Timeout`.
#[derive(Debug, PartialEq)]
pub struct Elapsed(());
impl<S: Stream> Timeout<S> {
pub(super) fn new(stream: S, duration: Duration) -> Self {
let next = Instant::now() + duration;
let deadline = tokio::time::sleep_until(next);
Timeout {
stream: Fuse::new(stream),
deadline,
duration,
poll_deadline: true,
}
}
}
impl<S: Stream> Stream for Timeout<S> {
type Item = Result<S::Item, Elapsed>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self.as_mut().project().stream.poll_next(cx) {
Poll::Ready(v) => {
if v.is_some() {
let next = Instant::now() + self.duration;
self.as_mut().project().deadline.reset(next);
*self.as_mut().project().poll_deadline = true;
}
return Poll::Ready(v.map(Ok));
}
Poll::Pending => {}
};
if self.poll_deadline {
ready!(Pin::new(self.as_mut().project().deadline).poll(cx));
*self.as_mut().project().poll_deadline = false;
return Poll::Ready(Some(Err(Elapsed::new())));
}
Poll::Pending
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
// ===== impl Elapsed =====
impl Elapsed {
pub(crate) fn new() -> Self {
Elapsed(())
}
}
impl fmt::Display for Elapsed {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
"deadline has elapsed".fmt(fmt)
}
}
impl std::error::Error for Elapsed {}
impl From<Elapsed> for std::io::Error {
fn from(_err: Elapsed) -> std::io::Error {
std::io::ErrorKind::TimedOut.into()
}
}
tokio-stream/src/try_next.rs
+38
View File
@@ -0,0 +1,38 @@
use crate::{Next, Stream};
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use pin_project_lite::pin_project;
pin_project! {
/// Future for the [`try_next`](super::StreamExt::try_next) method.
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct TryNext<'a, St: ?Sized> {
#[pin]
inner: Next<'a, St>,
// Make this future `!Unpin` for compatibility with async trait methods.
#[pin]
_pin: PhantomPinned,
}
}
impl<'a, St: ?Sized> TryNext<'a, St> {
pub(super) fn new(stream: &'a mut St) -> Self {
Self {
inner: Next::new(stream),
_pin: PhantomPinned,
}
}
}
impl<T, E, St: ?Sized + Stream<Item = Result<T, E>> + Unpin> Future for TryNext<'_, St> {
type Output = Result<Option<T>, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.project();
me.inner.poll(cx).map(Option::transpose)
}
}