KAFKA-20173: Ensure Metered session-stores pass headers correctly

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

@@ -18,7 +18,6 @@
 
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
-import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.serialization.Serde;
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.Time;
@@ -45,9 +44,6 @@
 import org.apache.kafka.streams.state.WindowStoreIterator;
 
 import java.util.Objects;
-import java.util.Set;
-import java.util.concurrent.atomic.LongAdder;
-import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import static org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl.maybeMeasureLatency;
@@ -190,10 +186,7 @@ private <R> QueryResult<R> runWindowKeyQuery(
                             return vth == null ? null : ValueAndTimestamp.make(vth.value(), vth.timestamp());
                         }
                     );
-
-                    final QueryResult<MeteredWindowStoreIterator<ValueAndTimestamp<V>>> typedQueryResult =
-                        InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, typedResult);
-                    queryResult = (QueryResult<R>) typedQueryResult;
+                    queryResult = (QueryResult<R>) InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, typedResult);
                 } else {
                     // For non-timestamped stores, return plain V
                     final MeteredWindowStoreIterator<V> typedResult = meteredIterator(
@@ -203,10 +196,7 @@ private <R> QueryResult<R> runWindowKeyQuery(
                             return vth == null ? null : vth.value();
                         }
                     );
-
-                    final QueryResult<MeteredWindowStoreIterator<V>> typedQueryResult =
-                        InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, typedResult);
-                    queryResult = (QueryResult<R>) typedQueryResult;
+                    queryResult = (QueryResult<R>) InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, typedResult);
                 }
             } else {
                 queryResult = (QueryResult<R>) rawResult;
@@ -430,59 +420,16 @@ private <ValueType> MeteredWindowedKeyValueIterator<K, ValueType> meteredWindowe
             rawResult.getResult(),
             fetchSensor,
             iteratorDurationSensor,
+            this::deserializeValue,
             this::deserializeKey,
+            ValueTimestampHeaders::headers,
             valueConverter,
             time,
             numOpenIterators,
             openIterators
         );
     }
 
