apollo-opentelemetry 0.8.0

use super::Clock;
use opentelemetry::KeyValue;
use opentelemetry::metrics::Histogram;
use quanta::Instant;

/// Records elapsed time in seconds on a [`Histogram<f64>`] on drop.
///
/// Attributes can be added or updated at any point before drop via [`set`](Self::set).
/// Recording can be suppressed via [`cancel`](Self::cancel).
///
/// # Example
///
/// ```
/// # use opentelemetry::metrics::Histogram;
/// # use opentelemetry::KeyValue;
/// # use apollo_opentelemetry::metrics::HistogramExt;
/// # async fn process() {}
/// async fn handle(histogram: Histogram<f64>) {
///     let mut guard = histogram.record_duration_on_drop([KeyValue::new("http.method", "GET")]);
///
///     process().await; // future may be cancelled here — guard still records duration on drop
///
///     guard.set(KeyValue::new("http.status", 200_i64));
/// }
/// # fn main() {}
/// ```
#[must_use]
pub struct RecordDurationGuard {
    histogram: Histogram<f64>,
    clock: Clock,
    start: Instant,
    attributes: Vec<KeyValue>,
    cancelled: bool,
}

impl RecordDurationGuard {
    /// Create a new timer using wall-clock time.
    pub fn new(histogram: Histogram<f64>, attributes: impl Into<Vec<KeyValue>>) -> Self {
        Self::with_clock(histogram, Clock::new(), attributes)
    }

    /// Create a new timer with a custom clock. Useful for testing.
    pub fn with_clock(
        histogram: Histogram<f64>,
        clock: Clock,
        attributes: impl Into<Vec<KeyValue>>,
    ) -> Self {
        let start = clock.now();
        Self {
            histogram,
            clock,
            start,
            attributes: attributes.into(),
            cancelled: false,
        }
    }

    /// Add or update an attribute that will be included when the duration is recorded.
    ///
    /// If an attribute with the same key already exists, it is replaced.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use opentelemetry::KeyValue;
    /// # use apollo_opentelemetry::metrics::HistogramExt;
    /// # let histogram = global::meter_provider().meter("example").f64_histogram("request.duration").build();
    /// let mut guard = histogram.record_duration_on_drop([]);
    /// guard.set(KeyValue::new("http.status", 200_i64));
    /// ```
    pub fn set(&mut self, kv: KeyValue) {
        if let Some(existing) = self.attributes.iter_mut().find(|a| a.key == kv.key) {
            *existing = kv;
        } else {
            self.attributes.push(kv);
        }
    }

    /// Suppress recording on drop. Use when a different recorder takes over.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::HistogramExt;
    /// # let histogram = global::meter_provider().meter("example").f64_histogram("request.duration").build();
    /// let mut guard = histogram.record_duration_on_drop([]);
    /// if true /* handled by a more specific timer */ {
    ///     guard.cancel();
    /// }
    /// ```
    pub fn cancel(&mut self) {
        self.cancelled = true;
    }
}

impl Drop for RecordDurationGuard {
    fn drop(&mut self) {
        if !self.cancelled {
            // Can't use `.elapsed()` because it will use the global clock
            let elapsed = (self.clock.now() - self.start).as_secs_f64();
            self.histogram.record(elapsed, &self.attributes);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use apollo_opentelemetry_test::{TelemetryContext, assert_metrics_snapshot};
    use opentelemetry::KeyValue;
    use opentelemetry::global;
    use quanta::Clock;

    use super::RecordDurationGuard;

    #[test]
    fn records_on_drop() {
        let ctx = TelemetryContext::new();
        let histogram = global::meter_provider()
            .meter("test")
            .f64_histogram("test.duration")
            .build();

        let (clock, mock) = Clock::mock();
        {
            let _timer = RecordDurationGuard::with_clock(histogram, clock, []);
            mock.increment(Duration::from_millis(1500));
        }

        assert_metrics_snapshot!(ctx, @r"
        - name: test.duration
          data:
            type: Histogram
            data_points:
              - count: 1
                sum: 1.5
                min: 1.5
                max: 1.5
                bounds:
                  - 0
                  - 5
                  - 10
                  - 25
                  - 50
                  - 75
                  - 100
                  - 250
                  - 500
                  - 750
                  - 1000
                  - 2500
                  - 5000
                  - 7500
                  - 10000
                bucket_counts:
                  - 0
                  - 1
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
            temporality: Cumulative
        ");
    }

    #[test]
    fn cancelled_does_not_record() {
        let ctx = TelemetryContext::new();
        let histogram = global::meter_provider()
            .meter("test")
            .f64_histogram("test.duration")
            .build();

        let (clock, mock) = Clock::mock();
        {
            let mut timer = RecordDurationGuard::with_clock(histogram, clock, []);
            mock.increment(Duration::from_millis(1500));
            timer.cancel();
        }

        assert_metrics_snapshot!(ctx, @"[]");
    }

    #[test]
    fn set_deduplicates_attribute_by_key() {
        let ctx = TelemetryContext::new();
        let histogram = global::meter_provider()
            .meter("test")
            .f64_histogram("test.duration")
            .build();

        let (clock, mock) = Clock::mock();
        {
            let mut timer = RecordDurationGuard::with_clock(histogram, clock, []);
            mock.increment(Duration::from_millis(1500));
            timer.set(KeyValue::new("http.method", "GET"));
            timer.set(KeyValue::new("http.method", "POST"));
        }

        assert_metrics_snapshot!(ctx, @r"
        - name: test.duration
          data:
            type: Histogram
            data_points:
              - attributes:
                  http.method: POST
                count: 1
                sum: 1.5
                min: 1.5
                max: 1.5
                bounds:
                  - 0
                  - 5
                  - 10
                  - 25
                  - 50
                  - 75
                  - 100
                  - 250
                  - 500
                  - 750
                  - 1000
                  - 2500
                  - 5000
                  - 7500
                  - 10000
                bucket_counts:
                  - 0
                  - 1
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
                  - 0
            temporality: Cumulative
        ");
    }
}