Skip to content

Commit

Permalink
fix-or-feat(iroh): Set MaybeFuture to None on Poll::Ready (#3090)
Browse files Browse the repository at this point in the history
## 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.
  • Loading branch information
flub authored Jan 3, 2025
1 parent d236e04 commit 6599ea6
Showing 1 changed file with 33 additions and 4 deletions.
37 changes: 33 additions & 4 deletions iroh/src/util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,11 +53,18 @@ impl<T> MaybeFuture<T> {
impl<T: Future> Future for MaybeFuture<T> {
type Output = T::Output;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
match this {
MaybeFutureProj::Some(t) => t.poll(cx),
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.as_mut().project();
let poll_res = match this {
MaybeFutureProj::Some(ref mut t) => t.as_mut().poll(cx),
MaybeFutureProj::None => Poll::Pending,
};
match poll_res {
Poll::Ready(val) => {
self.as_mut().project_replace(Self::None);
Poll::Ready(val)
}
Poll::Pending => Poll::Pending,
}
}
}
Expand All @@ -70,3 +77,25 @@ impl<T: Future> Future for MaybeFuture<T> {
pub(crate) fn relay_only_mode() -> bool {
std::option_env!("DEV_RELAY_ONLY").is_some()
}

#[cfg(test)]
mod tests {
use std::pin::pin;

use tokio::time::Duration;

use super::*;

#[tokio::test]
async fn test_maybefuture_poll_after_use() {
let fut = async move { "hello" };
let mut maybe_fut = pin!(MaybeFuture::Some(fut));
let res = (&mut maybe_fut).await;

assert_eq!(res, "hello");

// Now poll again
let res = tokio::time::timeout(Duration::from_millis(10), maybe_fut).await;
assert!(res.is_err());
}
}

0 comments on commit 6599ea6

Please sign in to comment.