Directly implement Measure trait for metric aggregates (#2371)

Co-authored-by: Cijo Thomas <[email protected]>

Author: fraillt

opentelemetry-sdk/src/metrics/internal/aggregate.rs

@@ -97,77 +97,77 @@ impl Default for AggregateTimeInitiator {
     }
 }
 
+type Filter = Arc<dyn Fn(&KeyValue) -> bool + Send + Sync>;
+
+/// Applies filter on provided attribute set
+/// No-op, if filter is not set
+#[derive(Clone)]
+pub(crate) struct AttributeSetFilter {
+    filter: Option<Filter>,
+}
+
+impl AttributeSetFilter {
+    pub(crate) fn new(filter: Option<Filter>) -> Self {
+        Self { filter }
+    }
+
+    pub(crate) fn apply(&self, attrs: &[KeyValue], run: impl FnOnce(&[KeyValue])) {
+        if let Some(filter) = &self.filter {
+            let filtered_attrs: Vec<KeyValue> =
+                attrs.iter().filter(|kv| filter(kv)).cloned().collect();
+            run(&filtered_attrs);
+        } else {
+            run(attrs);
+        };
+    }
+}
+
 /// Builds aggregate functions
 pub(crate) struct AggregateBuilder<T> {
     /// The temporality used for the returned aggregate functions.
     temporality: Temporality,
 
     /// The attribute filter the aggregate function will use on the input of
     /// measurements.
-    filter: Option<Filter>,
+    filter: AttributeSetFilter,
 
     _marker: marker::PhantomData<T>,
 }
 
-type Filter = Arc<dyn Fn(&KeyValue) -> bool + Send + Sync>;
-
 impl<T: Number> AggregateBuilder<T> {
     pub(crate) fn new(temporality: Temporality, filter: Option<Filter>) -> Self {
         AggregateBuilder {
             temporality,
-            filter,
+            filter: AttributeSetFilter::new(filter),
             _marker: marker::PhantomData,
         }
     }
 
-    /// Wraps the passed in measure with an attribute filtering function.
-    fn filter(&self, f: impl Measure<T>) -> impl Measure<T> {
-        let filter = self.filter.clone();
-        move |n, attrs: &[KeyValue]| {
-            if let Some(filter) = &filter {
-                let filtered_attrs: Vec<KeyValue> =
-                    attrs.iter().filter(|kv| filter(kv)).cloned().collect();
-                f.call(n, &filtered_attrs);
-            } else {
-                f.call(n, attrs);
-            };
-        }
-    }
-
     /// Builds a last-value aggregate function input and output.
     pub(crate) fn last_value(&self) -> (impl Measure<T>, impl ComputeAggregation) {
-        let lv = Arc::new(LastValue::new(self.temporality));
-        let agg_lv = Arc::clone(&lv);
-
-        (
-            self.filter(move |n, a: &[KeyValue]| lv.measure(n, a)),
-            agg_lv,
-        )
+        let lv = Arc::new(LastValue::new(self.temporality, self.filter.clone()));
+        (lv.clone(), lv)
     }
 
     /// Builds a precomputed sum aggregate function input and output.
     pub(crate) fn precomputed_sum(
         &self,
         monotonic: bool,
     ) -> (impl Measure<T>, impl ComputeAggregation) {
-        let s = Arc::new(PrecomputedSum::new(self.temporality, monotonic));
-        let agg_sum = Arc::clone(&s);
+        let s = Arc::new(PrecomputedSum::new(
+            self.temporality,
+            self.filter.clone(),
+            monotonic,
+        ));
 
-        (
-            self.filter(move |n, a: &[KeyValue]| s.measure(n, a)),
-            agg_sum,
-        )
+        (s.clone(), s)
     }
 
     /// Builds a sum aggregate function input and output.
     pub(crate) fn sum(&self, monotonic: bool) -> (impl Measure<T>, impl ComputeAggregation) {
-        let s = Arc::new(Sum::new(self.temporality, monotonic));
-        let agg_sum = Arc::clone(&s);
+        let s = Arc::new(Sum::new(self.temporality, self.filter.clone(), monotonic));
 
-        (
-            self.filter(move |n, a: &[KeyValue]| s.measure(n, a)),
-            agg_sum,
-        )
+        (s.clone(), s)
     }
 
     /// Builds a histogram aggregate function input and output.
@@ -179,13 +179,13 @@ impl<T: Number> AggregateBuilder<T> {
     ) -> (impl Measure<T>, impl ComputeAggregation) {
         let h = Arc::new(Histogram::new(
             self.temporality,
+            self.filter.clone(),
             boundaries,
             record_min_max,
             record_sum,
         ));
-        let agg_h = Arc::clone(&h);
 
-        (self.filter(move |n, a: &[KeyValue]| h.measure(n, a)), agg_h)
+        (h.clone(), h)
     }
 
     /// Builds an exponential histogram aggregate function input and output.
@@ -198,14 +198,14 @@ impl<T: Number> AggregateBuilder<T> {
     ) -> (impl Measure<T>, impl ComputeAggregation) {
         let h = Arc::new(ExpoHistogram::new(
             self.temporality,
+            self.filter.clone(),
             max_size,
             max_scale,
             record_min_max,
             record_sum,
         ));
-        let agg_h = Arc::clone(&h);
 
-        (self.filter(move |n, a: &[KeyValue]| h.measure(n, a)), agg_h)
+        (h.clone(), h)
     }
 }
 

