use std::time::{Duration, Instant};
use tonic::{Request, metadata::MetadataMap};
const GRPC_TIMEOUT: &str = "grpc-timeout";
#[derive(Debug, Clone)]
pub struct DeadlineBudget {
started_at: Instant,
timeout: Duration,
}
impl DeadlineBudget {
pub fn new(timeout: Duration) -> Self {
Self {
started_at: Instant::now(),
timeout,
}
}
pub fn remaining(&self) -> Duration {
self.timeout.saturating_sub(self.started_at.elapsed())
}
pub fn expired(&self) -> bool {
self.remaining().is_zero()
}
}
pub fn insert_grpc_timeout<T>(
request: &mut Request<T>,
duration: Duration,
) -> Result<(), tonic::metadata::errors::InvalidMetadataValue> {
request
.metadata_mut()
.insert(GRPC_TIMEOUT, encode_grpc_timeout(duration).parse()?);
Ok(())
}
pub fn remaining_timeout_from_metadata(metadata: &MetadataMap) -> Option<Duration> {
let value = metadata.get(GRPC_TIMEOUT)?.to_str().ok()?;
decode_grpc_timeout(value)
}
fn encode_grpc_timeout(duration: Duration) -> String {
let millis = duration.as_millis().clamp(1, 99_999_999);
format!("{millis}m")
}
fn decode_grpc_timeout(value: &str) -> Option<Duration> {
let (number, unit) = value.split_at(value.len().checked_sub(1)?);
let amount = number.parse::<u64>().ok()?;
match unit {
"H" => Some(Duration::from_secs(amount * 60 * 60)),
"M" => Some(Duration::from_secs(amount * 60)),
"S" => Some(Duration::from_secs(amount)),
"m" => Some(Duration::from_millis(amount)),
"u" => Some(Duration::from_micros(amount)),
"n" => Some(Duration::from_nanos(amount)),
_ => None,
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use tonic::Request;
use super::{insert_grpc_timeout, remaining_timeout_from_metadata};
#[test]
fn timeout_round_trips_metadata() {
let mut request = Request::new(());
insert_grpc_timeout(&mut request, Duration::from_millis(25)).expect("insert");
assert_eq!(
remaining_timeout_from_metadata(request.metadata()),
Some(Duration::from_millis(25))
);
}
}