MINOR: De-duplicate the metering lifecycle across Metered*WithHeaders iterators

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

@@ -0,0 +1,103 @@
+/*
+ * 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.state.internals;
+
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.streams.state.KeyValueIterator;
+
+import java.util.Set;
+import java.util.concurrent.atomic.LongAdder;
+
+/**
+ * Shared metering lifecycle for the metered iterators of the {@code Metered*WithHeaders} stores,
+ * whatever result type they yield: the {@code KeyValueIterator}s returned by the store's own range/
+ * fetch/find methods and the {@code ReadOnlyRecordIterator}s that back the headers-aware IQv2
+ * range/window/session query types.
+ *
+ * <p>Every such iterator opens over a raw {@code KeyValueIterator<RawKey, byte[]>} and needs the
+ * same bookkeeping: stamp the open time (for the {@code oldest-iterator-open-since-ms} metric),
+ * register in {@code numOpenIterators}/{@code openIterators}, and on {@link #close()} record the
+ * operation and iterator-duration sensors and deregister. This base is deliberately result-type
+ * agnostic -- it implements only {@link MeteredIterator} and does not bind the yielded key/value
+ * types -- so each subclass declares its own result interface (a {@code KeyValueIterator} or a
+ * {@code ReadOnlyRecordIterator}) and implements just the parts that genuinely differ: the
+ * deserializing {@code next()} (and, for the {@code KeyValueIterator}s, a peeking {@code hasNext()}
+ * and {@code peekNextKey()}).
+ *
+ * @param <RawKey> the raw iterator's key type
+ */
+abstract class AbstractMeteredIterator<RawKey> implements MeteredIterator {
+
+    final KeyValueIterator<RawKey, byte[]> iter;
+    private final Sensor operationSensor;
+    private final Sensor iteratorSensor;
+    private final Time time;
+    private final LongAdder numOpenIterators;
+    private final Set<MeteredIterator> openIterators;
+    private final long startNs;
+    private final long startTimestampMs;
+
+    AbstractMeteredIterator(final KeyValueIterator<RawKey, byte[]> iter,
+                            final Sensor operationSensor,
+                            final Sensor iteratorSensor,
+                            final Time time,
+                            final LongAdder numOpenIterators,
+                            final Set<MeteredIterator> openIterators) {
+        this.iter = iter;
+        this.operationSensor = operationSensor;
+        this.iteratorSensor = iteratorSensor;
+        this.time = time;
+        this.numOpenIterators = numOpenIterators;
+        this.openIterators = openIterators;
+        this.startNs = time.nanoseconds();
+        this.startTimestampMs = time.milliseconds();
+        numOpenIterators.increment();
+        openIterators.add(this);
+    }
+
+    // Final: the constructor's openIterators.add(this) sorts through this via the set's
+    // startTimestamp comparator, i.e. on a not-yet-fully-constructed object. Keeping it final stops a
+    // subclass from overriding it with something that reads its own not-yet-assigned state.
+    @Override
+    public final long startTimestamp() {
+        return startTimestampMs;
+    }
+
+    /**
+     * Delegates to the raw iterator. Subclasses that buffer a peeked element (the
+     * {@code KeyValueIterator}s) override this to also account for the buffered element.
+     */
+    public boolean hasNext() {
+        return iter.hasNext();
+    }
+
+    // Final: this owns the metering lifecycle (sensor recording, numOpenIterators decrement,
+    // openIterators deregistration). A subclass that overrode it and forgot super.close() would
+    // silently drop the decrement and deregistration. Subclasses vary only in next()/hasNext().
+    public final void close() {
+        try {
+            iter.close();
+        } finally {
+            final long duration = time.nanoseconds() - startNs;
+            operationSensor.record(duration);
+            iteratorSensor.record(duration);
+            numOpenIterators.decrement();
+            openIterators.remove(this);
+        }
+    }
+}

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

@@ -510,24 +510,13 @@ public KeyValueIterator<Windowed<K>, AggregationWithHeaders<AGG>> backwardFindSe
     }
 
     private class MeteredSessionStoreWithHeadersIterator
