Singleflight Async Request Coalescing

use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use tokio::sync::{broadcast, Mutex};

type SharedResult = Arc<Result<String, String>>;

#[derive(Clone, Default)]
struct SingleFlight {
    inflight: Arc<Mutex<HashMap<String, broadcast::Sender<SharedResult>>>>,
}

impl SingleFlight {
    fn new() -> Self {
        Self::default()
    }

    async fn run<F>(&self, key: &str, load: F) -> Result<String, String>
    where
        F: Future<Output = Result<String, String>>,
    {
        let follower = {
            let mut inflight = self.inflight.lock().await;
            match inflight.get(key) {
                Some(sender) => Some(sender.subscribe()),
                None => {
                    let (sender, _) = broadcast::channel(1);
                    inflight.insert(key.to_string(), sender);
                    None
                }
            }
        };

        if let Some(mut receiver) = follower {
            return match receiver.recv().await {
                Ok(shared) => (*shared).clone(),
                Err(_) => Err("leader cancelled before publishing".to_string()),
            };
        }

        let result = load.await;
        let shared = Arc::new(result.clone());
        let sender = {
            let mut inflight = self.inflight.lock().await;
            inflight.remove(key)
        };
        if let Some(sender) = sender {
            let _ = sender.send(shared);
        }
        result
    }
}

#[tokio::main]
async fn main() {
    let group = SingleFlight::new();
    let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let (first, second) = tokio::join!(
        group.run("user:1", async {
            calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            tokio::time::sleep(std::time::Duration::from_millis(30)).await;
            Ok::<String, String>("payload".to_string())
        }),
        group.run("user:1", async {
            calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            tokio::time::sleep(std::time::Duration::from_millis(30)).await;
            Ok::<String, String>("payload".to_string())
        })
    );
    println!(
        "{:?} {:?} loads={}",
        first,
        second,
        calls.load(std::sync::atomic::Ordering::SeqCst)
    );
}