Address review: include headers in toString, cache typed reference in KeyValueStoreWrapper, normalize null headers in ChangeLogging

Author: Shekharrajak

streams/src/main/java/org/apache/kafka/streams/state/VersionedKeyValueStoreWithHeaders.java

@@ -36,9 +36,10 @@ public interface VersionedKeyValueStoreWithHeaders<K, V> extends VersionedKeyVal
      * @param key       The key
      * @param value     The value, it can be {@code null}. {@code null} is interpreted as a delete.
      * @param timestamp The timestamp for this record version
-     * @param headers   The record headers
+     * @param headers   The record headers; must not be {@code null}
      * @return The validTo timestamp of the newly put record, or special values as defined
      *         in {@link VersionedKeyValueStore#put(Object, Object, long)}.
+     * @throws NullPointerException if {@code headers} is {@code null}
      */
     long put(K key, V value, long timestamp, Headers headers);
 }

streams/src/main/java/org/apache/kafka/streams/state/VersionedRecord.java

@@ -106,7 +106,7 @@ public Headers headers() {
 
     @Override
     public String toString() {
-        return "<" + value + "," + timestamp + "," + validTo + ">";
+        return "<" + value + "," + timestamp + "," + validTo + "," + headers + ">";
     }
 
     @Override

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

@@ -47,8 +47,9 @@ public long put(final Bytes key, final byte[] value, final long timestamp) {
 
     @Override
     public long put(final Bytes key, final byte[] value, final long timestamp, final Headers headers) {
-        final long validTo = inner.put(key, value, timestamp, headers);
-        log(key, value, timestamp, headers != null ? headers : new RecordHeaders());
+        final Headers nonNullHeaders = headers != null ? headers : new RecordHeaders();
+        final long validTo = inner.put(key, value, timestamp, nonNullHeaders);
+        log(key, value, timestamp, nonNullHeaders);
         return validTo;
     }
 

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

@@ -51,6 +51,7 @@ public class KeyValueStoreWrapper<K, V> implements StateStore {
 
     private TimestampedKeyValueStoreWithHeaders<K, V> headersStore = null;
     private VersionedKeyValueStore<K, V> versionedStore = null;
+    private VersionedKeyValueStoreWithHeaders<K, V> versionedStoreWithHeaders = null;
 
     // same as either timestampedStore or versionedStore above. kept merely as a convenience
     // to simplify implementation for methods which do not depend on store type.
@@ -79,6 +80,9 @@ public KeyValueStoreWrapper(final ProcessorContext<?, ?> context, final String s
         try {
             versionedStore = context.getStateStore(storeName);
             store = versionedStore;
+            if (versionedStore instanceof VersionedKeyValueStoreWithHeaders) {
+                versionedStoreWithHeaders = (VersionedKeyValueStoreWithHeaders<K, V>) versionedStore;
+            }
         } catch (final ClassCastException e) {
             store = context.getStateStore(storeName);
             final String storeType = store == null ? "null" : store.getClass().getName();
@@ -129,8 +133,8 @@ public long put(final K key, final V value, final long timestamp, final Headers
             return PUT_RETURN_CODE_IS_LATEST;
         }
         if (versionedStore != null) {
-            if (versionedStore instanceof VersionedKeyValueStoreWithHeaders) {
-                return ((VersionedKeyValueStoreWithHeaders<K, V>) versionedStore).put(key, value, timestamp, headers);
+            if (versionedStoreWithHeaders != null) {
+                return versionedStoreWithHeaders.put(key, value, timestamp, headers);
             }
             return versionedStore.put(key, value, timestamp);
         }

streams/src/test/java/org/apache/kafka/streams/state/VersionedRecordWithHeadersTest.java

@@ -118,6 +118,8 @@ public void shouldPreserveEqualityForRecordsWithDifferentHeaders() {
     @Test
     public void shouldReturnCorrectToString() {
         final VersionedRecord<String> record = new VersionedRecord<>("value", 42L);
-        assertEquals("<value,42,Optional.empty>", record.toString());
+        final String str = record.toString();
+        assertTrue(str.startsWith("<value,42,Optional.empty,"));
+        assertTrue(str.contains("RecordHeaders"));
     }
 }