servers/http/
memory_limit.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Middleware for limiting total memory usage of concurrent HTTP request bodies.
16
17use axum::extract::{Request, State};
18use axum::middleware::Next;
19use axum::response::{IntoResponse, Response};
20use http::StatusCode;
21
22use crate::metrics::{METRIC_HTTP_MEMORY_USAGE_BYTES, METRIC_HTTP_REQUESTS_REJECTED_TOTAL};
23use crate::request_limiter::RequestMemoryLimiter;
24
25pub async fn memory_limit_middleware(
26    State(limiter): State<RequestMemoryLimiter>,
27    req: Request,
28    next: Next,
29) -> Response {
30    let content_length = req
31        .headers()
32        .get(http::header::CONTENT_LENGTH)
33        .and_then(|v| v.to_str().ok())
34        .and_then(|v| v.parse::<usize>().ok())
35        .unwrap_or(0);
36
37    let _guard = match limiter.try_acquire(content_length) {
38        Ok(guard) => guard.inspect(|g| {
39            METRIC_HTTP_MEMORY_USAGE_BYTES.set(g.current_usage() as i64);
40        }),
41        Err(e) => {
42            METRIC_HTTP_REQUESTS_REJECTED_TOTAL.inc();
43            return (
44                StatusCode::TOO_MANY_REQUESTS,
45                format!("Request body memory limit exceeded: {}", e),
46            )
47                .into_response();
48        }
49    };
50
51    next.run(req).await
52}
servers/http/memory_limit.rs

servers/http/
memory_limit.rs
