apollo-opentelemetry 0.8.0

use opentelemetry::KeyValue;
use opentelemetry::metrics::UpDownCounter;

/// Increments an [`UpDownCounter<T>`] on construction and decrements it on drop.
///
/// Useful for tracking the number of active connections or in-flight requests.
/// Use [`set`](TrackGuard::set) to transition between named states (e.g. `active`
/// → `idle`) without losing track of which state to decrement on drop.
///
/// # Example
///
/// ```
/// # use opentelemetry::metrics::UpDownCounter;
/// # use apollo_opentelemetry::metrics::UpDownCounterExt;
/// # async fn process() {}
/// async fn handle(counter: UpDownCounter<i64>) {
///     let _guard = counter.track([]);
///
///     process().await; // future may be cancelled here — guard still decrements on drop
/// }
/// # fn main() {}
/// ```
#[must_use]
pub struct TrackGuard<T: From<i8>> {
    counter: UpDownCounter<T>,
    attributes: Vec<KeyValue>,
}

impl<T: From<i8>> TrackGuard<T> {
    /// Create a new guard, immediately incrementing the counter by 1.
    pub fn new(counter: UpDownCounter<T>, attributes: impl Into<Vec<KeyValue>>) -> Self {
        let attributes = attributes.into();
        counter.add(T::from(1_i8), &attributes);
        Self {
            counter,
            attributes,
        }
    }

    /// Add or update an attribute, transitioning the counter to the new attribute set.
    ///
    /// Decrements the counter with the current attributes, upserts the key-value
    /// pair, then increments with the updated attributes. The guard remembers the
    /// new attributes and will decrement them on drop.
    ///
    /// # Example
    ///
    /// ```
    /// # use opentelemetry::KeyValue;
    /// # use opentelemetry::global;
    /// # use apollo_opentelemetry::metrics::UpDownCounterExt;
    /// let counter = global::meter_provider()
    ///     .meter("example")
    ///     .i64_up_down_counter("http.client.open_connections")
    ///     .build();
    ///
    /// let mut guard = counter.track([KeyValue::new("state", "idle")]); // idle = 1
    /// guard.set(KeyValue::new("state", "active"));                     // idle = 0, active = 1
    /// guard.set(KeyValue::new("state", "idle"));                       // idle = 1, active = 0
    /// drop(guard);                                                      // idle = 0
    /// ```
    pub fn set(&mut self, kv: KeyValue) {
        self.counter.add(T::from(-1_i8), &self.attributes);
        if let Some(existing) = self.attributes.iter_mut().find(|a| a.key == kv.key) {
            *existing = kv;
        } else {
            self.attributes.push(kv);
        }
        self.counter.add(T::from(1_i8), &self.attributes);
    }
}

impl<T: From<i8>> Drop for TrackGuard<T> {
    fn drop(&mut self) {
        self.counter.add(T::from(-1_i8), &self.attributes);
    }
}

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

    use super::TrackGuard;

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

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

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

        drop(guard);

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

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

        let mut guard = TrackGuard::new(counter, [KeyValue::new("state", "idle")]);

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

        guard.set(KeyValue::new("state", "active"));

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

        guard.set(KeyValue::new("state", "idle"));

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

        drop(guard);

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

    /// Drop must decrement the last set attribute, not the original one.
    #[test]
    fn drop_decrements_last_set_attribute() {
        let ctx = TelemetryContext::new();
        let counter = global::meter_provider()
            .meter("test")
            .i64_up_down_counter("test.connections")
            .build();

        let mut guard = TrackGuard::new(counter, [KeyValue::new("state", "idle")]);
        guard.set(KeyValue::new("state", "active"));
        drop(guard); // must decrement active, not idle

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