Revert "Avoid race conditions between contributing and writing to a stat aggr…" (#9498)

This reverts commit fe3046f.

Author: amarziali

dd-trace-core/src/main/java/datadog/trace/common/metrics/Aggregator.java

@@ -4,7 +4,7 @@
 
 import datadog.trace.common.metrics.SignalItem.StopSignal;
 import datadog.trace.core.util.LRUCache;
-import java.util.HashSet;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
@@ -121,20 +121,12 @@ public void accept(InboxItem item) {
           signal.ignore();
         }
       } else if (item instanceof Batch && !stopped) {
-        final Batch batch = (Batch) item;
+        Batch batch = (Batch) item;
         MetricKey key = batch.getKey();
         // important that it is still *this* batch pending, must not remove otherwise
         pending.remove(key, batch);
-        // operations concerning the aggregates should be atomic not to potentially loose points.
-        aggregates.compute(
-            key,
-            (k, v) -> {
-              if (v == null) {
-                v = new AggregateMetric();
-              }
-              batch.contributeTo(v);
-              return v;
-            });
+        AggregateMetric aggregate = aggregates.computeIfAbsent(key, k -> new AggregateMetric());
+        batch.contributeTo(aggregate);
         dirty = true;
         // return the batch for reuse
         batchPool.offer(batch);
@@ -146,20 +138,13 @@ private void report(long when, SignalItem signal) {
     boolean skipped = true;
     if (dirty) {
       try {
-        final Set<MetricKey> validKeys = expungeStaleAggregates();
-        if (!validKeys.isEmpty()) {
+        expungeStaleAggregates();
+        if (!aggregates.isEmpty()) {
           skipped = false;
-          writer.startBucket(validKeys.size(), when, reportingIntervalNanos);
-          for (MetricKey key : validKeys) {
-            // operations concerning the aggregates should be atomic not to potentially loose
-            // points.
-            aggregates.computeIfPresent(
-                key,
-                (k, v) -> {
-                  writer.add(k, v);
-                  v.clear();
-                  return v;
-                });
+          writer.startBucket(aggregates.size(), when, reportingIntervalNanos);
+          for (Map.Entry<MetricKey, AggregateMetric> aggregate : aggregates.entrySet()) {
+            writer.add(aggregate.getKey(), aggregate.getValue());
+            aggregate.getValue().clear();
           }
           // note that this may do IO and block
           writer.finishBucket();
@@ -176,27 +161,16 @@ private void report(long when, SignalItem signal) {
     }
   }
 
-  /**
-   * Remove keys whose values have zeroed metrics.
-   *
-   * @return a set containing the keys still valid.
-   */
-  private Set<MetricKey> expungeStaleAggregates() {
-    final HashSet<MetricKey> ret = new HashSet<>();
-    for (MetricKey metricKey : new HashSet<>(aggregates.keySet())) {
-      // operations concerning the aggregates should be atomic not to potentially loose points.
-      aggregates.computeIfPresent(
-          metricKey,
-          (k, v) -> {
-            if (v.getHitCount() == 0) {
-              commonKeys.remove(k);
-              return null;
-            }
-            ret.add(k);
-            return v;
-          });
+  private void expungeStaleAggregates() {
+    Iterator<Map.Entry<MetricKey, AggregateMetric>> it = aggregates.entrySet().iterator();
+    while (it.hasNext()) {
+      Map.Entry<MetricKey, AggregateMetric> pair = it.next();
+      AggregateMetric metric = pair.getValue();
+      if (metric.getHitCount() == 0) {
+        it.remove();
+        commonKeys.remove(pair.getKey());
+      }
     }
-    return ret;
   }
 
   private long wallClockTime() {