KAFKA-19398: (De)Register oldest-iterator-open-since-ms metric dynamically

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

@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.internals.metrics;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
+import org.apache.kafka.streams.state.internals.MeteredIterator;
+import org.apache.kafka.streams.state.internals.metrics.StateStoreMetrics;
+
+import java.util.Comparator;
+import java.util.NavigableSet;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.atomic.LongAdder;
+
+public class OpenIterators {
+    private final TaskId taskId;
+    private final String metricsScope;
+    private final String name;
+    private final StreamsMetricsImpl streamsMetrics;
+
+    private final LongAdder numOpenIterators = new LongAdder();
+    private final NavigableSet<MeteredIterator> openIterators = new ConcurrentSkipListSet<>(Comparator.comparingLong(MeteredIterator::startTimestamp));
+
+    private MetricName metricName;
+
+    public OpenIterators(final TaskId taskId,
+                         final String metricsScope,
+                         final String name,
+                         final StreamsMetricsImpl streamsMetrics) {
+        this.taskId = taskId;
+        this.metricsScope = metricsScope;
+        this.name = name;
+        this.streamsMetrics = streamsMetrics;
+    }
+
+    public void add(final MeteredIterator iterator) {
+        openIterators.add(iterator);
+        numOpenIterators.increment();
+
+        if (numOpenIterators.intValue() == 1) {
+            metricName = StateStoreMetrics.addOldestOpenIteratorGauge(taskId.toString(), metricsScope, name, streamsMetrics,
+                (config, now) -> openIterators.first().startTimestamp()
+            );
+        }
+    }
+
+    public void remove(final MeteredIterator iterator) {
+        if (numOpenIterators.intValue() == 1) {
+            streamsMetrics.removeMetric(metricName);
+        }
+        numOpenIterators.decrement();
+        openIterators.remove(iterator);
+    }
+
+    public long sum() {
+        return numOpenIterators.sum();
+    }
+}

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

@@ -335,6 +335,10 @@ public final void removeAllThreadLevelMetrics(final String threadId) {
         }
     }
 
