Skip to main content

rs_zero/rpc/
deadline.rs

1use std::time::{Duration, Instant};
2
3use tonic::{Request, metadata::MetadataMap};
4
5const GRPC_TIMEOUT: &str = "grpc-timeout";
6
7/// Request deadline budget tracked across retries and outbound calls.
8#[derive(Debug, Clone)]
9pub struct DeadlineBudget {
10    started_at: Instant,
11    timeout: Duration,
12}
13
14impl DeadlineBudget {
15    /// Creates a budget from a timeout.
16    pub fn new(timeout: Duration) -> Self {
17        Self {
18            started_at: Instant::now(),
19            timeout,
20        }
21    }
22
23    /// Returns the remaining duration, saturating at zero.
24    pub fn remaining(&self) -> Duration {
25        self.timeout.saturating_sub(self.started_at.elapsed())
26    }
27
28    /// Returns whether the budget is exhausted.
29    pub fn expired(&self) -> bool {
30        self.remaining().is_zero()
31    }
32}
33
34/// Inserts a `grpc-timeout` metadata value into a tonic request.
35pub fn insert_grpc_timeout<T>(
36    request: &mut Request<T>,
37    duration: Duration,
38) -> Result<(), tonic::metadata::errors::InvalidMetadataValue> {
39    request
40        .metadata_mut()
41        .insert(GRPC_TIMEOUT, encode_grpc_timeout(duration).parse()?);
42    Ok(())
43}
44
45/// Reads a `grpc-timeout` value from metadata.
46pub fn remaining_timeout_from_metadata(metadata: &MetadataMap) -> Option<Duration> {
47    let value = metadata.get(GRPC_TIMEOUT)?.to_str().ok()?;
48    decode_grpc_timeout(value)
49}
50
51fn encode_grpc_timeout(duration: Duration) -> String {
52    let millis = duration.as_millis().clamp(1, 99_999_999);
53    format!("{millis}m")
54}
55
56fn decode_grpc_timeout(value: &str) -> Option<Duration> {
57    let (number, unit) = value.split_at(value.len().checked_sub(1)?);
58    let amount = number.parse::<u64>().ok()?;
59    match unit {
60        "H" => Some(Duration::from_secs(amount * 60 * 60)),
61        "M" => Some(Duration::from_secs(amount * 60)),
62        "S" => Some(Duration::from_secs(amount)),
63        "m" => Some(Duration::from_millis(amount)),
64        "u" => Some(Duration::from_micros(amount)),
65        "n" => Some(Duration::from_nanos(amount)),
66        _ => None,
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use std::time::Duration;
73
74    use tonic::Request;
75
76    use super::{insert_grpc_timeout, remaining_timeout_from_metadata};
77
78    #[test]
79    fn timeout_round_trips_metadata() {
80        let mut request = Request::new(());
81        insert_grpc_timeout(&mut request, Duration::from_millis(25)).expect("insert");
82
83        assert_eq!(
84            remaining_timeout_from_metadata(request.metadata()),
85            Some(Duration::from_millis(25))
86        );
87    }
88}