Skip to content

Commit 6599ea6

Browse files
authored
fix-or-feat(iroh): Set MaybeFuture to None on Poll::Ready (#3090)
## Description This removes a footgun where after polling the future it is easy to forget to call MaybeFuture::set_none, which can result in a panic if the inner future is not fused. ## Breaking Changes <!-- Optional, if there are any breaking changes document them, including how to migrate older code. --> ## Notes & open questions <!-- Any notes, remarks or open questions you have to make about the PR. --> ## Change checklist - [x] Self-review. - [x] Documentation updates following the [style guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text), if relevant. - [x] Tests if relevant. - [x] All breaking changes documented.
1 parent d236e04 commit 6599ea6

File tree

1 file changed

+33
-4
lines changed

1 file changed

+33
-4
lines changed

iroh/src/util.rs

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,11 +53,18 @@ impl<T> MaybeFuture<T> {
5353
impl<T: Future> Future for MaybeFuture<T> {
5454
type Output = T::Output;
5555

56-
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
57-
let this = self.project();
58-
match this {
59-
MaybeFutureProj::Some(t) => t.poll(cx),
56+
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
57+
let mut this = self.as_mut().project();
58+
let poll_res = match this {
59+
MaybeFutureProj::Some(ref mut t) => t.as_mut().poll(cx),
6060
MaybeFutureProj::None => Poll::Pending,
61+
};
62+
match poll_res {
63+
Poll::Ready(val) => {
64+
self.as_mut().project_replace(Self::None);
65+
Poll::Ready(val)
66+
}
67+
Poll::Pending => Poll::Pending,
6168
}
6269
}
6370
}
@@ -70,3 +77,25 @@ impl<T: Future> Future for MaybeFuture<T> {
7077
pub(crate) fn relay_only_mode() -> bool {
7178
std::option_env!("DEV_RELAY_ONLY").is_some()
7279
}
80+
81+
#[cfg(test)]
82+
mod tests {
83+
use std::pin::pin;
84+
85+
use tokio::time::Duration;
86+
87+
use super::*;
88+
89+
#[tokio::test]
90+
async fn test_maybefuture_poll_after_use() {
91+
let fut = async move { "hello" };
92+
let mut maybe_fut = pin!(MaybeFuture::Some(fut));
93+
let res = (&mut maybe_fut).await;
94+
95+
assert_eq!(res, "hello");
96+
97+
// Now poll again
98+
let res = tokio::time::timeout(Duration::from_millis(10), maybe_fut).await;
99+
assert!(res.is_err());
100+
}
101+
}

0 commit comments

Comments
 (0)