1use std::time::{Duration, Instant};
2
3use tonic::{Request, metadata::MetadataMap};
4
5const GRPC_TIMEOUT: &str = "grpc-timeout";
6
7#[derive(Debug, Clone)]
9pub struct DeadlineBudget {
10 started_at: Instant,
11 timeout: Duration,
12}
13
14impl DeadlineBudget {
15 pub fn new(timeout: Duration) -> Self {
17 Self {
18 started_at: Instant::now(),
19 timeout,
20 }
21 }
22
23 pub fn remaining(&self) -> Duration {
25 self.timeout.saturating_sub(self.started_at.elapsed())
26 }
27
28 pub fn expired(&self) -> bool {
30 self.remaining().is_zero()
31 }
32}
33
34pub 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
45pub 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}