kafka-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From j...@apache.org
Subject [2/2] kafka git commit: KAFKA-2273; Sticky partition assignment strategy (KIP-54)
Date Thu, 18 May 2017 03:17:07 GMT
KAFKA-2273; Sticky partition assignment strategy (KIP-54)

This PR implements a new partition assignment strategy called "sticky", and it's purpose is to balance partitions across consumers in a way that minimizes moving partitions around, or, in other words, preserves existing partition assignments as much as possible.

This patch is co-authored with rajinisivaram and edoardocomar.

Author: Vahid Hashemian <vahidhashemian@us.ibm.com>

Reviewers: Jason Gustafson <jason@confluent.io>

Closes #1020 from vahidhashemian/KAFKA-2273


Project: http://git-wip-us.apache.org/repos/asf/kafka/repo
Commit: http://git-wip-us.apache.org/repos/asf/kafka/commit/e1abf177
Tree: http://git-wip-us.apache.org/repos/asf/kafka/tree/e1abf177
Diff: http://git-wip-us.apache.org/repos/asf/kafka/diff/e1abf177

Branch: refs/heads/trunk
Commit: e1abf17708918b82d3974ea028a4d74e3892fa0f
Parents: 9815e18
Author: Vahid Hashemian <vahidhashemian@us.ibm.com>
Authored: Wed May 17 20:13:19 2017 -0700
Committer: Jason Gustafson <jason@confluent.io>
Committed: Wed May 17 20:15:17 2017 -0700

----------------------------------------------------------------------
 .../kafka/clients/consumer/StickyAssignor.java  | 933 +++++++++++++++++++
 .../internals/AbstractPartitionAssignor.java    |   8 +-
 .../consumer/internals/ConsumerProtocol.java    |   1 -
 .../consumer/internals/PartitionAssignor.java   |   2 +-
 .../org/apache/kafka/common/TopicPartition.java |   1 -
 .../clients/consumer/StickyAssignorTest.java    | 689 ++++++++++++++
 .../kafka/api/PlaintextConsumerTest.scala       |  58 +-
 7 files changed, 1685 insertions(+), 7 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
