apollo-opentelemetry 0.8.0

use opentelemetry::KeyValue;
use opentelemetry::metrics::Histogram;

/// Records a value on a [`Histogram<T>`] on drop.
///
/// Useful for metrics like response body size that are only known when a stream is fully consumed.
///
/// # Example
///
/// ```
/// # use opentelemetry::metrics::Histogram;
/// # use opentelemetry::KeyValue;
/// # use apollo_opentelemetry::metrics::HistogramExt;
/// # async fn stream_body() -> f64 { 0.0 }
/// async fn handle(histogram: Histogram<f64>) {
///     let mut guard = histogram.record_on_drop(0.0, [KeyValue::new("service", "api")]);
///
///     let bytes = stream_body().await; // future may be cancelled here — guard still records on drop
///
///     guard.value(bytes);
/// }
/// # fn main() {}
/// ```
#[must_use]
pub struct RecordValueGuard<T> {
    histogram: Histogram<T>,
    value: Option<T>,
    attributes: Vec<KeyValue>,
}

impl<T> RecordValueGuard<T> {
    /// Create a new guard. `value` will be recorded on drop.
    pub fn new(histogram: Histogram<T>, value: T, attributes: impl Into<Vec<KeyValue>>) -> Self {
        Self {
            histogram,
            value: Some(value),
            attributes: attributes.into(),
        }
    }

    /// Add or update an attribute that will be included when the value 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("response.size").build();
    /// let mut guard = histogram.record_on_drop(512.0, []);
    /// guard.set(KeyValue::new("encoding", "gzip"));
    /// ```
    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);
        }
    }

    /// Overwrite the value to record on drop.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::HistogramExt;
    /// # let histogram = global::meter_provider().meter("example").f64_histogram("response.size").build();
    /// let mut guard = histogram.record_on_drop(0.0, []);
    /// guard.value(4096.0);
    /// ```
    pub fn value(&mut self, value: T) {
        self.value = Some(value);
    }

    /// Cancel recording — nothing will be recorded on drop.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::HistogramExt;
    /// # let histogram = global::meter_provider().meter("example").f64_histogram("response.size").build();
    /// let mut guard = histogram.record_on_drop(0.0, []);
    /// if true /* some error condition */ {
    ///     guard.cancel();
    /// }
    /// ```
    pub fn cancel(&mut self) {
        self.value = None;
    }
}

impl<T> Drop for RecordValueGuard<T> {
    fn drop(&mut self) {
        if let Some(value) = self.value.take() {
            self.histogram.record(value, &self.attributes);
        }
    }
}

#[cfg(test)]
mod tests {
    use apollo_opentelemetry_test::{TelemetryContext, assert_metrics_snapshot};
    use opentelemetry::global;

    use super::RecordValueGuard;

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

        {
            let _recorder = RecordValueGuard::new(histogram, 1024.0, []);
        }

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

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

        {
            let mut recorder = RecordValueGuard::new(histogram, 512.0, []);
            recorder.set(opentelemetry::KeyValue::new("encoding", "gzip"));
            recorder.set(opentelemetry::KeyValue::new("encoding", "br"));
        }

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

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

        {
            let mut recorder = RecordValueGuard::new(histogram, 1024.0, []);
            recorder.cancel();
        }

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