如何用tokio::sync::Semaphore限制Rust异步请求调用速率

核心实现思路 #

• 针对每秒最多40次请求：使用tokio::sync::Semaphore控制并发请求数，初始许可设为40，每次请求前获取许可、请求完成后释放，确保同时运行的请求不超过阈值。
• 针对每120秒最多200次请求：维护线程安全的请求计数器和周期起始时间，每发送一个请求计数器加1；当计数器达到200时，计算当前周期剩余时间，若不足120秒则等待至周期结束，再重置计数器和起始时间继续发送。
• 响应收集：通过JoinHandle跟踪每个异步请求的执行结果，最后统一收集所有成功响应到Vec中。

完整代码实现 #

use reqwest::header::HeaderMap;
use std::sync::{Arc, Mutex};
use tokio::sync::Semaphore;
use tokio::time::{sleep, Duration, Instant};

// 封装速率限制配置，便于修改维护
struct RateLimitConfig {
    per_second: usize,
    per_period: usize,
    period_secs: u64,
}

async fn make_requests(
    headers: &HeaderMap,
    requests: &[String],
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let config = RateLimitConfig {
        per_second: 40,
        per_period: 200,
        period_secs: 120,
    };

    // 并发控制信号量：每秒最多40个请求
    let semaphore = Arc::new(Semaphore::new(config.per_second));
    // 周期请求跟踪：计数器+周期起始时间，用Mutex保证线程安全
    let rate_tracker = Arc::new(Mutex::new((0usize, Instant::now())));
    let client = reqwest::Client::new();
    let headers_cloned = headers.clone();
    let mut task_handles = Vec::new();

    for req in requests {
        let semaphore = Arc::clone(&semaphore);
        let rate_tracker = Arc::clone(&rate_tracker);
        let client = client.clone();
        let headers = headers_cloned.clone();
        let req = req.clone();

        // 启动异步任务处理单个请求
        let handle = tokio::spawn(async move {
            // 处理120秒周期限制
            let mut tracker = rate_tracker.lock().unwrap();
            let (count, start_time) = &mut *tracker;
            if *count >= config.per_period {
                let elapsed = start_time.elapsed().as_secs();
                if elapsed < config.period_secs {
                    sleep(Duration::from_secs(config.period_secs - elapsed)).await;
                }
                // 重置周期计数器和起始时间
                *count = 0;
                *start_time = Instant::now();
            }
            *count += 1;
            drop(tracker); // 提前释放锁，减少锁竞争

            // 处理每秒并发限制
            let permit = semaphore.acquire().await.unwrap();
            let resp = client
                .get(&req)
                .headers(headers)
                .send()
                .await?
                .text()
                .await?;
            drop(permit); // 释放许可，让后续请求可以获取

            Ok(resp)
        });

        task_handles.push(handle);
    }

    // 收集所有任务结果
    let mut responses = Vec::new();
    for handle in task_handles {
        match handle.await? {
            Ok(resp) => responses.push(resp),
            Err(e) => eprintln!("请求失败: {}", e),
        }
    }

    Ok(responses)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 示例请求列表（可替换为实际业务请求）
    let requests = vec![String::from("https://example.com") ; 250];

    // 构建请求头
    let mut headers = HeaderMap::new();
    headers.insert("Accept", "application/json".parse().unwrap());

    // 初始化Tokio运行时
    let rt = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .unwrap();

    let responses = rt.block_on(make_requests(&headers, &requests));
    println!("共收到 {} 个有效响应", responses?.len());

    Ok(())
}

代码重构建议 #

• 函数命名：移除原函数名前的下划线_make_requests，Rust中以下划线开头的函数通常标记为未使用，业务逻辑用常规命名更规范。
• 减少重复克隆：提前克隆headers到headers_cloned，避免循环内重复调用headers.to_owned()，降低不必要的内存开销。
• 结果跟踪：原代码直接tokio::spawn后未跟踪请求结果，改为存储JoinHandle并统一收集，确保所有请求执行完成后再返回。
• 锁粒度优化：处理周期计数时，完成检查和更新后立即释放锁，避免持有锁期间执行请求逻辑，减少线程间的锁竞争。
• 配置解耦：将速率限制参数封装到RateLimitConfig结构体中，后续修改阈值无需改动核心逻辑。

内容的提问来源于stack exchange，提问作者Coldchain9