rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use std::time::{Duration, Instant};

use tonic::{Request, metadata::MetadataMap};

const GRPC_TIMEOUT: &str = "grpc-timeout";

/// Request deadline budget tracked across retries and outbound calls.
#[derive(Debug, Clone)]
pub struct DeadlineBudget {
    started_at: Instant,
    timeout: Duration,
}

impl DeadlineBudget {
    /// Creates a budget from a timeout.
    pub fn new(timeout: Duration) -> Self {
        Self {
            started_at: Instant::now(),
            timeout,
        }
    }

    /// Returns the remaining duration, saturating at zero.
    pub fn remaining(&self) -> Duration {
        self.timeout.saturating_sub(self.started_at.elapsed())
    }

    /// Returns whether the budget is exhausted.
    pub fn expired(&self) -> bool {
        self.remaining().is_zero()
    }
}

/// Inserts a `grpc-timeout` metadata value into a tonic request.
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(())
}

/// Reads a `grpc-timeout` value from metadata.
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))
        );
    }
}