From 90d9a91cd1ab8d7d3eed7d2beb1f19e5979d9a1e Mon Sep 17 00:00:00 2001
From: Utkarsh Umesan Pillai <66651184+utpilla@users.noreply.github.com>
Date: Fri, 5 Apr 2024 13:20:53 -0700
Subject: [PATCH] Minor refactoring for AggregatorStore and MetricPoint (#5510)

---
 src/OpenTelemetry/Metrics/AggregatorStore.cs | 85 ++++++++++++--------
 src/OpenTelemetry/Metrics/MetricPoint.cs     |  5 +-
 2 files changed, 55 insertions(+), 35 deletions(-)

diff --git a/src/OpenTelemetry/Metrics/AggregatorStore.cs b/src/OpenTelemetry/Metrics/AggregatorStore.cs
index 5271a329e40..5022e48dee1 100644
--- a/src/OpenTelemetry/Metrics/AggregatorStore.cs
+++ b/src/OpenTelemetry/Metrics/AggregatorStore.cs
@@ -270,39 +270,7 @@ internal void SnapshotDeltaWithMetricPointReclaim()
                 // the MetricPoint can be reused for other tags.
                 if (metricPoint.LookupData != null && Interlocked.CompareExchange(ref metricPoint.ReferenceCount, int.MinValue, 0) == 0)
                 {
-                    var lookupData = metricPoint.LookupData;
-
-                    metricPoint.Reclaim();
-
-                    Debug.Assert(this.TagsToMetricPointIndexDictionaryDelta != null, "this.tagsToMetricPointIndexDictionaryDelta was null");
-
-                    lock (this.TagsToMetricPointIndexDictionaryDelta!)
-                    {
-                        LookupData? dictionaryValue;
-                        if (lookupData.SortedTags != Tags.EmptyTags)
-                        {
-                            // Check if no other thread added a new entry for the same Tags.
-                            // If no, then remove the existing entries.
-                            if (this.TagsToMetricPointIndexDictionaryDelta.TryGetValue(lookupData.SortedTags, out dictionaryValue) &&
-                                dictionaryValue == lookupData)
-                            {
-                                this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.SortedTags, out var _);
-                                this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.GivenTags, out var _);
-                            }
-                        }
-                        else
-                        {
-                            if (this.TagsToMetricPointIndexDictionaryDelta.TryGetValue(lookupData.GivenTags, out dictionaryValue) &&
-                                dictionaryValue == lookupData)
-                            {
-                                this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.GivenTags, out var _);
-                            }
-                        }
-
-                        Debug.Assert(this.availableMetricPoints != null, "this.availableMetricPoints was null");
-
-                        this.availableMetricPoints!.Enqueue(i);
-                    }
+                    this.ReclaimMetricPoint(ref metricPoint, i);
                 }
 
                 continue;
@@ -372,6 +340,57 @@ private void TakeMetricPointSnapshot(ref MetricPoint metricPoint, bool outputDel
         }
     }
 
+    private void ReclaimMetricPoint(ref MetricPoint metricPoint, int metricPointIndex)
+    {
+        /*
+         This method does three things:
+          1. Set `metricPoint.LookupData` and `metricPoint.mpComponents` to `null` to have them collected faster by GC.
+          2. Tries to remove the entry for this MetricPoint from the lookup dictionary. An update thread which retrieves this
+             MetricPoint would realize that the MetricPoint is not valid for use since its reference count would have been set to a negative number.
+             When that happens, the update thread would also try to remove the entry for this MetricPoint from the lookup dictionary.
+             We only care about the entry getting removed from the lookup dictionary and not about which thread removes it.
+          3. Put the array index of this MetricPoint to the queue of available metric points. This makes it available for update threads
+             to use this MetricPoint to track newer dimension combinations.
+        */
+
+        var lookupData = metricPoint.LookupData;
+
+        // This method is only called after checking that `metricPoint.LookupData` is not `null`.
+        Debug.Assert(lookupData != null, "LookupData for the provided MetricPoint was null");
+
+        metricPoint.NullifyMetricPointState();
+
+        Debug.Assert(this.TagsToMetricPointIndexDictionaryDelta != null, "this.tagsToMetricPointIndexDictionaryDelta was null");
+
+        lock (this.TagsToMetricPointIndexDictionaryDelta!)
+        {
+            LookupData? dictionaryValue;
+            if (lookupData!.SortedTags != Tags.EmptyTags)
+            {
+                // Check if no other thread added a new entry for the same Tags.
+                // If no, then remove the existing entries.
+                if (this.TagsToMetricPointIndexDictionaryDelta.TryGetValue(lookupData.SortedTags, out dictionaryValue) &&
+                    dictionaryValue == lookupData)
+                {
+                    this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.SortedTags, out var _);
+                    this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.GivenTags, out var _);
+                }
+            }
+            else
+            {
+                if (this.TagsToMetricPointIndexDictionaryDelta.TryGetValue(lookupData.GivenTags, out dictionaryValue) &&
+                    dictionaryValue == lookupData)
+                {
+                    this.TagsToMetricPointIndexDictionaryDelta.TryRemove(lookupData.GivenTags, out var _);
+                }
+            }
+
+            Debug.Assert(this.availableMetricPoints != null, "this.availableMetricPoints was null");
+
+            this.availableMetricPoints!.Enqueue(metricPointIndex);
+        }
+    }
+
     [MethodImpl(MethodImplOptions.AggressiveInlining)]
     private void InitializeZeroTagPointIfNotInitialized()
     {
diff --git a/src/OpenTelemetry/Metrics/MetricPoint.cs b/src/OpenTelemetry/Metrics/MetricPoint.cs
index e1fe8e1695c..aa3974e2537 100644
--- a/src/OpenTelemetry/Metrics/MetricPoint.cs
+++ b/src/OpenTelemetry/Metrics/MetricPoint.cs
@@ -1197,9 +1197,10 @@ internal void TakeSnapshotWithExemplar(bool outputDelta)
     }
 
     /// <summary>
-    /// Denote that this MetricPoint is reclaimed.
+    /// This method sets the member object references of MetricPoint to `null`.
+    /// This is done to have them collected faster by GC.
     /// </summary>
-    internal void Reclaim()
+    internal void NullifyMetricPointState()
     {
         this.LookupData = null;
         this.mpComponents = null;