Kafka小流量灰度&蓝绿

Kafka小流量灰度&蓝绿.md

在全链路灰度背景下，Kafka作为一些系统的流量触发入口（消息驱动），Kafka必须能做到准确无误的识别灰度标识。 在比较单独新建灰度主题/独立集群等业界常规方案后，发现该方案对下游业务改动较大（需消费新灰度主题/新灰度集群等），推动实施周期长难以执行。 最终决定在原有topic的分区上动手解决这个问题，比如每个主题的分区0作为小流量灰度的专用灰度分区，奇偶分区来作为蓝绿分区。

另：该方案仅需在kafka-client上扩展自定义策略即可，无需改动server端代码。

发送

1. 实现Partitioner接口，在partition方法内添加灰度逻辑；
2. 客户端初始化时指定partitioner.class为如上实现类；

properties.put(ProducerConfig.ACKS_CONFIG, String.valueOf(this.extraConfig.getRequestRequiredAck().getValue()));

消费

1. 实现ConsumerPartitionAssignor接口 在subscriptionUserData(Set topics)方法内添加当前节点环境信息Beta等，或者直接在GROUP_INSTANCE_ID上识别也可； 在assign(Cluster metadata, GroupSubscription groupSubscription)中实现灰度节点分配0分区；

2. 客户端初始化时指定partition.assignment.strategy为如上实现类；

        String randomStr = RandomStringUtils.randomAlphanumeric(6);
  
        properties.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, Arrays.asList(SFRangeAssignor.class));
        properties.put(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG,
                X_ENV_TAG_VALUE.envTagName() + SFMQGreyConstants.GROUP_INSTANCE_ID_SPLIT + NetUtil.getHostName()
                        + randomStr);

TODO:

1. 消费分区策略增加cooperative实现；
2. 灰度主题配置 & 开启开关有统一的控制面来配置；

---

---

SFDefaultPartitioner.java

/*
 * 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 com.sf.kafka.api.grey;

import static com.sf.kafka.api.grey.SFMQGreyConstants.KAFKA_GREY_TOPICS2;
import static com.sf.kafka.api.grey.SFMQGreyConstants.X_ENV_TAG_VALUE;
import static com.sf.kafka.util.KafkaUtil.CLUSTER_ID_ALIAS;
import com.google.common.collect.Sets;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

/**
 * The default partitioning strategy:
 * <ul>
 * <li>If a partition is specified in the record, use it
 * <li>If no partition is specified but a key is present choose a partition based on a hash of the key
 * <li>If no partition or key is present choose the sticky partition that changes when the batch is full.
 *
 * See KIP-480 for details about sticky partitioning.
 *
 * @author YANGLiiN
 */
public class SFDefaultPartitioner implements Partitioner {

    private final SFStickyPartitionCache stickyPartitionCache = new SFStickyPartitionCache();

    public void configure(Map<String, ?> configs) {
    }

    /**
     * Compute the partition for the given record.
     *
     * @param topic      The topic name
     * @param key        The key to partition on (or null if no key)
     * @param keyBytes   serialized key to partition on (or null if no key)
     * @param value      The value to partition on or null
     * @param valueBytes serialized value to partition on or null
     * @param cluster    The current cluster metadata
     */
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        return partition(topic, key, keyBytes, value, valueBytes, cluster, cluster.partitionsForTopic(topic).size());
    }


    /**
     * Compute the partition for the given record.
     *
     * @param topic         The topic name
     * @param numPartitions The number of partitions of the given {@code topic}
     * @param key           The key to partition on (or null if no key)
     * @param keyBytes      serialized key to partition on (or null if no key)
     * @param value         The value to partition on or null
     * @param valueBytes    serialized value to partition on or null
     * @param cluster       The current cluster metadata
     */
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster,
                         int numPartitions) {
        if (keyBytes == null) {
            return stickyPartitionCache.partition(topic, cluster);
        }

        // hash the keyBytes to choose a partition
        int keyHash = Utils.toPositive(Utils.murmur2(keyBytes));

        // 已配置灰度、蓝绿主题 && 开启灰度开关
        String clusterId = cluster.clusterResource().clusterId();
        String clusterAlias = CLUSTER_ID_ALIAS.getOrDefault(clusterId, "");
        if (KAFKA_GREY_TOPICS2.getOrDefault(clusterAlias, Sets.newHashSet()).contains(topic)
                && !X_ENV_TAG_VALUE.isUnknown()) {
            Predicate<? super PartitionInfo> predicate = X_ENV_TAG_VALUE.partitionInfoPredicate();
            return selectPartition(topic, keyHash, cluster, predicate);
        }

        return keyHash % numPartitions;
    }

    private final int selectPartition(String topic, int keyHash, Cluster cluster,
                                      Predicate<? super PartitionInfo> predicate) {

        List<PartitionInfo> all = cluster.partitionsForTopic(topic);
        List<PartitionInfo> partitions = all.stream().filter(predicate).collect(Collectors.toList());

        return partitions.get(keyHash % partitions.size()).partition();
    }

    public void close() {
    }

    /**
     * If a batch completed for the current sticky partition, change the sticky partition.
     * Alternately, if no sticky partition has been determined, set one.
     */
    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
        stickyPartitionCache.nextPartition(topic, cluster, prevPartition);
    }
}

