2
0
Fork 1
mirror of https://github.com/async-rs/async-std.git synced 2025-01-29 16:55:34 +00:00

Fix double drop in StreamExt::cycle

This commit is contained in:
Taiki Endo 2020-11-02 07:10:18 +09:00
parent 11196c853d
commit e8dc2c0571

View file

@ -1,14 +1,19 @@
use core::mem::ManuallyDrop;
use core::pin::Pin;
use futures_core::ready;
use pin_project_lite::pin_project;
use crate::stream::Stream;
use crate::task::{Context, Poll};
/// A stream that will repeatedly yield the same list of elements.
#[derive(Debug)]
pub struct Cycle<S> {
orig: S,
source: ManuallyDrop<S>,
pin_project! {
/// A stream that will repeatedly yield the same list of elements.
#[derive(Debug)]
pub struct Cycle<S> {
orig: S,
#[pin]
source: S,
}
}
impl<S> Cycle<S>
@ -18,15 +23,7 @@ where
pub(crate) fn new(source: S) -> Self {
Self {
orig: source.clone(),
source: ManuallyDrop::new(source),
}
}
}
impl<S> Drop for Cycle<S> {
fn drop(&mut self) {
unsafe {
ManuallyDrop::drop(&mut self.source);
source,
}
}
}
@ -38,17 +35,14 @@ where
type Item = S::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
unsafe {
let this = self.get_unchecked_mut();
let mut this = self.project();
match futures_core::ready!(Pin::new_unchecked(&mut *this.source).poll_next(cx)) {
Some(item) => Poll::Ready(Some(item)),
None => {
ManuallyDrop::drop(&mut this.source);
this.source = ManuallyDrop::new(this.orig.clone());
Pin::new_unchecked(&mut *this.source).poll_next(cx)
}
match ready!(this.source.as_mut().poll_next(cx)) {
None => {
this.source.set(this.orig.clone());
this.source.poll_next(cx)
}
item => Poll::Ready(item),
}
}
}