From 8e6b479684d550fb70b570d059080488b893aaf4 Mon Sep 17 00:00:00 2001
From: Utkarsh Umesan Pillai <66651184+utpilla@users.noreply.github.com>
Date: Tue, 26 Nov 2024 16:00:51 -0800
Subject: [PATCH] Avoid additional HashMap allocation for Cumulative
 aggregation (#2352)

---
 opentelemetry-sdk/src/metrics/internal/mod.rs | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/opentelemetry-sdk/src/metrics/internal/mod.rs b/opentelemetry-sdk/src/metrics/internal/mod.rs
index 49544948bb..fba67e07b4 100644
--- a/opentelemetry-sdk/src/metrics/internal/mod.rs
+++ b/opentelemetry-sdk/src/metrics/internal/mod.rs
@@ -10,7 +10,7 @@ use std::collections::{HashMap, HashSet};
 use std::mem::swap;
 use std::ops::{Add, AddAssign, DerefMut, Sub};
 use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering};
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, OnceLock, RwLock};
 
 use aggregate::{is_under_cardinality_limit, STREAM_CARDINALITY_LIMIT};
 pub(crate) use aggregate::{AggregateBuilder, ComputeAggregation, Measure};
@@ -52,9 +52,10 @@ where
     /// Trackers store the values associated with different attribute sets.
     trackers: RwLock<HashMap<Vec<KeyValue>, Arc<A>>>,
 
-    /// Used by collect exclusively. The data type must match the one used in
-    /// `trackers` to allow mem::swap.
-    trackers_for_collect: RwLock<HashMap<Vec<KeyValue>, Arc<A>>>,
+    /// Used ONLY by Delta collect. The data type must match the one used in
+    /// `trackers` to allow mem::swap. Wrapping the type in `OnceLock` to
+    /// avoid this allocation for Cumulative aggregation.
+    trackers_for_collect: OnceLock<RwLock<HashMap<Vec<KeyValue>, Arc<A>>>>,
 
     /// Number of different attribute set stored in the `trackers` map.
     count: AtomicUsize,
@@ -73,9 +74,7 @@ where
     fn new(config: A::InitConfig) -> Self {
         ValueMap {
             trackers: RwLock::new(HashMap::with_capacity(1 + STREAM_CARDINALITY_LIMIT)),
-            // TODO: For cumulative, this is not required, so avoid this
-            // pre-allocation.
-            trackers_for_collect: RwLock::new(HashMap::with_capacity(1 + STREAM_CARDINALITY_LIMIT)),
+            trackers_for_collect: OnceLock::new(),
             has_no_attribute_value: AtomicBool::new(false),
             no_attribute_tracker: A::create(&config),
             count: AtomicUsize::new(0),
@@ -83,6 +82,12 @@ where
         }
     }
 
+    #[inline]
+    fn trackers_for_collect(&self) -> &RwLock<HashMap<Vec<KeyValue>, Arc<A>>> {
+        self.trackers_for_collect
+            .get_or_init(|| RwLock::new(HashMap::with_capacity(1 + STREAM_CARDINALITY_LIMIT)))
+    }
+
     fn measure(&self, value: A::PreComputedValue, attributes: &[KeyValue]) {
         if attributes.is_empty() {
             self.no_attribute_tracker.update(value);
@@ -178,7 +183,7 @@ where
             ));
         }
 
-        if let Ok(mut trackers_collect) = self.trackers_for_collect.write() {
+        if let Ok(mut trackers_collect) = self.trackers_for_collect().write() {
             if let Ok(mut trackers_current) = self.trackers.write() {
                 swap(trackers_collect.deref_mut(), trackers_current.deref_mut());
                 self.count.store(0, Ordering::SeqCst);