SFMQEnvTagEnum.java

package com.sf.kafka.api.grey;

import java.util.Locale;
import java.util.function.Predicate;
import org.apache.commons.lang3.StringUtils;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public enum SFMQEnvTagEnum {

    UNKNOWN("unknown", pi -> true, tp -> true),

    GREY("grey", pi -> pi.partition() == 0, tp -> tp.partition() == 0),

    NORMAL("normal", pi -> pi.partition() != 0, tp -> tp.partition() != 0),

    BLUE("blue", pi -> pi.partition() % 2 == 0, tp -> tp.partition() % 2 == 0),

    GREEN("green", pi -> pi.partition() % 2 == 1, tp -> tp.partition() % 2 == 1);


    private static final Logger log = LoggerFactory.getLogger(SFMQEnvTagEnum.class);


    private final String envTagName;
    private final Predicate<? super PartitionInfo> partitionInfoPredicate;
    private final Predicate<? super TopicPartition> topicPartitionPredicate;

    SFMQEnvTagEnum(String envTagName,
                   Predicate<? super PartitionInfo> partitionInfoPredicate,
                   Predicate<? super TopicPartition> topicPartitionPredicate) {
        this.envTagName = envTagName;
        this.partitionInfoPredicate = partitionInfoPredicate;
        this.topicPartitionPredicate = topicPartitionPredicate;
    }

    public String envTagName() {
        return envTagName;
    }

    public Predicate<? super PartitionInfo> partitionInfoPredicate() {
        return partitionInfoPredicate;
    }

    public Predicate<? super TopicPartition> topicPartitionPredicate() {
        return topicPartitionPredicate;
    }

    public static SFMQEnvTagEnum toEnvTagEnum(String envTagName) {
        if (StringUtils.isBlank(envTagName)) {
            return UNKNOWN;
        }

        SFMQEnvTagEnum tag;
        try {
            tag = SFMQEnvTagEnum.valueOf(envTagName.toUpperCase(Locale.ROOT));
        } catch (IllegalArgumentException e) {
            log.warn("env tag 配置不对！！！ 按未配置 ENV 场景走");
            tag = UNKNOWN;
        }
        return tag;
    }

    public boolean isUnknown() {
        return UNKNOWN == this;
    }

    @Override
    public String toString() {
        return envTagName;
    }
}

SFMQGreyConstants.java

package com.sf.kafka.api.grey;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import org.apache.commons.lang3.StringUtils;

public class SFMQGreyConstants {

    // sf k8s env key
    public static final String X_ENV_TAG_KEY = "ENV_TAG";
    // 当前节点环境标识
    public static final SFMQEnvTagEnum X_ENV_TAG_VALUE;

    // grey topic
    public static final String X_KAFKA_TOPICS_KEY = "x-kafka-topics";

    public static final Map<String, Set<String>> KAFKA_GREY_TOPICS2 = Maps.newHashMap();

    // 蓝绿场景，关闭ENV_TAG，使用x-kafka-env-tag配置
    public static final String X_KAFKA_ENV_TAG_KEY = "x-kafka-env-tag";


    static {
        // env
        String envTag = System.getProperty(X_KAFKA_ENV_TAG_KEY, System.getenv(X_KAFKA_ENV_TAG_KEY));
        if (StringUtils.isBlank(envTag)) {
            envTag = System.getProperty(X_ENV_TAG_KEY, System.getenv(X_ENV_TAG_KEY));
        }
        X_ENV_TAG_VALUE = SFMQEnvTagEnum.toEnvTagEnum(envTag);

        // grey topic
        String topics = System.getProperty(X_KAFKA_TOPICS_KEY, System.getenv(X_KAFKA_TOPICS_KEY));
        if (StringUtils.isNotBlank(topics)) {
            String[] clusterTopics = StringUtils.split(topics, ";");
            for (String cs : clusterTopics) {
                if (StringUtils.isNotBlank(cs)) {
                    String[] ct = StringUtils.split(cs, ":");
                    if (ct != null && ct.length == 2) {
                        String clusterAlias = ct[0];
                        String topicStr = ct[1];
                        Set<String> ts = Arrays.stream(StringUtils.split(topicStr, ",")).collect(Collectors.toSet());
                        Set<String> old = KAFKA_GREY_TOPICS2.getOrDefault(clusterAlias, Sets.newHashSet());
                        old.addAll(ts);
                        KAFKA_GREY_TOPICS2.put(clusterAlias, old);
                    }

                }
            }
        }

    }

