This is an automated email from the ASF dual-hosted git repository.
mjsax pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new cba250ead18 KAFKA-20197: Headers aware StreamPartitioner (#22570)
cba250ead18 is described below
commit cba250ead18e8fa72dbf3b7539e1bed76fb6fce2
Author: Uladzislau Blok <[email protected]>
AuthorDate: Sun Jul 26 03:15:28 2026 +0200
KAFKA-20197: Headers aware StreamPartitioner (#22570)
Implementation for header-aware partitioners (KIP-1321).
This PR allows to pass in a records `Headers` object into
`StreamPartitioner`.
Reviewers: Alieh Saeedi <[email protected]>, Matthias J. Sax
<[email protected]>
---
.../org/apache/kafka/streams/KafkaStreams.java | 34 +++-
.../streams/kstream/internals/KTableImpl.java | 85 ++++++++-
.../internals/WindowedStreamPartitioner.java | 19 +-
.../kafka/streams/processor/StreamPartitioner.java | 17 +-
.../internals/DefaultStreamPartitioner.java | 9 +-
.../processor/internals/RecordCollectorImpl.java | 2 +-
.../processor/internals/StreamsMetadataState.java | 27 ++-
.../KafkaStreamsNamedTopologyWrapper.java | 6 +-
.../org/apache/kafka/streams/KafkaStreamsTest.java | 68 ++++++-
.../kstream/internals/KStreamRepartitionTest.java | 10 +-
.../KTableKTableForeignKeyJoinScenarioTest.java | 91 ++++++++++
.../internals/WindowedStreamPartitionerTest.java | 47 +++--
.../internals/DefaultStreamPartitionerTest.java | 86 +++++++++
.../processor/internals/ProcessorTopologyTest.java | 6 +
.../processor/internals/RecordCollectorTest.java | 65 ++++++-
.../internals/StreamsMetadataStateTest.java | 197 +++++++++++++++++----
16 files changed, 681 insertions(+), 88 deletions(-)
diff --git a/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
b/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
index ad9aada7c0d..aecdf3f67bb 100644
--- a/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
+++ b/streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java
@@ -30,6 +30,8 @@ import org.apache.kafka.common.annotation.InterfaceAudience;
import org.apache.kafka.common.annotation.InterfaceStability.Evolving;
import org.apache.kafka.common.annotation.SuppressKafkaInternalApiUsage;
import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.metrics.KafkaMetricsContext;
import org.apache.kafka.common.metrics.MetricConfig;
import org.apache.kafka.common.metrics.Metrics;
@@ -1784,11 +1786,24 @@ public class KafkaStreams implements AutoCloseable {
return streamsMetadataState.allMetadataForStore(storeName);
}
+ /**
+ * Finds the metadata containing the active hosts and standby hosts where
the key being queried would reside,
+ * without requiring record headers to be provided.
+ * <p>
+ * If your partitioner or serializer makes use of headers, use the {@link
#queryMetadataForKey(String, Object, Headers, Serializer) Headers} overload,
otherwise the returned metadata may not match where the key actually resides.
+ */
+ public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
+ final K key,
+ final Serializer<K>
keySerializer) {
+ return queryMetadataForKey(storeName, key, new RecordHeaders(),
keySerializer);
+ }
+
/**
* Finds the metadata containing the active hosts and standby hosts where
the key being queried would reside.
*
* @param storeName the {@code storeName} to find metadata for
* @param key the key to find metadata for
+ * @param headers the record headers
* @param keySerializer serializer for the key
* @param <K> key type
* Returns {@link KeyQueryMetadata} containing all metadata about hosting
the given key for the given store,
@@ -1796,9 +1811,22 @@ public class KafkaStreams implements AutoCloseable {
*/
public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
final K key,
+ final Headers headers,
final Serializer<K>
keySerializer) {
validateIsRunningOrRebalancing();
- return streamsMetadataState.keyQueryMetadataForKey(storeName, key,
keySerializer);
+ return streamsMetadataState.keyQueryMetadataForKey(storeName, key,
headers, keySerializer);
+ }
+
+ /**
+ * Finds the metadata containing the active hosts and standby hosts where
the key being queried would reside,
+ * using the supplied partitioner and without requiring record headers to
be provided.
+ * <p>
+ * If your partitioner or serializer makes use of headers, use the {@link
#queryMetadataForKey(String, Object, Headers, StreamPartitioner) Headers}
overload, otherwise the returned metadata may not match where the key actually
resides.
+ */
+ public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
+ final K key,
+ final StreamPartitioner<?
super K, ?> partitioner) {
+ return queryMetadataForKey(storeName, key, new RecordHeaders(),
partitioner);
}
/**
@@ -1806,6 +1834,7 @@ public class KafkaStreams implements AutoCloseable {
*
* @param storeName the {@code storeName} to find metadata for
* @param key the key to find metadata for
+ * @param headers the record headers
* @param partitioner the partitioner to be used to locate the host for
the key
* @param <K> key type
* Returns {@link KeyQueryMetadata} containing all metadata about hosting
the given key for the given store, using
@@ -1813,9 +1842,10 @@ public class KafkaStreams implements AutoCloseable {
*/
public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
final K key,
+ final Headers headers,
final StreamPartitioner<?
super K, ?> partitioner) {
validateIsRunningOrRebalancing();
- return streamsMetadataState.keyQueryMetadataForKey(storeName, key,
partitioner);
+ return streamsMetadataState.keyQueryMetadataForKey(storeName, key,
headers, partitioner);
}
/**
diff --git
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
index 5f0d043c7e4..1a190037141 100644
---
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
+++
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/KTableImpl.java
@@ -16,6 +16,7 @@
*/
package org.apache.kafka.streams.kstream.internals;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.Serde;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KeyValue;
@@ -1279,7 +1280,7 @@ public class KTableImpl<K, S, V> extends
AbstractStream<K, V> implements KTable<
final StreamPartitioner<KO, SubscriptionWrapper<K>>
subscriptionSinkPartitioner =
tableJoinedInternal.otherPartitioner() == null
? null
- : (topic, key, val, numPartitions) ->
getPartition.apply(tableJoinedInternal.otherPartitioner().partitions(topic,
key, null, numPartitions));
+ : new
SubscriptionSinkPartitioner<>(tableJoinedInternal.otherPartitioner(),
getPartition);
final StreamSinkNode<KO, SubscriptionWrapper<K>> subscriptionSink =
new StreamSinkNode<>(
renamed.suffixWithOrElseGet("-subscription-registration-sink",
builder, SINK_NAME),
@@ -1345,16 +1346,12 @@ public class KTableImpl<K, S, V> extends
AbstractStream<K, V> implements KTable<
final String finalRepartitionTopicName =
renamed.suffixWithOrElseGet("-subscription-response", builder,
SUBSCRIPTION_RESPONSE) + TOPIC_SUFFIX;
builder.internalTopologyBuilder.addInternalTopic(finalRepartitionTopicName,
InternalTopicProperties.empty());
- final StreamPartitioner<K, SubscriptionResponseWrapper<VO>>
defaultForeignResponseSinkPartitioner =
- (topic, key, subscriptionResponseWrapper, numPartitions) -> {
- final Integer partition =
subscriptionResponseWrapper.primaryPartition();
- return partition == null ? Optional.empty() :
Optional.of(Collections.singleton(partition));
- };
+ final StreamPartitioner<K, SubscriptionResponseWrapper<VO>>
defaultForeignResponseSinkPartitioner = new
DefaultForeignResponseSinkPartitioner<>();
final StreamPartitioner<K, SubscriptionResponseWrapper<VO>>
foreignResponseSinkPartitioner =
tableJoinedInternal.partitioner() == null
? defaultForeignResponseSinkPartitioner
- : (topic, key, val, numPartitions) ->
getPartition.apply(tableJoinedInternal.partitioner().partitions(topic, key,
null, numPartitions));
+ : new
ForeignResponseSinkPartitioner<>(tableJoinedInternal.partitioner(),
getPartition);
final StreamSinkNode<K, SubscriptionResponseWrapper<VO>>
foreignResponseSink =
new StreamSinkNode<>(
@@ -1438,4 +1435,78 @@ public class KTableImpl<K, S, V> extends
AbstractStream<K, V> implements KTable<
tableNode.setOutputVersioned(materializedInternal.storeSupplier()
instanceof VersionedBytesStoreSupplier);
}
}
+
+ private static class SubscriptionSinkPartitioner<K, KO> implements
StreamPartitioner<KO, SubscriptionWrapper<K>> {
+ private final StreamPartitioner<KO, Void> otherPartitioner;
+ private final Function<Optional<Set<Integer>>, Optional<Set<Integer>>>
getPartition;
+
+ private SubscriptionSinkPartitioner(
+ final StreamPartitioner<KO, Void> otherPartitioner,
+ final Function<Optional<Set<Integer>>, Optional<Set<Integer>>>
getPartition
+ ) {
+ this.otherPartitioner = otherPartitioner;
+ this.getPartition = getPartition;
+ }
+
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final KO
key, final SubscriptionWrapper<K> value, final int numPartitions) {
+ throw new UnsupportedOperationException("This method is deprecated
and should not be called.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic,
+ final KO key,
+ final SubscriptionWrapper<K>
value,
+ final Headers headers,
+ final int numPartitions) {
+ return getPartition.apply(otherPartitioner.partitions(topic, key,
null, headers, numPartitions));
+ }
+ }
+
+ private static class ForeignResponseSinkPartitioner<K, VO> implements
StreamPartitioner<K, SubscriptionResponseWrapper<VO>> {
+ private final StreamPartitioner<K, Void> partitioner;
+ private final Function<Optional<Set<Integer>>, Optional<Set<Integer>>>
getPartition;
+
+ private ForeignResponseSinkPartitioner(
+ final StreamPartitioner<K, Void> partitioner,
+ final Function<Optional<Set<Integer>>, Optional<Set<Integer>>>
getPartition
+ ) {
+ this.partitioner = partitioner;
+ this.getPartition = getPartition;
+ }
+
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final K
key, final SubscriptionResponseWrapper<VO> value, final int numPartitions) {
+ throw new UnsupportedOperationException("This method is deprecated
and should not be called.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic,
+ final K key,
+ final
SubscriptionResponseWrapper<VO> value,
+ final Headers headers,
+ final int numPartitions) {
+ return getPartition.apply(partitioner.partitions(topic, key, null,
headers, numPartitions));
+ }
+ }
+
+ private static class DefaultForeignResponseSinkPartitioner<K, VO>
implements StreamPartitioner<K, SubscriptionResponseWrapper<VO>> {
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final K
key, final SubscriptionResponseWrapper<VO> value, final int numPartitions) {
+ throw new UnsupportedOperationException("This method is deprecated
and should not be called.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic,
+ final K key,
+ final
SubscriptionResponseWrapper<VO> value,
+ final Headers headers,
+ final int numPartitions) {
+ final Integer partition = value.primaryPartition();
+ return partition == null ? Optional.empty() :
Optional.of(Collections.singleton(partition));
+ }
+ }
}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
index 54735b98bb1..b2386e31321 100644
---
a/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
+++
b/streams/src/main/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitioner.java
@@ -17,7 +17,7 @@
package org.apache.kafka.streams.kstream.internals;
import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
-import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.streams.kstream.Windowed;
import org.apache.kafka.streams.processor.StreamPartitioner;
@@ -33,21 +33,28 @@ public class WindowedStreamPartitioner<K, V> implements
StreamPartitioner<Window
this.serializer = serializer;
}
+ @SuppressWarnings({"removal"})
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
Windowed<K> windowedKey, final V value, final int numPartitions) {
+ throw new UnsupportedOperationException("This method is deprecated and
should not be called.");
+ }
+
/**
* WindowedStreamPartitioner determines the partition number for a record
with the given windowed key and value
* and the current number of partitions. The partition number id
determined by the original key of the windowed key
* using the same logic as DefaultPartitioner so that the topic is
partitioned by the original key.
*
- * @param topic the topic name this record is sent to
- * @param windowedKey the key of the record
- * @param value the value of the record
+ * @param topic the topic name this record is sent to
+ * @param windowedKey the key of the record
+ * @param value the value of the record
+ * @param headers the record headers
* @param numPartitions the total number of partitions
* @return an integer between 0 and {@code numPartitions-1}, or {@code
null} if the default partitioning logic should be used
*/
@Override
- public Optional<Set<Integer>> partitions(final String topic, final
Windowed<K> windowedKey, final V value, final int numPartitions) {
+ public Optional<Set<Integer>> partitions(final String topic, final
Windowed<K> windowedKey, final V value, final Headers headers, final int
numPartitions) {
// for windowed key, the key bytes should never be null
- final byte[] keyBytes = serializer.serializeBaseKey(topic, new
RecordHeaders(), windowedKey);
+ final byte[] keyBytes = serializer.serializeBaseKey(topic, headers,
windowedKey);
// stick with the same built-in partitioner util functions that
producer used
// to make sure its behavior is consistent with the producer
diff --git
a/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
b/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
index 20e6cfb5d81..bb81affa811 100644
---
a/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
+++
b/streams/src/main/java/org/apache/kafka/streams/processor/StreamPartitioner.java
@@ -17,6 +17,7 @@
package org.apache.kafka.streams.processor;
import org.apache.kafka.common.annotation.InterfaceAudience;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.streams.Topology;
import java.util.Optional;
@@ -44,7 +45,7 @@ import java.util.Set;
* for that topic.
* <p>
* All StreamPartitioner implementations should be stateless and a pure
function so they can be shared across topic and sink nodes.
- *
+ *
* @param <K> the type of keys
* @param <V> the type of values
* @see Topology#addSink(String, String,
org.apache.kafka.common.serialization.Serializer,
@@ -56,16 +57,26 @@ import java.util.Set;
public interface StreamPartitioner<K, V> {
/**
- * Determine the number(s) of the partition(s) to which a record with the
given key and value should be sent,
+ * @deprecated Since 4.4. Use {@link #partitions(String, Object, Object,
Headers, int)} instead.
+ * This method is planned to be removed in 5.0.
+ */
+ @Deprecated(since = "4.4", forRemoval = true)
+ Optional<Set<Integer>> partitions(String topic, K key, V value, int
numPartitions);
+
+ /**
+ * Determine the number(s) of the partition(s) to which a record with the
given key and value should be sent,
* for the given topic and current partition count
* @param topic the topic name this record is sent to
* @param key the key of the record
* @param value the value of the record
+ * @param headers the record headers
* @param numPartitions the total number of partitions
* @return an Optional of Set of integers between 0 and {@code
numPartitions-1},
* Empty optional means using default partitioner
* Optional of an empty set means the record won't be sent to any
partitions i.e drop it.
* Optional of Set of integers means the partitions to which the record
should be sent to.
* */
- Optional<Set<Integer>> partitions(String topic, K key, V value, int
numPartitions);
+ default Optional<Set<Integer>> partitions(final String topic, final K key,
final V value, final Headers headers, final int numPartitions) {
+ return partitions(topic, key, value, numPartitions);
+ }
}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
index 439eb379e00..868ed183236 100644
---
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
+++
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitioner.java
@@ -17,6 +17,7 @@
package org.apache.kafka.streams.processor.internals;
import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.streams.processor.StreamPartitioner;
@@ -32,9 +33,15 @@ public class DefaultStreamPartitioner<K, V> implements
StreamPartitioner<K, V> {
this.keySerializer = keySerializer;
}
+ @SuppressWarnings({"removal"})
@Override
public Optional<Set<Integer>> partitions(final String topic, final K key,
final V value, final int numPartitions) {
- final byte[] keyBytes = keySerializer.serialize(topic, key);
+ throw new UnsupportedOperationException("This method is deprecated and
should not be called.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final K key,
final V value, final Headers headers, final int numPartitions) {
+ final byte[] keyBytes = keySerializer.serialize(topic, headers, key);
// if the key bytes are not available, we just return empty optional
to let the producer decide
// which partition to send internally; otherwise stick with the same
built-in partitioner
diff --git
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
index dd9baeb6743..19f57dd7995 100644
---
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
+++
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/RecordCollectorImpl.java
@@ -159,7 +159,7 @@ public class RecordCollectorImpl implements RecordCollector
{
);
}
if (!partitions.isEmpty()) {
- final Optional<Set<Integer>> maybeMulticastPartitions =
partitioner.partitions(topic, key, value, partitions.size());
+ final Optional<Set<Integer>> maybeMulticastPartitions =
partitioner.partitions(topic, key, value, headers, partitions.size());
if (maybeMulticastPartitions.isEmpty()) {
// A null//empty partition indicates we should use the
default partitioner
send(topic, key, value, headers, null, timestamp,
keySerializer, valueSerializer, processorNodeId, context);
diff --git
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
index 5d84c599cf4..ea2bc70b336 100644
---
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
+++
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsMetadataState.java
@@ -18,6 +18,7 @@ package org.apache.kafka.streams.processor.internals;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.internals.LogContext;
import org.apache.kafka.streams.KafkaStreams;
@@ -190,7 +191,7 @@ public class StreamsMetadataState {
/**
* Find the {@link KeyQueryMetadata}s for a given storeName and key. This
method will use the
* {@link DefaultStreamPartitioner} to locate the store. If a custom
partitioner has been used
- * please use {@link StreamsMetadataState#keyQueryMetadataForKey(String,
Object, StreamPartitioner)} instead.
+ * please use {@link StreamsMetadataState#keyQueryMetadataForKey(String,
Object, Headers, StreamPartitioner)} instead.
*
* Note: the key may not exist in the {@link
org.apache.kafka.streams.processor.StateStore},
* this method provides a way of finding which {@link KeyQueryMetadata} it
would exist on.
@@ -205,6 +206,7 @@ public class StreamsMetadataState {
*/
public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final
String storeName,
final K
key,
+ final
Headers headers,
final
Serializer<K> keySerializer) {
Objects.requireNonNull(keySerializer, "keySerializer can't be null");
if (topologyMetadata.hasNamedTopologies()) {
@@ -214,19 +216,22 @@ public class StreamsMetadataState {
}
return keyQueryMetadataForKey(storeName,
key,
+ headers,
new
DefaultStreamPartitioner<>(keySerializer));
}
/**
- * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object,
Serializer)}
+ * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object,
Headers, Serializer)}
*/
public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final
String storeName,
final K
key,
+ final
Headers headers,
final
Serializer<K> keySerializer,
final
String topologyName) {
Objects.requireNonNull(keySerializer, "keySerializer can't be null");
return keyQueryMetadataForKey(storeName,
key,
+ headers,
new
DefaultStreamPartitioner<>(keySerializer),
topologyName);
}
@@ -239,6 +244,7 @@ public class StreamsMetadataState {
*
* @param storeName Name of the store
* @param key Key to use
+ * @param headers the record headers
* @param partitioner partitioner to use to find correct partition for key
* @param <K> key type
* @return The {@link KeyQueryMetadata} for the storeName and key or
{@link KeyQueryMetadata#NOT_AVAILABLE}
@@ -246,9 +252,11 @@ public class StreamsMetadataState {
*/
public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final
String storeName,
final K
key,
+ final
Headers headers,
final
StreamPartitioner<? super K, ?> partitioner) {
Objects.requireNonNull(storeName, "storeName can't be null");
Objects.requireNonNull(key, "key can't be null");
+ Objects.requireNonNull(headers, "headers can't be null");
Objects.requireNonNull(partitioner, "partitioner can't be null");
if (topologyMetadata.hasNamedTopologies()) {
throw new IllegalArgumentException("Cannot invoke the
keyQueryMetadataForKey(storeName, key, partitioner)"
@@ -273,22 +281,23 @@ public class StreamsMetadataState {
if (sourceTopicsInfo == null) {
return null;
}
- return keyQueryMetadataForKey(storeName, key, partitioner,
sourceTopicsInfo);
+ return keyQueryMetadataForKey(storeName, key, headers, partitioner,
sourceTopicsInfo);
}
/**
- * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object,
StreamPartitioner)}
+ * See {@link StreamsMetadataState#keyQueryMetadataForKey(String, Object,
Headers, StreamPartitioner)}
*/
public synchronized <K> KeyQueryMetadata keyQueryMetadataForKey(final
String storeName,
final K
key,
+ final
Headers headers,
final
StreamPartitioner<? super K, ?> partitioner,
final
String topologyName) {
Objects.requireNonNull(storeName, "storeName can't be null");
Objects.requireNonNull(key, "key can't be null");
+ Objects.requireNonNull(headers, "headers can't be null");
Objects.requireNonNull(partitioner, "partitioner can't be null");
Objects.requireNonNull(topologyName, "topologyName can't be null");
-
if (!isInitialized()) {
return KeyQueryMetadata.NOT_AVAILABLE;
}
@@ -297,7 +306,7 @@ public class StreamsMetadataState {
if (sourceTopicsInfo == null) {
return null;
}
- return keyQueryMetadataForKey(storeName, key, partitioner,
sourceTopicsInfo, topologyName);
+ return keyQueryMetadataForKey(storeName, key, headers, partitioner,
sourceTopicsInfo, topologyName);
}
/**
@@ -474,10 +483,11 @@ public class StreamsMetadataState {
private <K> KeyQueryMetadata keyQueryMetadataForKey(final String storeName,
final K key,
+ final Headers headers,
final
StreamPartitioner<? super K, ?> partitioner,
final SourceTopicsInfo
sourceTopicsInfo) {
- final Integer partition =
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
key, null, sourceTopicsInfo.maxPartitions));
+ final Integer partition =
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
key, null, headers, sourceTopicsInfo.maxPartitions));
final Set<TopicPartition> matchingPartitions = new HashSet<>();
for (final String sourceTopic : sourceTopicsInfo.sourceTopics) {
matchingPartitions.add(new TopicPartition(sourceTopic, partition));
@@ -507,11 +517,12 @@ public class StreamsMetadataState {
private <K> KeyQueryMetadata keyQueryMetadataForKey(final String storeName,
final K key,
+ final Headers headers,
final
StreamPartitioner<? super K, ?> partitioner,
final SourceTopicsInfo
sourceTopicsInfo,
final String
topologyName) {
Objects.requireNonNull(topologyName, "topology name must not be null");
- final Integer partition =
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
key, null, sourceTopicsInfo.maxPartitions));
+ final Integer partition =
getPartition.apply(partitioner.partitions(sourceTopicsInfo.topicWithMostPartitions,
key, null, headers, sourceTopicsInfo.maxPartitions));
final Set<TopicPartition> matchingPartitions = new HashSet<>();
for (final String sourceTopic : sourceTopicsInfo.sourceTopics) {
matchingPartitions.add(new TopicPartition(sourceTopic, partition));
diff --git
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
index b0a2a3ff07c..c9feee9e246 100644
---
a/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
+++
b/streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java
@@ -20,6 +20,8 @@ import
org.apache.kafka.clients.admin.DeleteConsumerGroupOffsetsResult;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.GroupIdNotFoundException;
import org.apache.kafka.common.errors.GroupSubscribedToTopicException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.internals.LogContext;
@@ -418,7 +420,7 @@ public class KafkaStreamsNamedTopologyWrapper extends
KafkaStreams {
}
/**
- * See {@link KafkaStreams#queryMetadataForKey(String, Object, Serializer)}
+ * See {@link KafkaStreams#queryMetadataForKey(String, Object, Headers,
Serializer)}
*/
public <K> KeyQueryMetadata queryMetadataForKey(final String storeName,
final K key,
@@ -426,7 +428,7 @@ public class KafkaStreamsNamedTopologyWrapper extends
KafkaStreams {
final String topologyName)
{
verifyTopologyStateStore(topologyName, storeName);
validateIsRunningOrRebalancing();
- return streamsMetadataState.keyQueryMetadataForKey(storeName, key,
keySerializer, topologyName);
+ return streamsMetadataState.keyQueryMetadataForKey(storeName, key, new
RecordHeaders(), keySerializer, topologyName);
}
/**
diff --git
a/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
b/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
index e1d8af5d318..8558aaf0d88 100644
--- a/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/KafkaStreamsTest.java
@@ -25,11 +25,14 @@ import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.Uuid;
import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.metrics.Metrics;
import org.apache.kafka.common.metrics.MetricsReporter;
import org.apache.kafka.common.metrics.Sensor.RecordingLevel;
import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.common.utils.LogCaptureAppender;
@@ -45,6 +48,7 @@ import
org.apache.kafka.streams.internals.metrics.ClientMetrics;
import org.apache.kafka.streams.kstream.Materialized;
import org.apache.kafka.streams.processor.StandbyUpdateListener;
import org.apache.kafka.streams.processor.StateRestoreListener;
+import org.apache.kafka.streams.processor.StreamPartitioner;
import org.apache.kafka.streams.processor.api.Processor;
import org.apache.kafka.streams.processor.api.ProcessorContext;
import org.apache.kafka.streams.processor.api.Record;
@@ -1109,13 +1113,60 @@ public class KafkaStreamsTest {
final AtomicReference<StreamThread.State> state2 =
prepareStreamThread(streamThreadTwo, 2);
prepareThreadState(streamThreadOne, state1);
prepareThreadState(streamThreadTwo, state2);
+ final StreamPartitioner<String, Object> simplePartitioner = new
SimplePartitioner();
try (final KafkaStreams streams = new
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
- assertThrows(StreamsNotStartedException.class, () ->
streams.queryMetadataForKey("store", "key", (topic, key, value, numPartitions)
-> Optional.of(Collections.singleton(0))));
+ assertThrows(StreamsNotStartedException.class, () ->
streams.queryMetadataForKey("store", "key", simplePartitioner));
streams.start();
waitForApplicationState(Collections.singletonList(streams),
KafkaStreams.State.RUNNING, DEFAULT_DURATION);
streams.close();
waitForApplicationState(Collections.singletonList(streams),
KafkaStreams.State.NOT_RUNNING, DEFAULT_DURATION);
- assertThrows(IllegalStateException.class, () ->
streams.queryMetadataForKey("store", "key", (topic, key, value, numPartitions)
-> Optional.of(Collections.singleton(0))));
+ assertThrows(IllegalStateException.class, () ->
streams.queryMetadataForKey("store", "key", simplePartitioner));
+ }
+ }
+
+ @Test
+ public void shouldPropagateSerializerAndHeadersToStreamsMetadataState() {
+ prepareStreams();
+ prepareStreamThread(streamThreadOne, 1);
+ prepareStreamThread(streamThreadTwo, 2);
+
+ try (final MockedConstruction<StreamsMetadataState>
metadataStateMockedConstruction = mockConstruction(StreamsMetadataState.class))
{
+ try (final KafkaStreams streams = new
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
+ streams.start();
+ final StreamsMetadataState mockMetadataState =
metadataStateMockedConstruction.constructed().get(0);
+
+ final Headers headers = new RecordHeaders();
+ headers.add("key", "value".getBytes());
+ final Serializer<String> serializer = new StringSerializer();
+
+ streams.queryMetadataForKey("store", "key", headers,
serializer);
+
+ verify(mockMetadataState).keyQueryMetadataForKey("store",
"key", headers, serializer);
+ }
+ }
+ }
+
+ @Test
+ public void shouldPropagatePartitionerAndHeadersToStreamsMetadataState() {
+ prepareStreams();
+ prepareStreamThread(streamThreadOne, 1);
+ prepareStreamThread(streamThreadTwo, 2);
+
+ try (final MockedConstruction<StreamsMetadataState>
metadataStateMockedConstruction = mockConstruction(StreamsMetadataState.class))
{
+ try (final KafkaStreams streams = new
KafkaStreams(getBuilderWithSource().build(), props, supplier, time)) {
+ streams.start();
+ final StreamsMetadataState mockMetadataState =
metadataStateMockedConstruction.constructed().get(0);
+
+ final Headers headers = new RecordHeaders();
+ headers.add("key", "value".getBytes());
+
+ @SuppressWarnings("unchecked")
+ final StreamPartitioner<String, Object> partitioner =
mock(StreamPartitioner.class);
+
+ streams.queryMetadataForKey("store", "key", headers,
partitioner);
+
+ verify(mockMetadataState).keyQueryMetadataForKey("store",
"key", headers, partitioner);
+ }
}
}
@@ -2056,4 +2107,17 @@ public class KafkaStreamsTest {
}
}
}
+
+ private static class SimplePartitioner implements
StreamPartitioner<String, Object> {
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions method
was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
+ return Optional.of(Collections.singleton(0));
+ }
+ }
}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
index 7a47632d449..ecb733a2a7d 100644
---
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KStreamRepartitionTest.java
@@ -16,6 +16,7 @@
*/
package org.apache.kafka.streams.kstream.internals;
+import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.serialization.IntegerDeserializer;
import org.apache.kafka.common.serialization.IntegerSerializer;
import org.apache.kafka.common.serialization.Serdes;
@@ -57,6 +58,7 @@ import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
@@ -85,8 +87,8 @@ public class KStreamRepartitionTest {
@SuppressWarnings("unchecked")
final StreamPartitioner<Integer, String> streamPartitionerMock =
mock(StreamPartitioner.class);
- when(streamPartitionerMock.partitions(anyString(), eq(0), eq("X0"),
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
- when(streamPartitionerMock.partitions(anyString(), eq(1), eq("X1"),
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
+ when(streamPartitionerMock.partitions(anyString(), eq(0), eq("X0"),
any(Headers.class),
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
+ when(streamPartitionerMock.partitions(anyString(), eq(1), eq("X1"),
any(Headers.class),
anyInt())).thenReturn(Optional.of(Collections.singleton(1)));
final String repartitionOperationName = "test";
final Repartitioned<Integer, String> repartitioned = Repartitioned
@@ -118,8 +120,8 @@ public class KStreamRepartitionTest {
assertTrue(testOutputTopic.readRecordsToList().isEmpty());
}
- verify(streamPartitionerMock).partitions(anyString(), eq(0), eq("X0"),
anyInt());
- verify(streamPartitionerMock).partitions(anyString(), eq(1), eq("X1"),
anyInt());
+ verify(streamPartitionerMock).partitions(anyString(), eq(0), eq("X0"),
any(Headers.class), anyInt());
+ verify(streamPartitionerMock).partitions(anyString(), eq(1), eq("X1"),
any(Headers.class), anyInt());
}
@Test
diff --git
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
index 1ccff8098c5..3f1ef503115 100644
---
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/KTableKTableForeignKeyJoinScenarioTest.java
@@ -16,6 +16,8 @@
*/
package org.apache.kafka.streams.kstream.internals;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.serialization.IntegerDeserializer;
import org.apache.kafka.common.serialization.IntegerSerializer;
import org.apache.kafka.common.serialization.Serdes;
@@ -33,25 +35,34 @@ import org.apache.kafka.streams.kstream.Consumed;
import org.apache.kafka.streams.kstream.KTable;
import org.apache.kafka.streams.kstream.Materialized;
import org.apache.kafka.streams.kstream.Produced;
+import org.apache.kafka.streams.kstream.TableJoined;
+import org.apache.kafka.streams.processor.StreamPartitioner;
import org.apache.kafka.streams.state.KeyValueStore;
+import org.apache.kafka.streams.test.TestRecord;
import org.apache.kafka.streams.utils.UniqueTopicSerdeScope;
import org.apache.kafka.test.StreamsTestUtils;
import org.apache.kafka.test.TestUtils;
+import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
+import java.util.Optional;
import java.util.Properties;
import java.util.Set;
+import java.util.concurrent.atomic.AtomicBoolean;
import static org.apache.kafka.common.utils.Utils.mkEntry;
import static org.apache.kafka.common.utils.Utils.mkMap;
import static org.apache.kafka.common.utils.Utils.mkProperties;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
public class KTableKTableForeignKeyJoinScenarioTest {
@@ -394,4 +405,84 @@ public class KTableKTableForeignKeyJoinScenarioTest {
assertThat(x, is(Collections.singletonMap(1,
"(999-alpha,(999-alpha,beta))")));
}
}
+
+ @Test
+ public void shouldPropagateHeadersToCustomForeignKeyJoinPartitioners() {
+ final StreamsBuilder builder = new StreamsBuilder();
+ final KTable<Integer, String> aTable = builder.table("A");
+ final KTable<Integer, String> bTable = builder.table("B");
+
+ final AtomicBoolean leftPartitionerCalled = new AtomicBoolean(false);
+ final AtomicBoolean rightPartitionerCalled = new AtomicBoolean(false);
+
+ final StreamPartitioner<Integer, Void> leftPartitioner = new
StreamPartitioner<>() {
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
Integer key, final Void value, final Headers headers, final int numPartitions) {
+ leftPartitionerCalled.set(true);
+ assertNotNull(headers);
+ assertNotNull(headers.lastHeader("test-header"));
+ assertEquals("test-value", new
String(headers.lastHeader("test-header").value()));
+ return Optional.of(Collections.singleton(0));
+ }
+
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
Integer key, final Void value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+ };
+
+ final StreamPartitioner<Integer, Void> rightPartitioner = new
StreamPartitioner<>() {
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
Integer key, final Void value, final Headers headers, final int numPartitions) {
+ rightPartitionerCalled.set(true);
+ assertNotNull(headers);
+ assertNotNull(headers.lastHeader("test-header"));
+ assertEquals("test-value", new
String(headers.lastHeader("test-header").value()));
+ return Optional.of(Collections.singleton(0));
+ }
+
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
Integer key, final Void value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+ };
+
+ final KTable<Integer, String> fkJoinResult = aTable.join(
+ bTable,
+ value -> Integer.parseInt(value.split("-")[0]),
+ (aVal, bVal) -> "(" + aVal + "," + bVal + ")",
+ TableJoined.with(leftPartitioner, rightPartitioner),
+ Materialized.as("asdf")
+ );
+
+ fkJoinResult.toStream().to("output");
+
+ final Properties config = new Properties();
+ config.setProperty(StreamsConfig.APPLICATION_ID_CONFIG, "test-app");
+ config.setProperty(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG,
"dummy:1234");
+ config.setProperty(StreamsConfig.STATE_DIR_CONFIG,
TestUtils.tempDirectory().getAbsolutePath());
+ config.setProperty(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG,
Serdes.Integer().getClass().getName());
+ config.setProperty(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG,
Serdes.String().getClass().getName());
+
+ try (final TopologyTestDriver driver = new
TopologyTestDriver(builder.build(), config)) {
+ final TestInputTopic<Integer, String> aTopic =
driver.createInputTopic("A", new IntegerSerializer(), new StringSerializer());
+ final TestInputTopic<Integer, String> bTopic =
driver.createInputTopic("B", new IntegerSerializer(), new StringSerializer());
+
+ final Headers headers = new RecordHeaders();
+ headers.add("test-header", "test-value".getBytes());
+
+ // Pipe input to A with headers. This triggers
SubscriptionSinkPartitioner,
+ // which delegates to rightPartitioner.
+ aTopic.pipeInput(new TestRecord<>(1, "999-alpha", headers));
+
+ // Pipe input to B with headers. This triggers
ForeignResponseSinkPartitioner,
+ // which delegates to leftPartitioner.
+ bTopic.pipeInput(new TestRecord<>(999, "beta", headers));
+ }
+
+ assertTrue(leftPartitionerCalled.get());
+ assertTrue(rightPartitionerCalled.get());
+ }
}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
index 25a15efdd4c..c8d27458290 100644
---
a/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/kstream/internals/WindowedStreamPartitionerTest.java
@@ -20,8 +20,9 @@ import
org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.serialization.IntegerSerializer;
-import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.streams.kstream.TimeWindowedSerializer;
import org.apache.kafka.streams.kstream.Windowed;
@@ -35,14 +36,17 @@ import java.util.Random;
import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
public class WindowedStreamPartitionerTest {
private final String topicName = "topic";
private final IntegerSerializer intSerializer = new IntegerSerializer();
- private final StringSerializer stringSerializer = new StringSerializer();
private final List<PartitionInfo> infos = Arrays.asList(
new PartitionInfo(topicName, 0, Node.noNode(), new Node[0], new
Node[0]),
@@ -57,28 +61,37 @@ public class WindowedStreamPartitionerTest {
Collections.emptySet(), Collections.emptySet());
@Test
- public void testCopartitioning() {
- final Random rand = new Random();
+ public void shouldThrowUnsupportedOperationExceptionForDeprecatedMethod() {
final WindowedSerializer<Integer> timeWindowedSerializer = new
TimeWindowedSerializer<>(intSerializer);
final WindowedStreamPartitioner<Integer, String> streamPartitioner =
new WindowedStreamPartitioner<>(timeWindowedSerializer);
+ final Windowed<Integer> windowedKey = new Windowed<>(1, new
TimeWindow(10, 20));
+
+ assertThrows(UnsupportedOperationException.class, () ->
streamPartitioner.partitions(topicName, windowedKey, "value", infos.size()));
+ }
- for (int k = 0; k < 10; k++) {
- final Integer key = rand.nextInt();
- final byte[] keyBytes = intSerializer.serialize(topicName, key);
+ @Test
+ public void testCopartitioningWithHeaders() {
+ final Random rand = new Random();
+ final Headers headers = new RecordHeaders();
+ headers.add("key", "value".getBytes());
- final String value = key.toString();
+ @SuppressWarnings("unchecked")
+ final WindowedSerializer<Integer> mockSerializer =
mock(WindowedSerializer.class);
+ final WindowedStreamPartitioner<Integer, String> streamPartitioner =
new WindowedStreamPartitioner<>(mockSerializer);
- final Set<Integer> expected =
Set.of(BuiltInPartitioner.partitionForKey(keyBytes,
cluster.partitionsForTopic(topicName).size()));
+ final Integer key = rand.nextInt();
+ final String value = key.toString();
+ final TimeWindow window = new TimeWindow(10, 20);
+ final Windowed<Integer> windowedKey = new Windowed<>(key, window);
+ final byte[] keyBytes = intSerializer.serialize(topicName, key);
+ final Set<Integer> expected =
Set.of(BuiltInPartitioner.partitionForKey(keyBytes,
cluster.partitionsForTopic(topicName).size()));
- for (int w = 1; w < 10; w++) {
- final TimeWindow window = new TimeWindow(10 * w, 20 * w);
+ when(mockSerializer.serializeBaseKey(topicName, headers,
windowedKey)).thenReturn(keyBytes);
- final Windowed<Integer> windowedKey = new Windowed<>(key,
window);
- final Optional<Set<Integer>> actual =
streamPartitioner.partitions(topicName, windowedKey, value, infos.size());
+ final Optional<Set<Integer>> actual =
streamPartitioner.partitions(topicName, windowedKey, value, headers,
infos.size());
- assertTrue(actual.isPresent());
- assertEquals(expected, actual.get());
- }
- }
+ assertTrue(actual.isPresent());
+ assertEquals(expected, actual.get());
+ verify(mockSerializer).serializeBaseKey(topicName, headers,
windowedKey);
}
}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
new file mode 100644
index 00000000000..57864c656cc
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/DefaultStreamPartitionerTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.processor.internals;
+
+import org.apache.kafka.clients.producer.internals.BuiltInPartitioner;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.serialization.Serializer;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.Optional;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class DefaultStreamPartitionerTest {
+
+ private static final String TOPIC = "topic";
+ private static final String KEY = "key";
+ private static final String VALUE = "value";
+ private static final int NUM_PARTITIONS = 5;
+
+ @SuppressWarnings("unchecked")
+ @Test
+ public void shouldPropagateHeadersToSerializer() {
+ final Serializer<String> keySerializer = mock(Serializer.class);
+ final DefaultStreamPartitioner<String, String>
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+ final Headers headers = new RecordHeaders();
+ headers.add("key", "value".getBytes());
+ final byte[] expectedKeyBytes = "serializedKey".getBytes();
+
+ when(keySerializer.serialize(TOPIC, headers,
KEY)).thenReturn(expectedKeyBytes);
+
+ final Optional<Set<Integer>> partition =
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, headers, NUM_PARTITIONS);
+
+ verify(keySerializer).serialize(TOPIC, headers, KEY);
+ assertTrue(partition.isPresent());
+
assertEquals(Collections.singleton(BuiltInPartitioner.partitionForKey(expectedKeyBytes,
NUM_PARTITIONS)), partition.get());
+ }
+
+ @SuppressWarnings({"unchecked"})
+ @Test
+ public void shouldThrowUnsupportedOperationExceptionForDeprecatedMethod() {
+ final Serializer<String> keySerializer = mock(Serializer.class);
+ final DefaultStreamPartitioner<String, String>
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+
+ assertThrows(UnsupportedOperationException.class, () ->
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, NUM_PARTITIONS));
+ }
+
+ @SuppressWarnings("unchecked")
+ @Test
+ public void shouldReturnEmptyOptionalWhenSerializedKeyIsNull() {
+ final Serializer<String> keySerializer = mock(Serializer.class);
+ final DefaultStreamPartitioner<String, String>
defaultStreamPartitioner = new DefaultStreamPartitioner<>(keySerializer);
+ final Headers headers = new RecordHeaders();
+
+ when(keySerializer.serialize(TOPIC, headers, KEY)).thenReturn(null);
+
+ final Optional<Set<Integer>> partition =
defaultStreamPartitioner.partitions(TOPIC, KEY, VALUE, headers, NUM_PARTITIONS);
+
+ verify(keySerializer).serialize(TOPIC, headers, KEY);
+ assertFalse(partition.isPresent());
+ }
+}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
index f39d818f2b8..55e8dddc2d3 100644
---
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/ProcessorTopologyTest.java
@@ -1103,8 +1103,14 @@ public class ProcessorTopologyTest {
}
static class DroppingPartitioner implements StreamPartitioner<String,
String> {
+ @SuppressWarnings("removal")
@Override
public Optional<Set<Integer>> partitions(final String topic, final
String key, final String value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions method
was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final String value, final Headers headers, final int numPartitions)
{
final Set<Integer> partitions = new HashSet<>();
for (int i = 1; i < numPartitions; i += 2) {
partitions.add(i);
diff --git
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
index 61c125058d5..b2cf4acdec2 100644
---
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/RecordCollectorTest.java
@@ -109,6 +109,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@ExtendWith(MockitoExtension.class)
@@ -137,8 +138,18 @@ public class RecordCollectorTest {
private final StringSerializer stringSerializer = new StringSerializer();
private final ByteArraySerializer byteArraySerializer = new
ByteArraySerializer();
- private final StreamPartitioner<String, Object> streamPartitioner =
- (topic, key, value, numPartitions) ->
Optional.of(Collections.singleton(Integer.parseInt(key) % numPartitions));
+ private final StreamPartitioner<String, Object> streamPartitioner = new
StreamPartitioner<String, Object>() {
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions method
was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
+ return Optional.of(Collections.singleton(Integer.parseInt(key) %
numPartitions));
+ }
+ };
private final MockProducer<byte[], byte[]> mockProducer
= new MockProducer<>(cluster, true, new
org.apache.kafka.clients.producer.RoundRobinPartitioner(), new
ByteArraySerializer(), new ByteArraySerializer());
@@ -304,8 +315,14 @@ public class RecordCollectorTest {
@Test
public void shouldSendOnlyToEvenPartitions() {
class EvenPartitioner implements StreamPartitioner<String, Object> {
+ @SuppressWarnings("removal")
@Override
public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
final Set<Integer> partitions = new HashSet<>();
for (int i = 0; i < numPartitions; i += 2) {
partitions.add(i);
@@ -369,8 +386,14 @@ public class RecordCollectorTest {
public void shouldBroadcastToAllPartitions() {
class BroadcastingPartitioner implements StreamPartitioner<String,
Object> {
+ @SuppressWarnings("removal")
@Override
public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
return Optional.of(IntStream.range(0,
numPartitions).boxed().collect(Collectors.toSet()));
}
}
@@ -430,8 +453,14 @@ public class RecordCollectorTest {
public void shouldDropAllRecords() {
class DroppingPartitioner implements StreamPartitioner<String, Object>
{
+ @SuppressWarnings("removal")
@Override
public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
return Optional.of(Collections.emptySet());
}
}
@@ -503,8 +532,14 @@ public class RecordCollectorTest {
public void shouldUseDefaultPartitionerViaPartitions() {
class DefaultPartitioner implements StreamPartitioner<String, Object> {
+ @SuppressWarnings("removal")
@Override
public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
return Optional.empty();
}
}
@@ -562,8 +597,18 @@ public class RecordCollectorTest {
@Test
public void shouldUseDefaultPartitionerAsPartitionReturnsEmptyOptional() {
- final StreamPartitioner<String, Object> streamPartitioner =
- (topic, key, value, numPartitions) -> Optional.empty();
+ final StreamPartitioner<String, Object> streamPartitioner = new
StreamPartitioner<String, Object>() {
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions
method was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
+ return Optional.empty();
+ }
+ };
final SinkNode<?, ?> sinkNode = new SinkNode<>(
sinkNodeName,
@@ -613,6 +658,18 @@ public class RecordCollectorTest {
assertEquals(9, mockProducer.history().size());
}
+ @Test
+ public void shouldPropagateHeadersToPartitioner() {
+ final Headers headers = new RecordHeaders(new Header[] {new
RecordHeader("key", "value".getBytes())});
+ @SuppressWarnings("unchecked")
+ final StreamPartitioner<String, Object> mockPartitioner =
mock(StreamPartitioner.class);
+ when(mockPartitioner.partitions(topic, "3", "0", headers,
3)).thenReturn(Optional.of(Collections.singleton(0)));
+
+ collector.send(topic, "3", "0", headers, null, stringSerializer,
stringSerializer, null, context, mockPartitioner);
+
+ verify(mockPartitioner).partitions(topic, "3", "0", headers, 3);
+ }
+
@Test
public void shouldUseDefaultPartitionerAsStreamPartitionerIsNull() {
diff --git
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
index 4ea73fcbf03..ae30308380d 100644
---
a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsMetadataStateTest.java
@@ -18,6 +18,8 @@ package org.apache.kafka.streams.processor.internals;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.serialization.Serdes;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.utils.internals.LogContext;
@@ -53,6 +55,11 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
public class StreamsMetadataStateTest {
@@ -129,20 +136,10 @@ public class StreamsMetadataStateTest {
topologyMetadata.buildAndRewriteTopology();
metadataState = new StreamsMetadataState(topologyMetadata, hostOne,
logContext);
metadataState.onChange(hostToActivePartitions,
hostToStandbyPartitions, partitionInfos);
- partitioner = (topic, key, value, numPartitions) ->
Optional.of(Collections.singleton(1));
+ partitioner = new FixedPartitionPartitioner(1);
storeNames = Set.of("table-one", "table-two", "merged-table",
globalTable);
}
- static class MultiValuedPartitioner implements StreamPartitioner<String,
Object> {
- @Override
- public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
- final Set<Integer> partitions = new HashSet<>();
- partitions.add(0);
- partitions.add(1);
- return Optional.of(partitions);
- }
- }
-
@Test
public void shouldNotThrowExceptionWhenOnChangeNotCalled() {
final Collection<StreamsMetadata> metadata = new StreamsMetadataState(
@@ -237,9 +234,13 @@ public class StreamsMetadataStateTest {
Collections.singletonMap(tp4, new PartitionInfo("topic-three", 1,
null, null, null)));
final KeyQueryMetadata expected = new KeyQueryMetadata(hostThree,
Set.of(hostTwo), 0);
- final KeyQueryMetadata actual =
metadataState.keyQueryMetadataForKey("table-three",
- "the-key",
-
Serdes.String().serializer());
+ final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+ "table-three",
+ "the-key",
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ );
+
assertEquals(expected, actual);
}
@@ -253,9 +254,12 @@ public class StreamsMetadataStateTest {
final KeyQueryMetadata expected = new KeyQueryMetadata(hostTwo,
Collections.emptySet(), 1);
- final KeyQueryMetadata actual =
metadataState.keyQueryMetadataForKey("table-three",
+ final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+ "table-three",
"the-key",
- partitioner);
+ new RecordHeaders(),
+ partitioner
+ );
assertEquals(expected, actual);
assertEquals(1, actual.partition());
}
@@ -269,15 +273,23 @@ public class StreamsMetadataStateTest {
Collections.singletonMap(tp4, new PartitionInfo("topic-three",
1, null, null, null)));
- assertThrows(IllegalArgumentException.class, () ->
metadataState.keyQueryMetadataForKey("table-three",
+ assertThrows(IllegalArgumentException.class, () ->
metadataState.keyQueryMetadataForKey(
+ "table-three",
"the-key",
- new MultiValuedPartitioner()));
+ new RecordHeaders(),
+ new MultiValuedPartitioner()
+ ));
}
@Test
public void shouldReturnNotAvailableWhenClusterIsEmpty() {
metadataState.onChange(Collections.emptyMap(), Collections.emptyMap(),
Collections.emptyMap());
- final KeyQueryMetadata result =
metadataState.keyQueryMetadataForKey("table-one", "a",
Serdes.String().serializer());
+ final KeyQueryMetadata result = metadataState.keyQueryMetadataForKey(
+ "table-one",
+ "a",
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ );
assertEquals(KeyQueryMetadata.NOT_AVAILABLE, result);
}
@@ -291,39 +303,76 @@ public class StreamsMetadataStateTest {
final KeyQueryMetadata expected = new KeyQueryMetadata(hostTwo,
Set.of(hostOne), 2);
- final KeyQueryMetadata actual =
metadataState.keyQueryMetadataForKey("merged-table", "the-key",
- (topic, key, value, numPartitions) ->
Optional.of(Collections.singleton(2)));
+ final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+ "merged-table",
+ "the-key",
+ new RecordHeaders(),
+ new FixedPartitionPartitioner(2)
+ );
assertEquals(expected, actual);
}
@Test
public void shouldReturnNullOnGetWithKeyWhenStoreDoesntExist() {
- final KeyQueryMetadata actual =
metadataState.keyQueryMetadataForKey("not-a-store",
+ final KeyQueryMetadata actual = metadataState.keyQueryMetadataForKey(
+ "not-a-store",
"key",
- Serdes.String().serializer());
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ );
assertNull(actual);
}
@Test
public void shouldThrowWhenKeyIsNull() {
- assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey("table-three", null,
Serdes.String().serializer()));
+ assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(
+ "table-three",
+ null,
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ ));
}
@Test
public void shouldThrowWhenSerializerIsNull() {
- assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey("table-three", "key", (Serializer<Object>)
null));
+ assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(
+ "table-three",
+ "key",
+ new RecordHeaders(),
+ (Serializer<Object>) null
+ ));
}
@Test
public void shouldThrowIfStoreNameIsNull() {
- assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(null, "key",
Serdes.String().serializer()));
+ assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(
+ null,
+ "key",
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ ));
}
@SuppressWarnings("unchecked")
@Test
public void shouldThrowIfStreamPartitionerIsNull() {
- assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(null, "key", (StreamPartitioner) null));
+ assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(
+ null,
+ "key",
+ new RecordHeaders(),
+ (StreamPartitioner) null
+ ));
+ }
+
+ @Test
+ public void shouldThrowIfHeadersIsNull() {
+ assertThrows(NullPointerException.class, () ->
metadataState.keyQueryMetadataForKey(
+ "table-three",
+ "key",
+ null,
+ Serdes.String().serializer()
+ ));
}
@Test
@@ -347,7 +396,12 @@ public class StreamsMetadataStateTest {
@Test
public void shouldGetQueryMetadataForGlobalStoreWithKey() {
- final KeyQueryMetadata metadata =
metadataState.keyQueryMetadataForKey(globalTable, "key",
Serdes.String().serializer());
+ final KeyQueryMetadata metadata = metadataState.keyQueryMetadataForKey(
+ globalTable,
+ "key",
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ );
assertEquals(hostOne, metadata.activeHost());
assertTrue(metadata.standbyHosts().isEmpty());
}
@@ -360,12 +414,22 @@ public class StreamsMetadataStateTest {
logContext
);
streamsMetadataState.onChange(hostToActivePartitions,
hostToStandbyPartitions, partitionInfos);
- assertNotNull(streamsMetadataState.keyQueryMetadataForKey(globalTable,
"key", Serdes.String().serializer()));
+ assertNotNull(streamsMetadataState.keyQueryMetadataForKey(
+ globalTable,
+ "key",
+ new RecordHeaders(),
+ Serdes.String().serializer()
+ ));
}
@Test
public void shouldGetQueryMetadataForGlobalStoreWithKeyAndPartitioner() {
- final KeyQueryMetadata metadata =
metadataState.keyQueryMetadataForKey(globalTable, "key", partitioner);
+ final KeyQueryMetadata metadata = metadataState.keyQueryMetadataForKey(
+ globalTable,
+ "key",
+ new RecordHeaders(),
+ partitioner
+ );
assertEquals(hostOne, metadata.activeHost());
assertTrue(metadata.standbyHosts().isEmpty());
}
@@ -378,7 +442,43 @@ public class StreamsMetadataStateTest {
logContext
);
streamsMetadataState.onChange(hostToActivePartitions,
hostToStandbyPartitions, partitionInfos);
- assertNotNull(streamsMetadataState.keyQueryMetadataForKey(globalTable,
"key", partitioner));
+ assertNotNull(streamsMetadataState.keyQueryMetadataForKey(
+ globalTable,
+ "key",
+ new RecordHeaders(),
+ partitioner)
+ );
+ }
+
+ @Test
+ public void shouldPropagateHeadersToPartitioner() {
+ final Headers headers = new RecordHeaders();
+ headers.add("key", "value".getBytes());
+
+ @SuppressWarnings("unchecked")
+ final StreamPartitioner<String, Object> mockPartitioner =
mock(StreamPartitioner.class);
+ when(mockPartitioner.partitions(
+ eq("topic-three"),
+ eq("the-key"),
+ eq(null),
+ eq(headers),
+ anyInt()
+ )).thenReturn(Optional.of(Collections.singleton(0)));
+
+ metadataState.keyQueryMetadataForKey(
+ "table-three",
+ "the-key",
+ headers,
+ mockPartitioner
+ );
+
+ verify(mockPartitioner).partitions(
+ "topic-three",
+ "the-key",
+ null,
+ headers,
+ 1
+ );
}
@Test
@@ -404,4 +504,39 @@ public class StreamsMetadataStateTest {
assertFalse(metadataState.allMetadata().isEmpty(), "encapsulation
broken");
}
+
+ private static class MultiValuedPartitioner implements
StreamPartitioner<String, Object> {
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions method
was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
+ final Set<Integer> partitions = new HashSet<>();
+ partitions.add(0);
+ partitions.add(1);
+ return Optional.of(partitions);
+ }
+ }
+
+ private static class FixedPartitionPartitioner implements
StreamPartitioner<String, Object> {
+ private final int partition;
+
+ FixedPartitionPartitioner(final int partition) {
+ this.partition = partition;
+ }
+
+ @SuppressWarnings("removal")
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final int numPartitions) {
+ throw new AssertionError("Deprecated 4-argument partitions method
was called instead of 5-argument method containing headers.");
+ }
+
+ @Override
+ public Optional<Set<Integer>> partitions(final String topic, final
String key, final Object value, final Headers headers, final int numPartitions)
{
+ return Optional.of(Collections.singleton(partition));
+ }
+ }
}