----------------------------------------------------------------------
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
new file mode 100644
index 0000000..58e5915
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java
@@ -0,0 +1,933 @@
+/*
+ * 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.clients.consumer;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.TreeSet;
+
+import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.protocol.types.ArrayOf;
+import org.apache.kafka.common.protocol.types.Field;
+import org.apache.kafka.common.protocol.types.Schema;
+import org.apache.kafka.common.protocol.types.Struct;
+import org.apache.kafka.common.protocol.types.Type;
+import org.apache.kafka.common.utils.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * The sticky assignor serves two purposes. First, it guarantees an assignment that is as balanced as possible, meaning either:
+ * - the numbers of topic partitions assigned to consumers differ by at most one; or
+ * - each consumer that has 2+ fewer topic partitions than some other consumer cannot get any of those topic partitions transferred to it.
+ * Second, it preserved as many existing assignment as possible when a reassignment occurs. This helps in saving some of the
+ * overhead processing when topic partitions move from one consumer to another.
+ *
+ * Starting fresh it would work by distributing the partitions over consumers as evenly as possible. Even though this may sound similar to
+ * how round robin assignor works, the second example below shows that it is not.
+ * During a reassignment it would perform the reassignment in such a way that in the new assignment
+ * 1. topic partitions are still distributed as evenly as possible, and
+ * 2. topic partitions stay with their previously assigned consumers as much as possible.
+ * Of course, the first goal above takes precedence over the second one.
+ *
+ * <b>Example 1.</b> Suppose there are three consumers <code>C0</code>, <code>C1</code>, <code>C2</code>,
+ * four topics <code>t0,</code> <code>t1</code>, <code>t2</code>, <code>t3</code>, and each topic has 2 partitions,
+ * resulting in partitions <code>t0p0</code>, <code>t0p1</code>, <code>t1p0</code>, <code>t1p1</code>, <code>t2p0</code>,
+ * <code>t2p1</code>, <code>t3p0</code>, <code>t3p1</code>. Each consumer is subscribed to all three topics.
+ *
+ * The assignment with both sticky and round robin assignors will be:
+ * <ul>
+ * <li><code>C0: [t0p0, t1p1, t3p0]<code></li>
+ * <li><code>C1: [t0p1, t2p0, t3p1]<code></li>
+ * <li><code>C2: [t1p0, t2p1]<code></li>
+ * </ul>
+ *
+ * Now, let's assume <code>C1</code> is removed and a reassignment is about to happen. The round robin assignor would produce:
+ * <ul>
+ * <li><code>C0: [t0p0, t1p0, t2p0, t3p0]</code></li>
+ * <li><code>C2: [t0p1, t1p1, t2p1, t3p1]</code></li>
+ * </ul>
+ *
+ * while the sticky assignor would result in:
+ * <ul>
+ * <li><code>C0 [t0p0, t1p1, t3p0, t2p0]</code></li>
+ * <li><code>C2 [t1p0, t2p1, t0p1, t3p1]</code></li>
+ * </ul>
+ * preserving all the previous assignments (unlike the round robin assignor).
+ *
+ * <b>Example 2.</b> There are three consumers <code>C0</code>, <code>C1</code>, <code>C2</code>,
+ * and three topics <code>t0</code>, <code>t1</code>, <code>t2</code>, with 1, 2, and 3 partitions respectively.
+ * Therefore, the partitions are <code>t0p0</code>, <code>t1p0</code>, <code>t1p1</code>, <code>t2p0</code>,
+ * <code>t2p1</code>, <code>t2p2</code>. <code>C0</code> is subscribed to <code>t0</code>; <code>C1</code> is subscribed to
+ * <code>t0</code>, <code>t1</code>; and <code>C2</code> is subscribed to <code>t0</code>, <code>t1</code>, <code>t2</code>.
+ *
+ * The round robin assignor would come up with the following assignment:
+ * <ul>
+ * <li><code>C0 [t0p0]</code></li>
+ * <li><code>C1 [t1p0]</code></li>
+ * <li><code>C2 [t1p1, t2p0, t2p1, t2p2]</code></li>
+ * </ul>
+ *
+ * which is not as balanced as the assignment suggested by sticky assignor:
+ * <ul>
+ * <li><code>C0 [t0p0]</code></li>
+ * <li><code>C1 [t1p0, t1p1]</code></li>
+ * <li><code>C2 [t2p0, t2p1, t2p2]</code></li>
+ * </ul>
+ *
+ * Now, if consumer <code>C0</code> is removed, these two assignors would produce the following assignments.
+ * Round Robin (preserves 3 partition assignments):
+ * <ul>
+ * <li><code>C1 [t0p0, t1p1]</code></li>
+ * <li><code>C2 [t1p0, t2p0, t2p1, t2p2]</code></li>
+ * </ul>
+ *
+ * Sticky (preserves 5 partition assignments):
+ * <ul>
+ * <li><code>C1 [t1p0, t1p1, t0p0]</code></li>
+ * <li><code>C2 [t2p0, t2p1, t2p2]</code></li>
+ * </ul>
+ *
+ * <h3>Impact on <code>ConsumerRebalanceListener</code></h3>
+ * The sticky assignment strategy can provide some optimization to those consumers that have some partition cleanup code
+ * in their <code>onPartitionsRevoked()</code> callback listeners. The cleanup code is placed in that callback listener
+ * because the consumer has no assumption or hope of preserving any of its assigned partitions after a rebalance when it
+ * is using range or round robin assignor. The listener code would look like this:
+ * <code>
+ * class TheOldRebalanceListener implements ConsumerRebalanceListener {
+ *
+ *   void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+ *     for (TopicPartition partition: partitions) {
+ *       commitOffsets(partition);
+ *       cleanupState(partition);
+ *     }
+ *   }
+ *
+ *   void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+ *     for (TopicPartition partition: partitions) {
+ *       initializeState(partition);
+ *       initializeOffset(partition);
+ *     }
+ *   }
+ * }
+ * </code>
+ *
+ * As mentioned above, one advantage of the sticky assignor is that, in general, it reduces the number of partitions that
+ * actually move from one consumer to another during a reassignment. Therefore, it allows consumers to do their cleanup
+ * more efficiently. Of course, they still can perform the partition cleanup in the <code>onPartitionsRevoked()</code>
+ * listener, but they can be more efficient and make a note of their partitions before and after the rebalance, and do the
+ * cleanup after the rebalance only on the partitions they have lost (which is normally not a lot). The code snippet below
+ * clarifies this point:
+ * <code>
+ * class TheNewRebalanceListener implements ConsumerRebalanceListener {
+ *   Collection<TopicPartition> lastAssignment = Collections.emptyList();
+ *
+ *   void onPartitionsRevoked(Collection<TopicPartition> partitions) {
+ *     for (TopicPartition partition: partitions)
+ *       commitOffsets(partition);
+ *   }
+ *
+ *   void onPartitionsAssigned(Collection<TopicPartition> assignment) {
+ *     for (TopicPartition partition: difference(lastAssignment, assignment))
+ *       cleanupState(partition);
+ *
+ *     for (TopicPartition partition: difference(assignment, lastAssignment))
+ *       initializeState(partition);
+ *
+ *     for (TopicPartition partition: assignment)
+ *       initializeOffset(partition);
+ *
+ *     this.lastAssignment = assignment;
+ *   }
+ * }
+ * </code>
+ *
+ * Any consumer that uses sticky assignment can leverage this listener like this:
+ * <code>consumer.subscribe(topics, new TheNewRebalanceListener());</code>
+ *
+ */
+public class StickyAssignor extends AbstractPartitionAssignor {
+    private static final Logger log = LoggerFactory.getLogger(StickyAssignor.class);
+
+    // these schemas are used for preserving consumer's previously assigned partitions
+    // list and sending it as user data to the leader during a rebalance
+    private static final String TOPIC_PARTITIONS_KEY_NAME = "previous_assignment";
+    private static final String TOPIC_KEY_NAME = "topic";
+    private static final String PARTITIONS_KEY_NAME = "partitions";
+    private static final Schema TOPIC_ASSIGNMENT = new Schema(
+            new Field(TOPIC_KEY_NAME, Type.STRING),
+            new Field(PARTITIONS_KEY_NAME, new ArrayOf(Type.INT32)));
+    private static final Schema STICKY_ASSIGNOR_USER_DATA = new Schema(
+            new Field(TOPIC_PARTITIONS_KEY_NAME, new ArrayOf(TOPIC_ASSIGNMENT)));
+
+    Map<String, List<TopicPartition>> currentAssignment = new HashMap<>();
+    private List<TopicPartition> memberAssignment = null;
+    private PartitionMovements partitionMovements;
+
+    public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsPerTopic,
+                                                    Map<String, List<String>> subscriptions) {
+        partitionMovements = new PartitionMovements();
+
+        prepopulateCurrentAssignments();
+        // make a deep copy of currentAssignment
+        Map<String, List<TopicPartition>> oldAssignment = deepCopy(currentAssignment);
+
+        // a mapping of all topic partitions to all consumers that can be assigned to them
+        final HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers = new HashMap<>();
+        // a mapping of all consumers to all potential topic partitions that can be assigned to them
+        final HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions = new HashMap<>();
+
+        // initialize partition2AllPotentialConsumers and consumer2AllPotentialPartitions in the following two for loops
+        for (Entry<String, Integer> entry: partitionsPerTopic.entrySet()) {
+            for (int i = 0; i < entry.getValue(); ++i)
+                partition2AllPotentialConsumers.put(new TopicPartition(entry.getKey(), i), new ArrayList<String>());
+        }
+
+        for (Entry<String, List<String>> entry: subscriptions.entrySet()) {
+            String consumer = entry.getKey();
+            consumer2AllPotentialPartitions.put(consumer, new ArrayList<TopicPartition>());
+            for (String topic: entry.getValue()) {
+                for (int i = 0; i < partitionsPerTopic.get(topic); ++i) {
+                    TopicPartition topicPartition = new TopicPartition(topic, i);
+                    consumer2AllPotentialPartitions.get(consumer).add(topicPartition);
+                    partition2AllPotentialConsumers.get(topicPartition).add(consumer);
+                }
+            }
+
+            // add this consumer to currentAssignment (with an empty topic partition assignment) if it does not already exist
+            if (!currentAssignment.containsKey(consumer))
+                currentAssignment.put(consumer, new ArrayList<TopicPartition>());
+        }
+
+        // a mapping of partition to current consumer
+        HashMap<TopicPartition, String> currentPartitionConsumer = new HashMap<>();
+        for (Map.Entry<String, List<TopicPartition>> entry: currentAssignment.entrySet())
+            for (TopicPartition topicPartition: entry.getValue())
+                currentPartitionConsumer.put(topicPartition, entry.getKey());
+
+        List<TopicPartition> sortedPartitions = sortPartitions(oldAssignment.isEmpty(), partition2AllPotentialConsumers, consumer2AllPotentialPartitions);
+
+        // all partitions that need to be assigned (initially set to all partitions but adjusted in the following loop)
+        List<TopicPartition> unassignedPartitions = new ArrayList<>(sortedPartitions);
+        for (Iterator<Map.Entry<String, List<TopicPartition>>> it = currentAssignment.entrySet().iterator(); it.hasNext();) {
+            Map.Entry<String, List<TopicPartition>> entry = it.next();
+            if (!subscriptions.containsKey(entry.getKey())) {
+                // if a consumer that existed before (and had some partition assignments) is now removed, remove it from currentAssignment
+                for (TopicPartition topicPartition: entry.getValue())
+                    currentPartitionConsumer.remove(topicPartition);
+                it.remove();
+            } else {
+                // otherwise (the consumer still exists)
+                for (Iterator<TopicPartition> partitionIter = entry.getValue().iterator(); partitionIter.hasNext();) {
+                    TopicPartition partition = partitionIter.next();
+                    if (!partition2AllPotentialConsumers.containsKey(partition)) {
+                        // if this topic partition of this consumer no longer exists remove it from currentAssignment of the consumer
+                        partitionIter.remove();
+                        currentPartitionConsumer.remove(partition);
+                    } else if (!subscriptions.get(entry.getKey()).contains(partition.topic())) {
+                        // if this partition cannot remain assigned to its current consumer because the consumer
+                        // is no longer subscribed to its topic remove it from currentAssignment of the consumer
+                        partitionIter.remove();
+                    } else
+                        // otherwise, remove the topic partition from those that need to be assigned only if
+                        // its current consumer is still subscribed to its topic (because it is already assigned
+                        // and we would want to preserve that assignment as much as possible)
+                        unassignedPartitions.remove(partition);
+                }
+            }
+        }
+        // at this point we have preserved all valid topic partition to consumer assignments and removed
+        // all invalid topic partitions and invalid consumers. Now we need to assign unassignedPartitions
+        // to consumers so that the topic partition assignments are as balanced as possible.
+
+        // an ascending sorted set of consumers based on how many topic partitions are already assigned to them
+        TreeSet<String> sortedCurrentSubscriptions = new TreeSet<>(new SubscriptionComparator(currentAssignment));
+        sortedCurrentSubscriptions.addAll(currentAssignment.keySet());
+
+        balance(sortedPartitions, unassignedPartitions, sortedCurrentSubscriptions, consumer2AllPotentialPartitions,
+                partition2AllPotentialConsumers, oldAssignment, currentPartitionConsumer);
+        return currentAssignment;
+    }
+
+    private void prepopulateCurrentAssignments() {
+        Map<String, Subscription> subscriptions = getSubscriptions();
+        if (subscriptions == null)
+            return;
+
+        currentAssignment.clear();
+        for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
+            ByteBuffer userData = subscriptionEntry.getValue().userData();
+            if (userData != null && userData.hasRemaining())
+                currentAssignment.put(subscriptionEntry.getKey(), deserializeTopicPartitionAssignment(userData));
+        }
+    }
+
+    @Override
+    public void onAssignment(Assignment assignment) {
+        memberAssignment = assignment.partitions();
+    }
+
+    @Override
+    public Subscription subscription(Set<String> topics) {
+        if (memberAssignment == null)
+            return new Subscription(new ArrayList<>(topics));
+
+        return new Subscription(new ArrayList<>(topics), serializeTopicPartitionAssignment(memberAssignment));
+    }
+
+    @Override
+    public String name() {
+        return "sticky";
+    }
+
+    /**
+     * determine if the current assignment is a balanced one
+     *
+     * @param sortedCurrentSubscriptions: an ascending sorted set of consumers based on how many topic partitions are already assigned to them
+     * @param allSubscriptions: a mapping of all consumers to all potential topic partitions that can be assigned to them
+     * @return
+     */
+    private boolean isBalanced(TreeSet<String> sortedCurrentSubscriptions, Map<String, List<TopicPartition>> allSubscriptions) {
+        int min = currentAssignment.get(sortedCurrentSubscriptions.first()).size();
+        int max = currentAssignment.get(sortedCurrentSubscriptions.last()).size();
+        if (min >= max - 1)
+            // if minimum and maximum numbers of partitions assigned to consumers differ by at most one return true
+            return true;
+
+        // create a mapping from partitions to the consumer assigned to them
+        final HashMap<TopicPartition, String> allPartitions = new HashMap<>();
+        Set<Entry<String, List<TopicPartition>>> assignments = currentAssignment.entrySet();
+        for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
+            List<TopicPartition> topicPartitions = entry.getValue();
+            for (TopicPartition topicPartition: topicPartitions) {
+                if (allPartitions.containsKey(topicPartition))
+                    log.error(topicPartition + " is assigned to more than one consumer.");
+                allPartitions.put(topicPartition, entry.getKey());
+            }
+        }
+
+        // for each consumer that does not have all the topic partitions it can get make sure none of the topic partitions it
+        // could but did not get cannot be moved to it (because that would break the balance)
+        for (String consumer: sortedCurrentSubscriptions) {
+            List<TopicPartition> consumerPartitions = currentAssignment.get(consumer);
+            int consumerPartitionCount = consumerPartitions.size();
+
+            // skip if this consumer already has all the topic partitions it can get
+            if (consumerPartitionCount == allSubscriptions.get(consumer).size())
+                continue;
+
+            // otherwise make sure it cannot get any more
+            List<TopicPartition> potentialTopicPartitions = allSubscriptions.get(consumer);
+            for (TopicPartition topicPartition: potentialTopicPartitions) {
+                if (!currentAssignment.get(consumer).contains(topicPartition)) {
+                    String otherConsumer = allPartitions.get(topicPartition);
+                    int otherConsumerPartitionCount = currentAssignment.get(otherConsumer).size();
+                    if (consumerPartitionCount < otherConsumerPartitionCount) {
+                        log.debug(topicPartition + " can be moved from consumer " + otherConsumer + " to consumer " + consumer + " for a more balanced assignment.");
+                        return false;
+                    }
+                }
+            }
+        }
+        return true;
+    }
+
+    /**
+     * @return the balance score of the given assignment, as the sum of assigned partitions size difference of all consumer pairs.
+     * A perfectly balanced assignment (with all consumers getting the same number of partitions) has a balance score of 0.
+     * Lower balance score indicates a more balanced assignment.
+     */
+    private int getBalanceScore(Map<String, List<TopicPartition>> assignment) {
+        int score = 0;
+
+        Map<String, Integer> consumer2AssignmentSize = new HashMap<>();
+        for (Entry<String, List<TopicPartition>> entry: assignment.entrySet())
+            consumer2AssignmentSize.put(entry.getKey(), entry.getValue().size());
+
+        Iterator<Entry<String, Integer>> it = consumer2AssignmentSize.entrySet().iterator();
+        while (it.hasNext()) {
+            Entry<String, Integer> entry = it.next();
+            int consumerAssignmentSize = entry.getValue();
+            it.remove();
+            for (Entry<String, Integer> otherEntry: consumer2AssignmentSize.entrySet())
+                score += Math.abs(consumerAssignmentSize - otherEntry.getValue());
+        }
+
+        return score;
+    }
+
+    /**
+     * Sort valid partitions so they are processed in the potential reassignment phase in the proper order
+     * that causes minimal partition movement among consumers (hence honoring maximal stickiness)
+     *
+     * @param isFreshAssignment whether this is a new assignment, or a reassignment of an existing one
+     * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers
+     * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from
+     * @return sorted list of valid partitions
+     */
+    private List<TopicPartition> sortPartitions(boolean isFreshAssignment,
+                                                HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers,
+                                                HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions) {
+        List<TopicPartition> sortedPartitions = new ArrayList<>();
+
+        if (!isFreshAssignment && areSubscriptionsIdentical(partition2AllPotentialConsumers, consumer2AllPotentialPartitions)) {
+            // if this is a reassignment and the subscriptions are identical (all consumers can consumer from all topics)
+            // then we just need to simply list partitions in a round robin fashion (from consumers with
+            // most assigned partitions to those with least)
+            Map<String, List<TopicPartition>> assignments = deepCopy(currentAssignment);
+            for (Entry<String, List<TopicPartition>> entry: assignments.entrySet()) {
+                List<TopicPartition> toRemove = new ArrayList<>();
+                for (TopicPartition partition: entry.getValue())
+                    if (!partition2AllPotentialConsumers.keySet().contains(partition))
+                        toRemove.add(partition);
+                for (TopicPartition partition: toRemove)
+                    entry.getValue().remove(partition);
+            }
+            TreeSet<String> sortedConsumers = new TreeSet<>(new SubscriptionComparator(assignments));
+            sortedConsumers.addAll(assignments.keySet());
+
+            while (!sortedConsumers.isEmpty()) {
+                String consumer = sortedConsumers.pollLast();
+                List<TopicPartition> remainingPartitions = assignments.get(consumer);
+                if (!remainingPartitions.isEmpty()) {
+                    sortedPartitions.add(remainingPartitions.remove(0));
+                    sortedConsumers.add(consumer);
+                }
+            }
+
+            for (TopicPartition partition: partition2AllPotentialConsumers.keySet()) {
+                if (!sortedPartitions.contains(partition))
+                    sortedPartitions.add(partition);
+            }
+
+        } else {
+            // an ascending sorted set of topic partitions based on how many consumers can potentially use them
+            TreeSet<TopicPartition> sortedAllPartitions = new TreeSet<>(new PartitionComparator(partition2AllPotentialConsumers));
+            sortedAllPartitions.addAll(partition2AllPotentialConsumers.keySet());
+
+            while (!sortedAllPartitions.isEmpty())
+                sortedPartitions.add(sortedAllPartitions.pollFirst());
+        }
+
+        return sortedPartitions;
+    }
+
+    /**
+     * @param partition2AllPotentialConsumers a mapping of partitions to their potential consumers
+     * @param consumer2AllPotentialPartitions a mapping of consumers to potential partitions they can consumer from
+     * @return true if potential consumers of partitions are the same, and potential partitions consumers can
+     * consumer from are the same too
+     */
+    private boolean areSubscriptionsIdentical(HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers,
+                                              HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions) {
+        if (!hasIdenticalListElements(partition2AllPotentialConsumers.values()))
+            return false;
+
+        if (!hasIdenticalListElements(consumer2AllPotentialPartitions.values()))
+            return false;
+
+        return true;
+    }
+
+    /**
+     * @param col a collection of elements of type list
+     * @return true if all lists in the collection have the same members; false otherwise
+     */
+    private <T> boolean hasIdenticalListElements(Collection<List<T>> col) {
+        Iterator<List<T>> it = col.iterator();
+        List<T> cur = it.next();
+        while (it.hasNext()) {
+            List<T> next = it.next();
+            if (!(cur.containsAll(next) && next.containsAll(cur)))
+                return false;
+            cur = next;
+        }
+        return true;
+    }
+
+    /**
+     * @return the consumer to which the given partition is assigned. The assignment should improve the overall balance
+     * of the partition assignments to consumers.
+     */
+    private String assignPartition(TopicPartition partition, TreeSet<String> sortedCurrentSubscriptions,
+            HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions, HashMap<TopicPartition, String> currentPartitionConsumer) {
+        for (String consumer: sortedCurrentSubscriptions) {
+            if (consumer2AllPotentialPartitions.get(consumer).contains(partition)) {
+                sortedCurrentSubscriptions.remove(consumer);
+                currentAssignment.get(consumer).add(partition);
+                currentPartitionConsumer.put(partition, consumer);
+                sortedCurrentSubscriptions.add(consumer);
+                return consumer;
+            }
+        }
+        return null;
+    }
+
+    private boolean canParticipateInReassignment(TopicPartition partition, HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers) {
+        // if a partition has two or more potential consumers it is subject to reassignment.
+        return partition2AllPotentialConsumers.get(partition).size() >= 2;
+    }
+
+    private boolean canParticipateInReassignment(String consumer,
+                                                 HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions,
+                                                 HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers) {
+        List<TopicPartition> currentPartitions = currentAssignment.get(consumer);
+        int currentAssignmentSize = currentPartitions.size();
+        int maxAssignmentSize = consumer2AllPotentialPartitions.get(consumer).size();
+        if (currentAssignmentSize > maxAssignmentSize)
+            log.error("The consumer " + consumer + " is assigned more partitions than the maximum possible.");
+
+        if (currentAssignmentSize < maxAssignmentSize)
+            // if a consumer is not assigned all its potential partitions it is subject to reassignment
+            return true;
+
+        for (TopicPartition partition: currentPartitions)
+            // if any of the partitions assigned to a consumer is subject to reassignment the consumer itself
+            // is subject to reassignment
+            if (canParticipateInReassignment(partition, partition2AllPotentialConsumers))
+                return true;
+
+        return false;
+    }
+
+    /**
+     * Balance the current assignment using the data structures created in the assign(...) method above.
+     */
+    private void balance(List<TopicPartition> sortedPartitions, List<TopicPartition> unassignedPartitions, TreeSet<String> sortedCurrentSubscriptions,
+                         HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions, HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers,
+                         Map<String, List<TopicPartition>> oldAssignment, HashMap<TopicPartition, String> currentPartitionConsumer) {
+        boolean initializing = currentAssignment.get(sortedCurrentSubscriptions.last()).isEmpty();
+        boolean reassignmentPerformed = false;
+
+        // assign all unassigned partitions
+        for (TopicPartition partition: unassignedPartitions) {
+            // skip if there is no potential consumer for the partition
+            if (partition2AllPotentialConsumers.get(partition).isEmpty())
+                continue;
+
+            assignPartition(partition, sortedCurrentSubscriptions, consumer2AllPotentialPartitions, currentPartitionConsumer);
+        }
+
+        // narrow down the reassignment scope to only those partitions that can actually be reassigned
+        Set<TopicPartition> fixedPartitions = new HashSet<>();
+        for (TopicPartition partition: partition2AllPotentialConsumers.keySet())
+            if (!canParticipateInReassignment(partition, partition2AllPotentialConsumers))
+                fixedPartitions.add(partition);
+        sortedPartitions.removeAll(fixedPartitions);
+
+        // narrow down the reassignment scope to only those consumers that are subject to reassignment
+        Map<String, List<TopicPartition>> fixedAssignments = new HashMap<>();
+        for (String consumer: consumer2AllPotentialPartitions.keySet())
+            if (!canParticipateInReassignment(consumer, consumer2AllPotentialPartitions, partition2AllPotentialConsumers)) {
+                sortedCurrentSubscriptions.remove(consumer);
+                fixedAssignments.put(consumer, currentAssignment.remove(consumer));
+            }
+
+        // create a deep copy of the current assignment so we can revert to it if we do not get a more balanced assignment later
+        Map<String, List<TopicPartition>> preBalanceAssignment = deepCopy(currentAssignment);
+        HashMap<TopicPartition, String> preBalancePartitionConsumers = new HashMap<>(currentPartitionConsumer);
+
+        reassignmentPerformed = performReassignments(sortedPartitions, sortedCurrentSubscriptions,
+                consumer2AllPotentialPartitions, partition2AllPotentialConsumers, currentPartitionConsumer);
+
+        // if we are not preserving existing assignments and we have made changes to the current assignment
+        // make sure we are getting a more balanced assignment; otherwise, revert to previous assignment
+        if (!initializing && reassignmentPerformed && getBalanceScore(currentAssignment) >= getBalanceScore(preBalanceAssignment)) {
+            deepCopy(preBalanceAssignment, currentAssignment);
+            currentPartitionConsumer.clear();
+            currentPartitionConsumer.putAll(preBalancePartitionConsumers);
+        }
+
+        // add the fixed assignments (those that could not change) back
+        for (Entry<String, List<TopicPartition>> entry: fixedAssignments.entrySet()) {
+            String consumer = entry.getKey();
+            currentAssignment.put(consumer, entry.getValue());
+            sortedCurrentSubscriptions.add(consumer);
+        }
+
+        fixedAssignments.clear();
+    }
+
+    private boolean performReassignments(List<TopicPartition> reassignablePartitions, TreeSet<String> sortedCurrentSubscriptions,
+                                         HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions,
+                                         HashMap<TopicPartition, List<String>> partition2AllPotentialConsumers,
+                                         HashMap<TopicPartition, String> currentPartitionConsumer) {
+        boolean reassignmentPerformed = false;
+        boolean modified;
+
+        // repeat reassignment until no partition can be moved to improve the balance
+        do {
+            modified = false;
+            // reassign all reassignable partitions (starting from the partition with least potential consumers and if needed)
+            // until the full list is processed or a balance is achieved
+            Iterator<TopicPartition> partitionIterator = reassignablePartitions.iterator();
+            while (partitionIterator.hasNext() && !isBalanced(sortedCurrentSubscriptions, consumer2AllPotentialPartitions)) {
+                TopicPartition partition = partitionIterator.next();
+
+                // the partition must have at least two consumers
+                if (partition2AllPotentialConsumers.get(partition).size() <= 1)
+                    log.error("Expected more than one potential consumer for partition '" + partition + "'");
+
+                // the partition must have a current consumer
+                String consumer = currentPartitionConsumer.get(partition);
+                if (consumer == null)
+                    log.error("Expected partition '" + partition + "' to be assigned to a consumer");
+
+                // check if a better-suited consumer exist for the partition; if so, reassign it
+                for (String otherConsumer: partition2AllPotentialConsumers.get(partition)) {
+                    if (currentAssignment.get(consumer).size() > currentAssignment.get(otherConsumer).size() + 1) {
+                        reassignPartition(partition, sortedCurrentSubscriptions, currentPartitionConsumer, consumer2AllPotentialPartitions);
+                        reassignmentPerformed = true;
+                        modified = true;
+                        break;
+                    }
+                }
+            }
+        } while (modified);
+
+        return reassignmentPerformed;
+    }
+
+    private void reassignPartition(TopicPartition partition, TreeSet<String> sortedCurrentSubscriptions,
+                                   HashMap<TopicPartition, String> currentPartitionConsumer,
+                                   HashMap<String, List<TopicPartition>> consumer2AllPotentialPartitions) {
+        String consumer = currentPartitionConsumer.get(partition);
+
+        // find the new consumer
+        String newConsumer = null;
+        for (String anotherConsumer: sortedCurrentSubscriptions) {
+            if (consumer2AllPotentialPartitions.get(anotherConsumer).contains(partition)) {
+                newConsumer = anotherConsumer;
+                break;
+            }
+        }
+
+        assert newConsumer != null;
+
+        // find the correct partition movement considering the stickiness requirement
+        TopicPartition partitionToBeMoved = partitionMovements.getTheActualPartitionToBeMoved(partition, consumer, newConsumer);
+        processPartitionMovement(partitionToBeMoved, newConsumer, sortedCurrentSubscriptions, currentPartitionConsumer);
+
+        return;
+    }
+
+    private void processPartitionMovement(TopicPartition partition, String newConsumer,
+                                          TreeSet<String> sortedCurrentSubscriptions,
+                                          HashMap<TopicPartition, String> currentPartitionConsumer) {
+        String oldConsumer = currentPartitionConsumer.get(partition);
+
+        sortedCurrentSubscriptions.remove(oldConsumer);
+        sortedCurrentSubscriptions.remove(newConsumer);
+
+        partitionMovements.movePartition(partition, oldConsumer, newConsumer);
+
+        currentAssignment.get(oldConsumer).remove(partition);
+        currentAssignment.get(newConsumer).add(partition);
+        currentPartitionConsumer.put(partition, newConsumer);
+        sortedCurrentSubscriptions.add(newConsumer);
+        sortedCurrentSubscriptions.add(oldConsumer);
+    }
+
+    boolean isSticky() {
+        return partitionMovements.isSticky();
+    }
+
+    private static ByteBuffer serializeTopicPartitionAssignment(List<TopicPartition> partitions) {
+        Struct struct = new Struct(STICKY_ASSIGNOR_USER_DATA);
+        List<Struct> topicAssignments = new ArrayList<>();
+        for (Map.Entry<String, List<Integer>> topicEntry : CollectionUtils.groupDataByTopic(partitions).entrySet()) {
+            Struct topicAssignment = new Struct(TOPIC_ASSIGNMENT);
+            topicAssignment.set(TOPIC_KEY_NAME, topicEntry.getKey());
+            topicAssignment.set(PARTITIONS_KEY_NAME, topicEntry.getValue().toArray());
+            topicAssignments.add(topicAssignment);
+        }
+        struct.set(TOPIC_PARTITIONS_KEY_NAME, topicAssignments.toArray());
+        ByteBuffer buffer = ByteBuffer.allocate(STICKY_ASSIGNOR_USER_DATA.sizeOf(struct));
+        STICKY_ASSIGNOR_USER_DATA.write(buffer, struct);
+        buffer.flip();
+        return buffer;
+    }
+
+    private static List<TopicPartition> deserializeTopicPartitionAssignment(ByteBuffer buffer) {
+        Struct struct = STICKY_ASSIGNOR_USER_DATA.read(buffer);
+        List<TopicPartition> partitions = new ArrayList<>();
+        for (Object structObj : struct.getArray(TOPIC_PARTITIONS_KEY_NAME)) {
+            Struct assignment = (Struct) structObj;
+            String topic = assignment.getString(TOPIC_KEY_NAME);
+            for (Object partitionObj : assignment.getArray(PARTITIONS_KEY_NAME)) {
+                Integer partition = (Integer) partitionObj;
+                partitions.add(new TopicPartition(topic, partition));
+            }
+        }
+        return partitions;
+    }
+
+    private void deepCopy(Map<String, List<TopicPartition>> source, Map<String, List<TopicPartition>> dest) {
+        dest.clear();
+        for (Entry<String, List<TopicPartition>> entry: source.entrySet())
+            dest.put(entry.getKey(), new ArrayList<>(entry.getValue()));
+    }
+
+    private Map<String, List<TopicPartition>> deepCopy(Map<String, List<TopicPartition>> assignment) {
+        Map<String, List<TopicPartition>> copy = new HashMap<>();
+        deepCopy(assignment, copy);
+        return copy;
+    }
+
+    private static class PartitionComparator implements Comparator<TopicPartition>, Serializable {
+        private static final long serialVersionUID = 1L;
+        private Map<TopicPartition, List<String>> map;
+
+        PartitionComparator(Map<TopicPartition, List<String>> map) {
+            this.map = map;
+        }
+
+        @Override
+        public int compare(TopicPartition o1, TopicPartition o2) {
+            int ret = map.get(o1).size() - map.get(o2).size();
+            if (ret == 0) {
+                ret = o1.topic().compareTo(o2.topic());
+                if (ret == 0)
+                    ret = o1.partition() - o2.partition();
+            }
+            return ret;
+        }
+    }
+
+    private static class SubscriptionComparator implements Comparator<String>, Serializable {
+        private static final long serialVersionUID = 1L;
+        private Map<String, List<TopicPartition>> map;
+
+        SubscriptionComparator(Map<String, List<TopicPartition>> map) {
+            this.map = map;
+        }
+
+        @Override
+        public int compare(String o1, String o2) {
+            int ret = map.get(o1).size() - map.get(o2).size();
+            if (ret == 0)
+                ret = o1.compareTo(o2);
+            return ret;
+        }
+    }
+
+    /**
+     * This class maintains some data structures to simplify lookup of partition movements among consumers. At each point of
+     * time during a partition rebalance it keeps track of partition movements corresponding to each topic, and also possible
+     * movement (in form a <code>ConsumerPair</code> object) for each partition.
+     */
+    private static class PartitionMovements {
+        private Map<String, Map<ConsumerPair, Set<TopicPartition>>> partitionMovementsByTopic = new HashMap<>();
+        private Map<TopicPartition, ConsumerPair> partitionMovements = new HashMap<>();
+
+        private ConsumerPair removeMovementRecordOfPartition(TopicPartition partition) {
+            ConsumerPair pair = partitionMovements.remove(partition);
+
+            String topic = partition.topic();
+            Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            partitionMovementsForThisTopic.get(pair).remove(partition);
+            if (partitionMovementsForThisTopic.get(pair).isEmpty())
+                partitionMovementsForThisTopic.remove(pair);
+            if (partitionMovementsByTopic.get(topic).isEmpty())
+                partitionMovementsByTopic.remove(topic);
+
+            return pair;
+        }
+
+        private void addPartitionMovementRecord(TopicPartition partition, ConsumerPair pair) {
+            partitionMovements.put(partition, pair);
+
+            String topic = partition.topic();
+            if (!partitionMovementsByTopic.containsKey(topic))
+                partitionMovementsByTopic.put(topic, new HashMap<ConsumerPair, Set<TopicPartition>>());
+
+            Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            if (!partitionMovementsForThisTopic.containsKey(pair))
+                partitionMovementsForThisTopic.put(pair, new HashSet<TopicPartition>());
+
+            partitionMovementsForThisTopic.get(pair).add(partition);
+        }
+
+        private void movePartition(TopicPartition partition, String oldConsumer, String newConsumer) {
+            ConsumerPair pair = new ConsumerPair(oldConsumer, newConsumer);
+
+            if (partitionMovements.containsKey(partition)) {
+                // this partition has previously moved
+                ConsumerPair existingPair = removeMovementRecordOfPartition(partition);
+                assert existingPair.dstMemberId.equals(oldConsumer);
+                if (!existingPair.srcMemberId.equals(newConsumer)) {
+                    // the partition is not moving back to its previous consumer
+                    // return new ConsumerPair2(existingPair.src, newConsumer);
+                    addPartitionMovementRecord(partition, new ConsumerPair(existingPair.srcMemberId, newConsumer));
+                }
+            } else
+                addPartitionMovementRecord(partition, pair);
+        }
+
+        private TopicPartition getTheActualPartitionToBeMoved(TopicPartition partition, String oldConsumer, String newConsumer) {
+            String topic = partition.topic();
+
+            if (!partitionMovementsByTopic.containsKey(topic))
+                return partition;
+
+            if (partitionMovements.containsKey(partition)) {
+                // this partition has previously moved
+                assert oldConsumer.equals(partitionMovements.get(partition).dstMemberId);
+                oldConsumer = partitionMovements.get(partition).srcMemberId;
+            }
+
+            Map<ConsumerPair, Set<TopicPartition>> partitionMovementsForThisTopic = partitionMovementsByTopic.get(topic);
+            ConsumerPair reversePair = new ConsumerPair(newConsumer, oldConsumer);
+            if (!partitionMovementsForThisTopic.containsKey(reversePair))
+                return partition;
+
+            return partitionMovementsForThisTopic.get(reversePair).iterator().next();
+        }
+
+        private boolean isLinked(String src, String dst, Set<ConsumerPair> pairs, List<String> currentPath) {
+            if (src.equals(dst))
+                return false;
+
+            if (pairs.isEmpty())
+                return false;
+
+            if (new ConsumerPair(src, dst).in(pairs)) {
+                currentPath.add(src);
+                currentPath.add(dst);
+                return true;
+            }
+
+            for (ConsumerPair pair: pairs)
+                if (pair.srcMemberId.equals(src)) {
+                    Set<ConsumerPair> reducedSet = new HashSet<>(pairs);
+                    reducedSet.remove(pair);
+                    currentPath.add(pair.srcMemberId);
+                    return isLinked(pair.dstMemberId, dst, reducedSet, currentPath);
+                }
+
+            return false;
+        }
+
+        private boolean in(List<String> cycle, Set<List<String>> cycles) {
+            List<String> superCycle = new ArrayList<>(cycle);
+            superCycle.remove(superCycle.size() - 1);
+            superCycle.addAll(cycle);
+            for (List<String> foundCycle: cycles) {
+                if (foundCycle.size() == cycle.size() && Collections.indexOfSubList(superCycle, foundCycle) != -1)
+                    return true;
+            }
+            return false;
+        }
+
+        private boolean hasCycles(Set<ConsumerPair> pairs) {
+            Set<List<String>> cycles = new HashSet<>();
+            for (ConsumerPair pair: pairs) {
+                Set<ConsumerPair> reducedPairs = new HashSet<>(pairs);
+                reducedPairs.remove(pair);
+                List<String> path = new ArrayList<>(Collections.singleton(pair.srcMemberId));
+                if (isLinked(pair.dstMemberId, pair.srcMemberId, reducedPairs, path) && !in(path, cycles)) {
+                    cycles.add(new ArrayList<>(path));
+                    log.error("A cycle of length " + (path.size() - 1) + " was found: " + path.toString());
+                }
+            }
+
+            // for now we want to make sure there is no partition movements of the same topic between a pair of consumers.
+            // the odds of finding a cycle among more than two consumers seem to be very low (according to various randomized
+            // tests with the given sticky algorithm) that it should not worth the added complexity of handling those cases.
+            for (List<String> cycle: cycles)
+                if (cycle.size() == 3) // indicates a cycle of length 2
+                    return true;
+            return false;
+        }
+
+        private boolean isSticky() {
+            for (Map.Entry<String, Map<ConsumerPair, Set<TopicPartition>>> topicMovements: this.partitionMovementsByTopic.entrySet()) {
+                Set<ConsumerPair> topicMovementPairs = topicMovements.getValue().keySet();
+                if (hasCycles(topicMovementPairs)) {
+                    log.error("Stickiness is violated for topic " + topicMovements.getKey()
+                            + "\nPartition movements for this topic occurred among the following consumer pairs:"
+                            + "\n" + topicMovements.getValue().toString());
+                    return false;
+                }
+            }
+
+            return true;
+        }
+    }
+
+    /**
+     * <code>ConsumerPair</code> represents a pair of Kafka consumer ids involved in a partition reassignment. Each
+     * <code>ConsumerPair</code> object, which contains a source (<code>src</code>) and a destination (<code>dst</code>)
+     * element, normally corresponds to a particular partition or topic, and indicates that the particular partition or some
+     * partition of the particular topic was moved from the source consumer to the destination consumer during the rebalance.
+     * This class is used, through the <code>PartitionMovements</code> class, by the sticky assignor and helps in determining
+     * whether a partition reassignment results in cycles among the generated graph of consumer pairs.
+     */
+    private static class ConsumerPair {
+        private final String srcMemberId;
+        private final String dstMemberId;
+
+        ConsumerPair(String srcMemberId, String dstMemberId) {
+            this.srcMemberId = srcMemberId;
+            this.dstMemberId = dstMemberId;
+        }
+
+        public String toString() {
+            return this.srcMemberId + "->" + this.dstMemberId;
+        }
+
+        @Override
+        public int hashCode() {
+            final int prime = 31;
+            int result = 1;
+            result = prime * result + ((this.srcMemberId == null) ? 0 : this.srcMemberId.hashCode());
+            result = prime * result + ((this.dstMemberId == null) ? 0 : this.dstMemberId.hashCode());
+            return result;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            if (obj == null)
+                return false;
+
+            if (!getClass().isInstance(obj))
+                return false;
+
+            ConsumerPair otherPair = (ConsumerPair) obj;
+            return this.srcMemberId.equals(otherPair.srcMemberId) && this.dstMemberId.equals(otherPair.dstMemberId);
+        }
+
+        private boolean in(Set<ConsumerPair> pairs) {
+            for (ConsumerPair pair: pairs)
+                if (this.equals(pair))
+                    return true;
+            return false;
+        }
+    }
+}