    // 消费组id分隔符
    public static final String GROUP_INSTANCE_ID_SPLIT = "_-_";

}

SFRangeAssignor.java

/*
 * 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 com.sf.kafka.api.grey;

import static com.sf.kafka.api.grey.SFMQEnvTagEnum.BLUE;
import static com.sf.kafka.api.grey.SFMQEnvTagEnum.GREEN;
import static com.sf.kafka.api.grey.SFMQEnvTagEnum.GREY;
import static com.sf.kafka.api.grey.SFMQEnvTagEnum.NORMAL;
import static com.sf.kafka.api.grey.SFMQGreyConstants.GROUP_INSTANCE_ID_SPLIT;
import static com.sf.kafka.api.grey.SFMQGreyConstants.KAFKA_GREY_TOPICS2;
import static com.sf.kafka.util.KafkaUtil.CLUSTER_ID_ALIAS;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.apache.commons.lang3.StringUtils;
import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 在 range assignor 分配策略基础上添加 小流量灰度&蓝绿 支持。
 *
 * <p>The sf range assignor works on a per-topic basis. For each topic, we lay out the available partitions in numeric
 * order
 * and the consumers in lexicographic order. We then divide the number of partitions by the total number of
 * consumers to determine the number of partitions to assign to each consumer. If it does not evenly
 * divide, then the first few consumers will have one extra partition.
 *
 * <p>For example, suppose there are two consumers <code>C0</code> and <code>C1</code>, two topics <code>t0</code> and
 * <code>t1</code>, and each topic has 3 partitions, resulting in partitions <code>t0p0</code>, <code>t0p1</code>,
 * <code>t0p2</code>, <code>t1p0</code>, <code>t1p1</code>, and <code>t1p2</code>.
 *
 * <p>The assignment will be:
 * <ul>
 * <li><code>C0: [t0p0, t0p1, t1p0, t1p1]</code></li>
 * <li><code>C1: [t0p2, t1p2]</code></li>
 * </ul>
 *
 * Since the introduction of static membership, we could leverage <code>group.instance.id</code> to make the
 * assignment behavior more sticky.
 * For the above example, after one rolling bounce, group coordinator will attempt to assign new <code>member
 * .id</code> towards consumers,
 * for example <code>C0</code> -&gt; <code>C3</code> <code>C1</code> -&gt; <code>C2</code>.
 *
 * <p>The assignment could be completely shuffled to:
 * <ul>
 * <li><code>C3 (was C0): [t0p2, t1p2] (before was [t0p0, t0p1, t1p0, t1p1])</code>
 * <li><code>C2 (was C1): [t0p0, t0p1, t1p0, t1p1] (before was [t0p2, t1p2])</code>
 * </ul>
 *
 * The assignment change was caused by the change of <code>member.id</code> relative order, and
 * can be avoided by setting the group.instance.id.
 * Consumers will have individual instance ids <code>I1</code>, <code>I2</code>. As long as
 * 1. Number of members remain the same across generation
 * 2. Static members' identities persist across generation
 * 3. Subscription pattern doesn't change for any member
 *
 * <p>The assignment will always be:
 * <ul>
 * <li><code>I0: [t0p0, t0p1, t1p0, t1p1]</code>
 * <li><code>I1: [t0p2, t1p2]</code>
 * </ul>
 *
 * @author YANGLiiN
 */
public class SFRangeAssignor extends AbstractPartitionAssignor {
    private static final Logger log = LoggerFactory.getLogger(SFRangeAssignor.class);

    private String clusterId = null;
    private String clusterAlias = null;

    public static final String SF_RANGE_ASSIGNOR_NAME = "sf-range";

    @Override
    public String name() {
        return SF_RANGE_ASSIGNOR_NAME;
    }

