apollo-opentelemetry 0.8.0

use opentelemetry::KeyValue;
use opentelemetry::metrics::Counter;

/// Adds a value to a [`Counter<T>`] on drop.
///
/// Ensures the counter is incremented at the end of an operation even if the future is cancelled.
/// The value can be updated via [`value`](Self::value) or suppressed via [`cancel`](Self::cancel).
/// Attributes can be added or updated via [`set`](Self::set).
///
/// # Example
///
/// ```
/// # use opentelemetry::metrics::Counter;
/// # use apollo_opentelemetry::metrics::CounterExt;
/// # async fn process() {}
/// async fn handle(counter: Counter<u64>) {
///     let guard = counter.add_on_drop(1, []);
///
///     process().await; // future may be cancelled here — guard still increments on drop
/// }
/// # fn main() {}
/// ```
#[must_use]
pub struct AddValueGuard<T> {
    counter: Counter<T>,
    value: Option<T>,
    attributes: Vec<KeyValue>,
}

impl<T> AddValueGuard<T> {
    /// Create a new guard. `value` will be added to the counter on drop.
    pub fn new(counter: Counter<T>, value: T, attributes: impl Into<Vec<KeyValue>>) -> Self {
        Self {
            counter,
            value: Some(value),
            attributes: attributes.into(),
        }
    }

    /// Add or update an attribute that will be included when the value is added.
    ///
    /// If an attribute with the same key already exists, it is replaced.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use opentelemetry::KeyValue;
    /// # use apollo_opentelemetry::metrics::CounterExt;
    /// # let counter = global::meter_provider().meter("example").u64_counter("bytes.sent").build();
    /// let mut guard = counter.add_on_drop(512, []);
    /// guard.set(KeyValue::new("route", "/api/v1/data"));
    /// ```
    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 that will be added on drop.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::CounterExt;
    /// # let counter = global::meter_provider().meter("example").u64_counter("bytes.sent").build();
    /// let mut guard = counter.add_on_drop(0, []);
    /// guard.value(512);
    /// ```
    pub fn value(&mut self, value: T) {
        self.value = Some(value);
    }

    /// Cancel the counter add — nothing will be added on drop.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::CounterExt;
    /// # let counter = global::meter_provider().meter("example").u64_counter("requests.completed").build();
    /// let mut guard = counter.add_on_drop(1, []);
    /// if true /* request should not be counted */ {
    ///     guard.cancel();
    /// }
    /// ```
    pub fn cancel(&mut self) {
        self.value = None;
    }
}

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

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

    use super::AddValueGuard;

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

        let guard = AddValueGuard::new(counter, 5, []);

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

        drop(guard);

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

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

        let mut guard = AddValueGuard::new(counter, 100, []);
        guard.value(42);
        drop(guard);

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

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

        let mut guard = AddValueGuard::new(counter, 1, []);
        guard.set(opentelemetry::KeyValue::new("route", "/a"));
        guard.set(opentelemetry::KeyValue::new("route", "/b"));
        drop(guard);

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

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

        let mut guard = AddValueGuard::new(counter, 5, []);
        guard.cancel();
        drop(guard);

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