http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
----------------------------------------------------------------------
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
index 5c97693..bc87ed0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractPartitionAssignor.java
@@ -34,6 +34,7 @@ import java.util.Set;
  */
 public abstract class AbstractPartitionAssignor implements PartitionAssignor {
     private static final Logger log = LoggerFactory.getLogger(AbstractPartitionAssignor.class);
+    private Map<String, Subscription> subscriptions = null;
 
     /**
      * Perform the group assignment given the partition counts and member subscriptions
@@ -52,6 +53,7 @@ public abstract class AbstractPartitionAssignor implements PartitionAssignor {
 
     @Override
     public Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions) {
+        this.subscriptions = new HashMap<>(subscriptions);
         Set<String> allSubscribedTopics = new HashSet<>();
         Map<String, List<String>> topicSubscriptions = new HashMap<>();
         for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
@@ -71,13 +73,17 @@ public abstract class AbstractPartitionAssignor implements PartitionAssignor {
 
         Map<String, List<TopicPartition>> rawAssignments = assign(partitionsPerTopic, topicSubscriptions);
 
-        // this class has maintains no user data, so just wrap the results
+        // this class maintains no user data, so just wrap the results
         Map<String, Assignment> assignments = new HashMap<>();
         for (Map.Entry<String, List<TopicPartition>> assignmentEntry : rawAssignments.entrySet())
             assignments.put(assignmentEntry.getKey(), new Assignment(assignmentEntry.getValue()));
         return assignments;
     }
 
+    protected Map<String, Subscription> getSubscriptions() {
+        return subscriptions;
+    }
+
     @Override
     public void onAssignment(Assignment assignment) {
         // this assignor maintains no internal state, so nothing to do

http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
----------------------------------------------------------------------
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
index 392e272..f8be9a0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerProtocol.java
@@ -146,7 +146,6 @@ public class ConsumerProtocol {
         // otherwise, assume versions can be parsed as V0
     }
 
-
     private static Map<String, List<Integer>> asMap(Collection<TopicPartition> partitions) {
         Map<String, List<Integer>> partitionMap = new HashMap<>();
         for (TopicPartition partition : partitions) {

http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
----------------------------------------------------------------------
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
index 86683a0..4a7c7a8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/PartitionAssignor.java
@@ -67,7 +67,7 @@ public interface PartitionAssignor {
 
 
     /**
-     * Unique name for this assignor (e.g. "range" or "roundrobin")
+     * Unique name for this assignor (e.g. "range" or "roundrobin" or "sticky")
      * @return non-null unique name
      */
     String name();

http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
----------------------------------------------------------------------
diff --git a/clients/src/main/java/org/apache/kafka/common/TopicPartition.java b/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
index 832bcd8..dc79c2e 100644
--- a/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
+++ b/clients/src/main/java/org/apache/kafka/common/TopicPartition.java
@@ -75,5 +75,4 @@ public final class TopicPartition implements Serializable {
     public String toString() {
         return topic + "-" + partition;
     }
-
 }

http://git-wip-us.apache.org/repos/asf/kafka/blob/e1abf177/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
----------------------------------------------------------------------
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
new file mode 100644
index 0000000..e9cc828
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/StickyAssignorTest.java
@@ -0,0 +1,689 @@
+/*
+ * 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.clients.consumer;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Random;
+import java.util.Set;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.CollectionUtils;
+import org.apache.kafka.common.utils.Utils;
+import org.junit.Test;
+
+public class StickyAssignorTest {
+
+    private StickyAssignor assignor = new StickyAssignor();
+
+    @Test
+    public void testOneConsumerNoTopic() {
+        String consumerId = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        Map<String, List<String>> subscriptions = Collections.singletonMap(consumerId, Collections.<String>emptyList());
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(assignment.get(consumerId).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerNonexistentTopic() {
+        String topic = "topic";
+        String consumerId = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 0);
+        Map<String, List<String>> subscriptions = Collections.singletonMap(consumerId, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+
+        assertEquals(Collections.singleton(consumerId), assignment.keySet());
+        assertTrue(assignment.get(consumerId).isEmpty());
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerOneTopic() {
+        String topic = "topic";
+        String consumerId = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        Map<String, List<String>> subscriptions = Collections.singletonMap(consumerId, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOnlyAssignsPartitionsFromSubscribedTopics() {
+        String topic = "topic";
+        String otherTopic = "other";
+        String consumerId = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        partitionsPerTopic.put(otherTopic, 3);
+        Map<String, List<String>> subscriptions = Collections.singletonMap(consumerId, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testOneConsumerMultipleTopics() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumerId = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 1);
+        partitionsPerTopic.put(topic2, 2);
+        Map<String, List<String>> subscriptions = Collections.singletonMap(consumerId, topics(topic1, topic2));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic1, 0), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumerId));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersOneTopicOnePartition() {
+        String topic = "topic";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 1);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic));
+        subscriptions.put(consumer2, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 0)), assignment.get(consumer1));
+        assertEquals(Collections.<TopicPartition>emptyList(), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersOneTopicTwoPartitions() {
+        String topic = "topic";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 2);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic));
+        subscriptions.put(consumer2, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 0)), assignment.get(consumer1));
+        assertEquals(Arrays.asList(tp(topic, 1)), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testMultipleConsumersMixedTopicSubscriptions() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+        String consumer3 = "consumer3";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 3);
+        partitionsPerTopic.put(topic2, 2);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic1));
+        subscriptions.put(consumer2, topics(topic1, topic2));
+        subscriptions.put(consumer3, topics(topic1));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
+        assertEquals(Arrays.asList(tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
+        assertEquals(Arrays.asList(tp(topic1, 1)), assignment.get(consumer3));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testTwoConsumersTwoTopicsSixPartitions() {
+        String topic1 = "topic1";
+        String topic2 = "topic2";
+        String consumer1 = "consumer1";
+        String consumer2 = "consumer2";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic1, 3);
+        partitionsPerTopic.put(topic2, 3);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic1, topic2));
+        subscriptions.put(consumer2, topics(topic1, topic2));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
+        assertEquals(Arrays.asList(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+    }
+
+    @Test
+    public void testAddRemoveConsumerOneTopic() {
+        String topic = "topic";
+        String consumer1 = "consumer";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+
+        String consumer2 = "consumer2";
+        subscriptions.put(consumer2, topics(topic));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertEquals(Arrays.asList(tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
+        assertEquals(Arrays.asList(tp(topic, 0)), assignment.get(consumer2));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+        assertTrue(assignor.isSticky());
+
+        subscriptions.remove(consumer1);
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        assertTrue(assignment.get(consumer2).contains(tp(topic, 0)));
+        assertTrue(assignment.get(consumer2).contains(tp(topic, 1)));
+        assertTrue(assignment.get(consumer2).contains(tp(topic, 2)));
+
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+        assertTrue(assignor.isSticky());
+    }
+
+    /**
+     * This unit test performs sticky assignment for a scenario that round robin assignor handles poorly.
+     * Topics (partitions per topic): topic1 (2), topic2 (1), topic3 (2), topic4 (1), topic5 (2)
+     * Subscriptions:
+     *  - consumer1: topic1, topic2, topic3, topic4, topic5
+     *  - consumer2: topic1, topic3, topic5
+     *  - consumer3: topic1, topic3, topic5
+     *  - consumer4: topic1, topic2, topic3, topic4, topic5
+     * Round Robin Assignment Result:
+     *  - consumer1: topic1-0, topic3-0, topic5-0
+     *  - consumer2: topic1-1, topic3-1, topic5-1
+     *  - consumer3:
+     *  - consumer4: topic2-0, topic4-0
+     * Sticky Assignment Result:
+     *  - consumer1: topic2-0, topic3-0
+     *  - consumer2: topic1-0, topic3-1
+     *  - consumer3: topic1-1, topic5-0
+     *  - consumer4: topic4-0, topic5-1
+     */
+    @Test
+    public void testPoorRoundRobinAssignmentScenario() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i <= 5; i++)
+            partitionsPerTopic.put(String.format("topic%d", i), (i % 2) + 1);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put("consumer1", Arrays.asList("topic1", "topic2", "topic3", "topic4", "topic5"));
+        subscriptions.put("consumer2", Arrays.asList("topic1", "topic3", "topic5"));
+        subscriptions.put("consumer3", Arrays.asList("topic1", "topic3", "topic5"));
+        subscriptions.put("consumer4", Arrays.asList("topic1", "topic2", "topic3", "topic4", "topic5"));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+    }
+
+    @Test
+    public void testAddRemoveTopicTwoConsumers() {
+        String topic = "topic";
+        String consumer1 = "consumer";
+        String consumer2 = "consumer2";
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put(topic, 3);
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put(consumer1, topics(topic));
+        subscriptions.put(consumer2, topics(topic));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        assertTrue(isFullyBalanced(assignment));
+        verifyValidityAndBalance(subscriptions, assignment);
+        // verify stickiness
+        List<TopicPartition> consumer1Assignment1 = assignment.get(consumer1);
+        List<TopicPartition> consumer2Assignment1 = assignment.get(consumer2);
+        assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) ||
+                   (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1));
+
+        String topic2 = "topic2";
+        partitionsPerTopic.put(topic2, 3);
+        subscriptions.put(consumer1, topics(topic, topic2));
+        subscriptions.put(consumer2, topics(topic, topic2));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+        // verify stickiness
+        List<TopicPartition> consumer1assignment = assignment.get(consumer1);
+        List<TopicPartition> consumer2assignment = assignment.get(consumer2);
+        assertTrue(consumer1assignment.size() == 3 && consumer2assignment.size() == 3);
+        assertTrue(consumer1assignment.containsAll(consumer1Assignment1));
+        assertTrue(consumer2assignment.containsAll(consumer2Assignment1));
+        assertTrue(assignor.isSticky());
+
+        partitionsPerTopic.remove(topic);
+        subscriptions.put(consumer1, topics(topic2));
+        subscriptions.put(consumer2, topics(topic2));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        // verify balance
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(isFullyBalanced(assignment));
+        // verify stickiness
+        List<TopicPartition> consumer1Assignment3 = assignment.get(consumer1);
+        List<TopicPartition> consumer2Assignment3 = assignment.get(consumer2);
+        assertTrue((consumer1Assignment3.size() == 1 && consumer2Assignment3.size() == 2) ||
+                   (consumer1Assignment3.size() == 2 && consumer2Assignment3.size() == 1));
+        assertTrue(consumer1assignment.containsAll(consumer1Assignment3));
+        assertTrue(consumer2assignment.containsAll(consumer2Assignment3));
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testReassignmentAfterOneConsumerLeaves() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 20; i++)
+            partitionsPerTopic.put(String.format("topic%02d", i), i);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        for (int i = 1; i < 20; i++) {
+            List<String> topics = new ArrayList<String>();
+            for (int j = 1; j <= i; j++)
+                topics.add(String.format("topic%02d", j));
+            subscriptions.put(String.format("consumer%02d", i), topics);
+        }
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+
+        subscriptions.remove("consumer10");
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testReassignmentAfterOneConsumerAdded() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put("topic", 20);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        for (int i = 1; i < 10; i++)
+            subscriptions.put(String.format("consumer%02d", i), Collections.singletonList("topic"));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+
+        subscriptions.put("consumer10", Collections.singletonList("topic"));
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testSameSubscriptions() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 15; i++)
+            partitionsPerTopic.put(String.format("topic%02d", i), i);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        for (int i = 1; i < 9; i++) {
+            List<String> topics = new ArrayList<String>();
+            for (int j = 1; j <= partitionsPerTopic.size(); j++)
+                topics.add(String.format("topic%02d", j));
+            subscriptions.put(String.format("consumer%02d", i), topics);
+        }
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+
+        subscriptions.remove("consumer05");
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testLargeAssignmentWithMultipleConsumersLeaving() {
+        Random rand = new Random();
+        int topicCount = 40;
+        int consumerCount = 200;
+
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 0; i < topicCount; i++)
+            partitionsPerTopic.put(getTopicName(i, topicCount), rand.nextInt(10) + 1);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        for (int i = 0; i < consumerCount; i++) {
+            List<String> topics = new ArrayList<String>();
+            for (int j = 0; j < rand.nextInt(20); j++)
+                topics.add(getTopicName(rand.nextInt(topicCount), topicCount));
+            subscriptions.put(getConsumerName(i, consumerCount), topics);
+        }
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+
+        for (int i = 0; i < 100; ++i) {
+            String c = getConsumerName(rand.nextInt(consumerCount), consumerCount);
+            subscriptions.remove(c);
+        }
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testNewSubscription() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i < 5; i++)
+            partitionsPerTopic.put(String.format("topic%02d", i), 1);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        for (int i = 0; i < 3; i++) {
+            List<String> topics = new ArrayList<String>();
+            for (int j = i; j <= 3 * i - 2; j++)
+                topics.add(String.format("topic%02d", j));
+            subscriptions.put(String.format("consumer%02d", i), topics);
+        }
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+
+        subscriptions.get("consumer00").add("topic01");
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+    }
+
+    @Test
+    public void testReassignmentWithRandomSubscriptionsAndChanges() {
+        final int minNumConsumers = 20;
+        final int maxNumConsumers = 40;
+        final int minNumTopics = 10;
+        final int maxNumTopics = 20;
+
+        for (int round = 1; round <= 100; ++round) {
+            int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics);
+
+            ArrayList<String> topics = new ArrayList<>();
+            for (int i = 0; i < numTopics; ++i)
+                topics.add(getTopicName(i, maxNumTopics));
+
+            Map<String, Integer> partitionsPerTopic = new HashMap<>();
+            for (int i = 0; i < numTopics; ++i)
+                partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1);
+
+            int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
+
+            Map<String, List<String>> subscriptions = new HashMap<>();
+            for (int i = 0; i < numConsumers; ++i) {
+                List<String> sub = Utils.sorted(getRandomSublist(topics));
+                subscriptions.put(getConsumerName(i, maxNumConsumers), sub);
+            }
+
+            StickyAssignor assignor = new StickyAssignor();
+
+            Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+            verifyValidityAndBalance(subscriptions, assignment);
+
+            subscriptions.clear();
+            for (int i = 0; i < numConsumers; ++i) {
+                List<String> sub = Utils.sorted(getRandomSublist(topics));
+                subscriptions.put(getConsumerName(i, maxNumConsumers), sub);
+            }
+
+            assignment = assignor.assign(partitionsPerTopic, subscriptions);
+            verifyValidityAndBalance(subscriptions, assignment);
+            assertTrue(assignor.isSticky());
+        }
+    }
+
+    @Test
+    public void testMoveExistingAssignments() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        for (int i = 1; i <= 6; i++)
+            partitionsPerTopic.put(String.format("topic%02d", i), 1);
+
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put("consumer01", topics("topic01", "topic02"));
+        subscriptions.put("consumer02", topics("topic01", "topic02", "topic03", "topic04"));
+        subscriptions.put("consumer03", topics("topic02", "topic03", "topic04", "topic05", "topic06"));
+
+        assignor.currentAssignment.put("consumer01", new ArrayList<>(Arrays.asList(tp("topic01", 0))));
+        assignor.currentAssignment.put("consumer02", new ArrayList<>(Arrays.asList(tp("topic02", 0), tp("topic03", 0))));
+        assignor.currentAssignment.put("consumer03", new ArrayList<>(Arrays.asList(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0))));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+    }
+
+    @Test
+    public void testStickiness() {
+        Map<String, Integer> partitionsPerTopic = new HashMap<>();
+        partitionsPerTopic.put("topic01", 3);
+        Map<String, List<String>> subscriptions = new HashMap<>();
+        subscriptions.put("consumer01", topics("topic01"));
+        subscriptions.put("consumer02", topics("topic01"));
+        subscriptions.put("consumer03", topics("topic01"));
+        subscriptions.put("consumer04", topics("topic01"));
+
+        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        Map<String, TopicPartition> partitionsAssigned = new HashMap<>();
+
+        Set<Entry<String, List<TopicPartition>>> assignments = assignment.entrySet();
+        for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
+            String consumer = entry.getKey();
+            List<TopicPartition> topicPartitions = entry.getValue();
+            int size = topicPartitions.size();
+            assertTrue("Consumer " + consumer + " is assigned more topic partitions than expected.", size <= 1);
+            if (size == 1)
+                partitionsAssigned.put(consumer, topicPartitions.get(0));
+        }
+
+        // removing the potential group leader
+        subscriptions.remove("consumer01");
+
+        assignment = assignor.assign(partitionsPerTopic, subscriptions);
+        verifyValidityAndBalance(subscriptions, assignment);
+        assertTrue(assignor.isSticky());
+
+        assignments = assignment.entrySet();
+        for (Map.Entry<String, List<TopicPartition>> entry: assignments) {
+            String consumer = entry.getKey();
+            List<TopicPartition> topicPartitions = entry.getValue();
+            assertEquals("Consumer " + consumer + " is assigned more topic partitions than expected.", 1, topicPartitions.size());
+            assertTrue("Stickiness was not honored for consumer " + consumer,
+                (!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))));
+        }
+    }
+
+    private String getTopicName(int i, int maxNum) {
+        return getCanonicalName("t", i, maxNum);
+    }
+
+    private String getConsumerName(int i, int maxNum) {
+        return getCanonicalName("c", i, maxNum);
+    }
+
+    private String getCanonicalName(String str, int i, int maxNum) {
+        return str + pad(i, Integer.toString(maxNum).length());
+    }
+
+    private String pad(int num, int digits) {
+        StringBuilder sb = new StringBuilder();
+        int iDigits = Integer.toString(num).length();
+
+        for (int i = 1; i <= digits - iDigits; ++i)
+            sb.append("0");
+
+        sb.append(num);
+        return sb.toString();
+    }
+
+    private static List<String> topics(String... topics) {
+        return Arrays.asList(topics);
+    }
+
+    private static TopicPartition tp(String topic, int partition) {
+        return new TopicPartition(topic, partition);
+    }
+
+    private static boolean isFullyBalanced(Map<String, List<TopicPartition>> assignment) {
+        int min = Integer.MAX_VALUE;
+        int max = Integer.MIN_VALUE;
+        for (List<TopicPartition> topicPartitions: assignment.values()) {
+            int size = topicPartitions.size();
+            if (size < min)
+                min = size;
+            if (size > max)
+                max = size;
+        }
+        return max - min <= 1;
+    }
+
+    private static List<String> getRandomSublist(ArrayList<String> list) {
+        List<String> selectedItems = new ArrayList<>(list);
+        int len = list.size();
+        Random random = new Random();
+        int howManyToRemove = random.nextInt(len);
+
+        for (int i = 1; i <= howManyToRemove; ++i)
+            selectedItems.remove(random.nextInt(selectedItems.size()));
+
+        return selectedItems;
+    }
+
+    /**
+     * Verifies that the given assignment is valid and balanced with respect to the given subscriptions
+     * Validity requirements:
+     * - each consumer is subscribed to topics of all partitions assigned to it, and
+     * - each partition is assigned to no more than one consumer
+     * Balance requirements:
+     * - the assignment is fully balanced (the numbers of topic partitions assigned to consumers differ by at most one), or
+     * - there is no topic partition that can be moved from one consumer to another with 2+ fewer topic partitions
+     *
+     * @param subscriptions: topic subscriptions of each consumer
+     * @param assignment: given assignment for balance check
+     */
+    private static void verifyValidityAndBalance(Map<String, List<String>> subscriptions, Map<String, List<TopicPartition>> assignments) {
+        int size = subscriptions.size();
+        assert size == assignments.size();
+
+        List<String> consumers = Utils.sorted(assignments.keySet());
+
+        for (int i = 0; i < size; ++i) {
+            String consumer = consumers.get(i);
+            List<TopicPartition> partitions = assignments.get(consumer);
+            for (TopicPartition partition: partitions)
+                assertTrue("Error: Partition " + partition + "is assigned to c" + i + ", but it is not subscribed to Topic t" + partition.topic()
+                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                        subscriptions.get(consumer).contains(partition.topic()));
+
+            if (i == size - 1)
+                continue;
+
+            for (int j = i + 1; j < size; ++j) {
+                String otherConsumer = consumers.get(j);
+                List<TopicPartition> otherPartitions = assignments.get(otherConsumer);
+
+                Set<TopicPartition> intersection = new HashSet<>(partitions);
+                intersection.retainAll(otherPartitions);
+                assertTrue("Error: Consumers c" + i + " and c" + j + " have common partitions assigned to them: " + intersection.toString()
+                        + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                        intersection.isEmpty());
+
+                int len = partitions.size();
+                int otherLen = otherPartitions.size();
+
+                if (Math.abs(len - otherLen) <= 1)
+                    continue;
+
+                Map<String, List<Integer>> map = CollectionUtils.groupDataByTopic(partitions);
+                Map<String, List<Integer>> otherMap = CollectionUtils.groupDataByTopic(otherPartitions);
+
+                if (len > otherLen) {
+                    for (String topic: map.keySet())
+                        assertTrue("Error: Some partitions can be moved from c" + i + " to c" + j + " to achieve a better balance"
+                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions."
+                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                                !otherMap.containsKey(topic));
+                }
+
+                if (otherLen > len) {
+                    for (String topic: otherMap.keySet())
+                        assertTrue("Error: Some partitions can be moved from c" + j + " to c" + i + " to achieve a better balance"
+                                + "\nc" + i + " has " + len + " partitions, and c" + j + " has " + otherLen + " partitions."
+                                + "\nSubscriptions: " + subscriptions.toString() + "\nAssignments: " + assignments.toString(),
+                                !map.containsKey(topic));
+                }
+            }
+        }
+    }
+}


Mime
View raw message