Implementation

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

pub struct RetryFuture<F, T, E> {
    op: F,
    max_retries: usize,
    attempts: usize,
}

impl<F, T, E> RetryFuture<F, T, E> {
    pub fn new(max_retries: usize, op: F) -> Self {
        Self {
            op,
            max_retries,
            attempts: 0,
        }
    }
}

impl<F, T, E> Future for RetryFuture<F, T, E>
where
    F: FnMut() -> Result<T, E> + Unpin,
{
    type Output = Result<T, E>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.attempts += 1;
        match (self.op)() {
            Ok(val) => Poll::Ready(Ok(val)),
            Err(err) => {
                if self.attempts >= self.max_retries {
                    Poll::Ready(Err(err))
                } else {
                    cx.waker().wake_by_ref();
                    Poll::Pending
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_retry_future_success() {
        let mut count = 0;
        let retry_fut = RetryFuture::new(3, move || {
            count += 1;
            if count == 2 {
                Ok("SUCCESS")
            } else {
                Err("TRANSIENT_ERR")
            }
        });

        let res = retry_fut.await;
        assert_eq!(res, Ok("SUCCESS"));
    }

    #[tokio::test]
    async fn test_retry_future_exhaustion() {
        let retry_fut = RetryFuture::new(2, || Err::<(), _>("PERMANENT_ERR"));
        let res = retry_fut.await;
        assert_eq!(res, Err("PERMANENT_ERR"));
    }
}

Technical Explanation

1. Manual Future Mechanics: Implementing Future directly requires defining type Output and writing the poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> state machine loop.
2. Waker Notification: When returning Poll::Pending, calling cx.waker().wake_by_ref() informs Tokio's executor to reschedule the future for polling again.

Exercise 2: Shared Single-Execution Asynchronous Cell (AsyncOnceCell)

Scenario: Database connection pools or configuration loaders use an AsyncOnceCell<T> to ensure an expensive initialization future runs only once. If multiple caller tasks .await the cell concurrently while initialization is in progress, all callers register their Wakers and receive the initialized result once complete.

Construct a shared single-execution cell using manual Future waker registration.

Requirements:

1. Define AsyncOnceCell<T> with shared Arc<Mutex<CellState<T>>>.
2. CellState<T> contains value: Option<T>, wakers: Vec<Waker>, and is_initializing: bool.
3. Implement get_or_init<F>(&self, init: F) returning a custom OnceCellFuture<T>.
4. Add unit tests asserting single initialization execution across multiple concurrent task callers.

[!check]- Answer

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};

struct CellState<T> {
    value: Option<T>,
    wakers: Vec<Waker>,
}

pub struct AsyncOnceCell<T> {
    state: Arc<Mutex<CellState<T>>>,
}

impl<T: Clone> AsyncOnceCell<T> {
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(CellState {
                value: None,
                wakers: Vec::new(),
            })),
        }
    }

    pub fn set(&self, val: T) {
        let mut guard = self.state.lock().unwrap();
        guard.value = Some(val);
        for waker in guard.wakers.drain(..) {
            waker.wake();
        }
    }

    pub fn get(&self) -> OnceCellFuture<T> {
        OnceCellFuture {
            state: Arc::clone(&self.state),
        }
    }
}

pub struct OnceCellFuture<T> {
    state: Arc<Mutex<CellState<T>>>,
}

impl<T: Clone> Future for OnceCellFuture<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut guard = self.state.lock().unwrap();
        if let Some(ref val) = guard.value {
            Poll::Ready(val.clone())
        } else {
            guard.wakers.push(cx.waker().clone());
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[tokio::test]
    async fn test_async_once_cell_multi_waiter() {
        let cell = Arc::new(AsyncOnceCell::<String>::new());
        let cell1 = Arc::clone(&cell);
        let cell2 = Arc::clone(&cell);

        let h1 = tokio::spawn(async move { cell1.get().await });
        let h2 = tokio::spawn(async move { cell2.get().await });

        tokio::time::sleep(Duration::from_millis(10)).await;
        cell.set("CONFIG_DATA".into());

        assert_eq!(h1.await.unwrap(), "CONFIG_DATA");
        assert_eq!(h2.await.unwrap(), "CONFIG_DATA");
    }
}

Technical Explanation

1. Multi-Waiter Waker Tracking: CellState maintains wakers: Vec<Waker>. Concurrent callers polling OnceCellFuture push their wakers into the list when data is pending.
2. Fan-Out Waker Notification: Calling cell.set(val) pops all stored wakers via guard.wakers.drain(..) and invokes .wake() on each, unblocking all .awaiting tasks simultaneously.

Exercise 3: Custom Cancellation-Safe Select2 Combinator with Safe Pinned Projection

Scenario: tokio::select! races two futures. To understand its internal mechanics, build a custom zero-allocation Select2<F1, F2> combinator implementing Future.

Construct a manual Select2 combinator with pinned projection.

Requirements:

1. Implement Select2<F1, F2> struct holding fut1: F1 and fut2: F2.
2. Implement Future returning Poll::Ready(Either<F1::Output, F2::Output>).
3. Perform structural pin projection using unsafe { Pin::new_unchecked(...) }.
4. Add unit tests asserting branch winning behavior and cancellation of the losing future.

[!check]- Answer

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

#[derive(Debug, PartialEq, Eq)]
pub enum Either<A, B> {
    Left(A),
    Right(B),
}

pub struct Select2<F1, F2> {
    fut1: F1,
    fut2: F2,
}

impl<F1, F2> Select2<F1, F2> {
    pub fn new(fut1: F1, fut2: F2) -> Self {
        Self { fut1, fut2 }
    }
}

impl<F1: Future, F2: Future> Future for Select2<F1, F2> {
    type Output = Either<F1::Output, F2::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // SAFETY: Structural pin projection for fut1 and fut2
        let (p1, p2) = unsafe {
            let this = self.get_unchecked_mut();
            (
                Pin::new_unchecked(&mut this.fut1),
                Pin::new_unchecked(&mut this.fut2),
            )
        };

        if let Poll::Ready(out1) = p1.poll(cx) {
            return Poll::Ready(Either::Left(out1));
        }

        if let Poll::Ready(out2) = p2.poll(cx) {
            return Poll::Ready(Either::Right(out2));
        }

        Poll::Pending
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[tokio::test]
    async fn test_select2_left_wins() {
        let f1 = async {
            tokio::time::sleep(Duration::from_millis(5)).await;
            "LEFT_WIN"
        };
        let f2 = async {
            tokio::time::sleep(Duration::from_millis(100)).await;
            "RIGHT_WIN"
        };

        let res = Select2::new(f1, f2).await;
        assert_eq!(res, Either::Left("LEFT_WIN"));
    }
}

Technical Explanation

1. Structural Pin Projection: Select2 projects Pin<&mut Select2<F1, F2>> into Pin<&mut F1> and Pin<&mut F2>.
2. First-Completion Resolution: Polling both projected futures inside poll returns Poll::Ready(Either::Left) or Poll::Ready(Either::Right) for whichever future completes first.