KAFKA-19748: fix metrics leak in Kafka Streams (#20633)

This PR fixes a leak in StreamsMetricImpl not removing a
store-level-metric correctly, and thus leaking objects.

Reviewers: Eduwer Camacaro <[email protected]>, Bill Bejeck
 <[email protected]>

Author: mjsax

streams/src/main/java/org/apache/kafka/streams/internals/metrics/OpenIterators.java

@@ -62,7 +62,7 @@ public void add(final MeteredIterator iterator) {
 
     public void remove(final MeteredIterator iterator) {
         if (openIterators.size() == 1) {
-            streamsMetrics.removeMetric(metricName);
+            streamsMetrics.removeStoreLevelMetric(metricName);
         }
         openIterators.remove(iterator);
         updateOldestStartTimestamp();

streams/src/main/java/org/apache/kafka/streams/processor/internals/metrics/StreamsMetricsImpl.java

@@ -41,12 +41,14 @@
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.LinkedList;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 public class StreamsMetricsImpl implements StreamsMetrics {
 
@@ -339,6 +341,30 @@ public void removeMetric(final MetricName metricName) {
         metrics.removeMetric(metricName);
     }
 
+    public void removeStoreLevelMetric(final MetricName metricName) {
+        metrics.removeMetric(metricName);
+
+        final List<String> metricsScopeCandidates = metricName.tags().keySet().stream()
+            .filter(tag -> !tag.equals(THREAD_ID_TAG) && !tag.equals(TASK_ID_TAG))
+            .collect(Collectors.toList());
+        if (metricsScopeCandidates.size() != 1) {
+            // should never happen
+            throw new IllegalStateException("Expected exactly one metric scope tag, but found " + metricsScopeCandidates);
+        }
+
+        final Deque<MetricName> metricsForStore = storeLevelMetrics.get(
+            storeSensorPrefix(
+                metricName.tags().get(THREAD_ID_TAG),
+                metricName.tags().get(TASK_ID_TAG),
+                metricName.tags().get(metricsScopeCandidates.get(0))
+            )
+        );
+
+        if (metricsForStore != null) {
+            metricsForStore.remove(metricName);
+        }
+    }
+
     public Map<String, String> taskLevelTagMap(final String threadId, final String taskId) {
         final Map<String, String> tagMap = new LinkedHashMap<>();
         tagMap.put(THREAD_ID_TAG, threadId);

streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStore.java

@@ -487,9 +487,11 @@ private class MeteredKeyValueTimestampedIterator implements KeyValueIterator<K,
         private final long startTimestamp;
         private final Function<byte[], V> valueDeserializer;
 
-        private MeteredKeyValueTimestampedIterator(final KeyValueIterator<Bytes, byte[]> iter,
-                                        final Sensor sensor,
-                                        final Function<byte[], V> valueDeserializer) {
+        private MeteredKeyValueTimestampedIterator(
+            final KeyValueIterator<Bytes, byte[]> iter,
+            final Sensor sensor,
+            final Function<byte[], V> valueDeserializer
+        ) {
             this.iter = iter;
             this.sensor = sensor;
             this.valueDeserializer = valueDeserializer;

streams/src/test/java/org/apache/kafka/streams/internals/metrics/OpenIteratorsTest.java

@@ -58,11 +58,11 @@ public void shouldCalculateOldestStartTimestampCorrectly() {
         assertThat(gauge.value(null, 0), is(2L));
 
         openIterators.remove(meteredIterator2);
-        verify(streamsMetrics, never()).removeMetric(any());
+        verify(streamsMetrics, never()).removeStoreLevelMetric(any());
         assertThat(gauge.value(null, 0), is(5L));
 
         openIterators.remove(meteredIterator1);
-        verify(streamsMetrics).removeMetric(any());
+        verify(streamsMetrics).removeStoreLevelMetric(any());
         assertThat(gauge.value(null, 0), is(5L));
 
         openIterators.add(meteredIterator3);