gcloud_sdk/
request_timeouts.rs

1use std::any::Any;
2use std::fmt;
3use std::future::Future;
4use std::pin::{pin, Pin};
5use std::sync::Arc;
6use std::task::{Context, Poll};
7use std::time::Duration;
8
9use bytes::Bytes;
10use futures::future::{select, Either};
11use hyper::body::{Body, Frame, SizeHint};
12use hyper::header::{HeaderMap, HeaderName, HeaderValue};
13use hyper::Uri;
14use tokio::time::{sleep, sleep_until, Instant, Sleep};
15use tonic::Status;
16use tracing::error;
17
18const GRPC_TIMEOUT: HeaderName = HeaderName::from_static("grpc-timeout");
19
20/// The largest `grpc-timeout` value: the gRPC spec allows at most eight digits.
21const MAX_GRPC_TIMEOUT_VALUE: u128 = 99_999_999;
22
23/// The `grpc-timeout` units from the finest, each with its length in nanoseconds.
24const GRPC_TIMEOUT_UNITS: [(char, u128); 6] = [
25    ('n', 1),
26    ('u', 1_000),
27    ('m', 1_000_000),
28    ('S', 1_000_000_000),
29    ('M', 60 * 1_000_000_000),
30    ('H', 60 * 60 * 1_000_000_000),
31];
32
33/// Client-side time limits for a gRPC call. Both are off by default.
34///
35/// Set them for every call of a client with
36/// [`GoogleApiClient::with_request_timeouts`](crate::GoogleApiClient::with_request_timeouts),
37/// or for one call by inserting a value into the request's extensions, which replaces both
38/// limits of the client for that call, so `RequestTimeouts::new()` switches them off:
39///
40/// ```rust,no_run
41/// # use std::time::Duration;
42/// # use gcloud_sdk::RequestTimeouts;
43/// # fn sketch<M>(message: M) {
44/// let mut request: tonic::Request<M> = tonic::Request::new(message);
45/// request
46///     .extensions_mut()
47///     .insert(RequestTimeouts::new().with_total(Duration::from_secs(10)));
48/// # }
49/// ```
50///
51/// A call that runs out of a limit fails with
52/// [`Code::DeadlineExceeded`](tonic::Code::DeadlineExceeded), and the status's `source()` is
53/// a [`RequestTimeoutElapsed`].
54///
55/// The limits bound the response body only when the service under the middleware answers
56/// with a [`tonic::body::Body`], as the channel of a
57/// [`GoogleAuthMiddleware`](crate::GoogleAuthMiddleware) does. Over any other service they
58/// bound the wait for the response headers alone.
59#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
60pub struct RequestTimeouts {
61    total: Option<Duration>,
62    idle: Option<Duration>,
63}
64
65impl RequestTimeouts {
66    /// No limits: a call waits for as long as its connection lives.
67    pub const fn new() -> Self {
68        Self {
69            total: None,
70            idle: None,
71        }
72    }
73
74    /// The longest the whole call may take, from the request, its authentication included,
75    /// to the last response message and the trailers.
76    ///
77    /// Sent to the server as `grpc-timeout`, so a server that honours it stops working on
78    /// the call too; a `grpc-timeout` the request already carries is kept. It ends a stream
79    /// however steadily the stream delivers, so it suits unary calls.
80    pub const fn with_total(self, total: Duration) -> Self {
81        Self {
82            total: Some(total),
83            ..self
84        }
85    }
86
87    /// The longest the caller waits for the response to start and for each next part of
88    /// the response, counted only while the caller is waiting for one. Not sent to the
89    /// server.
90    ///
91    /// It suits server streams. A stream that stays quiet until the client writes, such as a
92    /// bidirectional stream, fails on it.
93    pub const fn with_idle(self, idle: Duration) -> Self {
94        Self {
95            idle: Some(idle),
96            ..self
97        }
98    }
99
100    pub(crate) fn start(self) -> Option<CallClock> {
101        if self.total.is_none() && self.idle.is_none() {
102            return None;
103        }
104        let started = Instant::now();
105        Some(CallClock {
106            started,
107            total: self
108                .total
109                .and_then(|limit| Some((limit, started.checked_add(limit)?))),
110            idle: self.idle,
111        })
112    }
113}
114
115/// The limit of a [`RequestTimeouts`] that a call ran out of: the source of the
116/// [`Code::DeadlineExceeded`](tonic::Code::DeadlineExceeded) status the call fails with.
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118#[non_exhaustive]
119pub enum RequestTimeoutElapsed {
120    /// The total limit, set with [`RequestTimeouts::with_total`].
121    Total(Duration),
122    /// The idle limit, set with [`RequestTimeouts::with_idle`].
123    Idle(Duration),
124}
125
126impl fmt::Display for RequestTimeoutElapsed {
127    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
128        match self {
129            Self::Total(limit) => write!(f, "total request timeout of {limit:?} elapsed"),
130            Self::Idle(limit) => write!(f, "idle request timeout of {limit:?} elapsed"),
131        }
132    }
133}
134
135impl std::error::Error for RequestTimeoutElapsed {}
136
137impl From<RequestTimeoutElapsed> for Status {
138    fn from(elapsed: RequestTimeoutElapsed) -> Self {
139        let mut status = Status::deadline_exceeded(elapsed.to_string());
140        status.set_source(Arc::new(elapsed));
141        status
142    }
143}
144
145/// The limits of one call, measured from the moment it entered the middleware. A limit too
146/// large to add to that moment is no limit.
147#[derive(Clone, Copy)]
148pub(crate) struct CallClock {
149    started: Instant,
150    /// The total limit and the moment it runs out.
151    total: Option<(Duration, Instant)>,
152    idle: Option<Duration>,
153}
154
155impl CallClock {
156    /// Sends the time left of the total limit as `grpc-timeout`, unless the request carries
157    /// a `grpc-timeout` of its own, and fails the call when no time is left.
158    pub(crate) fn insert_grpc_timeout(&self, headers: &mut HeaderMap) -> Result<(), Status> {
159        let Some((limit, deadline)) = self.total else {
160            return Ok(());
161        };
162        let remaining = deadline.saturating_duration_since(Instant::now());
163        if remaining.is_zero() {
164            return Err(RequestTimeoutElapsed::Total(limit).into());
165        }
166        if let Some(value) = grpc_timeout_value(remaining) {
167            headers.entry(GRPC_TIMEOUT).or_insert(value);
168        }
169        Ok(())
170    }
171
172    /// Runs `send`, the call to `uri` up to its response headers, within the earlier of the
173    /// two limits, and then bounds the response body by both.
174    pub(crate) async fn limit<B: 'static, E: From<Status>>(
175        self,
176        uri: Uri,
177        send: impl Future<Output = Result<hyper::Response<B>, E>>,
178    ) -> Result<hyper::Response<B>, E> {
179        let total = self
180            .total
181            .map(|(limit, deadline)| (deadline, RequestTimeoutElapsed::Total(limit)));
182        let idle = self.idle.and_then(|limit| {
183            Some((
184                self.started.checked_add(limit)?,
185                RequestTimeoutElapsed::Idle(limit),
186            ))
187        });
188        let response = match total.into_iter().chain(idle).min_by_key(|(at, _)| *at) {
189            // The limit is polled first: with `total` set, tonic's channel races a deadline
190            // read from `grpc-timeout`, no earlier than this one, and fails the call
191            // `Cancelled`.
192            Some((deadline, elapsed)) => {
193                match select(pin!(sleep_until(deadline)), pin!(send)).await {
194                    Either::Left(_) => {
195                        error!(
196                            req_uri = %uri,
197                            "Err: took {}ms ({elapsed})",
198                            self.started.elapsed().as_millis()
199                        );
200                        return Err(Status::from(elapsed).into());
201                    }
202                    Either::Right((response, _)) => response?,
203                }
204            }
205            None => send.await?,
206        };
207        Ok(response.map(|body| self.limit_body(body)))
208    }
209
210    /// The response keeps the inner service's body type, so only the `tonic::body::Body` a
211    /// tonic channel answers with can be replaced by a limited one.
212    fn limit_body<B: 'static>(self, mut body: B) -> B {
213        if let Some(body) = (&mut body as &mut dyn Any).downcast_mut::<tonic::body::Body>() {
214            *body = tonic::body::Body::new(TimeoutBody {
215                inner: std::mem::take(body),
216                total: self
217                    .total
218                    .map(|(limit, deadline)| (limit, Box::pin(sleep_until(deadline)))),
219                idle: self.idle.map(|limit| IdleTimer {
220                    limit,
221                    sleep: Box::pin(sleep(limit)),
222                    waiting: false,
223                }),
224            });
225        }
226        body
227    }
228}
229
230/// `timeout` as a `grpc-timeout` value in the finest unit that fits, rounded up, so that a
231/// deadline read from the header never comes before the one it was written from. `None`
232/// past eight digits of hours.
233fn grpc_timeout_value(timeout: Duration) -> Option<HeaderValue> {
234    let nanos = timeout.as_nanos();
235    let (value, unit) = GRPC_TIMEOUT_UNITS
236        .iter()
237        .map(|&(unit, unit_nanos)| (nanos.div_ceil(unit_nanos), unit))
238        .find(|&(value, _)| value <= MAX_GRPC_TIMEOUT_VALUE)?;
239    HeaderValue::try_from(format!("{value}{unit}")).ok()
240}
241
242struct TimeoutBody {
243    inner: tonic::body::Body,
244    /// The total limit and the sleep until it runs out.
245    total: Option<(Duration, Pin<Box<Sleep>>)>,
246    idle: Option<IdleTimer>,
247}
248
249/// The idle limit of a response body, armed when the caller starts waiting for a frame and
250/// disarmed when one arrives, so that a caller busy between frames never trips it.
251struct IdleTimer {
252    limit: Duration,
253    sleep: Pin<Box<Sleep>>,
254    waiting: bool,
255}
256
257impl Body for TimeoutBody {
258    type Data = Bytes;
259    type Error = Status;
260
261    fn poll_frame(
262        self: Pin<&mut Self>,
263        cx: &mut Context<'_>,
264    ) -> Poll<Option<Result<Frame<Bytes>, Status>>> {
265        let this = self.get_mut();
266        if let Some((limit, deadline)) = &mut this.total {
267            if deadline.as_mut().poll(cx).is_ready() {
268                return Poll::Ready(Some(Err(RequestTimeoutElapsed::Total(*limit).into())));
269            }
270        }
271        if let Some(idle) = &mut this.idle {
272            if !idle.waiting {
273                // A limit too large to add to the clock leaves the sleep at tokio's far
274                // future, where `sleep(limit)` put it.
275                if let Some(deadline) = Instant::now().checked_add(idle.limit) {
276                    idle.sleep.as_mut().reset(deadline);
277                }
278                idle.waiting = true;
279            }
280        }
281        match Pin::new(&mut this.inner).poll_frame(cx) {
282            Poll::Ready(None) => {
283                // A body polled again after its end must not fail on a limit.
284                this.total = None;
285                this.idle = None;
286                Poll::Ready(None)
287            }
288            Poll::Ready(frame) => {
289                if let Some(idle) = &mut this.idle {
290                    idle.waiting = false;
291                }
292                Poll::Ready(frame)
293            }
294            Poll::Pending => {
295                if let Some(idle) = &mut this.idle {
296                    if idle.sleep.as_mut().poll(cx).is_ready() {
297                        return Poll::Ready(Some(Err(
298                            RequestTimeoutElapsed::Idle(idle.limit).into()
299                        )));
300                    }
301                }
302                Poll::Pending
303            }
304        }
305    }
306
307    fn is_end_stream(&self) -> bool {
308        self.inner.is_end_stream()
309    }
310
311    fn size_hint(&self) -> SizeHint {
312        self.inner.size_hint()
313    }
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319    use crate::middleware::GoogleAuthMiddlewareService;
320    use crate::test_support::{silent_endpoint, stalling_stream_channel, StubCredentials};
321    use crate::{GoogleApi, GoogleAuthHeaders, GoogleAuthMiddleware};
322    use google_cloud_auth::credentials::Credentials;
323    use hyper::http::uri::PathAndQuery;
324    use std::error::Error as _;
325    use std::future::poll_fn;
326    use tokio::time::Instant;
327    use tonic::client::Grpc;
328    use tonic::codec::Streaming;
329    use tonic::Code;
330    use tonic_prost::ProstCodec;
331
332    const LIMIT: Duration = Duration::from_secs(5);
333
334    async fn unary(
335        service: GoogleAuthMiddleware,
336        request: tonic::Request<()>,
337    ) -> Result<tonic::Response<()>, Status> {
338        let mut grpc = Grpc::new(service);
339        grpc.ready().await.unwrap();
340        grpc.unary(
341            request,
342            PathAndQuery::from_static("/test.Silent/Call"),
343            ProstCodec::default(),
344        )
345        .await
346    }
347
348    async fn open_stream(
349        service: GoogleAuthMiddleware,
350        request: tonic::Request<()>,
351    ) -> Streaming<()> {
352        let mut grpc = Grpc::new(service);
353        grpc.ready().await.unwrap();
354        grpc.server_streaming(
355            request,
356            PathAndQuery::from_static("/test.StallingStream/Stream"),
357            ProstCodec::default(),
358        )
359        .await
360        .unwrap()
361        .into_inner()
362    }
363
364    #[tokio::test(start_paused = true)]
365    async fn total_limit_fails_a_call_the_server_never_answers() {
366        let (url, _) = silent_endpoint().await;
367        let client: GoogleApi<GoogleAuthMiddleware> = GoogleApi::from_function_with_credentials(
368            |service| service,
369            url,
370            None,
371            Credentials::from(StubCredentials::bearer("token")),
372        )
373        .await
374        .unwrap()
375        .with_request_timeouts(RequestTimeouts::new().with_total(LIMIT));
376
377        let started = Instant::now();
378        let status = unary(client.get(), tonic::Request::new(()))
379            .await
380            .unwrap_err();
381
382        assert_eq!(started.elapsed(), LIMIT);
383        assert_eq!(status.code(), Code::DeadlineExceeded);
384        assert_eq!(
385            status.source().and_then(|source| source.downcast_ref()),
386            Some(&RequestTimeoutElapsed::Total(LIMIT))
387        );
388    }
389
390    /// Past 100 s tonic's own encoder drops the microseconds of `grpc-timeout`, and the
391    /// channel would end the call on that shorter deadline, `Cancelled`.
392    #[tokio::test(start_paused = true)]
393    async fn total_limit_past_milliseconds_fails_with_deadline_exceeded() {
394        let limit = Duration::from_micros(110_000_500);
395        let (url, _) = silent_endpoint().await;
396        let client: GoogleApi<GoogleAuthMiddleware> = GoogleApi::from_function_with_credentials(
397            |service| service,
398            url,
399            None,
400            Credentials::from(StubCredentials::bearer("token")),
401        )
402        .await
403        .unwrap()
404        .with_request_timeouts(RequestTimeouts::new().with_total(limit));
405
406        let started = Instant::now();
407        let status = unary(client.get(), tonic::Request::new(()))
408            .await
409            .unwrap_err();
410
411        // tokio's timers fire on whole milliseconds.
412        assert!((limit..limit + Duration::from_millis(1)).contains(&started.elapsed()));
413        assert_eq!(status.code(), Code::DeadlineExceeded);
414        assert_eq!(
415            status.source().and_then(|source| source.downcast_ref()),
416            Some(&RequestTimeoutElapsed::Total(limit))
417        );
418    }
419
420    #[tokio::test(start_paused = true)]
421    async fn total_limit_ends_a_stream_that_keeps_delivering() {
422        let service = GoogleAuthMiddlewareService::new(
423            stalling_stream_channel(vec![Duration::from_secs(2); 10]).await,
424            GoogleAuthHeaders::from(Credentials::from(StubCredentials::bearer("token"))),
425            None,
426        )
427        .unwrap();
428        let mut request = tonic::Request::new(());
429        request
430            .extensions_mut()
431            .insert(RequestTimeouts::new().with_total(LIMIT));
432
433        let started = Instant::now();
434        let mut messages = open_stream(service, request).await;
435        let status = tokio::time::timeout(LIMIT * 2, async {
436            while messages.message().await?.is_some() {}
437            Ok::<_, Status>(())
438        })
439        .await
440        .expect("the total limit ends the stream")
441        .unwrap_err();
442
443        assert_eq!(started.elapsed(), LIMIT);
444        assert_eq!(status.code(), Code::DeadlineExceeded);
445        assert_eq!(
446            status.source().and_then(|source| source.downcast_ref()),
447            Some(&RequestTimeoutElapsed::Total(LIMIT))
448        );
449    }
450
451    /// A body that has ended stays ended, however long after its limits it is polled.
452    #[tokio::test(start_paused = true)]
453    async fn limits_end_with_the_response_body() {
454        let clock = RequestTimeouts::new()
455            .with_total(LIMIT)
456            .with_idle(LIMIT)
457            .start()
458            .unwrap();
459        let mut body = clock.limit_body(tonic::body::Body::new(String::from("message")));
460
461        assert!(poll_fn(|cx| Pin::new(&mut body).poll_frame(cx))
462            .await
463            .is_some_and(|frame| frame.is_ok()));
464        assert!(poll_fn(|cx| Pin::new(&mut body).poll_frame(cx))
465            .await
466            .is_none());
467        tokio::time::sleep(LIMIT * 2).await;
468        assert!(poll_fn(|cx| Pin::new(&mut body).poll_frame(cx))
469            .await
470            .is_none());
471    }
472
473    #[tokio::test(start_paused = true)]
474    async fn idle_limit_fails_a_stream_that_stops_after_its_first_message() {
475        let service = GoogleAuthMiddlewareService::new(
476            stalling_stream_channel(vec![Duration::ZERO]).await,
477            GoogleAuthHeaders::from(Credentials::from(StubCredentials::bearer("token"))),
478            None,
479        )
480        .unwrap();
481        let mut request = tonic::Request::new(());
482        request
483            .extensions_mut()
484            .insert(RequestTimeouts::new().with_idle(LIMIT));
485        let mut messages = open_stream(service, request).await;
486        assert_eq!(messages.message().await.unwrap(), Some(()));
487
488        let started = Instant::now();
489        let status = tokio::time::timeout(LIMIT * 2, messages.message())
490            .await
491            .expect("the idle limit ends the wait")
492            .unwrap_err();
493
494        assert_eq!(started.elapsed(), LIMIT);
495        assert_eq!(status.code(), Code::DeadlineExceeded);
496        assert_eq!(
497            status.source().and_then(|source| source.downcast_ref()),
498            Some(&RequestTimeoutElapsed::Idle(LIMIT))
499        );
500    }
501
502    /// The idle clock runs only while the caller waits: no wait here is longer than 3 s,
503    /// although the stream runs for 13 s and the server is silent for 7 s while the reader
504    /// is busy. The request's limits replace the service's total limit.
505    #[tokio::test(start_paused = true)]
506    async fn idle_limit_spares_a_slow_stream_and_a_slow_reader() {
507        let gaps = [3, 3, 7].map(Duration::from_secs).to_vec();
508        let mut service = GoogleAuthMiddlewareService::new(
509            stalling_stream_channel(gaps).await,
510            GoogleAuthHeaders::from(Credentials::from(StubCredentials::bearer("token"))),
511            None,
512        )
513        .unwrap();
514        service.set_request_timeouts(RequestTimeouts::new().with_total(LIMIT));
515        let mut request = tonic::Request::new(());
516        request
517            .extensions_mut()
518            .insert(RequestTimeouts::new().with_idle(LIMIT));
519        let mut messages = open_stream(service, request).await;
520
521        assert_eq!(messages.message().await.unwrap(), Some(()));
522        assert_eq!(messages.message().await.unwrap(), Some(()));
523        tokio::time::sleep(Duration::from_secs(4)).await;
524        assert_eq!(messages.message().await.unwrap(), Some(()));
525    }
526
527    /// Nothing set, nothing limited: a call to a server that never answers is still waiting
528    /// a minute later, when only the connection's keep-alive can end it.
529    #[tokio::test(start_paused = true)]
530    async fn without_limits_a_call_waits_for_the_connection() {
531        let (url, _) = silent_endpoint().await;
532        let client: GoogleApi<GoogleAuthMiddleware> = GoogleApi::from_function_with_credentials(
533            |service| service,
534            url,
535            None,
536            Credentials::from(StubCredentials::bearer("token")),
537        )
538        .await
539        .unwrap();
540
541        let call = unary(client.get(), tonic::Request::new(()));
542
543        assert!(tokio::time::timeout(Duration::from_secs(60), call)
544            .await
545            .is_err());
546    }
547}
gcloud_sdk/request_timeouts.rs

gcloud_sdk/
request_timeouts.rs
