482 lines
16 KiB
Rust
482 lines
16 KiB
Rust
use std::sync::Arc;
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::{Request, State},
|
|
http::{HeaderValue, StatusCode, header::RETRY_AFTER},
|
|
middleware::Next,
|
|
response::Response,
|
|
};
|
|
use http_body_util::BodyExt;
|
|
use tokio::sync::{OwnedSemaphorePermit, TryAcquireError};
|
|
|
|
use crate::{
|
|
http_error::AppError,
|
|
request_context::RequestContext,
|
|
state::{BackpressureState, HttpRequestPermitPool, HttpRequestPermitPoolKind},
|
|
};
|
|
|
|
pub async fn limit_concurrent_requests(
|
|
State(state): State<BackpressureState>,
|
|
request: Request,
|
|
next: Next,
|
|
) -> Response {
|
|
if should_bypass_backpressure(&request) {
|
|
return next.run(request).await;
|
|
}
|
|
|
|
let requested_pool = classify_request_permit_pool(request.uri().path());
|
|
let Some((permit_pool_kind, permit_pool)) = state.request_permit_pool(requested_pool) else {
|
|
return next.run(request).await;
|
|
};
|
|
|
|
match acquire_http_request_permit(permit_pool_kind, permit_pool) {
|
|
Ok(permit) => hold_permit_until_response_body_dropped(next.run(request).await, permit),
|
|
Err(_) => reject_overloaded_request(&request),
|
|
}
|
|
}
|
|
|
|
fn acquire_http_request_permit(
|
|
permit_pool_kind: HttpRequestPermitPoolKind,
|
|
permit_pool: Arc<HttpRequestPermitPool>,
|
|
) -> Result<HttpRequestPermitGuard, TryAcquireError> {
|
|
match permit_pool.clone().try_acquire_owned() {
|
|
Ok(permit) => {
|
|
crate::telemetry::update_http_request_permits_available(
|
|
permit_pool_kind,
|
|
permit_pool.available_permits(),
|
|
);
|
|
Ok(HttpRequestPermitGuard {
|
|
permit_pool_kind,
|
|
permit: Some(permit),
|
|
permit_pool,
|
|
})
|
|
}
|
|
Err(error) => {
|
|
crate::telemetry::update_http_request_permits_available(
|
|
permit_pool_kind,
|
|
permit_pool.available_permits(),
|
|
);
|
|
Err(error)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn hold_permit_until_response_body_dropped(
|
|
response: Response,
|
|
permit: HttpRequestPermitGuard,
|
|
) -> Response {
|
|
response.map(|body| {
|
|
Body::new(body.map_frame(move |frame| {
|
|
let _permit_guard = &permit;
|
|
frame
|
|
}))
|
|
})
|
|
}
|
|
|
|
struct HttpRequestPermitGuard {
|
|
permit_pool_kind: HttpRequestPermitPoolKind,
|
|
permit: Option<OwnedSemaphorePermit>,
|
|
permit_pool: Arc<HttpRequestPermitPool>,
|
|
}
|
|
|
|
impl Drop for HttpRequestPermitGuard {
|
|
fn drop(&mut self) {
|
|
drop(self.permit.take());
|
|
crate::telemetry::update_http_request_permits_available(
|
|
self.permit_pool_kind,
|
|
self.permit_pool.available_permits(),
|
|
);
|
|
}
|
|
}
|
|
|
|
fn reject_overloaded_request(request: &Request<Body>) -> Response {
|
|
let request_context = request.extensions().get::<RequestContext>().cloned();
|
|
let mut response = AppError::from_status(StatusCode::TOO_MANY_REQUESTS)
|
|
.with_message("服务繁忙，请稍后重试")
|
|
.into_response_with_context(request_context.as_ref());
|
|
response
|
|
.headers_mut()
|
|
.insert(RETRY_AFTER, HeaderValue::from_static("1"));
|
|
response
|
|
}
|
|
|
|
fn should_bypass_backpressure(request: &Request<Body>) -> bool {
|
|
request.uri().path() == "/healthz"
|
|
}
|
|
|
|
fn classify_request_permit_pool(path: &str) -> HttpRequestPermitPoolKind {
|
|
if is_gallery_list_path(path) {
|
|
HttpRequestPermitPoolKind::Gallery
|
|
} else if is_gallery_detail_path(path) {
|
|
HttpRequestPermitPoolKind::Detail
|
|
} else if path.starts_with("/admin/api/") {
|
|
HttpRequestPermitPoolKind::Admin
|
|
} else {
|
|
HttpRequestPermitPoolKind::Default
|
|
}
|
|
}
|
|
|
|
fn is_gallery_list_path(path: &str) -> bool {
|
|
matches!(
|
|
path,
|
|
"/api/runtime/puzzle/gallery" | "/api/runtime/custom-world-gallery"
|
|
)
|
|
}
|
|
|
|
fn is_gallery_detail_path(path: &str) -> bool {
|
|
let puzzle_prefix = "/api/runtime/puzzle/gallery/";
|
|
if let Some(profile_id) = path.strip_prefix(puzzle_prefix) {
|
|
return !profile_id.is_empty() && !profile_id.contains('/');
|
|
}
|
|
|
|
let custom_world_prefix = "/api/runtime/custom-world-gallery/";
|
|
if let Some(remainder) = path.strip_prefix(custom_world_prefix) {
|
|
let mut segments = remainder.split('/');
|
|
return matches!(
|
|
(segments.next(), segments.next(), segments.next()),
|
|
(Some(owner_user_id), Some(profile_id), None)
|
|
if !owner_user_id.is_empty() && !profile_id.is_empty()
|
|
);
|
|
}
|
|
|
|
false
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::Arc;
|
|
|
|
use axum::{
|
|
Router,
|
|
body::Body,
|
|
extract::Extension,
|
|
http::{Request, StatusCode, header::RETRY_AFTER},
|
|
middleware,
|
|
routing::get,
|
|
};
|
|
use tokio::sync::Notify;
|
|
use tower::ServiceExt;
|
|
|
|
use axum::extract::FromRef;
|
|
|
|
use crate::{
|
|
config::AppConfig,
|
|
state::{AppState, BackpressureState},
|
|
};
|
|
|
|
use super::{classify_request_permit_pool, limit_concurrent_requests};
|
|
|
|
#[derive(Clone)]
|
|
struct HeldRequestGate {
|
|
entered: Arc<Notify>,
|
|
release: Arc<Notify>,
|
|
}
|
|
|
|
async fn held_request(Extension(gate): Extension<HeldRequestGate>) -> &'static str {
|
|
gate.entered.notify_one();
|
|
gate.release.notified().await;
|
|
"ok"
|
|
}
|
|
|
|
async fn fast_request() -> &'static str {
|
|
"ok"
|
|
}
|
|
|
|
fn test_request(path: &str) -> Request<Body> {
|
|
Request::builder()
|
|
.uri(path)
|
|
.body(Body::empty())
|
|
.expect("test request should build")
|
|
}
|
|
|
|
fn build_test_app(max_concurrent_requests: usize, gate: HeldRequestGate) -> Router {
|
|
let mut config = AppConfig::default();
|
|
config.max_concurrent_requests = Some(max_concurrent_requests);
|
|
let state = AppState::new(config).expect("state should build");
|
|
let backpressure_state = BackpressureState::from_ref(&state);
|
|
|
|
Router::new()
|
|
.route("/held", get(held_request))
|
|
.route("/fast", get(fast_request))
|
|
.route("/healthz", get(fast_request))
|
|
.layer(middleware::from_fn_with_state(
|
|
backpressure_state,
|
|
limit_concurrent_requests,
|
|
))
|
|
.layer(Extension(gate))
|
|
.with_state(state)
|
|
}
|
|
|
|
fn build_grouped_test_app(
|
|
default_max_concurrent_requests: usize,
|
|
gallery_max_concurrent_requests: usize,
|
|
admin_max_concurrent_requests: usize,
|
|
gate: HeldRequestGate,
|
|
) -> Router {
|
|
let mut config = AppConfig::default();
|
|
config.max_concurrent_requests = Some(default_max_concurrent_requests);
|
|
config.gallery_max_concurrent_requests = Some(gallery_max_concurrent_requests);
|
|
config.admin_max_concurrent_requests = Some(admin_max_concurrent_requests);
|
|
let state = AppState::new(config).expect("state should build");
|
|
let backpressure_state = BackpressureState::from_ref(&state);
|
|
|
|
Router::new()
|
|
.route("/held", get(held_request))
|
|
.route("/api/runtime/puzzle/gallery", get(held_request))
|
|
.route("/api/runtime/custom-world-gallery", get(held_request))
|
|
.route("/api/runtime/puzzle/gallery/profile-1", get(held_request))
|
|
.route(
|
|
"/api/runtime/puzzle/gallery/profile-1/like",
|
|
get(fast_request),
|
|
)
|
|
.route(
|
|
"/api/runtime/custom-world-gallery/user-1/profile-1",
|
|
get(held_request),
|
|
)
|
|
.route("/admin/api/overview", get(held_request))
|
|
.route("/fast", get(fast_request))
|
|
.layer(middleware::from_fn_with_state(
|
|
backpressure_state,
|
|
limit_concurrent_requests,
|
|
))
|
|
.layer(Extension(gate))
|
|
.with_state(state)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn returns_429_when_concurrency_permits_are_exhausted() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let app = build_test_app(1, gate.clone());
|
|
let entered = gate.entered.notified();
|
|
|
|
let held_response = tokio::spawn(app.clone().oneshot(test_request("/held")));
|
|
entered.await;
|
|
|
|
let rejected_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("rejected request should complete");
|
|
assert_eq!(rejected_response.status(), StatusCode::TOO_MANY_REQUESTS);
|
|
assert_eq!(
|
|
rejected_response
|
|
.headers()
|
|
.get(RETRY_AFTER)
|
|
.and_then(|value| value.to_str().ok()),
|
|
Some("1")
|
|
);
|
|
|
|
gate.release.notify_one();
|
|
let completed_response = held_response
|
|
.await
|
|
.expect("held request task should join")
|
|
.expect("held request should complete");
|
|
assert_eq!(completed_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn healthz_bypasses_concurrency_backpressure() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let app = build_test_app(1, gate.clone());
|
|
let entered = gate.entered.notified();
|
|
|
|
let held_response = tokio::spawn(app.clone().oneshot(test_request("/held")));
|
|
entered.await;
|
|
|
|
let health_response = app
|
|
.clone()
|
|
.oneshot(test_request("/healthz"))
|
|
.await
|
|
.expect("healthz request should complete");
|
|
assert_eq!(health_response.status(), StatusCode::OK);
|
|
|
|
gate.release.notify_one();
|
|
let completed_response = held_response
|
|
.await
|
|
.expect("held request task should join")
|
|
.expect("held request should complete");
|
|
assert_eq!(completed_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn permit_is_held_until_response_body_is_dropped() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let app = build_test_app(1, gate);
|
|
|
|
let first_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("first request should complete");
|
|
assert_eq!(first_response.status(), StatusCode::OK);
|
|
|
|
let rejected_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("second request should complete");
|
|
assert_eq!(rejected_response.status(), StatusCode::TOO_MANY_REQUESTS);
|
|
|
|
drop(first_response);
|
|
|
|
let accepted_response = app
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("third request should complete");
|
|
assert_eq!(accepted_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn gallery_pool_rejects_gallery_without_blocking_default_routes() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let app = build_grouped_test_app(2, 1, 1, gate.clone());
|
|
let entered = gate.entered.notified();
|
|
|
|
let held_response = tokio::spawn(
|
|
app.clone()
|
|
.oneshot(test_request("/api/runtime/puzzle/gallery")),
|
|
);
|
|
entered.await;
|
|
|
|
let rejected_gallery_response = app
|
|
.clone()
|
|
.oneshot(test_request("/api/runtime/custom-world-gallery"))
|
|
.await
|
|
.expect("rejected gallery request should complete");
|
|
assert_eq!(
|
|
rejected_gallery_response.status(),
|
|
StatusCode::TOO_MANY_REQUESTS
|
|
);
|
|
|
|
let accepted_default_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("default request should complete");
|
|
assert_eq!(accepted_default_response.status(), StatusCode::OK);
|
|
|
|
gate.release.notify_one();
|
|
let completed_response = held_response
|
|
.await
|
|
.expect("held request task should join")
|
|
.expect("held request should complete");
|
|
assert_eq!(completed_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn detail_pool_falls_back_to_default_when_unset() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let mut config = AppConfig::default();
|
|
config.max_concurrent_requests = Some(1);
|
|
config.detail_max_concurrent_requests = None;
|
|
let state = AppState::new(config).expect("state should build");
|
|
let backpressure_state = BackpressureState::from_ref(&state);
|
|
let app = Router::new()
|
|
.route("/api/runtime/puzzle/gallery/profile-1", get(held_request))
|
|
.route("/fast", get(fast_request))
|
|
.layer(middleware::from_fn_with_state(
|
|
backpressure_state,
|
|
limit_concurrent_requests,
|
|
))
|
|
.layer(Extension(gate.clone()))
|
|
.with_state(state);
|
|
let entered = gate.entered.notified();
|
|
|
|
let held_response = tokio::spawn(
|
|
app.clone()
|
|
.oneshot(test_request("/api/runtime/puzzle/gallery/profile-1")),
|
|
);
|
|
entered.await;
|
|
|
|
let rejected_default_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("default request should complete");
|
|
assert_eq!(
|
|
rejected_default_response.status(),
|
|
StatusCode::TOO_MANY_REQUESTS
|
|
);
|
|
|
|
gate.release.notify_one();
|
|
let completed_response = held_response
|
|
.await
|
|
.expect("held request task should join")
|
|
.expect("held request should complete");
|
|
assert_eq!(completed_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn admin_pool_is_isolated_from_default_routes() {
|
|
let gate = HeldRequestGate {
|
|
entered: Arc::new(Notify::new()),
|
|
release: Arc::new(Notify::new()),
|
|
};
|
|
let app = build_grouped_test_app(2, 1, 1, gate.clone());
|
|
let entered = gate.entered.notified();
|
|
|
|
let held_response = tokio::spawn(app.clone().oneshot(test_request("/admin/api/overview")));
|
|
entered.await;
|
|
|
|
let rejected_admin_response = app
|
|
.clone()
|
|
.oneshot(test_request("/admin/api/overview"))
|
|
.await
|
|
.expect("rejected admin request should complete");
|
|
assert_eq!(
|
|
rejected_admin_response.status(),
|
|
StatusCode::TOO_MANY_REQUESTS
|
|
);
|
|
|
|
let accepted_default_response = app
|
|
.clone()
|
|
.oneshot(test_request("/fast"))
|
|
.await
|
|
.expect("default request should complete");
|
|
assert_eq!(accepted_default_response.status(), StatusCode::OK);
|
|
|
|
gate.release.notify_one();
|
|
let completed_response = held_response
|
|
.await
|
|
.expect("held request task should join")
|
|
.expect("held request should complete");
|
|
assert_eq!(completed_response.status(), StatusCode::OK);
|
|
}
|
|
|
|
#[test]
|
|
fn classifies_only_exact_gallery_detail_paths_as_detail() {
|
|
assert_eq!(
|
|
classify_request_permit_pool("/api/runtime/puzzle/gallery/profile-1"),
|
|
crate::state::HttpRequestPermitPoolKind::Detail
|
|
);
|
|
assert_eq!(
|
|
classify_request_permit_pool("/api/runtime/puzzle/gallery/profile-1/like"),
|
|
crate::state::HttpRequestPermitPoolKind::Default
|
|
);
|
|
assert_eq!(
|
|
classify_request_permit_pool("/api/runtime/custom-world-gallery/user-1/profile-1"),
|
|
crate::state::HttpRequestPermitPoolKind::Detail
|
|
);
|
|
assert_eq!(
|
|
classify_request_permit_pool("/api/runtime/custom-world-gallery/user-1/profile-1/like"),
|
|
crate::state::HttpRequestPermitPoolKind::Default
|
|
);
|
|
}
|
|
}
|