LiteHub中间件之限流实现-EW帮帮网

限流实现

本项目我是基于令牌桶实现的访问限流，请看以下代码
LimitMiddleware.h头文件

class LimitMiddleware: public Middleware 
{
public:
    // 构造函数
    // rate：令牌生成速率（个/秒）
    // capacity：桶最大容量（最多存多少个令牌）
    LimitMiddleware(int rate, int capacity);

    // 在请求处理前调用，用于限流
    void before(HttpRequest& request) override;
    void after(HttpResponse& response) override {};
    double gettokens() const {return tokens_;}
private:
    // 补充令牌（根据时间推移）
    void refillTokens();

private:
    int rate_;            // 令牌生成速率（个/秒）
    int capacity_;        // 桶容量（最大令牌数）
    double tokens_;       // 当前可用令牌数（允许小数，更精确）
    std::chrono::steady_clock::time_point lastRefillTime_; // 上一次补充时间
    std::mutex mutex_;    // 保护多线程访问
};

LimitMiddleware.cc文件

using namespace std::chrono;

// 令牌桶限流中间件
// 构造函数：指定令牌产生速率 rate (个/秒) 和桶容量 capacity (最多可存储多少令牌)
LimitMiddleware::LimitMiddleware(int rate, int capacity)
    : rate_(rate),
      capacity_(capacity),
      tokens_(capacity),  // 初始化时桶是满的，令牌数等于容量
      lastRefillTime_(steady_clock::now())  // 记录上次补充令牌的时间
{
}

// 令牌补充逻辑
// 根据距离上次补充的时间，按速率补充新令牌
void LimitMiddleware::refillTokens()
{
    auto now = steady_clock::now();
    auto elapsedMs = duration_cast<milliseconds>(now - lastRefillTime_).count();

    if (elapsedMs > 0)
    {   
        // 按照速率计算可以补充的令牌数
        double newTokens = (elapsedMs / 1000.0) * rate_;
        // 桶中的令牌数不能超过容量上限
        tokens_ = std::min((double)capacity_, tokens_ + newTokens);
        // 更新时间戳
        lastRefillTime_ = now;
    }
}

// 请求前执行：判断是否有足够的令牌
//如果有足够的令牌，进行下一步操作
//如果没有，拒绝访问，返回状态码429
void LimitMiddleware::before(HttpRequest& request)
{   
    // 加锁保证多线程安全
    std::lock_guard<std::mutex> lock(mutex_);

    // 先补充令牌
    refillTokens();

    if (tokens_ >= 1.0)
    {
        tokens_ -= 1.0;
        // 有足够令牌，消费 1 个，允许请求通过
    }
    else
    {   
        // 没有足够令牌，请求被拒绝，抛出 429 响应
        HttpResponse resp;
        resp.setStatusLine(request.getVersion(), http::HttpResponse::k429TooManyRequests, "Too Many Requests");
        resp.setCloseConnection(true);
        resp.setContentType("application/json");
        resp.setContentLength(0);
        resp.setBody("Rate limit exceeded. Please try again later.");

        throw resp;
    }
}

核心逻辑是：

按固定速率补充令牌（refillTokens）
请求到来时消费令牌（before）
没有令牌可用时返回（限流了，返回状态码429 Too Many Requests）

限流测试

在WebApps/LiteHubServer/src/LiteHubServer.cpp中的initializeMiddleware函数中，定义了限流中间件

 limitMiddleware_ = std::make_shared<http::middleware::LimitMiddleware>(1,100); // 每秒最多100个请求
 httpServer_.addMiddleware(limitMiddleware_);

这里定义的是一秒不超过100个请求，如果通过手动点击，这个1秒内怎么也到不了100次请求；所以我通过python脚本代码模拟一次大量的访问，python代码如下，

import requests
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from collections import Counter

# -------------------------------
# 压测参数配置
# -------------------------------
TOTAL_REQUESTS = 150      # 总请求数
MAX_WORKERS = 3           # 并发线程数
REQUEST_INTERVAL = 0.02   # 相邻请求的间隔（秒），避免瞬间爆发
TARGET_URL = "http://47.122.77.97/"  # 目标 URL

# -------------------------------
# 单次请求任务
# index: 请求编号
# -------------------------------
def send_request(index):
    try:
        start_time = time.time()
        # 发起 GET 请求
        r = requests.get(TARGET_URL, timeout=3)
        elapsed = time.time() - start_time
        # 打印日志：时间戳 + 请求序号 + 响应码 + 耗时
        print(f"[{time.strftime('%H:%M:%S')}] 请求 {index + 1:02d} --> 状态码: {r.status_code} (耗时: {elapsed:.2f}s)")
        return r.status_code
    except Exception as e:
        # 异常时打印错误信息
        print(f"[{time.strftime('%H:%M:%S')}] 请求 {index + 1:02d} --> 失败: {str(e)}")
        return str(e)

# -------------------------------
# 主程序入口
# -------------------------------
def main():
    print(f"开始压测，总请求数：{TOTAL_REQUESTS}，最大并发数：{MAX_WORKERS}")
    results = []

    # 记录整个压测开始时间
    total_start_time = time.time()

    # 创建线程池
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = []
        # 提交任务
        for i in range(TOTAL_REQUESTS):
            futures.append(executor.submit(send_request, i))
            time.sleep(REQUEST_INTERVAL)  # 控制相邻请求的间隔

        # 等待所有任务完成
        for future in as_completed(futures):
            results.append(future.result())

    # 记录整个压测结束时间
    total_elapsed = time.time() - total_start_time

    # -------------------------------
    # 统计与输出结果
    # -------------------------------
    print("\n-----------------------------")
    print("请求状态统计结果:")
    counts = Counter(results)
    for key, count in counts.items():
        print(f"{key}: {count} 次")

    print(f"\n压测总耗时: {total_elapsed:.2f} 秒")
    print("-----------------------------")

if __name__ == "__main__":
    main()