Rust | 实现 API 限速操作 Example
在这篇文章中，我们将讨论如何在 Rust 中实现 API 限速。当涉及到生产中的服务时，是为了确保不良行为者不会滥用 API——这就是 API 限速的作用所在。
我们将实现 “滑动窗口” 算法，通过动态周期来检查请求历史，并使用基本的内存 hashmap 来存储用户 IP 及其请求时间。我们还将研究如何使用 tower-governor 库来实现限速。
实现一个简单的滑动窗口限速器
让我们从头开始编写一个简单的基于 IP 的滑动窗口限速器。
在项目中加入以下依赖项：
[dependencies]
chrono = { version = "0.4.34", features = ["serde", "clock"] }
serde = { version = "1.0.196", features = ["derive"] }
我们将声明一个新的结构体 RateLimiter，它保存 IP 地址作为键的 HashMap，值为 Vec<DateTime>(Utc 时区的时间戳)。
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use std::net::IpAddr;
use chrono::{DateTime, Utc};
// 這是用戶訪問端點的請求限制(每分鐘)
// 如果用戶試圖超過這個限制，返回一個錯誤
const REQUEST_LIMIT: usize = 120;
#[derive(Clone, Default)]
pub struct RateLimiter {
requests: Arc<Mutex<HashMap<IpAddr, Vec<DateTime<Utc>>>>>,
}
首先，我们想要通过使用.lock() 来锁定我们的 HashMap，这给了我们写访问权限。然后，我们需要检查hashmap 是否包含一个键，其中包含我们想要使用.entry() 函数检查的IP 地址，然后通过保留有效的时间戳来修改它，并根据长度是否在请求限制之下push 一个新记录。然后检查记录长度是否大于请求限制—如果是，则返回错误；如果不是，则返回 Ok(())。
impl RateLimiter {
fn check_if_rate_limited(&self, ip_addr: IpAddr) -> Result<(), String> {
// 我們只想保留60秒前的時間戳
let throttle_time_limit = Utc::now() - std::time::Duration::from_secs(60);
let mut requests_hashmap = self.requests.lock().unwrap();
let mut requests_for_ip = requests_hashmap
// 檢索記錄，並允許我們在適當的地方修改它
.entry(ip_addr)
// 如果記錄爲空，則插入帶有當前時間戳的vec
.or_default();
requests_for_ip.retain(|x| x.to_utc() > throttle_time_limit);
requests_for_ip.push(Utc::now());
if requests_for_ip.len() > REQUEST_LIMIT {
return Err("IP is rate limited :(".to_string());
}
Ok(())
}
}
测试如下：
fn main() {
let rate_limiter = RateLimiter::default();
let localhost_v4 = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
// 這裏我們請求120次
for _ in 1..80 {
assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
}
// 等待30秒
std::thread::sleep(std::time::Duration::from_secs(30));
// 在這裏再做40個請求
for _ in 1..40 {
assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
}
// 再等30秒
std::thread::sleep(std::time::Duration::from_secs(30));
// 現在我們可以再做80個請求
for _ in 1..80 {
assert!(rate_limiter.check_if_rate_limited(localhost_v4).is_ok())
}
}
然而，生产上的限速系统通常比这先进得多。我们将在下面讨论如何利用 tower-governor crate 进行速率限制。
使用 tower-governor 实现限速器
提供速率限制的 tower 服务和层，后端是 governo 库。 tower-governor 很大程度上基于 actix-governor 所做的工作，可以与 Axum, Hyper, Tonic 和其他任何基于 tower 组件库的框架一起使用！
tower-governor 的特点：
-
速率限制请求基于对端的 IP 地址，IP 地址头或通过自定义设置
-
自定义流量限制的标准是按每秒计算
-
使用简单
-
高可定制性
-
高性能
-
健壮而灵活的 API
它是如何工作的？
每个调控器中间件都有一个存储配额的配置，配额指定可以从一个 IP 地址发送多少请求，如果超过配额则阻止进一步请求。
例如，如果配额允许 10 个请求，客户机可以在中间件开始阻塞之前的短时间内发送 10 个请求。
一旦使用了至少一个配额元素，配额元素将在指定时间后补充。
例如，如果周期为 2 秒，并且配额为空，则需要 2 秒来补充配额的一个元素。这意味着可以平均每两秒钟发送一个请求。
如果配额允许在同一时间段内发送10 个请求，则客户端可以再次发送10 个请求的突发事件，然后必须等待2 秒才能发送进一步的请求，或者在完整配额被补充之前等待20 秒，他可以发送另一个突发事件。
下面我们来实现一个例子。
首先，在项目中加入以下依赖项：
建议版本一致
[dependencies]
tokio = {version = "1.36.0", features = ["full"] }
tower_governor = "0.3.2"
axum = "0.7"
tracing = {version="0.1.37", features=["attributes"]}
tracing-subscriber = "0.3"
tower = "0.4.13"
然后，在 main.rs 文件中写入以下代码：
use axum::{routing::get, Router};
use std::net::SocketAddr;
use std::time::Duration;
use tokio::net::TcpListener;
use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
async fn hello() -> &'static str {
"Hello world"
}
#[tokio::main]
async fn main() {
// 配置跟蹤
// 構造一個訂閱者，將格式化的跟蹤信息打印到標準輸出
let subscriber = tracing_subscriber::FmtSubscriber::new();
// 使用subscriber處理在此點之後發出的跟蹤
tracing::subscriber::set_global_default(subscriber).unwrap();
// 允許每個IP地址最多有五個請求，每兩秒鐘補充一個
// 我們將其裝箱是因爲Axum 0.6要求所有層都是克隆的，因此我們需要一個靜態引用
let governor_conf = Box::new(
GovernorConfigBuilder::default()
.per_second(2)
.burst_size(5)
.finish()
.unwrap(),
);
let governor_limiter = governor_conf.limiter().clone();
let interval = Duration::from_secs(60);
// 一個單獨的後臺任務
std::thread::spawn(move || {
loop {
std::thread::sleep(interval);
tracing::info!("rate limiting storage size: {}", governor_limiter.len());
governor_limiter.retain_recent();
}
});
// 構建路由
let app = Router::new()
// `GET /`
.route("/", get(hello))
.layer(GovernorLayer {
// 我們可以泄漏它，因爲它是一次性創建的
config: Box::leak(governor_conf),
});
let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
tracing::debug!("listening on {}", addr);
let listener = TcpListener::bind(addr).await.unwrap();
axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
.await
.unwrap();
}
运行服务器后，在浏览器中刷新请求，2 秒内超过 5 次会提示以下错误：
服务器后台跟踪日志如下：
2024-05-09T05:41:10.065323Z INFO limit_api2: rate limiting storage size: 1
2024-05-09T05:41:20.067175Z INFO limit_api2: rate limiting storage size: 1
2024-05-09T05:41:30.071745Z INFO limit_api2: rate limiting storage size: 1
2024-05-09T05:41:40.073187Z INFO limit_api2: rate limiting storage size: 1
2024-05-09T05:41:50.077540Z INFO limit_api2: rate limiting storage size: 0
2024-05-09T05:42:00.077448Z INFO limit_api2: rate limiting storage size: 0