opentelemetry-sdk/src/metrics/internal/exponential_histogram.rs

@@ -13,7 +13,10 @@ use crate::metrics::{
     Temporality,
 };
 
-use super::{aggregate::AggregateTimeInitiator, Aggregator, ComputeAggregation, Number, ValueMap};
+use super::{
+    aggregate::{AggregateTimeInitiator, AttributeSetFilter},
+    Aggregator, ComputeAggregation, Measure, Number, ValueMap,
+};
 
 pub(crate) const EXPO_MAX_SCALE: i8 = 20;
 pub(crate) const EXPO_MIN_SCALE: i8 = -10;
@@ -357,6 +360,7 @@ pub(crate) struct ExpoHistogram<T: Number> {
     value_map: ValueMap<Mutex<ExpoHistogramDataPoint<T>>>,
     init_time: AggregateTimeInitiator,
     temporality: Temporality,
+    filter: AttributeSetFilter,
     record_sum: bool,
     record_min_max: bool,
 }
@@ -365,6 +369,7 @@ impl<T: Number> ExpoHistogram<T> {
     /// Create a new exponential histogram.
     pub(crate) fn new(
         temporality: Temporality,
+        filter: AttributeSetFilter,
         max_size: u32,
         max_scale: i8,
         record_min_max: bool,
@@ -377,22 +382,12 @@ impl<T: Number> ExpoHistogram<T> {
             }),
             init_time: AggregateTimeInitiator::default(),
             temporality,
+            filter,
             record_sum,
             record_min_max,
         }
     }
 
-    pub(crate) fn measure(&self, value: T, attrs: &[KeyValue]) {
-        let f_value = value.into_float();
-        // Ignore NaN and infinity.
-        // Only makes sense if T is f64, maybe this could be no-op for other cases?
-        if !f_value.is_finite() {
-            return;
-        }
-
-        self.value_map.measure(value, attrs);
-    }
-
     fn delta(&self, dest: Option<&mut dyn Aggregation>) -> (usize, Option<Box<dyn Aggregation>>) {
         let time = self.init_time.delta();
 
@@ -505,6 +500,24 @@ impl<T: Number> ExpoHistogram<T> {
     }
 }
 
+impl<T> Measure<T> for Arc<ExpoHistogram<T>>
+where
+    T: Number,
+{
+    fn call(&self, measurement: T, attrs: &[KeyValue]) {
+        let f_value = measurement.into_float();
+        // Ignore NaN and infinity.
+        // Only makes sense if T is f64, maybe this could be no-op for other cases?
+        if !f_value.is_finite() {
+            return;
+        }
+
+        self.filter.apply(attrs, |filtered| {
+            self.value_map.measure(measurement, filtered);
+        })
+    }
+}
+
 impl<T> ComputeAggregation for Arc<ExpoHistogram<T>>
 where
     T: Number,
