servers/http/
timeout.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
15use std::future::Future;
16use std::pin::Pin;
17use std::task::{Context, Poll};
18use std::time::Duration;
19
20use axum::body::Body;
21use axum::http::Request;
22use axum::response::Response;
23use http::StatusCode;
24use pin_project::pin_project;
25use tokio::time::{Instant, Sleep};
26use tower::{Layer, Service};
27
28use crate::http::header::constants::GREPTIME_DB_HEADER_TIMEOUT;
29
30/// [`Timeout`] response future
31///
32/// [`Timeout`]: crate::timeout::Timeout
33///
34/// Modified from https://github.com/tower-rs/tower-http/blob/tower-http-0.5.2/tower-http/src/timeout/service.rs
35#[derive(Debug)]
36#[pin_project]
37pub struct ResponseFuture<T> {
38    #[pin]
39    inner: T,
40    #[pin]
41    sleep: Sleep,
42    status_code: StatusCode,
43}
44
45impl<T> ResponseFuture<T> {
46    pub(crate) fn new(inner: T, sleep: Sleep, status_code: StatusCode) -> Self {
47        ResponseFuture {
48            inner,
49            sleep,
50            status_code,
51        }
52    }
53}
54
55impl<F, E> Future for ResponseFuture<F>
56where
57    F: Future<Output = Result<Response, E>>,
58{
59    type Output = Result<Response, E>;
60
61    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
62        let this = self.project();
63
64        if this.sleep.poll(cx).is_ready() {
65            let mut res = Response::default();
66            *res.status_mut() = *this.status_code;
67            return Poll::Ready(Ok(res));
68        }
69
70        this.inner.poll(cx)
71    }
72}
73
74/// Applies a timeout to requests via the supplied inner service.
75///
76/// Modified from https://github.com/tower-rs/tower-http/blob/tower-http-0.5.2/tower-http/src/timeout/service.rs
77#[derive(Debug, Clone)]
78pub struct DynamicTimeoutLayer {
79    default_timeout: Duration,
80    status_code_fn: fn(&Request<Body>) -> StatusCode,
81}
82
83impl DynamicTimeoutLayer {
84    /// Create a timeout from a duration
85    pub fn new(default_timeout: Duration) -> Self {
86        DynamicTimeoutLayer {
87            default_timeout,
88            status_code_fn: |_| StatusCode::REQUEST_TIMEOUT,
89        }
90    }
91
92    /// Sets a function that selects the timeout response status for each request.
93    pub fn with_status_code_fn(mut self, status_code_fn: fn(&Request<Body>) -> StatusCode) -> Self {
94        self.status_code_fn = status_code_fn;
95        self
96    }
97}
98
99impl<S> Layer<S> for DynamicTimeoutLayer {
100    type Service = DynamicTimeout<S>;
101
102    fn layer(&self, service: S) -> Self::Service {
103        DynamicTimeout::new(service, self.default_timeout, self.status_code_fn)
104    }
105}
106
107/// Modified from https://github.com/tower-rs/tower-http/blob/tower-http-0.5.2/tower-http/src/timeout/service.rs
108#[derive(Clone)]
109pub struct DynamicTimeout<S> {
110    inner: S,
111    default_timeout: Duration,
112    status_code_fn: fn(&Request<Body>) -> StatusCode,
113}
114
115impl<S> DynamicTimeout<S> {
116    /// Create a new [`DynamicTimeout`] with the given timeout
117    pub fn new(
118        inner: S,
119        default_timeout: Duration,
120        status_code_fn: fn(&Request<Body>) -> StatusCode,
121    ) -> Self {
122        DynamicTimeout {
123            inner,
124            default_timeout,
125            status_code_fn,
126        }
127    }
128}
129
130impl<S> Service<Request<Body>> for DynamicTimeout<S>
131where
132    S: Service<Request<Body>, Response = Response> + Send + 'static,
133{
134    type Response = S::Response;
135    type Error = S::Error;
136    type Future = ResponseFuture<S::Future>;
137
138    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
139        match self.inner.poll_ready(cx) {
140            Poll::Pending => Poll::Pending,
141            Poll::Ready(r) => Poll::Ready(r),
142        }
143    }
144
145    fn call(&mut self, request: Request<Body>) -> Self::Future {
146        let status_code = (self.status_code_fn)(&request);
147        let timeout = request
148            .headers()
149            .get(GREPTIME_DB_HEADER_TIMEOUT)
150            .and_then(|value| {
151                value
152                    .to_str()
153                    .ok()
154                    .and_then(|value| humantime::parse_duration(value).ok())
155            })
156            .unwrap_or(self.default_timeout);
157        let response = self.inner.call(request);
158
159        if timeout.is_zero() {
160            // 30 years. See `Instant::far_future`.
161            let far_future = Instant::now() + Duration::from_secs(86400 * 365 * 30);
162            ResponseFuture::new(response, tokio::time::sleep_until(far_future), status_code)
163        } else {
164            let sleep = tokio::time::sleep(timeout);
165            ResponseFuture::new(response, sleep, status_code)
166        }
167    }
168}
servers/http/timeout.rs

servers/http/
timeout.rs
