KAFKA-20132: Implement TimestampedKeyValueStoreWithHeaders (4/N)

streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalProcessorContextImpl.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.streams.StreamsConfig;
@@ -148,6 +149,16 @@ public void logChange(final String storeName,
         throw new UnsupportedOperationException("this should not happen: logChange() not supported in global processor context.");
     }
 
+    @Override
+    public void logChange(final String storeName,
+                          final Bytes key,
+                          final byte[] value,
+                          final long timestamp,
+                          final Headers headers,
+                          final Position position) {
+        throw new UnsupportedOperationException("this should not happen: logChange() not supported in global processor context.");
+    }
+
     @Override
     public void transitionToActive(final StreamTask streamTask, final RecordCollector recordCollector, final ThreadCache newCache) {
         throw new UnsupportedOperationException("this should not happen: transitionToActive() not supported in global processor context.");

streams/src/main/java/org/apache/kafka/streams/processor/internals/InternalProcessorContext.java

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.serialization.BytesSerializer;
 import org.apache.kafka.common.utils.Bytes;
@@ -122,6 +123,13 @@ void logChange(final String storeName,
                    final long timestamp,
                    final Position position);
 
+    void logChange(final String storeName,
+                   final Bytes key,
+                   final byte[] value,
+                   final long timestamp,
+                   final Headers headers,
+                   final Position position);
+
     String changelogFor(final String storeName);
 
     void addProcessorMetadataKeyValue(final String key, final long value);

streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorContextImpl.java

@@ -133,11 +133,7 @@ public void logChange(final String storeName,
         if (!consistencyEnabled) {
             headers = null;
         } else {
-            // Add the vector clock to the header part of every record
-            headers = new RecordHeaders();
-            headers.add(ChangelogRecordDeserializationHelper.CHANGELOG_VERSION_HEADER_RECORD_CONSISTENCY);
-            headers.add(new RecordHeader(ChangelogRecordDeserializationHelper.CHANGELOG_POSITION_HEADER_KEY,
-                    PositionSerde.serialize(position).array()));
+            addVectorClockToHeaders(headers = new RecordHeaders(), position);
         }
 
         collector.send(
@@ -153,6 +149,40 @@ public void logChange(final String storeName,
             null);
     }
 
+    @Override
+    public void logChange(final String storeName,
+                          final Bytes key,
+                          final byte[] value,
+                          final long timestamp,
+                          final Headers headers,
+                          final Position position) {
+        throwUnsupportedOperationExceptionIfStandby("logChange");
+
+        final TopicPartition changelogPartition = stateManager().registeredChangelogPartitionFor(storeName);
+
+        if (consistencyEnabled) {
+            addVectorClockToHeaders(headers, position);
+        }
+
+        collector.send(
+            changelogPartition.topic(),
+            key,
+            value,
+            headers,
+            changelogPartition.partition(),
+            timestamp,
+            BYTES_KEY_SERIALIZER,
+            BYTEARRAY_VALUE_SERIALIZER,
+            null,
+            null);
+    }
+
+    private void addVectorClockToHeaders(Headers headers, Position position) {
+        headers.add(ChangelogRecordDeserializationHelper.CHANGELOG_VERSION_HEADER_RECORD_CONSISTENCY);
+        headers.add(new RecordHeader(ChangelogRecordDeserializationHelper.CHANGELOG_POSITION_HEADER_KEY,
+            PositionSerde.serialize(position).array()));
+    }
+
     /**
      * @throws StreamsException if an attempt is made to access this state store from an unknown node
      * @throws UnsupportedOperationException if the current streamTask type is standby

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

@@ -0,0 +1,99 @@
+/*
+ * 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.utils.Bytes;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.state.KeyValueStore;
+
+import java.util.List;
+
+import static org.apache.kafka.streams.state.internals.ValueTimestampHeadersDeserializer.headers;
+import static org.apache.kafka.streams.state.internals.ValueTimestampHeadersDeserializer.rawValue;
+import static org.apache.kafka.streams.state.internals.ValueTimestampHeadersDeserializer.timestamp;
+
+/**
+ * Change-logging wrapper for a timestamped key-value bytes store whose values also carry headers.
+ * <p>
+ * the header-aware serialized value format produced by {@link ValueTimestampHeadersSerializer}.
+ * <p>
+ * Semantics:
+ *  - The inner store value format is:
+ *        [ varint header_length ][ header_bytes ][ 8-byte timestamp ][ value_bytes ]
+ *  - The changelog record value logged via {@code log(...)} remains just {@code value_bytes}
+ *    (no timestamp, no headers), and the timestamp is logged separately.
+ */
+public class ChangeLoggingTimestampedKeyValueBytesStoreWithHeaders
+    extends ChangeLoggingKeyValueBytesStore {
+
+    ChangeLoggingTimestampedKeyValueBytesStoreWithHeaders(final KeyValueStore<Bytes, byte[]> inner) {
+        super(inner);
+    }
+
+    @Override
+    public void put(final Bytes key,
+                    final byte[] valueTimestampHeaders) {
+        wrapped().put(key, valueTimestampHeaders);
+        log(
+            key,
+            rawValue(valueTimestampHeaders),
+            valueTimestampHeaders == null
+                ? internalContext.recordContext().timestamp()
+                : timestamp(valueTimestampHeaders),
+            headers(valueTimestampHeaders)
+        );
+    }
+
+    @Override
+    public byte[] putIfAbsent(final Bytes key,
+                              final byte[] valueTimestampHeaders) {
+        final byte[] previous = wrapped().putIfAbsent(key, valueTimestampHeaders);
+        if (previous == null) {
+            // then it was absent
+            log(
+                key,
+                rawValue(valueTimestampHeaders),
+                valueTimestampHeaders == null
+                    ? internalContext.recordContext().timestamp()
+                    : timestamp(valueTimestampHeaders),
+                headers(valueTimestampHeaders)
+            );
+        }
+        return previous;
+    }
+
+    @Override
+    public void putAll(final List<KeyValue<Bytes, byte[]>> entries) {
+        wrapped().putAll(entries);
+        for (final KeyValue<Bytes, byte[]> entry : entries) {
+            final byte[] valueTimestampHeaders = entry.value;
+            log(
+                entry.key,
+                rawValue(valueTimestampHeaders),
+                valueTimestampHeaders == null
+                    ? internalContext.recordContext().timestamp()
+                    : timestamp(valueTimestampHeaders),
+                headers(valueTimestampHeaders)
+            );
+        }
+    }
+
+    void log(final Bytes key, final byte[] value, final long timestamp, final Headers headers) {
+        internalContext.logChange(name(), key, value, timestamp, headers, wrapped().getPosition());
+    }
+}

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

@@ -164,4 +164,17 @@ static Headers headers(final byte[] rawValueTimestampHeaders) {
         final byte[] rawHeaders = readBytes(buffer, headersSize);
         return HEADERS_DESERIALIZER.deserialize("", rawHeaders);
     }
+    /**
+     * Extract raw value from serialized ValueTimestampHeaders.
+     */
+    static byte[] rawValue(final byte[] rawValueTimestampHeaders) {
+        if (rawValueTimestampHeaders == null) {
+            return null;
+        }
+
+        final ByteBuffer buffer = ByteBuffer.wrap(rawValueTimestampHeaders);
+        final int headersSize = ByteUtils.readVarint(buffer);
+        buffer.position(buffer.position() + headersSize + Long.BYTES);
+        return readBytes(buffer, buffer.remaining());
+    }
 }

streams/src/test/java/org/apache/kafka/streams/processor/internals/AbstractProcessorContextTest.java

@@ -258,6 +258,15 @@ public void logChange(final String storeName,
                               final Position position) {
         }
 
+        @Override
+        public void logChange(final String storeName,
+                              final Bytes key,
+                              final byte[] value,
+                              final long timestamp,
+                              final Headers headers,
+                              final Position position) {
+        }
+
         @Override
         public void transitionToActive(final StreamTask streamTask, final RecordCollector recordCollector, final ThreadCache newCache) {
         }