diff --git a/src/future_stopper.rs b/src/future_stopper.rs index 6601a03..4d0d09e 100644 --- a/src/future_stopper.rs +++ b/src/future_stopper.rs @@ -40,12 +40,12 @@ impl Future for FutureStopper { } match Pin::new(&mut *this.event_listener).poll(cx) { - Poll::Ready(()) => continue, + Poll::Ready(()) => { + *this.event_listener = this.stopper.0.event.listen(); + continue; + } Poll::Pending => { - return match this.future.poll(cx) { - Poll::Ready(output) => Poll::Ready(Some(output)), - Poll::Pending => Poll::Pending, - } + return this.future.poll(cx).map(Some); } }; } diff --git a/src/stopped.rs b/src/stopped.rs index 9e69967..15eb4eb 100644 --- a/src/stopped.rs +++ b/src/stopped.rs @@ -39,7 +39,10 @@ impl Future for Stopped { } match Pin::new(&mut *event_listener).poll(cx) { - Poll::Ready(()) => continue, + Poll::Ready(()) => { + *event_listener = stopper.0.event.listen(); + continue; + } Poll::Pending => return Poll::Pending, }; } diff --git a/src/stream_stopper.rs b/src/stream_stopper.rs index ffad82c..2020cb2 100644 --- a/src/stream_stopper.rs +++ b/src/stream_stopper.rs @@ -66,7 +66,10 @@ impl Stream for StreamStopper { } match Pin::new(&mut *this.event_listener).poll(cx) { - Poll::Ready(()) => continue, + Poll::Ready(()) => { + *this.event_listener = this.stopper.0.event.listen(); + continue; + } Poll::Pending => return this.stream.poll_next(cx), }; }