-        implements KeyValueIterator<Windowed<K>, AggregationWithHeaders<AGG>>, MeteredIterator {
+        extends AbstractMeteredIterator<Windowed<Bytes>>
+        implements KeyValueIterator<Windowed<K>, AggregationWithHeaders<AGG>> {
 
-        private final KeyValueIterator<Windowed<Bytes>, byte[]> iter;
-        private final long startNs;
-        private final long startTimestampMs;
         private KeyValue<Windowed<K>, AggregationWithHeaders<AGG>> cachedNext;
 
         private MeteredSessionStoreWithHeadersIterator(final KeyValueIterator<Windowed<Bytes>, byte[]> iter) {
-            this.iter = iter;
-            this.startNs = time.nanoseconds();
-            this.startTimestampMs = time.milliseconds();
-            numOpenIterators.increment();
-            openIterators.add(this);
-        }
-
-        @Override
-        public long startTimestamp() {
-            return startTimestampMs;
+            super(iter, fetchSensor, iteratorDurationSensor, time, numOpenIterators, openIterators);
         }
 
         @Override
@@ -552,19 +541,6 @@ public KeyValue<Windowed<K>, AggregationWithHeaders<AGG>> next() {
             return KeyValue.pair(windowedKey, value);
         }
 
-        @Override
-        public void close() {
-            try {
-                iter.close();
-            } finally {
-                final long duration = time.nanoseconds() - startNs;
-                fetchSensor.record(duration);
-                iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
-                openIterators.remove(this);
-            }
-        }
-
         @Override
         public Windowed<K> peekNextKey() {
             if (cachedNext == null) {
@@ -585,30 +561,13 @@ public Windowed<K> peekNextKey() {
      * non-negative when the window is constructed, so this iterator's {@code next()} can never throw.
      */
     private class MeteredSessionWithHeadersReadOnlyRecordIterator
-        implements ReadOnlyRecordIterator<Windowed<K>, AGG>, MeteredIterator {
-
-        private final KeyValueIterator<Windowed<Bytes>, byte[]> iter;
-        private final long startNs;
-        private final long startTimestampMs;
+        extends AbstractMeteredIterator<Windowed<Bytes>>
+        implements ReadOnlyRecordIterator<Windowed<K>, AGG> {
 
         private MeteredSessionWithHeadersReadOnlyRecordIterator(
             final KeyValueIterator<Windowed<Bytes>, byte[]> iter
         ) {
-            this.iter = iter;
-            this.startNs = time.nanoseconds();
-            this.startTimestampMs = time.milliseconds();
-            numOpenIterators.increment();
-            openIterators.add(this);
-        }
-
-        @Override
-        public long startTimestamp() {
-            return startTimestampMs;
-        }
-
-        @Override
-        public boolean hasNext() {
-            return iter.hasNext();
+            super(iter, fetchSensor, iteratorDurationSensor, time, numOpenIterators, openIterators);
         }
 
         @Override
@@ -626,18 +585,5 @@ public ReadOnlyRecord<Windowed<K>, AGG> next() {
             ((RecordHeaders) record.headers()).setReadOnly();
             return record;
         }
-
-        @Override
-        public void close() {
-            try {
-                iter.close();
-            } finally {
-                final long duration = time.nanoseconds() - startNs;
-                fetchSensor.record(duration);
-                iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
-                openIterators.remove(this);
-            }
-        }
     }
 }

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

@@ -657,48 +657,9 @@ private KeyValueIterator<K, ValueTimestampHeaders<V>> reverseAllInternal(
         return new MeteredTimestampedKeyValueStoreWithHeadersIterator(store.reverseAll(), allSensor);
     }
 
-    /**
-     * Shared scaffolding for the metered iterators below: tracks {@code num-open-iterators},
-     * {@code oldest-iterator-open-since-ms}, and per-operation iterator duration, and delegates
-     * closing the wrapped raw iterator. Subclasses only need to implement the deserializing
-     * {@code next()}/{@code hasNext()} (and, where applicable, {@code peekNextKey()}).
-     */
-    private abstract class AbstractMeteredIterator implements MeteredIterator {
-
-        final KeyValueIterator<Bytes, byte[]> iter;
-        private final Sensor sensor;
-        private final long startNs;
-        private final long startTimestampMs;
-
-        AbstractMeteredIterator(final KeyValueIterator<Bytes, byte[]> iter, final Sensor sensor) {
-            this.iter = iter;
-            this.sensor = sensor;
-            this.startNs = time.nanoseconds();
-            this.startTimestampMs = time.milliseconds();
-            numOpenIterators.increment();
-            openIterators.add(this);
-        }
-
-        @Override
-        public long startTimestamp() {
-            return startTimestampMs;
-        }
-
-        public void close() {
-            try {
-                iter.close();
-            } finally {
-                final long duration = time.nanoseconds() - startNs;
-                sensor.record(duration);
-                iteratorDurationSensor.record(duration);
-                numOpenIterators.decrement();
-                openIterators.remove(this);
-            }
-        }
-    }
-
     @SuppressWarnings("unchecked")
-    private class MeteredTimestampedKeyValueStoreWithHeadersQueryIterator extends AbstractMeteredIterator implements KeyValueIterator<K, V> {
+    private class MeteredTimestampedKeyValueStoreWithHeadersQueryIterator
+        extends AbstractMeteredIterator<Bytes> implements KeyValueIterator<K, V> {
 
         private final Function<byte[], ValueTimestampHeaders<V>> valueTimestampHeadersDeserializer;
 
@@ -711,7 +672,7 @@ private MeteredTimestampedKeyValueStoreWithHeadersQueryIterator(
             final Function<byte[], ValueTimestampHeaders<V>> valueTimestampHeadersDeserializer,
             final boolean returnPlainValue
         ) {
-            super(iter, sensor);
+            super(iter, sensor, iteratorDurationSensor, time, numOpenIterators, openIterators);
             this.valueTimestampHeadersDeserializer = valueTimestampHeadersDeserializer;
             this.returnPlainValue = returnPlainValue;
         }
@@ -777,7 +738,7 @@ public K peekNextKey() {
      * {@code next()} throws.
      */
     private class MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator
-        extends AbstractMeteredIterator implements ReadOnlyRecordIterator<K, V> {
+        extends AbstractMeteredIterator<Bytes> implements ReadOnlyRecordIterator<K, V> {
 
         private final Function<byte[], ValueTimestampHeaders<V>> valueTimestampHeadersDeserializer;
 
@@ -786,15 +747,10 @@ private MeteredTimestampedKeyValueStoreWithHeadersReadOnlyRecordIterator(
             final Sensor sensor,
             final Function<byte[], ValueTimestampHeaders<V>> valueTimestampHeadersDeserializer
         ) {
-            super(iter, sensor);
+            super(iter, sensor, iteratorDurationSensor, time, numOpenIterators, openIterators);
             this.valueTimestampHeadersDeserializer = valueTimestampHeadersDeserializer;
         }
 
-        @Override
-        public boolean hasNext() {
-            return iter.hasNext();
-        }
-
         @Override
         public ReadOnlyRecord<K, V> next() {
             final KeyValue<Bytes, byte[]> keyValue = iter.next();
@@ -817,15 +773,15 @@ public ReadOnlyRecord<K, V> next() {
     }
 
     private class MeteredTimestampedKeyValueStoreWithHeadersIterator
-        extends AbstractMeteredIterator implements KeyValueIterator<K, ValueTimestampHeaders<V>> {
+        extends AbstractMeteredIterator<Bytes> implements KeyValueIterator<K, ValueTimestampHeaders<V>> {
 
         private KeyValue<K, ValueTimestampHeaders<V>> cachedNext;
 
         private MeteredTimestampedKeyValueStoreWithHeadersIterator(
             final KeyValueIterator<Bytes, byte[]> iter,
             final Sensor sensor
         ) {
-            super(iter, sensor);
+            super(iter, sensor, iteratorDurationSensor, time, numOpenIterators, openIterators);
         }
 
         @Override