Code shared from the Rust Playground

playground.rs

use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::Mutex;

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

/// Like OnceCell, but returns an error on cancellation instead of retrying.
pub struct OnceCompute<T> {
    state: Mutex<Result<T, Cancelled>>,
    started: AtomicBool,
}

impl<T: Clone> OnceCompute<T> {
    pub fn new() -> Self {
        Self {
            // We start with state preset to Cancelled.
            // So that when the mutex is dropped without finishing,
            // consumers see this error instead of retrying.
            state: Mutex::new(Err(Cancelled)),
            started: AtomicBool::new(false),
        }
    }

    pub async fn get_or_init<F: Future<Output = T>>(&self, fut: F) -> Result<T, Cancelled> {
        if !self.started.swap(true, Ordering::AcqRel) {
            // We're the runner. Hold the lock across .await — waiters queue on it.
            let mut guard = self.state.lock().await;
            let value = fut.await; // if cancelled here, guard drops, Err(Cancelled) stays
            *guard = Ok(value.clone());
            return Ok(value);
        }
        // Waiter: blocks until runner drops the lock, then clones whatever's there
        self.state.lock().await.clone()
    }
}