From 0f46509e86fec949498fd06641f67777338da6ab Mon Sep 17 00:00:00 2001
From: dino2895 <dino2895un@gmai.com>
Date: Sat, 11 Jul 2026 06:04:04 +0800
Subject: [PATCH] KAFKA-20796: Fix open iterator metric for timestamped range
 queries

---
 .../MeteredTimestampedKeyValueStore.java      |  2 +
 .../MeteredTimestampedKeyValueStoreTest.java  | 56 +++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStore.java b/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStore.java
index 72caa56d88e35..e44f2af701fc6 100644
--- a/streams/src/main/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStore.java
+++ b/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);
             }
         }
diff --git a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreTest.java b/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreTest.java
index 6e05513ec9c7c..b5061b352b7c9 100644
--- a/streams/src/test/java/org/apache/kafka/streams/state/internals/MeteredTimestampedKeyValueStoreTest.java
+++ b/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() {