-    private final class MeteredWindowedKeyValueWithHeadersIterator<ValueType> extends MeteredWindowedKeyValueIterator<K, ValueType> {
-        private final BiFunction<byte[], Headers, K> deserializeKey;
-        private final Function<ValueTimestampHeaders<V>, ValueType> valueConverter;
-
-        MeteredWindowedKeyValueWithHeadersIterator(
-            final KeyValueIterator<Windowed<Bytes>, byte[]> iter,
-            final Sensor operationSensor,
-            final Sensor iteratorSensor,
-            final BiFunction<byte[], Headers, K> deserializeKey,
-            final Function<ValueTimestampHeaders<V>, ValueType> valueConverter,
-            final Time time,
-            final LongAdder numOpenIterators,
-            final Set<MeteredIterator> openIterators
-        ) {
-            super(
-                iter,
-                operationSensor,
-                iteratorSensor,
-                null, // should not be used in super-class
-                null, // should not be used in super-class
-                time,
-                numOpenIterators,
-                openIterators
-            );
-
-            this.deserializeKey = deserializeKey;
-            this.valueConverter = valueConverter;
-        }
-
-        @Override
-        public KeyValue<Windowed<K>, ValueType> next() {
-            final KeyValue<Windowed<Bytes>, byte[]> next = iter.next();
-            final ValueTimestampHeaders<V> valueTimestampHeaders = deserializeValue(next.value);
-            return KeyValue.pair(
-                windowedKey(next.key, valueTimestampHeaders.headers()),
-                valueConverter.apply(valueTimestampHeaders)
-            );
-        }
-
-        private Windowed<K> windowedKey(final Windowed<Bytes> bytesKey, final Headers headers) {
-            final K key = deserializeKey.apply(bytesKey.key().get(), headers);
-            return new Windowed<>(key, bytesKey.window());
-        }
-    }
-
     private boolean isUnderlyingStoreTimestamped() {
         StateStore store = wrapped();
         do {
@@ -507,16 +454,11 @@ private boolean isUnderlyingStoreTimestamped() {
         return false;
     }
 
-    protected Bytes serializeKey(final K key, final Headers headers) {
+    private Bytes serializeKey(final K key, final Headers headers) {
         return Bytes.wrap(serdes.rawKey(key, headers));
     }
 
-    @Override
-    protected K deserializeKey(final byte[] rawKey) {
-        throw new UnsupportedOperationException("MeteredTimestampedWindowStoreWithHeaders required to pass in Headers when deserializing a key.");
-    }
-
-    protected K deserializeKey(final byte[] rawKey, final Headers headers) {
+    private K deserializeKey(final byte[] rawKey, final Headers headers) {
         return serdes.keyFrom(rawKey, headers);
     }
 

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

@@ -541,7 +541,7 @@ private Bytes serializeKey(final K key) {
         return Bytes.wrap(serdes.rawKey(key, internalContext.headers()));
     }
 
-    protected K deserializeKey(final byte[] rawKey) {
+    private K deserializeKey(final byte[] rawKey) {
         return serdes.keyFrom(rawKey, internalContext.headers());
     }
 

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

@@ -0,0 +1,81 @@
+/*
+ * 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.header.Headers;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.utils.Bytes;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.state.KeyValueIterator;
+
+import java.util.Set;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+
+final class MeteredWindowedKeyValueWithHeadersIterator<K, VInner, VOuter> extends MeteredWindowedKeyValueIterator<K, VOuter> {
+    private final Function<byte[], VInner> deserializeValue;
+    private final BiFunction<byte[], Headers, K> deserializeKey;
+    private final Function<VInner, Headers> headersExtractor;
+    private final Function<VInner, VOuter> valueConverter;
+
+    MeteredWindowedKeyValueWithHeadersIterator(
+        final KeyValueIterator<Windowed<Bytes>, byte[]> iter,
+        final Sensor operationSensor,
+        final Sensor iteratorSensor,
+        final Function<byte[], VInner> deserializeValue,
+        final BiFunction<byte[], Headers, K> deserializeKey,
+        final Function<VInner, Headers> headersExtractor,
+        final Function<VInner, VOuter> valueConverter,
+        final Time time,
+        final LongAdder numOpenIterators,
+        final Set<MeteredIterator> openIterators
+    ) {
+        super(
+            iter,
+            operationSensor,
+            iteratorSensor,
+            null, // should not be used in super-class
+            null, // should not be used in super-class
+            time,
+            numOpenIterators,
+            openIterators
+        );
+
+        this.deserializeValue = deserializeValue;
+        this.deserializeKey = deserializeKey;
+        this.headersExtractor = headersExtractor;
+        this.valueConverter = valueConverter;
+    }
+
+    @Override
+    public KeyValue<Windowed<K>, VOuter> next() {
+        final KeyValue<Windowed<Bytes>, byte[]> next = iter.next();
+        final VInner valueTimestampHeaders = deserializeValue.apply(next.value);
+        return KeyValue.pair(
+            windowedKey(next.key, headersExtractor.apply(valueTimestampHeaders)),
+            valueConverter.apply(valueTimestampHeaders)
+        );
+    }
+
+    private Windowed<K> windowedKey(final Windowed<Bytes> bytesKey, final Headers headers) {
+        final K key = deserializeKey.apply(bytesKey.key().get(), headers);
+        return new Windowed<>(key, bytesKey.window());
+    }
+}

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

@@ -20,9 +20,7 @@
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.serialization.LongDeserializer;
-import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.internals.ByteUtils;
-import org.apache.kafka.streams.kstream.Windowed;
 import org.apache.kafka.streams.state.StateSerdes;
 
 import java.nio.ByteBuffer;
@@ -54,38 +52,6 @@ public static Headers headers(final byte[] valueWithHeaders) {
         return readHeaders(buffer);
     }
 
-    /**
-     * Serialize the key with headers into bytes
-     * @param key the key to serialize
-     * @param headers the Headers as context
-     * @param serdes the StateSerdes as serializer
-     * @return the Bytes of the key
-     */
-    public static <K> Bytes keyBytes(final K key, final Headers headers, final StateSerdes<K, ?> serdes) {
-        return Bytes.wrap(serdes.rawKey(key, headers));
-    }
-
-    /**
-     * Serialize the key into bytes
-     * @param key the key to serialize
-     * @param serdes the StateSerdes as serializer
-     * @return the Bytes of the key
-     */
-    static <K> Bytes keyBytes(final K key, final StateSerdes<K, ?> serdes) {
-        return keyBytes(key, new RecordHeaders(), serdes);
-    }
-
-    /**
-     * Serialize the session key with headers into bytes
-     * @param sessionKey the Windowed session key to serialize
-     * @param headers the Headers as context
-     * @param serdes the StateSerdes as serializer
-     * @return the Bytes of the key
-     */
-    static <K> Bytes keyBytes(final Windowed<K> sessionKey, final Headers headers, final StateSerdes<K, ?> serdes) {
-        return keyBytes(sessionKey.key(), headers, serdes);
-    }
-
     /**
      * Extract the raw aggregation bytes from serialized AggregationWithHeaders,
      * stripping the headers prefix.

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

@@ -103,6 +103,7 @@ public class MeteredSessionStoreTest {
     private static final Windowed<Bytes> WINDOWED_KEY_BYTES = new Windowed<>(KEY_BYTES, new SessionWindow(0, 0));
     private static final String VALUE = "value";
     private static final byte[] VALUE_BYTES = VALUE.getBytes();
+    private static final Headers HEADERS = new RecordHeaders();
     private static final long START_TIMESTAMP = 24L;
     private static final long END_TIMESTAMP = 42L;
     private static final int RETENTION_PERIOD = 100;
@@ -139,8 +140,7 @@ public void setUp() {
         setUpWithoutContext();
         metrics.config().recordLevel(Sensor.RecordingLevel.DEBUG);
         when(context.applicationId()).thenReturn(APPLICATION_ID);
-        when(context.metrics())
-                .thenReturn(new StreamsMetricsImpl(metrics, "test", mockTime));
+        when(context.metrics()).thenReturn(new StreamsMetricsImpl(metrics, "test", mockTime));
         when(context.taskId()).thenReturn(taskId);
         when(context.changelogFor(STORE_NAME)).thenReturn(CHANGELOG_TOPIC);
         when(innerStore.name()).thenReturn(STORE_NAME);
@@ -187,12 +187,13 @@ private void doShouldPassChangelogTopicNameToStateStoreSerde(final String topic)
         final Deserializer<String> valueDeserializer = mock(Deserializer.class);
         final Serializer<String> valueSerializer = mock(Serializer.class);
         when(keySerde.serializer()).thenReturn(keySerializer);
-        when(keySerializer.serialize(topic, new RecordHeaders(), KEY)).thenReturn(KEY.getBytes());
+        when(keySerializer.serialize(topic, HEADERS, KEY)).thenReturn(KEY.getBytes());
         when(valueSerde.deserializer()).thenReturn(valueDeserializer);
-        when(valueDeserializer.deserialize(topic, new RecordHeaders(), VALUE_BYTES)).thenReturn(VALUE);
+        when(valueDeserializer.deserialize(topic, HEADERS, VALUE_BYTES)).thenReturn(VALUE);
         when(valueSerde.serializer()).thenReturn(valueSerializer);
-        when(valueSerializer.serialize(topic, new RecordHeaders(), VALUE)).thenReturn(VALUE_BYTES);
+        when(valueSerializer.serialize(topic, HEADERS, VALUE)).thenReturn(VALUE_BYTES);
         when(innerStore.fetchSession(KEY_BYTES, START_TIMESTAMP, END_TIMESTAMP)).thenReturn(VALUE_BYTES);
+        when(context.headers()).thenReturn(HEADERS);
         store = new MeteredSessionStore<>(
             innerStore,
             STORE_TYPE,

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

@@ -841,6 +841,8 @@ private MeteredSessionStoreWithHeaders<String, String> createStoreWithMockSerdes
         lenient().when(keyDeserializer.deserialize(any(), eq(HEADERS), eq(KEY.getBytes())))
             .thenReturn(KEY);
 
+        when(context.headers()).thenReturn(new RecordHeaders());
+
         final MeteredSessionStoreWithHeaders<String, String> mockStore = new MeteredSessionStoreWithHeaders<>(
             innerStore,
             STORE_TYPE,