Best way to timeout across multiple async calls?

This is not asking how to wrap multiple async calls made at the same time in a single timeout like this:

tokio::time::timeout(
    Duration::from_millis(10), 
    tokio::join!(call1, call2, ...)
).await

The context is I have an async fn consisting of a loop which makes a series of async calls and it can be in one of two states based on the return values of those calls: A and not-A.

async fn foo() {
    // boolean (true when in state A, false otherwise)
    let mut state_a = true;
    loop {
        state_a = call_1().await;  // for simplicity the results are all bools
        state_a = call_2().await;
        ...
        state_a = call_n().await;
    }
}

When in state A, everything is stable. But when in state not-A, there is a deadline (e.g. 5 seconds) to get back into state A, and if the time elapses then we must continue (return to start of loop).

What's the cleanest way to add timer logic to this code to achieve the above? All I can think of is setting a tokio::time::Timeout 5 seconds in the future which every call is wrapped in. But then I have to initiate it every time state not-A is entered and also check whether we're in state A or not at every call to decide whether to wrap or not, the code quickly gets surprisingly messy. Is there a cleaner way?

Here's one that uses regular control flow. The state is stored in timeout, with None representing A and Some representing not-A. (playground)

use core::future::Future;
use core::time::Duration;
use tokio::time::timeout_at;
use tokio::time::Instant;

pub async fn foo() {
    let mut timeout: Option<Instant> = None;
    loop {
        if !run_and_update_timeout(&mut timeout, call_1()).await {
            continue;
        }
        if !run_and_update_timeout(&mut timeout, call_2()).await {
            continue;
        }
        if !run_and_update_timeout(&mut timeout, call_n()).await {
            continue;
        }
    }
}

/// Returns true when the next future should run, returns false when the loop should be continued.
async fn run_and_update_timeout<F>(timeout: &mut Option<Instant>, future: F) -> bool
where
    F: Future<Output = bool>,
{
    if let &mut Some(t) = timeout {
        // In state !A, timeout active
        let call_or_timeout = timeout_at(t, future).await;
        match call_or_timeout {
            // Future completed first and succeeded, clear timeout
            Ok(true) => *timeout = None,
            // Future completed and failed, move on to next future
            Ok(false) => (),
            // Timeout finished first, continue the loop
            Err(_) => {
                *timeout = None;
                return false;
            }
        }
    } else {
        // In state A, no timeout active
        if !future.await {
            // If this future fails, activate timeout
            *timeout = Some(Instant::now() + Duration::from_secs(5));
        }
    }
    true
}

Compared to Chayim's, this one has the advantage of only checking the state once per future instead of on every Pending, but in exchange, the state can't be changed while inside one of the futures.

One thing that wasn't specified in your question is what should happen to the state when continuing. I wrote this to reset the state to A, but you might've wanted to keep not-A and just reset the timer.

You can do almost anything by the magic of writing a future wrapper:

use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Poll};
use std::time::Duration;

use tokio::time::{Instant, Sleep};

pin_project_lite::pin_project! {
    pub struct NonATimeout<'a, Fut> {
        #[pin]
        fut: Fut,
        #[pin]
        sleep: Sleep,
        state: &'a AtomicBool,
        was_in_a: bool,
    }
}

impl<'a, Fut> NonATimeout<'a, Fut> {
    pub fn new(fut: Fut, state: &'a AtomicBool) -> Self {
        Self {
            fut,
            sleep: tokio::time::sleep(Duration::ZERO),
            state,
            was_in_a: true,
        }
    }
}

impl<Fut: Future<Output = ()>> Future for NonATimeout<'_, Fut> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();

        if this.fut.poll(cx).is_ready() {
            return Poll::Ready(());
        }

        if !this.state.load(Ordering::Relaxed) {
            // Not in A.
            if *this.was_in_a {
                this.sleep.as_mut().reset(Instant::now() + Duration::from_secs(5));
                *this.was_in_a = false;
            }
            if this.sleep.poll(cx).is_ready() {
                // Cancel the work.
                return Poll::Ready(());
            }
        }

        Poll::Pending
    }
}

Then using it like:

async fn foo() {
    let state_a = AtomicBool::new(true);
    loop {
        let do_work = async {
            state_a.store(call_1().await, Ordering::Relaxed);
            state_a.store(call_2().await, Ordering::Relaxed);
            // ...
            state_a.store(call_n().await, Ordering::Relaxed);
        };
        let timeout = NonATimeout::new(do_work, &state_a);
        timeout.await;
    }
}

If the state type is more complicated, you can use a Mutex. It will never be contended because we never acquire it in parallel. Technically, you can use your own non-atomic Cell wrapper for that, but that requires unsafe code.