apollo-opentelemetry 0.8.0

use super::Clock;
use opentelemetry::KeyValue;
use opentelemetry::metrics::{Counter, Histogram, UpDownCounter};

use super::{AddValueGuard, RecordDurationGuard, RecordValueGuard, TrackGuard};

/// Extension methods on [`Histogram<T>`] for creating RAII recording guards.
///
/// Import this trait to use [`record_on_drop`](HistogramExt::record_on_drop) and
/// [`record_duration_on_drop`](HistogramExt::record_duration_on_drop) directly on histogram instruments.
pub trait HistogramExt<T> {
    /// Records `value` on drop. The value can be updated via [`RecordValueGuard::value`]
    /// or suppressed via [`RecordValueGuard::cancel`].
    fn record_on_drop(&self, value: T, attributes: impl Into<Vec<KeyValue>>)
    -> RecordValueGuard<T>;

    /// Records elapsed time in seconds on drop. Attributes can be added via
    /// [`RecordDurationGuard::set`]; recording can be suppressed via [`RecordDurationGuard::cancel`].
    fn record_duration_on_drop(&self, attributes: impl Into<Vec<KeyValue>>) -> RecordDurationGuard
    where
        Histogram<T>: Into<Histogram<f64>>;

    /// Records elapsed time in seconds on drop using a custom clock. Useful for testing.
    fn record_duration_on_drop_with_clock(
        &self,
        clock: Clock,
        attributes: impl Into<Vec<KeyValue>>,
    ) -> RecordDurationGuard
    where
        Histogram<T>: Into<Histogram<f64>>;
}

impl<T> HistogramExt<T> for Histogram<T>
where
    Self: Clone,
{
    fn record_on_drop(
        &self,
        value: T,
        attributes: impl Into<Vec<KeyValue>>,
    ) -> RecordValueGuard<T> {
        RecordValueGuard::new(self.clone(), value, attributes)
    }

    fn record_duration_on_drop(&self, attributes: impl Into<Vec<KeyValue>>) -> RecordDurationGuard
    where
        Histogram<T>: Into<Histogram<f64>>,
    {
        RecordDurationGuard::new(self.clone().into(), attributes)
    }

    fn record_duration_on_drop_with_clock(
        &self,
        clock: Clock,
        attributes: impl Into<Vec<KeyValue>>,
    ) -> RecordDurationGuard
    where
        Histogram<T>: Into<Histogram<f64>>,
    {
        RecordDurationGuard::with_clock(self.clone().into(), clock, attributes)
    }
}

/// Extension methods on [`Counter<T>`] for creating RAII counter guards.
///
/// Import this trait to use [`add_on_drop`](CounterExt::add_on_drop) directly on counter instruments.
pub trait CounterExt<T> {
    /// Adds `value` to the counter on drop. The value can be updated via [`AddValueGuard::value`],
    /// attributes can be added or updated via [`AddValueGuard::set`], and recording can be
    /// suppressed via [`AddValueGuard::cancel`].
    fn add_on_drop(&self, value: T, attributes: impl Into<Vec<KeyValue>>) -> AddValueGuard<T>;
}

impl<T> CounterExt<T> for Counter<T>
where
    Self: Clone,
{
    fn add_on_drop(&self, value: T, attributes: impl Into<Vec<KeyValue>>) -> AddValueGuard<T> {
        AddValueGuard::new(self.clone(), value, attributes)
    }
}

/// Extension methods on [`UpDownCounter<T>`] for creating RAII active guards.
///
/// Import this trait to use [`track`](UpDownCounterExt::track) directly on up-down counter instruments.
pub trait UpDownCounterExt<T: From<i8>> {
    /// Increments the counter immediately and decrements it on drop.
    fn track(&self, attributes: impl Into<Vec<KeyValue>>) -> TrackGuard<T>;
}

impl<T: From<i8>> UpDownCounterExt<T> for UpDownCounter<T>
where
    Self: Clone,
{
    fn track(&self, attributes: impl Into<Vec<KeyValue>>) -> TrackGuard<T> {
        TrackGuard::new(self.clone(), attributes)
    }
}

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

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

    use super::{CounterExt, HistogramExt, UpDownCounterExt};

    #[test]
    fn counter_add_on_drop() {
        let ctx = TelemetryContext::new();
        let counter = global::meter_provider()
            .meter("test")
            .u64_counter("test.requests")
            .build();

        let guard = counter.add_on_drop(1, []);
        counter.add(1, &[]);
        drop(guard);

        assert_metrics_snapshot!(ctx, @r"
        - name: test.requests
          data:
            type: Sum
            data_points:
              - value: 2
            is_monotonic: true
            temporality: Cumulative
        ");
    }

    #[test]
    fn up_down_counter_track() {
        let ctx = TelemetryContext::new();
        let counter = global::meter_provider()
            .meter("test")
            .i64_up_down_counter("test.active")
            .build();

        let guard = counter.track([KeyValue::new("state", "active")]);
        let _guard2 = counter.track([KeyValue::new("state", "idle")]);
        drop(guard);

        assert_metrics_snapshot!(ctx, @r"
        - name: test.active
          data:
            type: Sum
            data_points:
              - attributes:
                  state: active
                value: 0
              - attributes:
                  state: idle
                value: 1
            is_monotonic: false
            temporality: Cumulative
        ");
    }

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

        let guard = histogram.record_on_drop(100.0, []);
        histogram.record(200.0, &[]);
        drop(guard);

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

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

        let (clock, mock) = quanta::Clock::mock();
        let guard = histogram.record_duration_on_drop_with_clock(clock, []);
        histogram.record(0.0, &[]);
        mock.increment(Duration::from_millis(500));
        drop(guard);

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