KAFKA-6718 / A Rack awareness for Kafka Streams

clients/src/main/java/org/apache/kafka/common/config/ConfigDef.java

@@ -16,8 +16,6 @@
  */
 package org.apache.kafka.common.config;
 
-import java.util.function.Function;
-import java.util.stream.Collectors;
 import org.apache.kafka.common.config.types.Password;
 import org.apache.kafka.common.utils.Utils;
 
@@ -33,8 +31,10 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.regex.Pattern;
+import java.util.stream.Collectors;
 
 /**
  * This class is used for specifying the set of expected configurations. For each configuration, you can specify
@@ -943,6 +943,27 @@ public String toString() {
         }
     }
 
+    public static class ListSize implements Validator {
+        final int maxSize;
+
+        private ListSize(final int maxSize) {
+            this.maxSize = maxSize;
+        }
+
+        public static ListSize max(final int maxSize) {
+            return new ListSize(maxSize);
+        }
+
+        @Override
+        public void ensureValid(final String name, final Object value) {
+            @SuppressWarnings("unchecked")
+            List<String> values = (List<String>) value;
+            if (values.size() > maxSize) {
+                throw new ConfigException(name, value, "exceeds maximum list size of [" + maxSize + "].");
+            }
+        }
+    }
+
     public static class ValidString implements Validator {
         final List<String> validStrings;
 

clients/src/test/java/org/apache/kafka/common/config/ConfigDefTest.java

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.config.ConfigDef.CaseInsensitiveValidString;
 import org.apache.kafka.common.config.ConfigDef.Importance;
+import org.apache.kafka.common.config.ConfigDef.ListSize;
 import org.apache.kafka.common.config.ConfigDef.Range;
 import org.apache.kafka.common.config.ConfigDef.Type;
 import org.apache.kafka.common.config.ConfigDef.ValidString;
@@ -711,4 +712,24 @@ public void testNiceTimeUnits() {
         assertEquals(" (365 days)", ConfigDef.niceTimeUnits(Duration.ofDays(365).toMillis()));
     }
 
+    @Test
+    public void testThrowsExceptionWhenListSizeExceedsLimit() {
+        assertThrows(ConfigException.class, () -> new ConfigDef().define("lst",
+                                                                         Type.LIST,
+                                                                         asList("a", "b"),
+                                                                         ListSize.max(1),
+                                                                         Importance.HIGH,
+                                                                         "lst doc"));
+    }
+
+    @Test
+    public void testNoExceptionIsThrownWhenListSizeIsWithinTheLimit() {
+        new ConfigDef().define("lst",
+                               Type.LIST,
+                               asList("a", "b"),
+                               ListSize.max(2),
+                               Importance.HIGH,
+                               "lst doc");
+    }
+
 }

streams/src/main/java/org/apache/kafka/streams/StreamsConfig.java

@@ -49,12 +49,16 @@
 import java.time.Duration;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.apache.kafka.common.IsolationLevel.READ_COMMITTED;
+import static org.apache.kafka.common.config.ConfigDef.ListSize.max;
 import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
 import static org.apache.kafka.common.config.ConfigDef.Range.between;
 import static org.apache.kafka.common.config.ConfigDef.ValidString.in;
@@ -146,9 +150,11 @@ public class StreamsConfig extends AbstractConfig {
     private static final long DEFAULT_COMMIT_INTERVAL_MS = 30000L;
     private static final long EOS_DEFAULT_COMMIT_INTERVAL_MS = 100L;
     private static final int DEFAULT_TRANSACTION_TIMEOUT = 10000;
+    private static final short DEFAULT_MAX_CLIENT_TAG_KEY_VALUE_LENGTH = 40;
 
     public static final int DUMMY_THREAD_INDEX = 1;
     public static final long MAX_TASK_IDLE_MS_DISABLED = -1;
+    public static final int MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE = 5;
 
     /**
      * Prefix used to provide default topic configs to be applied when creating internal topics.
@@ -214,6 +220,15 @@ public class StreamsConfig extends AbstractConfig {
     @SuppressWarnings("WeakerAccess")
     public static final String ADMIN_CLIENT_PREFIX = "admin.";
 
+    /**
+     * Prefix used to add arbitrary tags to a Kafka Stream's instance as key-value pairs.
+     * Example:
+     * client.tag.zone=zone1
+     * client.tag.cluster=cluster1
+     */
+    @SuppressWarnings("WeakerAccess")
+    public static final String CLIENT_TAG_PREFIX = "client.tag.";
+
     /**
      * Config value for parameter {@link #TOPOLOGY_OPTIMIZATION_CONFIG "topology.optimization"} for disabling topology optimization
      */
