Tokio Semaphore Bounded Concurrency

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, Semaphore};

struct Gatekeeper {
    permits: Arc<Semaphore>,
    key_states: Mutex<HashMap<String, usize>>,
    per_key: usize,
}

impl Gatekeeper {
    fn new(total: usize, per_key: usize) -> Self {
        Self {
            permits: Arc::new(Semaphore::new(total)),
            key_states: Mutex::new(HashMap::new()),
            per_key,
        }
    }

    async fn acquire(&self, key: &str) -> tokio::sync::OwnedSemaphorePermit {
        // 第一段：等按键配额有余量
        loop {
            let active = *self.key_states.lock().await
                .get(key).unwrap_or(&0);
            if active < self.per_key {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        }

        // 第二段：获取全局许可（挂起直到总量有空位）
        let permit = self
            .permits
            .clone()
            .acquire_owned()
            .await
            .expect("semaphore never closes");

        // 第三段：登记按键计数（真实实现应在 Drop 包装中递减）
        *self.key_states.lock().await
            .entry(key.to_string()).or_insert(0) += 1;
        permit
    }
}

#[tokio::main]
async fn main() {
    let keeper = Arc::new(Gatekeeper::new(4, 2));

    let mut handles = Vec::new();
    for i in 0..6 {
        let gate = keeper.clone();
        handles.push(tokio::spawn(async move {
            let _permit = gate.acquire(if i % 2 == 0 { "api" } else { "sync" }).await;
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
            i
        }));
    }
    for h in handles {
        println!("task {} finished", h.await.unwrap());
    }
}