    private Map<String, List<MemberInfo>> consumersPerTopic(Map<String, Subscription> consumerMetadata) {
        Map<String, List<MemberInfo>> topicToConsumers = new HashMap<>();
        for (Map.Entry<String, Subscription> subscriptionEntry : consumerMetadata.entrySet()) {
            String consumerId = subscriptionEntry.getKey();
            MemberInfo memberInfo = new MemberInfo(consumerId, subscriptionEntry.getValue().groupInstanceId());
            for (String topic : subscriptionEntry.getValue().topics()) {
                put(topicToConsumers, topic, memberInfo);
            }
        }
        return topicToConsumers;
    }


    @Override
    public GroupAssignment assign(Cluster metadata, GroupSubscription groupSubscription) {
        if (clusterId == null) {
            clusterId = metadata.clusterResource().clusterId();
            clusterAlias = CLUSTER_ID_ALIAS.getOrDefault(clusterId, "");

        }
        return super.assign(metadata, groupSubscription);
    }

    @Override
    public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsPerTopic,
                                                    Map<String, Subscription> subscriptions) {
        Map<String, List<MemberInfo>> consumersPerTopic = consumersPerTopic(subscriptions);

        Map<String, List<TopicPartition>> assignment = new HashMap<>();
        for (String memberId : subscriptions.keySet()) {
            assignment.put(memberId, new ArrayList<>());
        }

        for (Map.Entry<String, List<MemberInfo>> topicEntry : consumersPerTopic.entrySet()) {
            String topic = topicEntry.getKey();
            List<MemberInfo> consumersForTopic = topicEntry.getValue();

            Integer numPartitionsForTopic = partitionsPerTopic.get(topic);
            if (numPartitionsForTopic == null) {
                continue;
            }

            Collections.sort(consumersForTopic);

            List<TopicPartition> partitions = AbstractPartitionAssignor.partitions(topic, numPartitionsForTopic);

            boolean isGreyNormal = false;
            boolean isBlueGreen = false;

            List<MemberInfo> consumersForTopicGrey = Lists.newArrayList();
            List<MemberInfo> consumersForTopicNormal = Lists.newArrayList();
            List<MemberInfo> consumersForTopicBlue = Lists.newArrayList();
            List<MemberInfo> consumersForTopicGreen = Lists.newArrayList();

            if (KAFKA_GREY_TOPICS2.getOrDefault(clusterAlias, Sets.newHashSet()).contains(topic)) {
                for (MemberInfo memberInfo : consumersForTopic) {
                    String[] ids = StringUtils.split(memberInfo.memberId, GROUP_INSTANCE_ID_SPLIT);
                    if (ids.length >= 2) {
                        String envStr = ids[0];
                        SFMQEnvTagEnum envTag = SFMQEnvTagEnum.toEnvTagEnum(envStr);
                        if (!envTag.isUnknown()) {
                            switch (envTag) {
                                case GREY:
                                    isGreyNormal = true;
                                    consumersForTopicGrey.add(memberInfo);
                                    break;
                                case NORMAL:
                                    isGreyNormal = true;
                                    consumersForTopicNormal.add(memberInfo);
                                    break;

                                case BLUE:
                                    isBlueGreen = true;
                                    consumersForTopicBlue.add(memberInfo);
                                    break;
                                case GREEN:
                                    isBlueGreen = true;
                                    consumersForTopicGreen.add(memberInfo);
                                    break;
                            }
                        }
                    }

                }


                // 灰度(小流量)
                if (isGreyNormal) {

                    List<TopicPartition> partitionsGrey =
                            partitions.stream().filter(GREY.topicPartitionPredicate()).collect(Collectors.toList());
                    doAssign(consumersForTopicGrey, partitionsGrey, assignment);


                    List<TopicPartition> partitionsNormal =
                            partitions.stream().filter(NORMAL.topicPartitionPredicate()).collect(Collectors.toList());
                    doAssign(consumersForTopicNormal, partitionsNormal, assignment);

                }

                // 蓝绿
                if (isBlueGreen) {

                    List<TopicPartition> partitionsBlue =
                            partitions.stream().filter(BLUE.topicPartitionPredicate()).collect(Collectors.toList());
                    doAssign(consumersForTopicBlue, partitionsBlue, assignment);


                    List<TopicPartition> partitionsGreen =
                            partitions.stream().filter(GREEN.topicPartitionPredicate()).collect(Collectors.toList());
                    doAssign(consumersForTopicGreen, partitionsGreen, assignment);


                }

                // 只配置有灰度蓝绿主题，但是没配置环境标签(或配置错误)
                if (!isGreyNormal && !isBlueGreen) {
                    log.warn("配置了灰度蓝绿主题，但没正确配置[灰度]或[蓝绿]环境标！");
                    doAssign(consumersForTopic, partitions, assignment);
                }

                // 同时配置了灰度、蓝绿
                if (isGreyNormal && isBlueGreen) {
                    log.warn("请忽同时开启[灰度]和[蓝绿]！");
                }


            } else {
                doAssign(consumersForTopic, partitions, assignment);
            }


        }
        return assignment;
    }

    private void doAssign(List<MemberInfo> consumersForTopic, List<TopicPartition> partitions,
                          Map<String, List<TopicPartition>> assignment) {
        if (partitions != null && partitions.size() > 0 && consumersForTopic != null && consumersForTopic.size() > 0) {
            int numPartitionsPerConsumer = partitions.size() / consumersForTopic.size();
            int consumersWithExtraPartition = partitions.size() % consumersForTopic.size();

            for (int i = 0, n = consumersForTopic.size(); i < n; i++) {
                int start = numPartitionsPerConsumer * i + Math.min(i, consumersWithExtraPartition);
                int length = numPartitionsPerConsumer + (i + 1 > consumersWithExtraPartition ? 0 : 1);
                assignment.get(consumersForTopic.get(i).memberId).addAll(partitions.subList(start, start + length));
            }
        }
    }
}