@@ -427,6 +442,10 @@ public class StreamsConfig extends AbstractConfig {
     public static final String DEFAULT_TIMESTAMP_EXTRACTOR_CLASS_CONFIG = "default.timestamp.extractor";
     private static final String DEFAULT_TIMESTAMP_EXTRACTOR_CLASS_DOC = "Default timestamp extractor class that implements the <code>org.apache.kafka.streams.processor.TimestampExtractor</code> interface.";
 
+    /** {@code max.client.tag.length} */
+    public static final String MAX_CLIENT_TAG_KEY_VALUE_LENGTH_CONFIG = "max.client.tag.key.value.length";
+    private static final String MAX_CLIENT_TAG_KEY_VALUE_LENGTH_DOC = "The maximum length of key-value pairs set via " + CLIENT_TAG_PREFIX + " configuration prefix.";
+
     /** {@code max.task.idle.ms} */
     public static final String MAX_TASK_IDLE_MS_CONFIG = "max.task.idle.ms";
     private static final String MAX_TASK_IDLE_MS_DOC = "Maximum amount of time in milliseconds a stream task will stay idle when not all of its partition buffers contain records," +
@@ -494,6 +513,13 @@ public class StreamsConfig extends AbstractConfig {
     @SuppressWarnings("WeakerAccess")
     public static final String RECEIVE_BUFFER_CONFIG = CommonClientConfigs.RECEIVE_BUFFER_CONFIG;
 
+    /** {@code rack.aware.assignment.tags} */
+    @SuppressWarnings("WeakerAccess")
+    public static final String RACK_AWARE_ASSIGNMENT_TAGS_CONFIG = "rack.aware.assignment.tags";
+    private static final String RACK_AWARE_ASSIGNMENT_TAGS_DOC = "List of client tag keys used to distribute standby replicas across Kafka Streams instances." +
+                                                                 " When configured, Kafka Streams will make a best-effort to distribute" +
+                                                                 " the standby tasks over each client tag dimension.";
+
     /** {@code reconnect.backoff.ms} */
     @SuppressWarnings("WeakerAccess")
     public static final String RECONNECT_BACKOFF_MS_CONFIG = CommonClientConfigs.RECONNECT_BACKOFF_MS_CONFIG;
@@ -599,6 +625,7 @@ public class StreamsConfig extends AbstractConfig {
             ProducerConfig.TRANSACTIONAL_ID_CONFIG
         };
 
+
     static {
         CONFIG = new ConfigDef()
 
@@ -714,6 +741,12 @@ public class StreamsConfig extends AbstractConfig {
                     -1,
                     Importance.MEDIUM,
                     REPLICATION_FACTOR_DOC)
+            .define(RACK_AWARE_ASSIGNMENT_TAGS_CONFIG,
+                    Type.LIST,
+                    Collections.emptyList(),
+                    max(MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE),
+                    Importance.MEDIUM,
+                    RACK_AWARE_ASSIGNMENT_TAGS_DOC)
             .define(SECURITY_PROTOCOL_CONFIG,
                     Type.STRING,
                     CommonClientConfigs.DEFAULT_SECURITY_PROTOCOL,
@@ -763,6 +796,12 @@ public class StreamsConfig extends AbstractConfig {
                     9 * 60 * 1000L,
                     ConfigDef.Importance.LOW,
                     CommonClientConfigs.CONNECTIONS_MAX_IDLE_MS_DOC)
+            .define(MAX_CLIENT_TAG_KEY_VALUE_LENGTH_CONFIG,
+                    Type.SHORT,
+                    DEFAULT_MAX_CLIENT_TAG_KEY_VALUE_LENGTH,
+                    between(1, 200),
+                    Importance.LOW,
+                    MAX_CLIENT_TAG_KEY_VALUE_LENGTH_DOC)
             .define(METADATA_MAX_AGE_CONFIG,
                     ConfigDef.Type.LONG,
                     5 * 60 * 1000L,
@@ -1024,6 +1063,16 @@ public static String adminClientPrefix(final String adminClientProp) {
         return ADMIN_CLIENT_PREFIX + adminClientProp;
     }
 
+    /**
+     * Prefix a client tag key with {@link #CLIENT_TAG_PREFIX}.
+     *
+     * @param clientTagKey client tag key
+     * @return {@link #CLIENT_TAG_PREFIX} + {@code clientTagKey}
+     */
+    public static String clientTagPrefix(final String clientTagKey) {
+        return CLIENT_TAG_PREFIX + clientTagKey;
+    }
+
     /**
      * Prefix a property with {@link #TOPIC_PREFIX}
      * used to provide default topic configs to be applied when creating internal topics.
@@ -1107,6 +1156,38 @@ protected Map<String, Object> postProcessParsedConfig(final Map<String, Object>
             configUpdates.put(COMMIT_INTERVAL_MS_CONFIG, EOS_DEFAULT_COMMIT_INTERVAL_MS);
         }
 
+        // validate rack awareness related configuration
+        final short maxClientTagKeyValueLength = getShort(MAX_CLIENT_TAG_KEY_VALUE_LENGTH_CONFIG);
+        final List<String> rackAwareAssignmentTags = getList(RACK_AWARE_ASSIGNMENT_TAGS_CONFIG);
+        final Map<String, String> clientTags = getClientTags();
+
+        if (clientTags.size() > MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE) {
+            throw new ConfigException("At most " + MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE + " client tags " +
+                                      "can be specified using " + CLIENT_TAG_PREFIX + " prefix.");
+        }
+
+        for (final String rackAwareAssignmentTag : rackAwareAssignmentTags) {
+            if (!clientTags.containsKey(rackAwareAssignmentTag)) {
+                throw new ConfigException(RACK_AWARE_ASSIGNMENT_TAGS_CONFIG,
+                                          rackAwareAssignmentTags,
+                                          "Contains invalid value [" + rackAwareAssignmentTag + "] " +
+                                          "which doesn't have corresponding tag set via [" + CLIENT_TAG_PREFIX + "] prefix.");
+            }
+        }
+
+        clientTags.forEach((tagKey, tagValue) -> {
+            if (tagKey.length() > maxClientTagKeyValueLength) {
+                throw new ConfigException(CLIENT_TAG_PREFIX + tagKey,
+                                          tagValue,
+                                          "Key exceeds maximum length of " + maxClientTagKeyValueLength + ".");
+            }
+            if (tagValue.length() > maxClientTagKeyValueLength) {
+                throw new ConfigException(CLIENT_TAG_PREFIX + tagKey,
+                                          tagValue,
+                                          "Value exceeds max length of " + maxClientTagKeyValueLength + ".");
+            }
+        });
+
         return configUpdates;
     }
 
@@ -1238,6 +1319,7 @@ public Map<String, Object> getMainConsumerConfigs(final String groupId, final St
         consumerProps.put(REPLICATION_FACTOR_CONFIG, getInt(REPLICATION_FACTOR_CONFIG));
         consumerProps.put(APPLICATION_SERVER_CONFIG, getString(APPLICATION_SERVER_CONFIG));
         consumerProps.put(NUM_STANDBY_REPLICAS_CONFIG, getInt(NUM_STANDBY_REPLICAS_CONFIG));
+        consumerProps.put(RACK_AWARE_ASSIGNMENT_TAGS_CONFIG, getList(RACK_AWARE_ASSIGNMENT_TAGS_CONFIG));
         consumerProps.put(ACCEPTABLE_RECOVERY_LAG_CONFIG, getLong(ACCEPTABLE_RECOVERY_LAG_CONFIG));
         consumerProps.put(MAX_WARMUP_REPLICAS_CONFIG, getInt(MAX_WARMUP_REPLICAS_CONFIG));
         consumerProps.put(PROBING_REBALANCE_INTERVAL_MS_CONFIG, getLong(PROBING_REBALANCE_INTERVAL_MS_CONFIG));
@@ -1389,6 +1471,23 @@ public Map<String, Object> getAdminConfigs(final String clientId) {
         return props;
     }
 
+    /**
+     * Get the configured client tags set with {@link #CLIENT_TAG_PREFIX} prefix.
+     * @return Map of the client tags.
+     */
+    @SuppressWarnings("WeakerAccess")
+    public Map<String, String> getClientTags() {
+        return originalsWithPrefix(CLIENT_TAG_PREFIX)
+            .entrySet()
+            .stream()
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    tagEntry -> Objects.toString(tagEntry.getValue())
+                )
+            );
+    }
+
     private Map<String, Object> getClientPropsWithPrefix(final String prefix,
                                                          final Set<String> configNames) {
         final Map<String, Object> props = clientProps(configNames, originals());

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

@@ -330,6 +330,7 @@ public static StreamThread create(final InternalTopologyBuilder builder,
         referenceContainer.adminClient = adminClient;
         referenceContainer.streamsMetadataState = streamsMetadataState;
         referenceContainer.time = time;
+        referenceContainer.clientTags = config.getClientTags();
 
         log.info("Creating restore consumer client");
         final Map<String, Object> restoreConsumerConfigs = config.getRestoreConsumerConfigs(getRestoreConsumerClientId(threadId));

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

@@ -128,19 +128,20 @@ private static class ClientMetadata {
         private final ClientState state;
         private final SortedSet<String> consumers;
 
-        ClientMetadata(final String endPoint) {
+        ClientMetadata(final String endPoint, final Map<String, String> clientTags) {
 
             // get the host info, or null if no endpoint is configured (ie endPoint == null)
             hostInfo = HostInfo.buildFromEndpoint(endPoint);
 
             // initialize the consumer memberIds
             consumers = new TreeSet<>();
 
-            // initialize the client state
-            state = new ClientState();
+            // initialize the client state with client tags
+            state = new ClientState(clientTags);
         }
 
-        void addConsumer(final String consumerMemberId, final List<TopicPartition> ownedPartitions) {
+        void addConsumer(final String consumerMemberId,
+                         final List<TopicPartition> ownedPartitions) {
             consumers.add(consumerMemberId);
             state.incrementCapacity();
             state.addOwnedPartitions(ownedPartitions, consumerMemberId);
@@ -188,6 +189,7 @@ public String toString() {
 
     private Supplier<TaskAssignor> taskAssignorSupplier;
     private byte uniqueField;
+    private Map<String, String> clientTags;
 
     /**
      * We need to have the PartitionAssignor and its StreamThread to be mutually accessible since the former needs
@@ -221,6 +223,7 @@ public void configure(final Map<String, ?> configs) {
         taskAssignorSupplier = assignorConfiguration::taskAssignor;
         assignmentListener = assignorConfiguration.assignmentListener();
         uniqueField = 0;
+        clientTags = referenceContainer.clientTags;
     }
 
     @Override
@@ -255,7 +258,8 @@ public ByteBuffer subscriptionUserData(final Set<String> topics) {
             userEndPoint,
             taskManager.getTaskOffsetSums(),
             uniqueField,
-            assignmentErrorCode.get()
+            assignmentErrorCode.get(),
+            clientTags
         ).encode();
     }
 
@@ -327,7 +331,7 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr
                 futureMetadataVersion = usedVersion;
                 processId = FUTURE_ID;
                 if (!clientMetadataMap.containsKey(FUTURE_ID)) {
-                    clientMetadataMap.put(FUTURE_ID, new ClientMetadata(null));
+                    clientMetadataMap.put(FUTURE_ID, new ClientMetadata(null, Collections.emptyMap()));
                 }
             } else {
                 processId = info.processId();
@@ -337,7 +341,7 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr
 
             // create the new client metadata if necessary
             if (clientMetadata == null) {
-                clientMetadata = new ClientMetadata(info.userEndPoint());
+                clientMetadata = new ClientMetadata(info.userEndPoint(), info.clientTags());
                 clientMetadataMap.put(info.processId(), clientMetadata);
             }
 
@@ -1255,6 +1259,7 @@ public void onAssignment(final Assignment assignment, final ConsumerGroupMetadat
             case 7:
             case 8:
             case 9:
+            case 10:
                 validateActiveTaskEncoding(partitions, info, logPrefix);
 
                 activeTasks = getActiveTasks(partitions, info);

streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentInfo.java

@@ -184,6 +184,7 @@ public ByteBuffer encode() {
                 case 7:
                 case 8:
                 case 9:
+                case 10:
                     out.writeInt(usedVersion);
                     out.writeInt(commonlySupportedVersion);
                     encodeActiveAndStandbyTaskAssignment(out);
@@ -359,6 +360,7 @@ public static AssignmentInfo decode(final ByteBuffer data) {
                 case 7:
                 case 8:
                 case 9:
+                case 10:
                     commonlySupportedVersion = in.readInt();
                     assignmentInfo = new AssignmentInfo(usedVersion, commonlySupportedVersion);
                     decodeActiveTasks(assignmentInfo, in);