servers/grpc/
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//! Aggregate memory admission for gRPC services.
16//!
17//! Tonic decompresses and decodes a complete request message (up to
18//! `max_recv_message_size`, 512 MiB by default) *before* the typed handler is
19//! invoked, so a handler that charges the [`ServerMemoryLimiter`] only sees
20//! the request after the memory is already allocated. For transport-compressed
21//! messages (`grpc-encoding: gzip/zstd`) a few KiB on the wire can expand to
22//! hundreds of MiB during that window, which would bypass any aggregate quota.
23//!
24//! [`MemoryLimiterExtensionLayer`] therefore reserves the configured maximum
25//! decoded message size for every compressed request *before* the inner
26//! (tonic) service runs, and holds the reservation for the whole request.
27//! Handlers can detect the reservation via [`PreDecodeMemoryReservation`] in
28//! the request extensions and skip their own post-decode charge to avoid
29//! double accounting.
30//!
31//! The reservation is intentionally conservative (it upper-bounds the decoded
32//! size before it is known); with the default unlimited limiter it is a no-op.
33
34use std::convert::Infallible;
35use std::sync::Arc;
36use std::task::{Context, Poll};
37
38use axum::response::IntoResponse;
39use common_memory_manager::MemoryGuard;
40use futures::future::BoxFuture;
41use http::Request;
42use tonic::Status;
43use tonic::codegen::Service;
44use tonic::server::NamedService;
45use tower::Layer;
46
47use crate::request_memory_limiter::ServerMemoryLimiter;
48use crate::request_memory_metrics::RequestMemoryMetrics;
49
50/// The gRPC request compression selector header.
51const GRPC_ENCODING_HEADER: &str = "grpc-encoding";
52/// The "no compression" encoding value.
53const IDENTITY_ENCODING: &str = "identity";
54
55/// Present in the request extensions when memory for the (compressed)
56/// request's decoded message was reserved before tonic decompression.
57///
58/// Handlers that charge the [`ServerMemoryLimiter`] after decoding should skip
59/// that charge when this marker is present: the reservation already covers the
60/// peak decoding memory and stays alive for the whole request.
61#[derive(Clone)]
62pub(crate) struct PreDecodeMemoryReservation {
63    _guard: Arc<MemoryGuard<RequestMemoryMetrics>>,
64}
65
66impl PreDecodeMemoryReservation {
67    pub(crate) fn reserved_bytes(&self) -> u64 {
68        self._guard.granted_bytes()
69    }
70}
71
72#[derive(Clone)]
73pub struct MemoryLimiterExtensionLayer {
74    limiter: ServerMemoryLimiter,
75    max_decoding_message_size: usize,
76}
77
78impl MemoryLimiterExtensionLayer {
79    pub fn new(limiter: ServerMemoryLimiter, max_decoding_message_size: usize) -> Self {
80        Self {
81            limiter,
82            max_decoding_message_size,
83        }
84    }
85}
86
87impl<S> Layer<S> for MemoryLimiterExtensionLayer {
88    type Service = MemoryLimiterExtensionService<S>;
89
90    fn layer(&self, service: S) -> Self::Service {
91        MemoryLimiterExtensionService {
92            inner: service,
93            limiter: self.limiter.clone(),
94            max_decoding_message_size: self.max_decoding_message_size,
95        }
96    }
97}
98
99#[derive(Clone)]
100pub struct MemoryLimiterExtensionService<S> {
101    inner: S,
102    limiter: ServerMemoryLimiter,
103    max_decoding_message_size: usize,
104}
105
106impl<S: NamedService> NamedService for MemoryLimiterExtensionService<S> {
107    const NAME: &'static str = S::NAME;
108}
109
110impl<S, ReqBody> Service<Request<ReqBody>> for MemoryLimiterExtensionService<S>
111where
112    S: Service<Request<ReqBody>, Error = Infallible> + Clone + Send + 'static,
113    S::Response: axum::response::IntoResponse,
114    S::Future: Send + 'static,
115    ReqBody: Send + 'static,
116{
117    type Response = axum::response::Response;
118    type Error = Infallible;
119    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
120
121    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
122        self.inner.poll_ready(cx)
123    }
124
125    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
126        // Own a clone of the service so the returned future does not borrow
127        // `self` (tower's load-shedding pattern).
128        let mut this = self.clone();
129        Box::pin(async move {
130            req.extensions_mut().insert(this.limiter.clone());
131
132            let compressed = req
133                .headers()
134                .get(GRPC_ENCODING_HEADER)
135                .and_then(|value| value.to_str().ok())
136                .is_some_and(|value| !value.eq_ignore_ascii_case(IDENTITY_ENCODING));
137
138            if compressed {
139                // A compressed message can expand up to the decoding limit
140                // inside tonic, before any handler runs. Reserve that worst
141                // case against the aggregate quota so the decoding phase is
142                // admitted as well. No-op for an unlimited limiter.
143                let reservation = this.max_decoding_message_size as u64;
144                match this.limiter.acquire(reservation).await {
145                    Ok(guard) => {
146                        req.extensions_mut().insert(PreDecodeMemoryReservation {
147                            _guard: Arc::new(guard),
148                        });
149                    }
150                    Err(e) => {
151                        return Ok(Status::resource_exhausted(format!(
152                            "request memory limit exceeded: {e}"
153                        ))
154                        .into_http::<axum::body::Body>()
155                        .into_response());
156                    }
157                }
158            }
159
160            match this.inner.call(req).await {
161                Ok(response) => Ok(response.into_response()),
162                Err(e) => match e {},
163            }
164        })
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use std::convert::Infallible;
171
172    use axum::body::Body;
173    use axum::response::IntoResponse;
174    use common_memory_manager::OnExhaustedPolicy;
175    use futures_util::future::BoxFuture;
176    use http::{HeaderValue, StatusCode};
177    use tower::ServiceExt;
178
179    use super::*;
180    use crate::grpc::memory_limit::MemoryLimiterExtensionLayer;
181
182    #[derive(Clone)]
183    struct EchoService;
184
185    impl Service<Request<Body>> for EchoService {
186        type Response = axum::response::Response;
187        type Error = Infallible;
188        type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
189
190        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
191            Poll::Ready(Ok(()))
192        }
193
194        fn call(&mut self, req: Request<Body>) -> Self::Future {
195            let saw_reservation = req
196                .extensions()
197                .get::<PreDecodeMemoryReservation>()
198                .is_some();
199            let saw_limiter = req.extensions().get::<ServerMemoryLimiter>().is_some();
200            Box::pin(async move {
201                Ok((
202                    StatusCode::OK,
203                    format!("reservation={saw_reservation} limiter={saw_limiter}"),
204                )
205                    .into_response())
206            })
207        }
208    }
209
210    impl NamedService for EchoService {
211        const NAME: &'static str = "test.Echo";
212    }
213
214    #[tokio::test]
215    async fn test_inserts_limiter_for_uncompressed_requests() {
216        let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail);
217        let mut svc =
218            MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService);
219
220        let req = Request::builder().body(Body::empty()).unwrap();
221        let res = svc.ready().await.unwrap().call(req).await.unwrap();
222        assert_eq!(res.status(), StatusCode::OK);
223        let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap();
224        assert_eq!(
225            &body[..],
226            &b"reservation=false limiter=true"[..],
227            "uncompressed requests must not be pre-reserved"
228        );
229        assert_eq!(0, limiter.used_bytes());
230    }
231
232    #[tokio::test]
233    async fn test_reserves_max_message_size_for_compressed_requests() {
234        let max_size = 512 * 1024 * 1024;
235        let limiter = ServerMemoryLimiter::new(1024 * 1024 * 1024, OnExhaustedPolicy::Fail);
236        let mut svc =
237            MemoryLimiterExtensionLayer::new(limiter.clone(), max_size).layer(EchoService);
238
239        // The inner service observes the reservation; use a oneshot call and
240        // check the limiter was charged while the call is in flight via the
241        // response extensions observed by the echo service.
242        let req = Request::builder()
243            .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("gzip"))
244            .body(Body::empty())
245            .unwrap();
246        let res = svc.ready().await.unwrap().call(req).await.unwrap();
247        assert_eq!(res.status(), StatusCode::OK);
248        let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap();
249        assert_eq!(
250            &body[..],
251            &b"reservation=true limiter=true"[..],
252            "compressed requests must carry a pre-decode reservation"
253        );
254        // The reservation is released once the inner call (request) finishes.
255        assert_eq!(0, limiter.used_bytes());
256    }
257
258    #[tokio::test]
259    async fn test_rejects_compressed_request_when_quota_exhausted() {
260        // Quota smaller than the decoding limit: the reservation cannot fit.
261        let limiter = ServerMemoryLimiter::new(1024, OnExhaustedPolicy::Fail);
262        let mut svc =
263            MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService);
264
265        let req = Request::builder()
266            .header(GRPC_ENCODING_HEADER, HeaderValue::from_static("zstd"))
267            .body(Body::empty())
268            .unwrap();
269        let res = svc.ready().await.unwrap().call(req).await.unwrap();
270        assert_eq!(res.status(), StatusCode::OK);
271        // gRPC errors are HTTP 200 with a grpc-status header.
272        assert_eq!(
273            res.headers()
274                .get("grpc-status")
275                .and_then(|v| v.to_str().ok()),
276            Some("8"),
277            "expected RESOURCE_EXHAUSTED grpc status"
278        );
279    }
280
281    #[tokio::test]
282    async fn test_identity_encoding_is_not_reserved() {
283        let limiter = ServerMemoryLimiter::new(1024 * 1024, OnExhaustedPolicy::Fail);
284        let mut svc =
285            MemoryLimiterExtensionLayer::new(limiter.clone(), 512 * 1024 * 1024).layer(EchoService);
286
287        let req = Request::builder()
288            .header(
289                GRPC_ENCODING_HEADER,
290                HeaderValue::from_static(IDENTITY_ENCODING),
291            )
292            .body(Body::empty())
293            .unwrap();
294        let res = svc.ready().await.unwrap().call(req).await.unwrap();
295        let body = axum::body::to_bytes(res.into_body(), 1024).await.unwrap();
296        assert_eq!(&body[..], &b"reservation=false limiter=true"[..]);
297    }
298}
servers/grpc/memory_limit.rs

servers/grpc/
memory_limit.rs