SFStickyPartitionCache.java

/*
 * 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 com.sf.kafka.api.grey;

import static com.sf.kafka.api.grey.SFMQEnvTagEnum.UNKNOWN;
import static com.sf.kafka.api.grey.SFMQGreyConstants.KAFKA_GREY_TOPICS2;
import static com.sf.kafka.api.grey.SFMQGreyConstants.X_ENV_TAG_VALUE;
import static com.sf.kafka.util.KafkaUtil.CLUSTER_ID_ALIAS;
import com.google.common.collect.Sets;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

/**
 * An internal class that implements a cache used for sticky partitioning behavior. The cache tracks the current sticky
 * partition for any given topic. This class should not be used externally.
 */
public class SFStickyPartitionCache {
    private final ConcurrentMap<String, Integer> indexCache;

    public SFStickyPartitionCache() {
        this.indexCache = new ConcurrentHashMap<>();
    }

    public int partition(String topic, Cluster cluster) {
        Integer part = indexCache.get(topic);
        if (part == null) {
            return nextPartition(topic, cluster, -1);
        }
        return part;
    }

    public int nextPartition(String topic, Cluster cluster, int prevPartition) {
        List<PartitionInfo> allPartitions = cluster.partitionsForTopic(topic);
        List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);

        Integer oldPart = indexCache.get(topic);
        Integer newPart = oldPart;
        // Check that the current sticky partition for the topic is either not set or that the partition that
        // triggered the new batch matches the sticky partition that needs to be changed.
        if (oldPart == null || oldPart == prevPartition) {

            Predicate<? super PartitionInfo> predicate = UNKNOWN.partitionInfoPredicate();

            // 已配置灰度、蓝绿主题 && 开启灰度开关
            String clusterId = cluster.clusterResource().clusterId();
            String clusterAlias = CLUSTER_ID_ALIAS.getOrDefault(clusterId, "");
            if (KAFKA_GREY_TOPICS2.getOrDefault(clusterAlias, Sets.newHashSet()).contains(topic)
                    && !X_ENV_TAG_VALUE.isUnknown()) {
                predicate = X_ENV_TAG_VALUE.partitionInfoPredicate();
            }

            newPart = selectNextPartition(allPartitions, availablePartitions, oldPart, predicate);

            // Only change the sticky partition if it is null or prevPartition matches the current sticky partition.
            if (oldPart == null) {
                indexCache.putIfAbsent(topic, newPart);
            } else {
                indexCache.replace(topic, prevPartition, newPart);
            }
            return indexCache.get(topic);
        }
        return indexCache.get(topic);
    }


    private static Integer selectNextPartition(List<PartitionInfo> allPartitions,
                                               List<PartitionInfo> availablePartitions,
                                               Integer oldPart, Predicate<? super PartitionInfo> predicate) {
        Integer newPart = null;

        List<PartitionInfo> partitions = availablePartitions.stream().filter(predicate).collect(Collectors.toList());

        if (partitions.size() < 1) {
            List<PartitionInfo> tmp =
                    allPartitions.stream().filter(predicate).collect(Collectors.toList());
            Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
            newPart = tmp.get(random % tmp.size()).partition();
        } else if (partitions.size() == 1) {
            newPart = partitions.get(0).partition();
        } else {
            while (newPart == null || newPart.equals(oldPart)) {
                int random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
                newPart = partitions.get(random % partitions.size()).partition();
            }
        }

        return newPart;
    }

}