yumechi-no-kuni-proxy-worker/src/timing.rs

use axum::{
    http::{HeaderName, HeaderValue},
    response::IntoResponse,
};

/// A cross-platform `Instant` implementation for measuring time
#[cfg(not(target_arch = "wasm32"))]
pub struct Instant(std::time::Instant);

impl Instant {
    /// Create a new `Instant` from the current time
    #[cfg(not(target_arch = "wasm32"))]
    #[must_use] pub fn now() -> Self {
        Self(std::time::Instant::now())
    }

    /// Get the elapsed time since the instant was created
    #[cfg(not(target_arch = "wasm32"))]
    #[must_use] pub fn elapsed(&self) -> std::time::Duration {
        self.0.elapsed()
    }
}

#[cfg(all(target_arch = "wasm32", not(feature = "cf-worker")))]
mod js_performance {
    use wasm_bindgen::prelude::*;
    #[wasm_bindgen]
    extern "C" {
        #[wasm_bindgen(js_namespace = performance)]
        pub fn now() -> f64;
    }
}

/// A cross-platform `Instant` implementation for measuring time
#[cfg(target_arch = "wasm32")]
pub struct Instant(f64);

#[cfg(all(target_arch = "wasm32", not(feature = "cf-worker")))]
impl Instant {
    /// Create a new `Instant` from the current time
    pub fn now() -> Self {
        Self(js_performance::now())
    }

    /// Get the elapsed time since the instant was created
    pub fn elapsed(&self) -> std::time::Duration {
        std::time::Duration::from_secs_f64((js_performance::now() - self.0) / 1000.0)
    }
}

#[cfg(all(target_arch = "wasm32", feature = "cf-worker"))]
impl Instant {
    /// Create a new `Instant` from the current time
    pub fn now() -> Self {
        Self(worker::Date::now().as_millis() as f64)
    }

    /// Get the elapsed time since the instant was created
    pub fn elapsed(&self) -> std::time::Duration {
        std::time::Duration::from_secs_f64(
            (worker::Date::now().as_millis() as f64 - self.0) / 1000.0,
        )
    }
}

/// A response that includes timing information
pub struct WithTimingInfo<T> {
    pub(crate) key: &'static str,
    pub(crate) inner: T,
    pub(crate) dur: std::time::Duration,
}

impl<T> WithTimingInfo<T> {
    pub(crate) fn new(key: &'static str, inner: T, dur: std::time::Duration) -> Self {
        Self { key, inner, dur }
    }

    /// Get the duration of the response
    pub fn duration(&self) -> std::time::Duration {
        self.dur
    }
}

impl<T: IntoResponse> IntoResponse for WithTimingInfo<T> {
    fn into_response(self) -> axum::response::Response {
        let mut res = self.inner.into_response();

        res.headers_mut().insert(
            HeaderName::from_lowercase(format!("x-timing-{}-ms", self.key).as_bytes()).unwrap(),
            if cfg!(feature = "cf-worker") && self.dur.is_zero() {
                HeaderValue::from_static("N/A") // Cloudflare redacts some timing information
            } else {
                self.dur.as_millis().to_string().parse().unwrap()
            },
        );

        res
    }
}

/// A response that includes timing information, if available
pub struct WithMaybeTimingInfo<T> {
    pub(crate) key: &'static str,
    pub(crate) inner: T,
    pub(crate) dur: Option<std::time::Duration>,
}

impl<T> WithMaybeTimingInfo<T> {
    pub(crate) fn new(key: &'static str, inner: T, dur: Option<std::time::Duration>) -> Self {
        Self { key, inner, dur }
    }

    /// Get the duration of the response
    pub fn duration(&self) -> Option<std::time::Duration> {
        self.dur
    }
}

impl<T: IntoResponse> IntoResponse for WithMaybeTimingInfo<T> {
    fn into_response(self) -> axum::response::Response {
        let mut res = self.inner.into_response();

        if let Some(dur) = self.dur {
            res.headers_mut().insert(
                HeaderName::from_lowercase(format!("x-timing-{}-ms", self.key).as_bytes()).unwrap(),
                if cfg!(feature = "cf-worker") && dur.is_zero() {
                    HeaderValue::from_static("N/A") // Cloudflare redacts some timing information
                } else {
                    dur.as_millis().to_string().parse().unwrap()
                },
            );
        }

        res
    }
}

pub(crate) trait IntoResponseExt: Sized {
    fn with_opt_timing_info(
        self,
        key: &'static str,
        dur: Option<std::time::Duration>,
    ) -> WithMaybeTimingInfo<Self>;
    fn with_timing_info(self, key: &'static str, dur: std::time::Duration) -> WithTimingInfo<Self>;
}

impl<T: IntoResponse> IntoResponseExt for T {
    fn with_opt_timing_info(
        self,
        key: &'static str,
        dur: Option<std::time::Duration>,
    ) -> WithMaybeTimingInfo<Self> {
        WithMaybeTimingInfo::new(key, self, dur)
    }
    fn with_timing_info(self, key: &'static str, dur: std::time::Duration) -> WithTimingInfo<Self> {
        WithTimingInfo::new(key, self, dur)
    }
}