@@ -682,9 +695,16 @@ mod tests {
         ];
 
         for test in test_cases {
-            let h = ExpoHistogram::new(Temporality::Cumulative, 4, 20, true, true);
+            let h = Arc::new(ExpoHistogram::new(
+                Temporality::Cumulative,
+                AttributeSetFilter::new(None),
+                4,
+                20,
+                true,
+                true,
+            ));
             for v in test.values {
-                h.measure(v, &[]);
+                Measure::call(&h, v, &[]);
             }
             let dp = h.value_map.no_attribute_tracker.lock().unwrap();
 
@@ -731,9 +751,16 @@ mod tests {
         ];
 
         for test in test_cases {
-            let h = ExpoHistogram::new(Temporality::Cumulative, 4, 20, true, true);
+            let h = Arc::new(ExpoHistogram::new(
+                Temporality::Cumulative,
+                AttributeSetFilter::new(None),
+                4,
+                20,
+                true,
+                true,
+            ));
             for v in test.values {
-                h.measure(v, &[]);
+                Measure::call(&h, v, &[]);
             }
             let dp = h.value_map.no_attribute_tracker.lock().unwrap();
 

opentelemetry-sdk/src/metrics/internal/histogram.rs

@@ -9,7 +9,9 @@ use crate::metrics::Temporality;
 use opentelemetry::KeyValue;
 
 use super::aggregate::AggregateTimeInitiator;
+use super::aggregate::AttributeSetFilter;
 use super::ComputeAggregation;
+use super::Measure;
 use super::ValueMap;
 use super::{Aggregator, Number};
 
@@ -72,6 +74,7 @@ pub(crate) struct Histogram<T: Number> {
     value_map: ValueMap<Mutex<Buckets<T>>>,
     init_time: AggregateTimeInitiator,
     temporality: Temporality,
+    filter: AttributeSetFilter,
     bounds: Vec<f64>,
     record_min_max: bool,
     record_sum: bool,
@@ -81,6 +84,7 @@ impl<T: Number> Histogram<T> {
     #[allow(unused_mut)]
     pub(crate) fn new(
         temporality: Temporality,
+        filter: AttributeSetFilter,
         mut bounds: Vec<f64>,
         record_min_max: bool,
         record_sum: bool,
@@ -97,24 +101,13 @@ impl<T: Number> Histogram<T> {
             value_map: ValueMap::new(buckets_count),
             init_time: AggregateTimeInitiator::default(),
             temporality,
+            filter,
             bounds,
             record_min_max,
             record_sum,
         }
     }
 
-    pub(crate) fn measure(&self, measurement: T, attrs: &[KeyValue]) {
-        let f = measurement.into_float();
-        // This search will return an index in the range `[0, bounds.len()]`, where
-        // it will return `bounds.len()` if value is greater than the last element
-        // of `bounds`. This aligns with the buckets in that the length of buckets
-        // is `bounds.len()+1`, with the last bucket representing:
-        // `(bounds[bounds.len()-1], +∞)`.
-        let index = self.bounds.partition_point(|&x| x < f);
-
-        self.value_map.measure((measurement, index), attrs);
-    }
-
     fn delta(&self, dest: Option<&mut dyn Aggregation>) -> (usize, Option<Box<dyn Aggregation>>) {
         let time = self.init_time.delta();
 
@@ -216,6 +209,25 @@ impl<T: Number> Histogram<T> {
     }
 }
 
+impl<T> Measure<T> for Arc<Histogram<T>>
+where
+    T: Number,
+{
+    fn call(&self, measurement: T, attrs: &[KeyValue]) {
+        let f = measurement.into_float();
+        // This search will return an index in the range `[0, bounds.len()]`, where
+        // it will return `bounds.len()` if value is greater than the last element
+        // of `bounds`. This aligns with the buckets in that the length of buckets
+        // is `bounds.len()+1`, with the last bucket representing:
+        // `(bounds[bounds.len()-1], +∞)`.
+        let index = self.bounds.partition_point(|&x| x < f);
+
+        self.filter.apply(attrs, |filtered| {
+            self.value_map.measure((measurement, index), filtered);
+        })
+    }
+}
+
 impl<T> ComputeAggregation for Arc<Histogram<T>>
 where
     T: Number,
@@ -236,12 +248,13 @@ mod tests {
     fn check_buckets_are_selected_correctly() {
         let hist = Arc::new(Histogram::<i64>::new(
             Temporality::Cumulative,
+            AttributeSetFilter::new(None),
             vec![1.0, 3.0, 6.0],
             false,
             false,
         ));
         for v in 1..11 {
-            hist.measure(v, &[]);
+            Measure::call(&hist, v, &[]);
         }
         let (count, dp) = ComputeAggregation::call(&hist, None);
         let dp = dp.unwrap();