Tokio 信号量有界并发

示例说明

信号量是并发限流的正统原语：许可数即并发上限，acquire 在许可耗尽时挂起等待，Drop 归还时自动唤醒等待者——Rust 的 RAII 语义让归还不可遗忘，任务 panic 或提前返回都会释放许可。acquire_owned 返回的所有权许可可以跨任务移动，配合按键配额可以实现总量与调用方双重限制的公平调度。

技术要点

• 许可的 Drop 即归还不可遗忘
• acquire_owned 许可可跨任务移动
• 按键配额实现调用方级公平

代码解读

Gatekeeper 维护全局信号量与按键计数两张表；acquire 循环里先查该键的活跃数——超过 per_key 就退避重试，未超则获取全局许可并计数加一。KeyedPermit 演示许可与状态表的组合结构，into_permit 展示所有权转移。main 并发六个任务验证总并发不超过信号量上限。

使用建议

退避重试可用等待者队列替代轮询。键计数要防泄漏——真实实现应把计数清理放进 Drop。许可不可克隆，转移即唯一。

代码示例

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());
    }
}