+    public void removeMetric(final MetricName metricName) {
+        metrics.removeMetric(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);
@@ -517,13 +521,13 @@ public final Sensor storeLevelSensor(final String taskId,
         return getSensors(storeLevelSensors, sensorSuffix, sensorPrefix, recordingLevel, parents);
     }
 
-    public <T> void addStoreLevelMutableMetric(final String taskId,
-                                               final String metricsScope,
-                                               final String storeName,
-                                               final String name,
-                                               final String description,
-                                               final RecordingLevel recordingLevel,
-                                               final Gauge<T> valueProvider) {
+    public <T> MetricName addStoreLevelMutableMetric(final String taskId,
+                                                     final String metricsScope,
+                                                     final String storeName,
+                                                     final String name,
+                                                     final String description,
+                                                     final RecordingLevel recordingLevel,
+                                                     final Gauge<T> valueProvider) {
         final MetricName metricName = metrics.metricName(
             name,
             STATE_STORE_LEVEL_GROUP,
@@ -535,6 +539,8 @@ public <T> void addStoreLevelMutableMetric(final String taskId,
             final String key = storeSensorPrefix(Thread.currentThread().getName(), taskId, storeName);
             storeLevelMetrics.computeIfAbsent(key, ignored -> new LinkedList<>()).push(metricName);
         }
+
+        return metricName;
     }
 
     public final void removeAllStoreLevelSensorsAndMetrics(final String taskId,

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

@@ -23,6 +23,7 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.errors.ProcessorStateException;
+import org.apache.kafka.streams.internals.metrics.OpenIterators;
 import org.apache.kafka.streams.kstream.internals.Change;
 import org.apache.kafka.streams.kstream.internals.WrappingNullableUtils;
 import org.apache.kafka.streams.processor.StateStore;
@@ -48,14 +49,9 @@
 import org.apache.kafka.streams.state.internals.metrics.StateStoreMetrics;
 
 import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
-import java.util.NavigableSet;
 import java.util.Objects;
-import java.util.concurrent.ConcurrentSkipListSet;
-import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 
 import static org.apache.kafka.common.utils.Utils.mkEntry;
@@ -96,8 +92,7 @@ public class MeteredKeyValueStore<K, V>
     private StreamsMetricsImpl streamsMetrics;
     private TaskId taskId;
 
-    protected LongAdder numOpenIterators = new LongAdder();
-    protected NavigableSet<MeteredIterator> openIterators = new ConcurrentSkipListSet<>(Comparator.comparingLong(MeteredIterator::startTimestamp));
+    protected OpenIterators openIterators;
 
     @SuppressWarnings("rawtypes")
     private final Map<Class, QueryHandler> queryHandlers =
@@ -153,13 +148,8 @@ private void registerMetrics() {
         e2eLatencySensor = StateStoreMetrics.e2ELatencySensor(taskId.toString(), metricsScope, name(), streamsMetrics);
         iteratorDurationSensor = StateStoreMetrics.iteratorDurationSensor(taskId.toString(), metricsScope, name(), streamsMetrics);
         StateStoreMetrics.addNumOpenIteratorsGauge(taskId.toString(), metricsScope, name(), streamsMetrics,
-                (config, now) -> numOpenIterators.sum());
-        StateStoreMetrics.addOldestOpenIteratorGauge(taskId.toString(), metricsScope, name(), streamsMetrics,
-                (config, now) -> {
-                    final Iterator<MeteredIterator> openIteratorsIterator = openIterators.iterator();
-                    return openIteratorsIterator.hasNext() ? openIteratorsIterator.next().startTimestamp() : null;
-                }
-        );
+                (config, now) -> openIterators.sum());
+        openIterators = new OpenIterators(taskId, metricsScope, name(), streamsMetrics);
     }
 
     protected Serde<V> prepareValueSerdeForStore(final Serde<V> valueSerde, final SerdeGetter getter) {
@@ -445,7 +435,6 @@ private MeteredKeyValueIterator(final KeyValueIterator<Bytes, byte[]> iter,
             this.sensor = sensor;
             this.startTimestamp = time.milliseconds();
             this.startNs = time.nanoseconds();
-            numOpenIterators.increment();
             openIterators.add(this);
         }
 
@@ -475,7 +464,6 @@ public void close() {
                 final long duration = time.nanoseconds() - startNs;
                 sensor.record(duration);
                 iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
                 openIterators.remove(this);
             }
         }
@@ -502,7 +490,6 @@ private MeteredKeyValueTimestampedIterator(final KeyValueIterator<Bytes, byte[]>
             this.valueDeserializer = valueDeserializer;
             this.startTimestamp = time.milliseconds();
             this.startNs = time.nanoseconds();
-            numOpenIterators.increment();
             openIterators.add(this);
         }
 
@@ -532,7 +519,6 @@ public void close() {
                 final long duration = time.nanoseconds() - startNs;
                 sensor.record(duration);
                 iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
                 openIterators.remove(this);
             }
         }

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

@@ -18,39 +18,34 @@
 
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.streams.internals.metrics.OpenIterators;
 import org.apache.kafka.streams.state.VersionedRecord;
 import org.apache.kafka.streams.state.VersionedRecordIterator;
 
-import java.util.Set;
-import java.util.concurrent.atomic.LongAdder;
 import java.util.function.Function;
 
 class MeteredMultiVersionedKeyQueryIterator<V> implements VersionedRecordIterator<V>, MeteredIterator {
 
     private final VersionedRecordIterator<byte[]> iterator;
     private final Function<VersionedRecord<byte[]>, VersionedRecord<V>> deserializeValue;
-    private final LongAdder numOpenIterators;
     private final Sensor sensor;
     private final Time time;
     private final long startNs;
     private final long startTimestampMs;
-    private final Set<MeteredIterator> openIterators;
+    private final OpenIterators openIterators;
 
     public MeteredMultiVersionedKeyQueryIterator(final VersionedRecordIterator<byte[]> iterator,
                                                  final Sensor sensor,
                                                  final Time time,
                                                  final Function<VersionedRecord<byte[]>, VersionedRecord<V>> deserializeValue,
-                                                 final LongAdder numOpenIterators,
-                                                 final Set<MeteredIterator> openIterators) {
+                                                 final OpenIterators openIterators) {
         this.iterator = iterator;
         this.deserializeValue = deserializeValue;
-        this.numOpenIterators = numOpenIterators;
         this.openIterators = openIterators;
         this.sensor = sensor;
         this.time = time;
         this.startNs = time.nanoseconds();
         this.startTimestampMs = time.milliseconds();
-        numOpenIterators.increment();
         openIterators.add(this);
     }
 
@@ -65,7 +60,6 @@ public void close() {
             iterator.close();
         } finally {
             sensor.record(time.nanoseconds() - startNs);
-            numOpenIterators.decrement();
             openIterators.remove(this);
         }
     }

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

@@ -326,7 +326,6 @@ private MeteredTimestampedKeyValueStoreIterator(final KeyValueIterator<Bytes, by
             this.startNs = time.nanoseconds();
             this.startTimestampMs = time.milliseconds();
             this.returnPlainValue = returnPlainValue;
-            numOpenIterators.increment();
             openIterators.add(this);
         }
 
@@ -360,7 +359,6 @@ public void close() {
                 final long duration = time.nanoseconds() - startNs;
                 sensor.record(duration);
                 iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
                 openIterators.remove(this);
             }
         }

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

@@ -269,7 +269,6 @@ private <R> QueryResult<R> runMultiVersionedKeyQuery(final Query<R> query, final
                             iteratorDurationSensor,
                             time,
                             StoreQueryUtils.deserializeValue(plainValueSerdes),
-                            numOpenIterators,
                             openIterators
                         );
                 final QueryResult<MeteredMultiVersionedKeyQueryIterator<V>> typedQueryResult =

streams/src/main/java/org/apache/kafka/streams/state/internals/metrics/StateStoreMetrics.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.state.internals.metrics;
 
+import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.metrics.Gauge;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.Sensor.RecordingLevel;
@@ -455,12 +456,12 @@ public static void addNumOpenIteratorsGauge(final String taskId,
 
     }
 
-    public static void addOldestOpenIteratorGauge(final String taskId,
-                                                  final String storeType,
-                                                  final String storeName,
-                                                  final StreamsMetricsImpl streamsMetrics,
-                                                  final Gauge<Long> oldestOpenIteratorGauge) {
-        streamsMetrics.addStoreLevelMutableMetric(
+    public static MetricName addOldestOpenIteratorGauge(final String taskId,
+                                                        final String storeType,
+                                                        final String storeName,
+                                                        final StreamsMetricsImpl streamsMetrics,
+                                                        final Gauge<Long> oldestOpenIteratorGauge) {
+        return streamsMetrics.addStoreLevelMutableMetric(
                 taskId,
                 storeType,
                 storeName,

streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredKeyValueStoreTest.java

@@ -525,15 +525,16 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
         when(inner.all()).thenReturn(KeyValueIterators.emptyIterator());
         init();
 
-        final KafkaMetric oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
-        assertThat(oldestIteratorTimestampMetric, not(nullValue()));
-
-        assertThat(oldestIteratorTimestampMetric.metricValue(), nullValue());
+        KafkaMetric oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
 
         KeyValueIterator<String, String> second = null;
         final long secondTimestamp;
         try {
             try (final KeyValueIterator<String, String> unused = metered.all()) {
+                oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+                assertThat(oldestIteratorTimestampMetric, not(nullValue()));
+
                 final long oldestTimestamp = mockTime.milliseconds();
                 assertThat((Long) oldestIteratorTimestampMetric.metricValue(), equalTo(oldestTimestamp));
                 mockTime.sleep(100);
@@ -553,7 +554,8 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
             }
         }
 
-        assertThat((Integer) oldestIteratorTimestampMetric.metricValue(), nullValue());
+        oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
     }
 
     private KafkaMetric metric(final MetricName metricName) {

streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreTest.java

@@ -503,15 +503,16 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
         when(inner.all()).thenReturn(KeyValueIterators.emptyIterator());
         init();
 
-        final KafkaMetric oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
-        assertThat(oldestIteratorTimestampMetric, not(nullValue()));
-
-        assertThat(oldestIteratorTimestampMetric.metricValue(), nullValue());
+        KafkaMetric oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
 
         KeyValueIterator<String, ValueAndTimestamp<String>> second = null;
         final long secondTimestamp;
         try {
             try (final KeyValueIterator<String, ValueAndTimestamp<String>> unused = metered.all()) {
+                oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+                assertThat(oldestIteratorTimestampMetric, not(nullValue()));
+
                 final long oldestTimestamp = mockTime.milliseconds();
                 assertThat((Long) oldestIteratorTimestampMetric.metricValue(), equalTo(oldestTimestamp));
                 mockTime.sleep(100);
@@ -531,6 +532,7 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
             }
         }
 
-        assertThat((Integer) oldestIteratorTimestampMetric.metricValue(), nullValue());
+        oldestIteratorTimestampMetric = metric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
     }
 }

streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredVersionedKeyValueStoreTest.java

@@ -426,16 +426,17 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
         when(inner.query(any(), any(), any())).thenReturn(
                 QueryResult.forResult(new LogicalSegmentIterator(Collections.emptyListIterator(), RAW_KEY, 0L, 0L, ResultOrder.ANY)));
 
-        final KafkaMetric oldestIteratorTimestampMetric = getMetric("oldest-iterator-open-since-ms");
-        assertThat(oldestIteratorTimestampMetric, not(nullValue()));
-
-        assertThat(oldestIteratorTimestampMetric.metricValue(), nullValue());
+        KafkaMetric oldestIteratorTimestampMetric = getMetric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
 
         final QueryResult<VersionedRecordIterator<String>> first = store.query(query, bound, config);
         VersionedRecordIterator<String> secondIterator = null;
         final long secondTime;
         try {
             try (final VersionedRecordIterator<String> unused = first.getResult()) {
+                oldestIteratorTimestampMetric = getMetric("oldest-iterator-open-since-ms");
+                assertThat(oldestIteratorTimestampMetric, not(nullValue()));
+
                 final long oldestTimestamp = mockTime.milliseconds();
                 assertThat((Long) oldestIteratorTimestampMetric.metricValue(), equalTo(oldestTimestamp));
                 mockTime.sleep(100);
@@ -457,7 +458,8 @@ public void shouldTrackOldestOpenIteratorTimestamp() {
             }
         }
 
-        assertThat((Integer) oldestIteratorTimestampMetric.metricValue(), nullValue());
+        oldestIteratorTimestampMetric = getMetric("oldest-iterator-open-since-ms");
+        assertThat(oldestIteratorTimestampMetric, nullValue());
     }
 
     private KafkaMetric getMetric(final String name) {