Code shared from the Rust Playground

playground.rs

use std::sync::Mutex;
use tokio::sync::Notify;

enum State<T> {
    Empty,
    Running,
    Ready(T),
    Cancelled,
}

/// Like OnceCell, but returns an error on cancellation instead of retrying.
///
/// Same as before, but uses std mutex instead of tokio.
/// We need a manual state machine here.
pub struct OnceCompute<T> {
    state: Mutex<State<T>>,
    notify: Notify,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Cancelled;

impl<T: Clone> OnceCompute<T> {
    pub fn new() -> Self {
        Self {
            state: Mutex::new(State::Empty),
            notify: Notify::new(),
        }
    }

    pub async fn get_or_init<F: Future<Output = T>>(&self, fut: F) -> Result<T, Cancelled> {
        let action = {
            // lock held briefly. No await across lock guard.
            let mut state = self.state.lock().unwrap();
            match &*state {
                // If we already see data or error, return it.
                State::Ready(v) => return Ok(v.clone()),
                State::Cancelled => return Err(Cancelled),
                // Already running in another future, wait for it.
                State::Running => false,
                // Not running, this future becomes the runner.
                State::Empty => {
                    *state = State::Running;
                    true
                }
            }
        };

        if action {
            self.run(fut).await
        } else {
            self.wait().await
        }
    }

    async fn wait(&self) -> Result<T, Cancelled> {
        loop {
            // wait for notification in a loop.
            let notified = self.notify.notified();
            let resolved = {
                let state = self.state.lock().unwrap();
                match &*state {
                    // if we see data, return.
                    // otherwise, wait for another notification.
                    State::Ready(v) => Some(Ok(v.clone())),
                    State::Cancelled => Some(Err(Cancelled)),
                    _ => None,
                }
            };

            match resolved {
                Some(result) => return result,
                None => notified.await,
            }
        }
    }

    async fn run<F: Future<Output = T>>(&self, fut: F) -> Result<T, Cancelled> {
        // If this future is dropped (cancelled), the guard fires.
        let guard = CancelGuard { cell: self };
        let value = fut.await;
        // Future is done, so let's make sure Drop of CancelGuard doesn't run.
        std::mem::forget(guard);

        *self.state.lock().unwrap() = State::Ready(value.clone());
        self.notify.notify_waiters();
        Ok(value)
    }
}

/// Makes sure that the Cancelled state is set if the future gets cancelled early.
struct CancelGuard<'a, T> {
    cell: &'a OnceCompute<T>,
}

impl<T> Drop for CancelGuard<'_, T> {
    fn drop(&mut self) {
        *self.cell.state.lock().unwrap() = State::Cancelled;
        self.cell.notify.notify_waiters();
    }
}