Updates per comments

Author: bbejeck

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

@@ -23,7 +23,6 @@
 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;
@@ -49,9 +48,14 @@
 import org.apache.kafka.streams.state.internals.metrics.StateStoreMetrics;
 
 import java.util.ArrayList;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.NavigableSet;
+import java.util.NoSuchElementException;
 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;
@@ -93,7 +97,9 @@ public class MeteredKeyValueStore<K, V>
     private TaskId taskId;
     private Sensor restoreSensor;
 
-    protected OpenIterators openIterators;
+    protected LongAdder numOpenIterators = new LongAdder();
+    protected NavigableSet<MeteredIterator> openIterators = new ConcurrentSkipListSet<>(Comparator.comparingLong(MeteredIterator::startTimestamp));
+
 
     @SuppressWarnings("rawtypes")
     private final Map<Class, QueryHandler> queryHandlers =
@@ -148,10 +154,15 @@ 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) -> openIterators.sum());
-        openIterators = new OpenIterators();
+                (config, now) -> numOpenIterators.sum());
         StateStoreMetrics.addOldestOpenIteratorGauge(taskId.toString(), metricsScope, name(), streamsMetrics,
-            (config, now) -> openIterators.oldestStartTimestamp());
+            (config, now) -> {
+                try {
+                    return openIterators.isEmpty() ? 0L : openIterators.first().startTimestamp();
+                } catch (final NoSuchElementException ignore) {
+                    return 0L;
+                }
+            });
     }
 
     @Override
@@ -442,6 +453,7 @@ private MeteredKeyValueIterator(final KeyValueIterator<Bytes, byte[]> iter,
             this.sensor = sensor;
             this.startTimestamp = time.milliseconds();
             this.startNs = time.nanoseconds();
+            numOpenIterators.increment();
             openIterators.add(this);
         }
 
@@ -471,6 +483,7 @@ public void close() {
                 final long duration = time.nanoseconds() - startNs;
                 sensor.record(duration);
                 iteratorDurationSensor.record(duration);
+                numOpenIterators.decrement();
                 openIterators.remove(this);
             }
         }
@@ -499,6 +512,7 @@ private MeteredKeyValueTimestampedIterator(
             this.valueDeserializer = valueDeserializer;
             this.startTimestamp = time.milliseconds();
             this.startNs = time.nanoseconds();
+            numOpenIterators.increment();
             openIterators.add(this);
         }
 
@@ -528,6 +542,7 @@ 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/MeteredKeyValueStore.java

@@ -18,10 +18,11 @@
 
 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 {
@@ -32,20 +33,24 @@ class MeteredMultiVersionedKeyQueryIterator<V> implements VersionedRecordIterato
     private final Time time;
     private final long startNs;
     private final long startTimestampMs;
-    private final OpenIterators openIterators;
+    private final Set<MeteredIterator> openIterators;
+    private final LongAdder numOpenIterators;
 
     public MeteredMultiVersionedKeyQueryIterator(final VersionedRecordIterator<byte[]> iterator,
                                                  final Sensor sensor,
                                                  final Time time,
                                                  final Function<VersionedRecord<byte[]>, VersionedRecord<V>> deserializeValue,
-                                                 final OpenIterators openIterators) {
+                                                 final LongAdder numOpenIterators,
+                                                 final Set<MeteredIterator> 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);
     }
 
@@ -60,6 +65,7 @@ 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/MeteredMultiVersionedKeyQueryIterator.java

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