KAFKA-20796: Fix open iterator metric for timestamped range queries

Author: dino2895

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

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

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

@@ -37,6 +37,12 @@
 import org.apache.kafka.streams.processor.internals.InternalProcessorContext;
 import org.apache.kafka.streams.processor.internals.ProcessorStateManager;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
+import org.apache.kafka.streams.query.PositionBound;
+import org.apache.kafka.streams.query.Query;
+import org.apache.kafka.streams.query.QueryConfig;
+import org.apache.kafka.streams.query.QueryResult;
+import org.apache.kafka.streams.query.RangeQuery;
+import org.apache.kafka.streams.query.TimestampedRangeQuery;
 import org.apache.kafka.streams.state.KeyValueIterator;
 import org.apache.kafka.streams.state.KeyValueStore;
 import org.apache.kafka.streams.state.ValueAndTimestamp;
@@ -451,6 +457,56 @@ public void shouldTrackOpenIteratorsMetric() {
         assertThat((Long) openIteratorsMetric.metricValue(), equalTo(0L));
     }
 
+    @Test
+    public void shouldTrackOpenIteratorsMetricForRangeQuery() {
+        assertTracksOpenIteratorsMetricForQuery(
+            RangeQuery.<String, String>withNoBounds()
+        );
+    }
+
+    @Test
+    public void shouldTrackOpenIteratorsMetricForTimestampedRangeQuery() {
+        assertTracksOpenIteratorsMetricForQuery(
+            TimestampedRangeQuery.<String, String>withNoBounds()
+        );
+    }
+
+    @SuppressWarnings({"rawtypes", "unchecked"})
+    private void assertTracksOpenIteratorsMetricForQuery(final Query<?> query) {
+        setUp();
+
+        final QueryResult<KeyValueIterator<Bytes, byte[]>> rawResult =
+            QueryResult.forResult(KeyValueIterators.emptyIterator());
+
+        when(inner.query(
+            any(),
+            any(PositionBound.class),
+            any(QueryConfig.class)
+        )).thenReturn((QueryResult) rawResult);
+
+        init();
+
+        final KafkaMetric openIteratorsMetric = metric("num-open-iterators");
+        assertThat(openIteratorsMetric, not(nullValue()));
+        assertThat((Long) openIteratorsMetric.metricValue(), equalTo(0L));
+
+        final QueryResult<?> result = metered.query(
+            query,
+            PositionBound.unbounded(),
+            new QueryConfig(false)
+        );
+
+        assertTrue(result.isSuccess());
+
+        try (final KeyValueIterator<?, ?> unused =
+                (KeyValueIterator<?, ?>) result.getResult()) {
+            assertThat((Long) openIteratorsMetric.metricValue(), equalTo(1L));
+        }
+
+        assertThat((Long) openIteratorsMetric.metricValue(), equalTo(0L));
+    }
+    
+
     @SuppressWarnings("unused")
     @Test
     public void shouldTimeIteratorDuration() {