rdkafka_sticky_assignor.c

/*
 * librdkafka - The Apache Kafka C/C++ library
 *
 * Copyright (c) 2020 Magnus Edenhill
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * 1. Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright notice,
 *    this list of conditions and the following disclaimer in the documentation
 *    and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */


#include "rdkafka_int.h"
#include "rdkafka_assignor.h"
#include "rdkafka_request.h"
#include "rdmap.h"
#include "rdunittest.h"

#include <stdarg.h>
#include <math.h> /* abs() */

/**
 * @name KIP-54 and KIP-341 Sticky assignor.
 *
 * Closely mimicking the official Apache Kafka AbstractStickyAssignor
 * implementation.
 */

/** FIXME
 * Remaining:
 *  isSticky() -- used by tests
 */


/** @brief Assignor state from last rebalance */
typedef struct rd_kafka_sticky_assignor_state_s {
        rd_kafka_topic_partition_list_t *prev_assignment;
        int32_t generation_id;
} rd_kafka_sticky_assignor_state_t;



/**
 * Auxilliary glue types
 */

/**
 * @struct ConsumerPair_t represents a pair of consumer member ids involved in
 *         a partition reassignment, indicating a source consumer a partition
 *         is moving from and a destination partition the same partition is
 *         moving to.
 *
 * @sa PartitionMovements_t
 */
typedef struct ConsumerPair_s {
        const char *src; /**< Source member id */
        const char *dst; /**< Destination member id */
} ConsumerPair_t;


static ConsumerPair_t *ConsumerPair_new(const char *src, const char *dst) {
        ConsumerPair_t *cpair;

        cpair      = rd_malloc(sizeof(*cpair));
        cpair->src = src ? rd_strdup(src) : NULL;
        cpair->dst = dst ? rd_strdup(dst) : NULL;

        return cpair;
}


static void ConsumerPair_free(void *p) {
        ConsumerPair_t *cpair = p;
        if (cpair->src)
                rd_free((void *)cpair->src);
        if (cpair->dst)
                rd_free((void *)cpair->dst);
        rd_free(cpair);
}

static int ConsumerPair_cmp(const void *_a, const void *_b) {
        const ConsumerPair_t *a = _a, *b = _b;
        int r = strcmp(a->src ? a->src : "", b->src ? b->src : "");
        if (r)
                return r;
        return strcmp(a->dst ? a->dst : "", b->dst ? b->dst : "");
}


static unsigned int ConsumerPair_hash(const void *_a) {
        const ConsumerPair_t *a = _a;
        return 31 * (a->src ? rd_map_str_hash(a->src) : 1) +
               (a->dst ? rd_map_str_hash(a->dst) : 1);
}



typedef struct ConsumerGenerationPair_s {
        const char *consumer; /**< Memory owned by caller */
        int generation;
} ConsumerGenerationPair_t;

static void ConsumerGenerationPair_destroy(void *ptr) {
        ConsumerGenerationPair_t *cgpair = ptr;
        rd_free(cgpair);
}

/**
 * @param consumer This memory will be referenced, not copied, and thus must
 *                 outlive the ConsumerGenerationPair_t object.
 */
static ConsumerGenerationPair_t *
ConsumerGenerationPair_new(const char *consumer, int generation) {
        ConsumerGenerationPair_t *cgpair = rd_malloc(sizeof(*cgpair));
        cgpair->consumer                 = consumer;
        cgpair->generation               = generation;
        return cgpair;
}

static int ConsumerGenerationPair_cmp_generation(const void *_a,
                                                 const void *_b) {
        const ConsumerGenerationPair_t *a = _a, *b = _b;
        return a->generation - b->generation;
}



/**
 * Hash map types.
 *
 * Naming convention is:
 * map_<keytype>_<valuetype>_t
 *
 * Where the keytype and valuetype are spoken names of the types and
 * not the specific C types (since that'd be too long).
 */
typedef RD_MAP_TYPE(const char *,
                    rd_kafka_topic_partition_list_t *) map_str_toppar_list_t;

typedef RD_MAP_TYPE(const rd_kafka_topic_partition_t *,
                    const char *) map_toppar_str_t;

typedef RD_MAP_TYPE(const rd_kafka_topic_partition_t *,
                    rd_list_t *) map_toppar_list_t;

typedef RD_MAP_TYPE(const rd_kafka_topic_partition_t *,
                    ConsumerGenerationPair_t *) map_toppar_cgpair_t;

typedef RD_MAP_TYPE(const rd_kafka_topic_partition_t *,
                    ConsumerPair_t *) map_toppar_cpair_t;

typedef RD_MAP_TYPE(const ConsumerPair_t *,
                    rd_kafka_topic_partition_list_t *) map_cpair_toppar_list_t;

/* map<string, map<ConsumerPair*, topic_partition_list_t*>> */
typedef RD_MAP_TYPE(const char *,
                    map_cpair_toppar_list_t *) map_str_map_cpair_toppar_list_t;



/** Glue type helpers */

static map_cpair_toppar_list_t *map_cpair_toppar_list_t_new(void) {
        map_cpair_toppar_list_t *map = rd_calloc(1, sizeof(*map));

        RD_MAP_INIT(map, 0, ConsumerPair_cmp, ConsumerPair_hash, NULL,
                    rd_kafka_topic_partition_list_destroy_free);

        return map;
}

static void map_cpair_toppar_list_t_free(void *ptr) {
        map_cpair_toppar_list_t *map = ptr;
        RD_MAP_DESTROY(map);
        rd_free(map);
}



/**
 * @struct Provides current state of partition movements between consumers
 *         for each topic, and possible movements for each partition.
 */
typedef struct PartitionMovements_s {
        map_toppar_cpair_t partitionMovements;
        map_str_map_cpair_toppar_list_t partitionMovementsByTopic;
} PartitionMovements_t;


static void PartitionMovements_init(PartitionMovements_t *pmov,
                                    size_t topic_cnt) {
        RD_MAP_INIT(&pmov->partitionMovements, topic_cnt * 3,
                    rd_kafka_topic_partition_cmp, rd_kafka_topic_partition_hash,
                    NULL, ConsumerPair_free);

        RD_MAP_INIT(&pmov->partitionMovementsByTopic, topic_cnt, rd_map_str_cmp,
                    rd_map_str_hash, NULL, map_cpair_toppar_list_t_free);
}

static void PartitionMovements_destroy(PartitionMovements_t *pmov) {
        RD_MAP_DESTROY(&pmov->partitionMovementsByTopic);
        RD_MAP_DESTROY(&pmov->partitionMovements);
}


static ConsumerPair_t *PartitionMovements_removeMovementRecordOfPartition(
    PartitionMovements_t *pmov,
    const rd_kafka_topic_partition_t *toppar) {

        ConsumerPair_t *cpair;
        map_cpair_toppar_list_t *partitionMovementsForThisTopic;
        rd_kafka_topic_partition_list_t *plist;

        cpair = RD_MAP_GET(&pmov->partitionMovements, toppar);
        rd_assert(cpair);

        partitionMovementsForThisTopic =
            RD_MAP_GET(&pmov->partitionMovementsByTopic, toppar->topic);

        plist = RD_MAP_GET(partitionMovementsForThisTopic, cpair);
        rd_assert(plist);

        rd_kafka_topic_partition_list_del(plist, toppar->topic,
                                          toppar->partition);
        if (plist->cnt == 0)
                RD_MAP_DELETE(partitionMovementsForThisTopic, cpair);
        if (RD_MAP_IS_EMPTY(partitionMovementsForThisTopic))
                RD_MAP_DELETE(&pmov->partitionMovementsByTopic, toppar->topic);

        return cpair;
}

static void PartitionMovements_addPartitionMovementRecord(
    PartitionMovements_t *pmov,
    const rd_kafka_topic_partition_t *toppar,
    ConsumerPair_t *cpair) {
        map_cpair_toppar_list_t *partitionMovementsForThisTopic;
        rd_kafka_topic_partition_list_t *plist;

        RD_MAP_SET(&pmov->partitionMovements, toppar, cpair);

        partitionMovementsForThisTopic =
            RD_MAP_GET_OR_SET(&pmov->partitionMovementsByTopic, toppar->topic,
                              map_cpair_toppar_list_t_new());

        plist = RD_MAP_GET_OR_SET(partitionMovementsForThisTopic, cpair,
                                  rd_kafka_topic_partition_list_new(16));

        rd_kafka_topic_partition_list_add(plist, toppar->topic,
                                          toppar->partition);
}

static void
PartitionMovements_movePartition(PartitionMovements_t *pmov,
                                 const rd_kafka_topic_partition_t *toppar,
                                 const char *old_consumer,
                                 const char *new_consumer) {

        if (RD_MAP_GET(&pmov->partitionMovements, toppar)) {
                /* This partition has previously moved */
                ConsumerPair_t *existing_cpair;

                existing_cpair =
                    PartitionMovements_removeMovementRecordOfPartition(pmov,
                                                                       toppar);

                rd_assert(!rd_strcmp(existing_cpair->dst, old_consumer));

                if (rd_strcmp(existing_cpair->src, new_consumer)) {
                        /* Partition is not moving back to its
                         * previous consumer */
                        PartitionMovements_addPartitionMovementRecord(
                            pmov, toppar,
                            ConsumerPair_new(existing_cpair->src,
                                             new_consumer));
                }
        } else {
                PartitionMovements_addPartitionMovementRecord(
                    pmov, toppar, ConsumerPair_new(old_consumer, new_consumer));
        }
}

static const rd_kafka_topic_partition_t *
PartitionMovements_getTheActualPartitionToBeMoved(
    PartitionMovements_t *pmov,
    const rd_kafka_topic_partition_t *toppar,
    const char *oldConsumer,
    const char *newConsumer) {

        ConsumerPair_t *cpair;
        ConsumerPair_t reverse_cpair = {.src = newConsumer, .dst = oldConsumer};
        map_cpair_toppar_list_t *partitionMovementsForThisTopic;
        rd_kafka_topic_partition_list_t *plist;

        if (!RD_MAP_GET(&pmov->partitionMovementsByTopic, toppar->topic))
                return toppar;

        cpair = RD_MAP_GET(&pmov->partitionMovements, toppar);
        if (cpair) {
                /* This partition has previously moved */
                rd_assert(!rd_strcmp(oldConsumer, cpair->dst));

                oldConsumer = cpair->src;
        }

        partitionMovementsForThisTopic =
            RD_MAP_GET(&pmov->partitionMovementsByTopic, toppar->topic);

        plist = RD_MAP_GET(partitionMovementsForThisTopic, &reverse_cpair);
        if (!plist)
                return toppar;

        return &plist->elems[0];
}

#if FIXME

static rd_bool_t hasCycles(map_cpair_toppar_list_t *pairs) {
        return rd_true;  // FIXME
}

/**
 * @remark This method is only used by the AbstractStickyAssignorTest
 *         in the Java client.
 */
static rd_bool_t PartitionMovements_isSticky(rd_kafka_t *rk,
                                             PartitionMovements_t *pmov) {
        const char *topic;
        map_cpair_toppar_list_t *topicMovementPairs;

        RD_MAP_FOREACH(topic, topicMovementPairs,
                       &pmov->partitionMovementsByTopic) {
                if (hasCycles(topicMovementPairs)) {
                        const ConsumerPair_t *cpair;
                        const rd_kafka_topic_partition_list_t *partitions;

                        rd_kafka_log(
                            rk, LOG_ERR, "STICKY",
                            "Sticky assignor: Stickiness is violated for "
                            "topic %s: partition movements for this topic "
                            "occurred among the following consumers: ",
                            topic);
                        RD_MAP_FOREACH(cpair, partitions, topicMovementPairs) {
                                rd_kafka_log(rk, LOG_ERR, "STICKY", " %s -> %s",
                                             cpair->src, cpair->dst);
                        }

                        if (partitions)
                                ; /* Avoid unused warning */

                        return rd_false;
                }
        }

        return rd_true;
}
#endif


/**
 * @brief Comparator to sort ascendingly by rd_map_elem_t object value as
 *        topic partition list count, or by member id if the list count is
 *        identical.
 *        Used to sort sortedCurrentSubscriptions list.
 *
 * elem.key is the consumer member id string,
 * elem.value is the partition list.
 */
static int sort_by_map_elem_val_toppar_list_cnt(const void *_a,
                                                const void *_b) {
        const rd_map_elem_t *a = _a, *b = _b;
        const rd_kafka_topic_partition_list_t *al = a->value, *bl = b->value;
        int r = al->cnt - bl->cnt;
        if (r)
                return r;
        return strcmp((const char *)a->key, (const char *)b->key);
}


/**
 * @brief Assign partition to the most eligible consumer.
 *
 * The assignment should improve the overall balance of the partition
 * assignments to consumers.
 */
static void
assignPartition(const rd_kafka_topic_partition_t *partition,
                rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
                map_str_toppar_list_t *currentAssignment,
                map_str_toppar_list_t *consumer2AllPotentialPartitions,
                map_toppar_str_t *currentPartitionConsumer) {
        const rd_map_elem_t *elem;
        int i;

        RD_LIST_FOREACH(elem, sortedCurrentSubscriptions, i) {
                const char *consumer = (const char *)elem->key;
                const rd_kafka_topic_partition_list_t *partitions;

                partitions =
                    RD_MAP_GET(consumer2AllPotentialPartitions, consumer);
                if (!rd_kafka_topic_partition_list_find(
                        partitions, partition->topic, partition->partition))
                        continue;

                rd_kafka_topic_partition_list_add(
                    RD_MAP_GET(currentAssignment, consumer), partition->topic,
                    partition->partition);

                RD_MAP_SET(currentPartitionConsumer,
                           rd_kafka_topic_partition_copy(partition), consumer);

                /* Re-sort sortedCurrentSubscriptions since this consumer's
                 * assignment count has increased.
                 * This is an O(N) operation since it is a single shuffle. */
                rd_list_sort(sortedCurrentSubscriptions,
                             sort_by_map_elem_val_toppar_list_cnt);
                return;
        }
}

/**
 * @returns true if the partition has two or more potential consumers.
 */
static RD_INLINE rd_bool_t partitionCanParticipateInReassignment(
    const rd_kafka_topic_partition_t *partition,
    map_toppar_list_t *partition2AllPotentialConsumers) {
        rd_list_t *consumers;

        if (!(consumers =
                  RD_MAP_GET(partition2AllPotentialConsumers, partition)))
                return rd_false;

        return rd_list_cnt(consumers) >= 2;
}


/**
 * @returns true if consumer can participate in reassignment based on
 *          its current assignment.
 */
static RD_INLINE rd_bool_t consumerCanParticipateInReassignment(
    rd_kafka_t *rk,
    const char *consumer,
    map_str_toppar_list_t *currentAssignment,
    map_str_toppar_list_t *consumer2AllPotentialPartitions,
    map_toppar_list_t *partition2AllPotentialConsumers) {
        const rd_kafka_topic_partition_list_t *currentPartitions =
            RD_MAP_GET(currentAssignment, consumer);
        int currentAssignmentSize = currentPartitions->cnt;
        int maxAssignmentSize =
            RD_MAP_GET(consumer2AllPotentialPartitions, consumer)->cnt;
        int i;

        /* FIXME: And then what? Is this a local error? If so, assert. */
        if (currentAssignmentSize > maxAssignmentSize)
                rd_kafka_log(rk, LOG_ERR, "STICKY",
                             "Sticky assignor error: "
                             "Consumer %s is assigned more partitions (%d) "
                             "than the maximum possible (%d)",
                             consumer, currentAssignmentSize,
                             maxAssignmentSize);

        /* If a consumer is not assigned all its potential partitions it is
         * subject to reassignment. */
        if (currentAssignmentSize < maxAssignmentSize)
                return rd_true;

        /* If any of the partitions assigned to a consumer is subject to
         * reassignment the consumer itself is subject to reassignment. */
        for (i = 0; i < currentPartitions->cnt; i++) {
                const rd_kafka_topic_partition_t *partition =
                    &currentPartitions->elems[i];

                if (partitionCanParticipateInReassignment(
                        partition, partition2AllPotentialConsumers))
                        return rd_true;
        }

        return rd_false;
}


/**
 * @brief Process moving partition from old consumer to new consumer.
 */
static void processPartitionMovement(
    rd_kafka_t *rk,
    PartitionMovements_t *partitionMovements,
    const rd_kafka_topic_partition_t *partition,
    const char *newConsumer,
    map_str_toppar_list_t *currentAssignment,
    rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
    map_toppar_str_t *currentPartitionConsumer) {

        const char *oldConsumer =
            RD_MAP_GET(currentPartitionConsumer, partition);

        PartitionMovements_movePartition(partitionMovements, partition,
                                         oldConsumer, newConsumer);

        rd_kafka_topic_partition_list_add(
            RD_MAP_GET(currentAssignment, newConsumer), partition->topic,
            partition->partition);

        rd_kafka_topic_partition_list_del(
            RD_MAP_GET(currentAssignment, oldConsumer), partition->topic,
            partition->partition);

        RD_MAP_SET(currentPartitionConsumer,
                   rd_kafka_topic_partition_copy(partition), newConsumer);

        /* Re-sort after assignment count has changed. */
        rd_list_sort(sortedCurrentSubscriptions,
                     sort_by_map_elem_val_toppar_list_cnt);

        rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                     "%s [%" PRId32 "] %sassigned to %s (from %s)",
                     partition->topic, partition->partition,
                     oldConsumer ? "re" : "", newConsumer,
                     oldConsumer ? oldConsumer : "(none)");
}


/**
 * @brief Reassign \p partition to \p newConsumer
 */
static void reassignPartitionToConsumer(
    rd_kafka_t *rk,
    PartitionMovements_t *partitionMovements,
    const rd_kafka_topic_partition_t *partition,
    map_str_toppar_list_t *currentAssignment,
    rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
    map_toppar_str_t *currentPartitionConsumer,
    const char *newConsumer) {

        const char *consumer = RD_MAP_GET(currentPartitionConsumer, partition);
        const rd_kafka_topic_partition_t *partitionToBeMoved;

        /* Find the correct partition movement considering
         * the stickiness requirement. */
        partitionToBeMoved = PartitionMovements_getTheActualPartitionToBeMoved(
            partitionMovements, partition, consumer, newConsumer);

        processPartitionMovement(rk, partitionMovements, partitionToBeMoved,
                                 newConsumer, currentAssignment,
                                 sortedCurrentSubscriptions,
                                 currentPartitionConsumer);
}

/**
 * @brief Reassign \p partition to an eligible new consumer.
 */
static void
reassignPartition(rd_kafka_t *rk,
                  PartitionMovements_t *partitionMovements,
                  const rd_kafka_topic_partition_t *partition,
                  map_str_toppar_list_t *currentAssignment,
                  rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
                  map_toppar_str_t *currentPartitionConsumer,
                  map_str_toppar_list_t *consumer2AllPotentialPartitions) {

        const rd_map_elem_t *elem;
        int i;

        /* Find the new consumer */
        RD_LIST_FOREACH(elem, sortedCurrentSubscriptions, i) {
                const char *newConsumer = (const char *)elem->key;

                if (rd_kafka_topic_partition_list_find(
                        RD_MAP_GET(consumer2AllPotentialPartitions,
                                   newConsumer),
                        partition->topic, partition->partition)) {
                        reassignPartitionToConsumer(
                            rk, partitionMovements, partition,
                            currentAssignment, sortedCurrentSubscriptions,
                            currentPartitionConsumer, newConsumer);

                        return;
                }
        }

        rd_assert(!*"reassignPartition(): no new consumer found");
}



/**
 * @brief Determine if the current assignment is balanced.
 *
 * @param currentAssignment the assignment whose balance needs to be checked
 * @param sortedCurrentSubscriptions an ascending sorted set of consumers based
 *                                   on how many topic partitions are already
 *                                   assigned to them
 * @param consumer2AllPotentialPartitions a mapping of all consumers to all
 *                                        potential topic partitions that can be
 *                                        assigned to them.
 *                                        This parameter is called
 *                                        allSubscriptions in the Java
 *                                        implementation, but we choose this
 *                                        name to be more consistent with its
 *                                        use elsewhere in the code.
 * @param partition2AllPotentialConsumers a mapping of all partitions to
 *                                        all potential consumers.
 *
 * @returns true if the given assignment is balanced; false otherwise
 */
static rd_bool_t
isBalanced(rd_kafka_t *rk,
           map_str_toppar_list_t *currentAssignment,
           const rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
           map_str_toppar_list_t *consumer2AllPotentialPartitions,
           map_toppar_list_t *partition2AllPotentialConsumers) {

        int minimum = ((const rd_kafka_topic_partition_list_t
                            *)((const rd_map_elem_t *)rd_list_first(
                                   sortedCurrentSubscriptions))
                           ->value)
                          ->cnt;
        int maximum = ((const rd_kafka_topic_partition_list_t
                            *)((const rd_map_elem_t *)rd_list_last(
                                   sortedCurrentSubscriptions))
                           ->value)
                          ->cnt;

        /* Mapping from partitions to the consumer assigned to them */
        // FIXME: don't create prior to min/max check below */
        map_toppar_str_t allPartitions = RD_MAP_INITIALIZER(
            RD_MAP_CNT(partition2AllPotentialConsumers),
            rd_kafka_topic_partition_cmp, rd_kafka_topic_partition_hash,
            NULL /* references currentAssignment */,
            NULL /* references currentAssignment */);

        /* Iterators */
        const rd_kafka_topic_partition_list_t *partitions;
        const char *consumer;
        const rd_map_elem_t *elem;
        int i;

        /* The assignment is balanced if minimum and maximum numbers of
         * partitions assigned to consumers differ by at most one. */
        if (minimum >= maximum - 1) {
                rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                             "Assignment is balanced: "
                             "minimum %d and maximum %d partitions assigned "
                             "to each consumer",
                             minimum, maximum);
                RD_MAP_DESTROY(&allPartitions);
                return rd_true;
        }

        /* Create a mapping from partitions to the consumer assigned to them */
        RD_MAP_FOREACH(consumer, partitions, currentAssignment) {

                for (i = 0; i < partitions->cnt; i++) {
                        const rd_kafka_topic_partition_t *partition =
                            &partitions->elems[i];
                        const char *existing;
                        if ((existing = RD_MAP_GET(&allPartitions, partition)))
                                rd_kafka_log(rk, LOG_ERR, "STICKY",
                                             "Sticky assignor: %s [%" PRId32
                                             "] "
                                             "is assigned to more than one "
                                             "consumer (%s and %s)",
                                             partition->topic,
                                             partition->partition, existing,
                                             consumer);

                        RD_MAP_SET(&allPartitions, partition, consumer);
                }
        }


        /* 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.
         *
         * Note: Since sortedCurrentSubscriptions elements are pointers to
         *       currentAssignment's element we get both the consumer
         *       and partition list in elem here. */
        RD_LIST_FOREACH(elem, sortedCurrentSubscriptions, i) {
                const char *consumer = (const char *)elem->key;
                const rd_kafka_topic_partition_list_t *potentialTopicPartitions;
                const rd_kafka_topic_partition_list_t *consumerPartitions;

                consumerPartitions =
                    (const rd_kafka_topic_partition_list_t *)elem->value;

                potentialTopicPartitions =
                    RD_MAP_GET(consumer2AllPotentialPartitions, consumer);

                /* Skip if this consumer already has all the topic partitions
                 * it can get. */
                if (consumerPartitions->cnt == potentialTopicPartitions->cnt)
                        continue;

                /* Otherwise make sure it can't get any more partitions */

                for (i = 0; i < potentialTopicPartitions->cnt; i++) {
                        const rd_kafka_topic_partition_t *partition =
                            &potentialTopicPartitions->elems[i];
                        const char *otherConsumer;
                        int otherConsumerPartitionCount;

                        if (rd_kafka_topic_partition_list_find(
                                consumerPartitions, partition->topic,
                                partition->partition))
                                continue;

                        otherConsumer = RD_MAP_GET(&allPartitions, partition);
                        otherConsumerPartitionCount =
                            RD_MAP_GET(currentAssignment, otherConsumer)->cnt;

                        if (consumerPartitions->cnt <
                            otherConsumerPartitionCount) {
                                rd_kafka_dbg(
                                    rk, ASSIGNOR, "STICKY",
                                    "%s [%" PRId32
                                    "] can be moved from "
                                    "consumer %s (%d partition(s)) to "
                                    "consumer %s (%d partition(s)) "
                                    "for a more balanced assignment",
                                    partition->topic, partition->partition,
                                    otherConsumer, otherConsumerPartitionCount,
                                    consumer, consumerPartitions->cnt);
                                RD_MAP_DESTROY(&allPartitions);
                                return rd_false;
                        }
                }
        }

        RD_MAP_DESTROY(&allPartitions);
        return rd_true;
}


/**
 * @brief Perform reassignment.
 *
 * @returns true if reassignment was performed.
 */
static rd_bool_t
performReassignments(rd_kafka_t *rk,
                     PartitionMovements_t *partitionMovements,
                     rd_kafka_topic_partition_list_t *reassignablePartitions,
                     map_str_toppar_list_t *currentAssignment,
                     map_toppar_cgpair_t *prevAssignment,
                     rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
                     map_str_toppar_list_t *consumer2AllPotentialPartitions,
                     map_toppar_list_t *partition2AllPotentialConsumers,
                     map_toppar_str_t *currentPartitionConsumer) {
        rd_bool_t reassignmentPerformed = rd_false;
        rd_bool_t modified, saveIsBalanced = rd_false;
        int iterations = 0;

        /* Repeat reassignment until no partition can be moved to
         * improve the balance. */
        do {
                int i;

                iterations++;

                modified = rd_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. */

                for (i = 0; i < reassignablePartitions->cnt &&
                            !isBalanced(rk, currentAssignment,
                                        sortedCurrentSubscriptions,
                                        consumer2AllPotentialPartitions,
                                        partition2AllPotentialConsumers);
                     i++) {
                        const rd_kafka_topic_partition_t *partition =
                            &reassignablePartitions->elems[i];
                        const rd_list_t *consumers = RD_MAP_GET(
                            partition2AllPotentialConsumers, partition);
                        const char *consumer, *otherConsumer;
                        const ConsumerGenerationPair_t *prevcgp;
                        const rd_kafka_topic_partition_list_t *currAssignment;
                        int j;

                        /* FIXME: Is this a local error/bug? If so, assert */
                        if (rd_list_cnt(consumers) <= 1)
                                rd_kafka_log(
                                    rk, LOG_ERR, "STICKY",
                                    "Sticky assignor: expected more than "
                                    "one potential consumer for partition "
                                    "%s [%" PRId32 "]",
                                    partition->topic, partition->partition);

                        /* The partition must have a current consumer */
                        consumer =
                            RD_MAP_GET(currentPartitionConsumer, partition);
                        rd_assert(consumer);

                        currAssignment =
                            RD_MAP_GET(currentAssignment, consumer);
                        prevcgp = RD_MAP_GET(prevAssignment, partition);

                        if (prevcgp &&
                            currAssignment->cnt >
                                RD_MAP_GET(currentAssignment, prevcgp->consumer)
                                        ->cnt +
                                    1) {
                                reassignPartitionToConsumer(
                                    rk, partitionMovements, partition,
                                    currentAssignment,
                                    sortedCurrentSubscriptions,
                                    currentPartitionConsumer,
                                    prevcgp->consumer);
                                reassignmentPerformed = rd_true;
                                modified              = rd_true;
                                continue;
                        }

                        /* Check if a better-suited consumer exists for the
                         * partition; if so, reassign it. */
                        RD_LIST_FOREACH(otherConsumer, consumers, j) {
                                if (consumer == otherConsumer)
                                        continue;

                                if (currAssignment->cnt <=
                                    RD_MAP_GET(currentAssignment, otherConsumer)
                                            ->cnt +
                                        1)
                                        continue;

                                reassignPartition(
                                    rk, partitionMovements, partition,
                                    currentAssignment,
                                    sortedCurrentSubscriptions,
                                    currentPartitionConsumer,
                                    consumer2AllPotentialPartitions);

                                reassignmentPerformed = rd_true;
                                modified              = rd_true;
                                break;
                        }
                }

                if (i < reassignablePartitions->cnt)
                        saveIsBalanced = rd_true;

        } while (modified);

        rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                     "Reassignment %sperformed after %d iteration(s) of %d "
                     "reassignable partition(s)%s",
                     reassignmentPerformed ? "" : "not ", iterations,
                     reassignablePartitions->cnt,
                     saveIsBalanced ? ": assignment is balanced" : "");

        return reassignmentPerformed;
}


/**
 * @returns 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.
 * FIXME: should be called imbalance score then?
 */
static int getBalanceScore(map_str_toppar_list_t *assignment) {
        const char *consumer;
        const rd_kafka_topic_partition_list_t *partitions;
        int *sizes;
        int cnt   = 0;
        int score = 0;
        int i, next;

        /* If there is just a single consumer the assignment will be balanced */
        if (RD_MAP_CNT(assignment) < 2)
                return 0;

        sizes = rd_malloc(sizeof(*sizes) * RD_MAP_CNT(assignment));

        RD_MAP_FOREACH(consumer, partitions, assignment)
        sizes[cnt++] = partitions->cnt;

        for (next = 0; next < cnt; next++)
                for (i = next + 1; i < cnt; i++)
                        score += abs(sizes[next] - sizes[i]);

        rd_free(sizes);

        if (consumer)
                ; /* Avoid unused warning */

        return score;
}



/**
 * @brief Balance the current assignment using the data structures
 *        created in assign_cb(). */
static void balance(rd_kafka_t *rk,
                    PartitionMovements_t *partitionMovements,
                    map_str_toppar_list_t *currentAssignment,
                    map_toppar_cgpair_t *prevAssignment,
                    rd_kafka_topic_partition_list_t *sortedPartitions,
                    rd_kafka_topic_partition_list_t *unassignedPartitions,
                    rd_list_t *sortedCurrentSubscriptions /*rd_map_elem_t*/,
                    map_str_toppar_list_t *consumer2AllPotentialPartitions,
                    map_toppar_list_t *partition2AllPotentialConsumers,
                    map_toppar_str_t *currentPartitionConsumer,
                    rd_bool_t revocationRequired) {

        /* If the consumer with most assignments (thus the last element
         * in the ascendingly ordered sortedCurrentSubscriptions list) has
         * zero partitions assigned it means there is no current assignment
         * for any consumer and the group is thus initializing for the first
         * time. */
        rd_bool_t initializing = ((const rd_kafka_topic_partition_list_t
                                       *)((const rd_map_elem_t *)rd_list_last(
                                              sortedCurrentSubscriptions))
                                      ->value)
                                     ->cnt == 0;
        rd_bool_t reassignmentPerformed = rd_false;

        map_str_toppar_list_t fixedAssignments =
                RD_MAP_INITIALIZER(RD_MAP_CNT(partition2AllPotentialConsumers),
                                   rd_map_str_cmp,
                                   rd_map_str_hash,
                                   NULL,
                                   NULL /* Will transfer ownership of the list
                                         * to currentAssignment at the end of
                                         * this function. */);

        map_str_toppar_list_t preBalanceAssignment = RD_MAP_INITIALIZER(
            RD_MAP_CNT(currentAssignment), rd_map_str_cmp, rd_map_str_hash,
            NULL /* references currentAssignment */,
            rd_kafka_topic_partition_list_destroy_free);
        map_toppar_str_t preBalancePartitionConsumers = RD_MAP_INITIALIZER(
            RD_MAP_CNT(partition2AllPotentialConsumers),
            rd_kafka_topic_partition_cmp, rd_kafka_topic_partition_hash,
            rd_kafka_topic_partition_destroy_free,
            NULL /* refs currentPartitionConsumer */);
        int newScore, oldScore;
        /* Iterator variables */
        const rd_kafka_topic_partition_t *partition;
        const void *ignore;
        const rd_map_elem_t *elem;
        int i;

        /* Assign all unassigned partitions */
        for (i = 0; i < unassignedPartitions->cnt; i++) {
                partition = &unassignedPartitions->elems[i];

                /* Skip if there is no potential consumer for the partition.
                 * FIXME: How could this be? */
                if (rd_list_empty(RD_MAP_GET(partition2AllPotentialConsumers,
                                             partition))) {
                        rd_dassert(!*"sticky assignor bug");
                        continue;
                }

                assignPartition(
                    partition, sortedCurrentSubscriptions, currentAssignment,
                    consumer2AllPotentialPartitions, currentPartitionConsumer);
        }


        /* Narrow down the reassignment scope to only those partitions that can
         * actually be reassigned. */
        RD_MAP_FOREACH(partition, ignore, partition2AllPotentialConsumers) {
                if (partitionCanParticipateInReassignment(
                        partition, partition2AllPotentialConsumers))
                        continue;

                rd_kafka_topic_partition_list_del(
                    sortedPartitions, partition->topic, partition->partition);
                rd_kafka_topic_partition_list_del(unassignedPartitions,
                                                  partition->topic,
                                                  partition->partition);
        }

        if (ignore)
                ; /* Avoid unused warning */


        /* Narrow down the reassignment scope to only those consumers that are
         * subject to reassignment. */
        RD_LIST_FOREACH(elem, sortedCurrentSubscriptions, i) {
                const char *consumer = (const char *)elem->key;
                rd_kafka_topic_partition_list_t *partitions;

                if (consumerCanParticipateInReassignment(
                        rk, consumer, currentAssignment,
                        consumer2AllPotentialPartitions,
                        partition2AllPotentialConsumers))
                        continue;

                rd_list_remove_elem(sortedCurrentSubscriptions, i);
                i--; /* Since the current element is removed we need
                      * to rewind the iterator. */

                partitions = rd_kafka_topic_partition_list_copy(
                    RD_MAP_GET(currentAssignment, consumer));
                RD_MAP_DELETE(currentAssignment, consumer);

                RD_MAP_SET(&fixedAssignments, consumer, partitions);
        }


        rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                     "Prepared balanced reassignment for %d consumers, "
                     "%d available partition(s) where of %d are unassigned "
                     "(initializing=%s, revocationRequired=%s, "
                     "%d fixed assignments)",
                     (int)RD_MAP_CNT(consumer2AllPotentialPartitions),
                     sortedPartitions->cnt, unassignedPartitions->cnt,
                     initializing ? "true" : "false",
                     revocationRequired ? "true" : "false",
                     (int)RD_MAP_CNT(&fixedAssignments));

        /* Create a deep copy of the current assignment so we can revert to it
         * if we do not get a more balanced assignment later. */
        RD_MAP_COPY(&preBalanceAssignment, currentAssignment,
                    NULL /* just reference the key */,
                    (rd_map_copy_t *)rd_kafka_topic_partition_list_copy);
        RD_MAP_COPY(&preBalancePartitionConsumers, currentPartitionConsumer,
                    rd_kafka_topic_partition_copy_void,
                    NULL /* references assign_cb(members) fields */);


        /* If we don't already need to revoke something due to subscription
         * changes, first try to balance by only moving newly added partitions.
         */
        if (!revocationRequired && unassignedPartitions->cnt > 0)
                performReassignments(
                    rk, partitionMovements, unassignedPartitions,
                    currentAssignment, prevAssignment,
                    sortedCurrentSubscriptions, consumer2AllPotentialPartitions,
                    partition2AllPotentialConsumers, currentPartitionConsumer);

        reassignmentPerformed = performReassignments(
            rk, partitionMovements, sortedPartitions, currentAssignment,
            prevAssignment, 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 &&
            (newScore = getBalanceScore(currentAssignment)) >=
                (oldScore = getBalanceScore(&preBalanceAssignment))) {

                rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                             "Reassignment performed but keeping previous "
                             "assignment since balance score did not improve: "
                             "new score %d (%d consumers) vs "
                             "old score %d (%d consumers): "
                             "lower score is better",
                             newScore, (int)RD_MAP_CNT(currentAssignment),
                             oldScore, (int)RD_MAP_CNT(&preBalanceAssignment));

                RD_MAP_COPY(
                    currentAssignment, &preBalanceAssignment,
                    NULL /* just reference the key */,
                    (rd_map_copy_t *)rd_kafka_topic_partition_list_copy);

                RD_MAP_CLEAR(currentPartitionConsumer);
                RD_MAP_COPY(currentPartitionConsumer,
                            &preBalancePartitionConsumers,
                            rd_kafka_topic_partition_copy_void,
                            NULL /* references assign_cb(members) fields */);
        }

        RD_MAP_DESTROY(&preBalancePartitionConsumers);
        RD_MAP_DESTROY(&preBalanceAssignment);

        /* Add the fixed assignments (those that could not change) back. */
        if (!RD_MAP_IS_EMPTY(&fixedAssignments)) {
                const rd_map_elem_t *elem;

                RD_MAP_FOREACH_ELEM(elem, &fixedAssignments.rmap) {
                        const char *consumer = elem->key;
                        rd_kafka_topic_partition_list_t *partitions =
                            (rd_kafka_topic_partition_list_t *)elem->value;

                        RD_MAP_SET(currentAssignment, consumer, partitions);

                        rd_list_add(sortedCurrentSubscriptions, (void *)elem);
                }

                /* Re-sort */
                rd_list_sort(sortedCurrentSubscriptions,
                             sort_by_map_elem_val_toppar_list_cnt);
        }

        RD_MAP_DESTROY(&fixedAssignments);
}



/**
 * @brief Populate subscriptions, current and previous assignments based on the
 *        \p members assignments.
 */
static void prepopulateCurrentAssignments(
    rd_kafka_t *rk,
    rd_kafka_group_member_t *members,
    size_t member_cnt,
    map_str_toppar_list_t *subscriptions,
    map_str_toppar_list_t *currentAssignment,
    map_toppar_cgpair_t *prevAssignment,
    map_toppar_str_t *currentPartitionConsumer,
    map_str_toppar_list_t *consumer2AllPotentialPartitions,
    size_t estimated_partition_cnt) {

        /* We need to process subscriptions' user data with each consumer's
         * reported generation in mind.
         * Higher generations overwrite lower generations in case of a conflict.
         * Conflicts will only exist if user data is for different generations.
         */

        /* For each partition we create a sorted list (by generation) of
         * its consumers. */
        RD_MAP_LOCAL_INITIALIZER(
            sortedPartitionConsumersByGeneration, member_cnt * 10 /* FIXME */,
            const rd_kafka_topic_partition_t *,
            /* List of ConsumerGenerationPair_t */
            rd_list_t *, rd_kafka_topic_partition_cmp,
            rd_kafka_topic_partition_hash, NULL, rd_list_destroy_free);
        const rd_kafka_topic_partition_t *partition;
        rd_list_t *consumers;
        int i;

        /* For each partition that is currently assigned to the group members
         * add the member and its generation to
         * sortedPartitionConsumersByGeneration (which is sorted afterwards)
         * indexed by the partition. */
        for (i = 0; i < (int)member_cnt; i++) {
                rd_kafka_group_member_t *consumer = &members[i];
                int j;

                RD_MAP_SET(subscriptions, consumer->rkgm_member_id->str,
                           consumer->rkgm_subscription);

                RD_MAP_SET(currentAssignment, consumer->rkgm_member_id->str,
                           rd_kafka_topic_partition_list_new(10));

                RD_MAP_SET(consumer2AllPotentialPartitions,
                           consumer->rkgm_member_id->str,
                           rd_kafka_topic_partition_list_new(
                               (int)estimated_partition_cnt));

                if (!consumer->rkgm_owned)
                        continue;

                for (j = 0; j < (int)consumer->rkgm_owned->cnt; j++) {
                        partition = &consumer->rkgm_owned->elems[j];

                        consumers = RD_MAP_GET_OR_SET(
                            &sortedPartitionConsumersByGeneration, partition,
                            rd_list_new(10, ConsumerGenerationPair_destroy));

                        if (consumer->rkgm_generation != -1 &&
                            rd_list_find(
                                consumers, &consumer->rkgm_generation,
                                ConsumerGenerationPair_cmp_generation)) {
                                rd_kafka_log(
                                    rk, LOG_WARNING, "STICKY",
                                    "Sticky assignor: "
                                    "%s [%" PRId32
                                    "] is assigned to "
                                    "multiple consumers with same "
                                    "generation %d: "
                                    "skipping member %.*s",
                                    partition->topic, partition->partition,
                                    consumer->rkgm_generation,
                                    RD_KAFKAP_STR_PR(consumer->rkgm_member_id));
                                continue;
                        }

                        rd_list_add(consumers,
                                    ConsumerGenerationPair_new(
                                        consumer->rkgm_member_id->str,
                                        consumer->rkgm_generation));

                        RD_MAP_SET(currentPartitionConsumer,
                                   rd_kafka_topic_partition_copy(partition),
                                   consumer->rkgm_member_id->str);
                }
        }

        /* Populate currentAssignment and prevAssignment.
         * prevAssignment holds the prior ConsumerGenerationPair_t
         * (before current) of each partition. */
        RD_MAP_FOREACH(partition, consumers,
                       &sortedPartitionConsumersByGeneration) {
                /* current and previous are the last two consumers
                 * of each partition. */
                ConsumerGenerationPair_t *current, *previous;
                rd_kafka_topic_partition_list_t *partitions;

                /* Sort the per-partition consumers list by generation */
                rd_list_sort(consumers, ConsumerGenerationPair_cmp_generation);

                /* Add current (highest generation) consumer
                 * to currentAssignment. */
                current    = rd_list_elem(consumers, 0);
                partitions = RD_MAP_GET(currentAssignment, current->consumer);
                rd_kafka_topic_partition_list_add(partitions, partition->topic,
                                                  partition->partition);

                /* Add previous (next highest generation) consumer, if any,
                 * to prevAssignment. */
                previous = rd_list_elem(consumers, 1);
                if (previous)
                        RD_MAP_SET(
                            prevAssignment,
                            rd_kafka_topic_partition_copy(partition),
                            ConsumerGenerationPair_new(previous->consumer,
                                                       previous->generation));
        }

        RD_MAP_DESTROY(&sortedPartitionConsumersByGeneration);
}


/**
 * @brief Populate maps for potential partitions per consumer and vice-versa.
 */
static void
populatePotentialMaps(const rd_kafka_assignor_topic_t *atopic,
                      map_toppar_list_t *partition2AllPotentialConsumers,
                      map_str_toppar_list_t *consumer2AllPotentialPartitions,
                      size_t estimated_partition_cnt) {
        int i;
        const rd_kafka_group_member_t *rkgm;

        /* for each eligible (subscribed and available) topic (\p atopic):
         *   for each member subscribing to that topic:
         *     and for each partition of that topic:
         *        add consumer and partition to:
         *          partition2AllPotentialConsumers
         *          consumer2AllPotentialPartitions
         */

        RD_LIST_FOREACH(rkgm, &atopic->members, i) {
                const char *consumer = rkgm->rkgm_member_id->str;
                rd_kafka_topic_partition_list_t *partitions =
                    RD_MAP_GET(consumer2AllPotentialPartitions, consumer);
                int j;

                rd_assert(partitions != NULL);

                for (j = 0; j < atopic->metadata->partition_cnt; j++) {
                        rd_kafka_topic_partition_t *partition;
                        rd_list_t *consumers;

                        /* consumer2AllPotentialPartitions[consumer] += part */
                        partition = rd_kafka_topic_partition_list_add(
                            partitions, atopic->metadata->topic,
                            atopic->metadata->partitions[j].id);

                        /* partition2AllPotentialConsumers[part] += consumer */
                        if (!(consumers =
                                  RD_MAP_GET(partition2AllPotentialConsumers,
                                             partition))) {
                                consumers = rd_list_new(
                                    RD_MAX(2, (int)estimated_partition_cnt / 2),
                                    NULL);
                                RD_MAP_SET(
                                    partition2AllPotentialConsumers,
                                    rd_kafka_topic_partition_copy(partition),
                                    consumers);
                        }
                        rd_list_add(consumers, (void *)consumer);
                }
        }
}


/**
 * @returns true if all consumers have identical subscriptions based on
 *          the currently available topics and partitions.
 *
 * @remark The Java code checks both partition2AllPotentialConsumers and
 *         and consumer2AllPotentialPartitions but since these maps
 *         are symmetrical we only check one of them.
 *         ^ FIXME, but we do.
 */
static rd_bool_t areSubscriptionsIdentical(
    map_toppar_list_t *partition2AllPotentialConsumers,
    map_str_toppar_list_t *consumer2AllPotentialPartitions) {
        const void *ignore;
        const rd_list_t *lcurr, *lprev                       = NULL;
        const rd_kafka_topic_partition_list_t *pcurr, *pprev = NULL;

        RD_MAP_FOREACH(ignore, lcurr, partition2AllPotentialConsumers) {
                if (lprev && rd_list_cmp(lcurr, lprev, rd_map_str_cmp))
                        return rd_false;
                lprev = lcurr;
        }

        RD_MAP_FOREACH(ignore, pcurr, consumer2AllPotentialPartitions) {
                if (pprev && rd_kafka_topic_partition_list_cmp(
                                 pcurr, pprev, rd_kafka_topic_partition_cmp))
                        return rd_false;
                pprev = pcurr;
        }

        if (ignore) /* Avoid unused warning */
                ;

        return rd_true;
}


/**
 * @brief Comparator to sort an rd_kafka_topic_partition_list_t in ascending
 *        order by the number of list elements in the .opaque field, or
 *        secondarily by the topic name.
 *        Used by sortPartitions().
 */
static int
toppar_sort_by_list_cnt(const void *_a, const void *_b, void *opaque) {
        const rd_kafka_topic_partition_t *a = _a, *b = _b;
        const rd_list_t *al = a->opaque, *bl = b->opaque;
        int r = rd_list_cnt(al) - rd_list_cnt(bl); /* ascending order */
        if (r)
                return r;
        return rd_kafka_topic_partition_cmp(a, b);
}


/**
 * @brief Sort valid partitions so they are processed in the potential
 *        reassignment phase in the proper order that causes minimal partition
 *        movement among consumers (hence honouring maximal stickiness).
 *
 * @returns The result of the partitions sort.
 */
static rd_kafka_topic_partition_list_t *
sortPartitions(rd_kafka_t *rk,
               map_str_toppar_list_t *currentAssignment,
               map_toppar_cgpair_t *prevAssignment,
               rd_bool_t isFreshAssignment,
               map_toppar_list_t *partition2AllPotentialConsumers,
               map_str_toppar_list_t *consumer2AllPotentialPartitions) {

        rd_kafka_topic_partition_list_t *sortedPartitions;
        map_str_toppar_list_t assignments = RD_MAP_INITIALIZER(
            RD_MAP_CNT(currentAssignment), rd_map_str_cmp, rd_map_str_hash,
            NULL, rd_kafka_topic_partition_list_destroy_free);
        rd_kafka_topic_partition_list_t *partitions;
        const rd_kafka_topic_partition_t *partition;
        const rd_list_t *consumers;
        const char *consumer;
        rd_list_t sortedConsumers; /* element is the (rd_map_elem_t *) from
                                    * assignments. */
        const rd_map_elem_t *elem;
        rd_bool_t wasEmpty;
        int i;

        sortedPartitions = rd_kafka_topic_partition_list_new(
            (int)RD_MAP_CNT(partition2AllPotentialConsumers));
        ;

        rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                     "Sort %d partitions in %s assignment",
                     (int)RD_MAP_CNT(partition2AllPotentialConsumers),
                     isFreshAssignment ? "fresh" : "existing");

        if (isFreshAssignment ||
            !areSubscriptionsIdentical(partition2AllPotentialConsumers,
                                       consumer2AllPotentialPartitions)) {
                /* Create an ascending sorted list of partitions based on
                 * how many consumers can potentially use them. */
                RD_MAP_FOREACH(partition, consumers,
                               partition2AllPotentialConsumers) {
                        rd_kafka_topic_partition_list_add(sortedPartitions,
                                                          partition->topic,
                                                          partition->partition)
                            ->opaque = (void *)consumers;
                }

                rd_kafka_topic_partition_list_sort(
                    sortedPartitions, toppar_sort_by_list_cnt, NULL);

                RD_MAP_DESTROY(&assignments);

                return sortedPartitions;
        }

        /* If this is a reassignment and the subscriptions are identical
         * then we just need to list partitions in a round robin fashion
         * (from consumers with most assigned partitions to those
         * with least assigned partitions). */

        /* Create an ascending sorted list of consumers by valid
         * partition count. The list element is the `rd_map_elem_t *`
         * of the assignments map. This allows us to get a sorted list
         * of consumers without too much data duplication. */
        rd_list_init(&sortedConsumers, (int)RD_MAP_CNT(currentAssignment),
                     NULL);

        RD_MAP_FOREACH(consumer, partitions, currentAssignment) {
                rd_kafka_topic_partition_list_t *partitions2;

                /* Sort assigned partitions for consistency (during tests) */
                rd_kafka_topic_partition_list_sort(partitions, NULL, NULL);

                partitions2 =
                    rd_kafka_topic_partition_list_new(partitions->cnt);

                for (i = 0; i < partitions->cnt; i++) {
                        partition = &partitions->elems[i];

                        /* Only add partitions from the current assignment
                         * that still exist. */
                        if (RD_MAP_GET(partition2AllPotentialConsumers,
                                       partition))
                                rd_kafka_topic_partition_list_add(
                                    partitions2, partition->topic,
                                    partition->partition);
                }

                if (partitions2->cnt > 0) {
                        elem = RD_MAP_SET(&assignments, consumer, partitions2);
                        rd_list_add(&sortedConsumers, (void *)elem);
                } else
                        rd_kafka_topic_partition_list_destroy(partitions2);
        }

        /* Sort consumers */
        rd_list_sort(&sortedConsumers, sort_by_map_elem_val_toppar_list_cnt);

        /* At this point sortedConsumers contains an ascending-sorted list
         * of consumers based on how many valid partitions are currently
         * assigned to them. */

        while (!rd_list_empty(&sortedConsumers)) {
                /* Take consumer with most partitions */
                const rd_map_elem_t *elem = rd_list_last(&sortedConsumers);
                const char *consumer      = (const char *)elem->key;
                /* Currently assigned partitions to this consumer */
                rd_kafka_topic_partition_list_t *remainingPartitions =
                    RD_MAP_GET(&assignments, consumer);
                /* Partitions that were assigned to a different consumer
                 * last time */
                rd_kafka_topic_partition_list_t *prevPartitions =
                    rd_kafka_topic_partition_list_new(
                        (int)RD_MAP_CNT(prevAssignment));
                rd_bool_t reSort = rd_true;

                /* From the partitions that had a different consumer before,
                 * keep only those that are assigned to this consumer now. */
                for (i = 0; i < remainingPartitions->cnt; i++) {
                        partition = &remainingPartitions->elems[i];
                        if (RD_MAP_GET(prevAssignment, partition))
                                rd_kafka_topic_partition_list_add(
                                    prevPartitions, partition->topic,
                                    partition->partition);
                }

                if (prevPartitions->cnt > 0) {
                        /* If there is a partition of this consumer that was
                         * assigned to another consumer before, then mark
                         * it as a good option for reassignment. */
                        partition = &prevPartitions->elems[0];

                        rd_kafka_topic_partition_list_del(remainingPartitions,
                                                          partition->topic,
                                                          partition->partition);

                        rd_kafka_topic_partition_list_add(sortedPartitions,
                                                          partition->topic,
                                                          partition->partition);

                        rd_kafka_topic_partition_list_del_by_idx(prevPartitions,
                                                                 0);

                } else if (remainingPartitions->cnt > 0) {
                        /* Otherwise mark any other one of the current
                         * partitions as a reassignment candidate. */
                        partition = &remainingPartitions->elems[0];

                        rd_kafka_topic_partition_list_add(sortedPartitions,
                                                          partition->topic,
                                                          partition->partition);

                        rd_kafka_topic_partition_list_del_by_idx(
                            remainingPartitions, 0);
                } else {
                        rd_list_remove_elem(&sortedConsumers,
                                            rd_list_cnt(&sortedConsumers) - 1);
                        /* No need to re-sort the list (below) */
                        reSort = rd_false;
                }

                rd_kafka_topic_partition_list_destroy(prevPartitions);

                if (reSort) {
                        /* Re-sort the list to keep the consumer with the most
                         * partitions at the end of the list.
                         * This should be an O(N) operation given it is at most
                         * a single shuffle. */
                        rd_list_sort(&sortedConsumers,
                                     sort_by_map_elem_val_toppar_list_cnt);
                }
        }


        wasEmpty = !sortedPartitions->cnt;

        RD_MAP_FOREACH(partition, consumers, partition2AllPotentialConsumers)
        rd_kafka_topic_partition_list_upsert(sortedPartitions, partition->topic,
                                             partition->partition);

        /* If all partitions were added in the foreach loop just above
         * it means there is no order to retain from the sorderConsumer loop
         * below and we sort the partitions according to their topic+partition
         * to get consistent results (mainly in tests). */
        if (wasEmpty)
                rd_kafka_topic_partition_list_sort(sortedPartitions, NULL,
                                                   NULL);

        rd_list_destroy(&sortedConsumers);
        RD_MAP_DESTROY(&assignments);

        return sortedPartitions;
}


/**
 * @brief Transfer currentAssignment to members array.
 */
static void assignToMembers(map_str_toppar_list_t *currentAssignment,
                            rd_kafka_group_member_t *members,
                            size_t member_cnt) {
        size_t i;

        for (i = 0; i < member_cnt; i++) {
                rd_kafka_group_member_t *rkgm = &members[i];
                const rd_kafka_topic_partition_list_t *partitions =
                    RD_MAP_GET(currentAssignment, rkgm->rkgm_member_id->str);
                if (rkgm->rkgm_assignment)
                        rd_kafka_topic_partition_list_destroy(
                            rkgm->rkgm_assignment);
                rkgm->rkgm_assignment =
                    rd_kafka_topic_partition_list_copy(partitions);
        }
}


/**
 * @brief KIP-54 and KIP-341/FIXME sticky assignor.
 *
 * This code is closely mimicking the AK Java AbstractStickyAssignor.assign().
 */
rd_kafka_resp_err_t
rd_kafka_sticky_assignor_assign_cb(rd_kafka_t *rk,
                                   const rd_kafka_assignor_t *rkas,
                                   const char *member_id,
                                   const rd_kafka_metadata_t *metadata,
                                   rd_kafka_group_member_t *members,
                                   size_t member_cnt,
                                   rd_kafka_assignor_topic_t **eligible_topics,
                                   size_t eligible_topic_cnt,
                                   char *errstr,
                                   size_t errstr_size,
                                   void *opaque) {
        /* FIXME: Let the cgrp pass the actual eligible partition count */
        size_t partition_cnt = member_cnt * 10; /* FIXME */

        /* Map of subscriptions. This is \p member turned into a map. */
        map_str_toppar_list_t subscriptions =
            RD_MAP_INITIALIZER(member_cnt, rd_map_str_cmp, rd_map_str_hash,
                               NULL /* refs members.rkgm_member_id */,
                               NULL /* refs members.rkgm_subscription */);

        /* Map member to current assignment */
        map_str_toppar_list_t currentAssignment =
            RD_MAP_INITIALIZER(member_cnt, rd_map_str_cmp, rd_map_str_hash,
                               NULL /* refs members.rkgm_member_id */,
                               rd_kafka_topic_partition_list_destroy_free);

        /* Map partition to ConsumerGenerationPair */
        map_toppar_cgpair_t prevAssignment =
            RD_MAP_INITIALIZER(partition_cnt, rd_kafka_topic_partition_cmp,
                               rd_kafka_topic_partition_hash,
                               rd_kafka_topic_partition_destroy_free,
                               ConsumerGenerationPair_destroy);

        /* Partition assignment movements between consumers */
        PartitionMovements_t partitionMovements;

        rd_bool_t isFreshAssignment;

        /* Mapping of all topic partitions to all consumers that can be
         * assigned to them.
         * Value is an rd_list_t* with elements referencing the \p members
         * \c rkgm_member_id->str. */
        map_toppar_list_t partition2AllPotentialConsumers = RD_MAP_INITIALIZER(
            partition_cnt, rd_kafka_topic_partition_cmp,
            rd_kafka_topic_partition_hash,
            rd_kafka_topic_partition_destroy_free, rd_list_destroy_free);

        /* Mapping of all consumers to all potential topic partitions that
         * can be assigned to them. */
        map_str_toppar_list_t consumer2AllPotentialPartitions =
            RD_MAP_INITIALIZER(member_cnt, rd_map_str_cmp, rd_map_str_hash,
                               NULL,
                               rd_kafka_topic_partition_list_destroy_free);

        /* Mapping of partition to current consumer. */
        map_toppar_str_t currentPartitionConsumer =
            RD_MAP_INITIALIZER(partition_cnt, rd_kafka_topic_partition_cmp,
                               rd_kafka_topic_partition_hash,
                               rd_kafka_topic_partition_destroy_free,
                               NULL /* refs members.rkgm_member_id->str */);

        rd_kafka_topic_partition_list_t *sortedPartitions;
        rd_kafka_topic_partition_list_t *unassignedPartitions;
        rd_list_t sortedCurrentSubscriptions;

        rd_bool_t revocationRequired = rd_false;

        /* Iteration variables */
        const char *consumer;
        rd_kafka_topic_partition_list_t *partitions;
        const rd_map_elem_t *elem;
        int i;

        /* Initialize PartitionMovements */
        PartitionMovements_init(&partitionMovements, eligible_topic_cnt);

        /* Prepopulate current and previous assignments */
        prepopulateCurrentAssignments(
            rk, members, member_cnt, &subscriptions, &currentAssignment,
            &prevAssignment, &currentPartitionConsumer,
            &consumer2AllPotentialPartitions, partition_cnt);

        isFreshAssignment = RD_MAP_IS_EMPTY(&currentAssignment);

        /* Populate partition2AllPotentialConsumers and
         * consumer2AllPotentialPartitions maps by each eligible topic. */
        for (i = 0; i < (int)eligible_topic_cnt; i++)
                populatePotentialMaps(
                    eligible_topics[i], &partition2AllPotentialConsumers,
                    &consumer2AllPotentialPartitions, partition_cnt);


        /* Sort valid partitions to minimize partition movements. */
        sortedPartitions = sortPartitions(
            rk, &currentAssignment, &prevAssignment, isFreshAssignment,
            &partition2AllPotentialConsumers, &consumer2AllPotentialPartitions);


        /* All partitions that need to be assigned (initially set to all
         * partitions but adjusted in the following loop) */
        unassignedPartitions =
            rd_kafka_topic_partition_list_copy(sortedPartitions);

        RD_MAP_FOREACH(consumer, partitions, &currentAssignment) {
                if (!RD_MAP_GET(&subscriptions, consumer)) {
                        /* If a consumer that existed before
                         * (and had some partition assignments) is now removed,
                         * remove it from currentAssignment and its
                         * partitions from currentPartitionConsumer */

                        rd_kafka_dbg(rk, ASSIGNOR, "STICKY",
                                     "Removing now non-existent consumer %s "
                                     "with %d previously assigned partitions",
                                     consumer, partitions->cnt);


                        for (i = 0; i < partitions->cnt; i++) {
                                const rd_kafka_topic_partition_t *partition =
                                    &partitions->elems[i];
                                RD_MAP_DELETE(&currentPartitionConsumer,
                                              partition);
                        }

                        /* FIXME: The delete could be optimized by passing the
                         *        underlying elem_t. */
                        RD_MAP_DELETE(&currentAssignment, consumer);

                } else {
                        /* Otherwise (the consumer still exists) */

                        for (i = 0; i < partitions->cnt; i++) {
                                const rd_kafka_topic_partition_t *partition =
                                    &partitions->elems[i];
                                rd_bool_t remove_part = rd_false;

                                if (!RD_MAP_GET(
                                        &partition2AllPotentialConsumers,
                                        partition)) {
                                        /* If this partition of this consumer
                                         * no longer exists remove it from
                                         * currentAssignment of the consumer */
                                        remove_part = rd_true;
                                        RD_MAP_DELETE(&currentPartitionConsumer,
                                                      partition);

                                } else if (!rd_kafka_topic_partition_list_find(
                                               RD_MAP_GET(&subscriptions,
                                                          consumer),
                                               partition->topic,
                                               RD_KAFKA_PARTITION_UA)) {
                                        /* If this partition cannot remain
                                         * assigned to its current consumer
                                         * because the consumer is no longer
                                         * subscribed to its topic, remove it
                                         * from the currentAssignment of the
                                         * consumer. */
                                        remove_part        = rd_true;
                                        revocationRequired = rd_true;
                                } 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). */
                                        rd_kafka_topic_partition_list_del(
                                            unassignedPartitions,
                                            partition->topic,
                                            partition->partition);
                                }

                                if (remove_part) {
                                        rd_kafka_topic_partition_list_del_by_idx(
                                            partitions, i);
                                        i--; /* Since the current element was
                                              * removed we need the next for
                                              * loop iteration to stay at the
                                              * same index. */
                                }
                        }
                }
        }


        /* 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 list of consumers based on how many topic
         * partitions are already assigned to them. The list element is
         * referencing the rd_map_elem_t* from the currentAssignment map. */
        rd_list_init(&sortedCurrentSubscriptions,
                     (int)RD_MAP_CNT(&currentAssignment), NULL);

        RD_MAP_FOREACH_ELEM(elem, &currentAssignment.rmap)
        rd_list_add(&sortedCurrentSubscriptions, (void *)elem);

        rd_list_sort(&sortedCurrentSubscriptions,
                     sort_by_map_elem_val_toppar_list_cnt);

        /* Balance the available partitions across consumers */
        balance(rk, &partitionMovements, &currentAssignment, &prevAssignment,
                sortedPartitions, unassignedPartitions,
                &sortedCurrentSubscriptions, &consumer2AllPotentialPartitions,
                &partition2AllPotentialConsumers, &currentPartitionConsumer,
                revocationRequired);

        /* Transfer currentAssignment (now updated) to each member's
         * assignment. */
        assignToMembers(&currentAssignment, members, member_cnt);


        rd_list_destroy(&sortedCurrentSubscriptions);

        PartitionMovements_destroy(&partitionMovements);

        rd_kafka_topic_partition_list_destroy(unassignedPartitions);
        rd_kafka_topic_partition_list_destroy(sortedPartitions);

        RD_MAP_DESTROY(&currentPartitionConsumer);
        RD_MAP_DESTROY(&consumer2AllPotentialPartitions);
        RD_MAP_DESTROY(&partition2AllPotentialConsumers);
        RD_MAP_DESTROY(&prevAssignment);
        RD_MAP_DESTROY(&currentAssignment);
        RD_MAP_DESTROY(&subscriptions);

        return RD_KAFKA_RESP_ERR_NO_ERROR;
}



/** @brief FIXME docstring */
static void rd_kafka_sticky_assignor_on_assignment_cb(
    const rd_kafka_assignor_t *rkas,
    void **assignor_state,
    const rd_kafka_topic_partition_list_t *partitions,
    const rd_kafkap_bytes_t *assignment_userdata,
    const rd_kafka_consumer_group_metadata_t *rkcgm) {
        rd_kafka_sticky_assignor_state_t *state =
            (rd_kafka_sticky_assignor_state_t *)*assignor_state;

        if (!state)
                state = rd_calloc(1, sizeof(*state));
        else
                rd_kafka_topic_partition_list_destroy(state->prev_assignment);

        state->prev_assignment = rd_kafka_topic_partition_list_copy(partitions);
        state->generation_id   = rkcgm->generation_id;

        *assignor_state = state;
}

/** @brief FIXME docstring */
static rd_kafkap_bytes_t *rd_kafka_sticky_assignor_get_metadata(
    const rd_kafka_assignor_t *rkas,
    void *assignor_state,
    const rd_list_t *topics,
    const rd_kafka_topic_partition_list_t *owned_partitions,
    const rd_kafkap_str_t *rack_id) {
        rd_kafka_sticky_assignor_state_t *state;
        rd_kafka_buf_t *rkbuf;
        rd_kafkap_bytes_t *metadata;
        rd_kafkap_bytes_t *kbytes;
        size_t len;

        /*
         * UserData (Version: 1) => [previous_assignment] generation
         *   previous_assignment => topic [partitions]
         *     topic => STRING
         *     partitions => partition
         *       partition => INT32
         *   generation => INT32
         *
         * If there is no previous assignment, UserData is NULL.
         */


        if (!assignor_state) {
                return rd_kafka_consumer_protocol_member_metadata_new(
                    topics, NULL, 0, owned_partitions, -1 /* generation */,
                    rack_id);
        }

        state = (rd_kafka_sticky_assignor_state_t *)assignor_state;

        rkbuf = rd_kafka_buf_new(1, 100);
        rd_assert(state->prev_assignment != NULL);
        const rd_kafka_topic_partition_field_t fields[] = {
            RD_KAFKA_TOPIC_PARTITION_FIELD_PARTITION,
            RD_KAFKA_TOPIC_PARTITION_FIELD_END};
        rd_kafka_buf_write_topic_partitions(rkbuf, state->prev_assignment,
                                            rd_false /*skip invalid offsets*/,
                                            rd_false /*any offset*/, fields);
        rd_kafka_buf_write_i32(rkbuf, state->generation_id);

        /* Get binary buffer and allocate a new Kafka Bytes with a copy. */
        rd_slice_init_full(&rkbuf->rkbuf_reader, &rkbuf->rkbuf_buf);
        len    = rd_slice_remains(&rkbuf->rkbuf_reader);
        kbytes = rd_kafkap_bytes_new(NULL, (int32_t)len);
        rd_slice_read(&rkbuf->rkbuf_reader, (void *)kbytes->data, len);
        rd_kafka_buf_destroy(rkbuf);

        metadata = rd_kafka_consumer_protocol_member_metadata_new(
            topics, kbytes->data, kbytes->len, owned_partitions,
            state->generation_id, rack_id);

        rd_kafkap_bytes_destroy(kbytes);

        return metadata;
}


/**
 * @brief Destroy assignor state
 */
static void rd_kafka_sticky_assignor_state_destroy(void *assignor_state) {
        rd_kafka_sticky_assignor_state_t *state =
            (rd_kafka_sticky_assignor_state_t *)assignor_state;

        rd_assert(assignor_state);

        rd_kafka_topic_partition_list_destroy(state->prev_assignment);
        rd_free(state);
}



/**
 * @name Sticky assignor unit tests
 *
 *
 * These are based on AbstractStickyAssignorTest.java
 *
 *
 *
 */



/**
 * @brief Set a member's owned partitions based on its assignment.
 *
 * For use between assignor_run(). This is mimicing a consumer receiving
 * its new assignment and including it in the next rebalance as its
 * owned-partitions.
 */
static void ut_set_owned(rd_kafka_group_member_t *rkgm) {
        if (rkgm->rkgm_owned)
                rd_kafka_topic_partition_list_destroy(rkgm->rkgm_owned);

        rkgm->rkgm_owned =
            rd_kafka_topic_partition_list_copy(rkgm->rkgm_assignment);
}


/**
 * @brief Verify assignment validity and balance.
 *
 * @remark Also updates the members owned partitions to the assignment.
 */

static int verifyValidityAndBalance0(const char *func,
                                     int line,
                                     rd_kafka_group_member_t *members,
                                     size_t member_cnt,
                                     const rd_kafka_metadata_t *metadata) {
        int fails = 0;
        int i;
        rd_bool_t verbose = rd_false; /* Enable for troubleshooting */

        RD_UT_SAY("%s:%d: verifying assignment for %d member(s):", func, line,
                  (int)member_cnt);

        for (i = 0; i < (int)member_cnt; i++) {
                const char *consumer = members[i].rkgm_member_id->str;
                const rd_kafka_topic_partition_list_t *partitions =
                    members[i].rkgm_assignment;
                int p, j;

                if (verbose)
                        RD_UT_SAY(
                            "%s:%d:   "
                            "consumer \"%s\", %d subscribed topic(s), "
                            "%d assigned partition(s):",
                            func, line, consumer,
                            members[i].rkgm_subscription->cnt, partitions->cnt);

                for (p = 0; p < partitions->cnt; p++) {
                        const rd_kafka_topic_partition_t *partition =
                            &partitions->elems[p];

                        if (verbose)
                                RD_UT_SAY("%s:%d:     %s [%" PRId32 "]", func,
                                          line, partition->topic,
                                          partition->partition);

                        if (!rd_kafka_topic_partition_list_find(
                                members[i].rkgm_subscription, partition->topic,
                                RD_KAFKA_PARTITION_UA)) {
                                RD_UT_WARN("%s [%" PRId32
                                           "] is assigned to "
                                           "%s but it is not subscribed to "
                                           "that topic",
                                           partition->topic,
                                           partition->partition, consumer);
                                fails++;
                        }
                }

                /* Update the member's owned partitions to match
                 * the assignment. */
                ut_set_owned(&members[i]);

                if (i == (int)member_cnt - 1)
                        continue;

                for (j = i + 1; j < (int)member_cnt; j++) {
                        const char *otherConsumer =
                            members[j].rkgm_member_id->str;
                        const rd_kafka_topic_partition_list_t *otherPartitions =
                            members[j].rkgm_assignment;
                        rd_bool_t balanced =
                            abs(partitions->cnt - otherPartitions->cnt) <= 1;

                        for (p = 0; p < partitions->cnt; p++) {
                                const rd_kafka_topic_partition_t *partition =
                                    &partitions->elems[p];

                                if (rd_kafka_topic_partition_list_find(
                                        otherPartitions, partition->topic,
                                        partition->partition)) {
                                        RD_UT_WARN(
                                            "Consumer %s and %s are both "
                                            "assigned %s [%" PRId32 "]",
                                            consumer, otherConsumer,
                                            partition->topic,
                                            partition->partition);
                                        fails++;
                                }


                                /* If assignment is imbalanced and this topic
                                 * is also subscribed by the other consumer
                                 * it means the assignment strategy failed to
                                 * properly balance the partitions. */
                                if (!balanced &&
                                    rd_kafka_topic_partition_list_find_topic(
                                        otherPartitions, partition->topic)) {
                                        RD_UT_WARN(
                                            "Some %s partition(s) can be "
                                            "moved from "
                                            "%s (%d partition(s)) to "
                                            "%s (%d partition(s)) to "
                                            "achieve a better balance",
                                            partition->topic, consumer,
                                            partitions->cnt, otherConsumer,
                                            otherPartitions->cnt);
                                        fails++;
                                }
                        }
                }
        }

        RD_UT_ASSERT(!fails, "%s:%d: See %d previous errors", func, line,
                     fails);

        return 0;
}


#define verifyValidityAndBalance(members, member_cnt, metadata)                \
        do {                                                                   \
                if (verifyValidityAndBalance0(__FUNCTION__, __LINE__, members, \
                                              member_cnt, metadata))           \
                        return 1;                                              \
        } while (0)


/**
 * @brief Checks that all assigned partitions are fully balanced.
 *
 * Only works for symmetrical subscriptions.
 */
static int isFullyBalanced0(const char *function,
                            int line,
                            const rd_kafka_group_member_t *members,
                            size_t member_cnt) {
        int min_assignment = INT_MAX;
        int max_assignment = -1;
        size_t i;

        for (i = 0; i < member_cnt; i++) {
                int size = members[i].rkgm_assignment->cnt;
                if (size < min_assignment)
                        min_assignment = size;
                if (size > max_assignment)
                        max_assignment = size;
        }

        RD_UT_ASSERT(max_assignment - min_assignment <= 1,
                     "%s:%d: Assignment not balanced: min %d, max %d", function,
                     line, min_assignment, max_assignment);

        return 0;
}

#define isFullyBalanced(members, member_cnt)                                   \
        do {                                                                   \
                if (isFullyBalanced0(__FUNCTION__, __LINE__, members,          \
                                     member_cnt))                              \
                        return 1;                                              \
        } while (0)


static void
ut_print_toppar_list(const rd_kafka_topic_partition_list_t *partitions) {
        int i;

        for (i = 0; i < partitions->cnt; i++)
                RD_UT_SAY(" %s [%" PRId32 "]", partitions->elems[i].topic,
                          partitions->elems[i].partition);
}



/**
 * @brief Verify that member's assignment matches the expected partitions.
 *
 * The va-list is a NULL-terminated list of (const char *topic, int partition)
 * tuples.
 *
 * @returns 0 on success, else raises a unittest error and returns 1.
 */
static int verifyAssignment0(const char *function,
                             int line,
                             rd_kafka_group_member_t *rkgm,
                             ...) {
        va_list ap;
        int cnt = 0;
        const char *topic;
        int fails = 0;

        va_start(ap, rkgm);
        while ((topic = va_arg(ap, const char *))) {
                int partition = va_arg(ap, int);
                cnt++;

                if (!rd_kafka_topic_partition_list_find(rkgm->rkgm_assignment,
                                                        topic, partition)) {
                        RD_UT_WARN(
                            "%s:%d: Expected %s [%d] not found in %s's "
                            "assignment (%d partition(s))",
                            function, line, topic, partition,
                            rkgm->rkgm_member_id->str,
                            rkgm->rkgm_assignment->cnt);
                        fails++;
                }
        }
        va_end(ap);

        if (cnt != rkgm->rkgm_assignment->cnt) {
                RD_UT_WARN(
                    "%s:%d: "
                    "Expected %d assigned partition(s) for %s, not %d",
                    function, line, cnt, rkgm->rkgm_member_id->str,
                    rkgm->rkgm_assignment->cnt);
                fails++;
        }

        if (fails)
                ut_print_toppar_list(rkgm->rkgm_assignment);

        RD_UT_ASSERT(!fails, "%s:%d: See previous errors", function, line);

        return 0;
}

#define verifyAssignment(rkgm, ...)                                            \
        do {                                                                   \
                if (verifyAssignment0(__FUNCTION__, __LINE__, rkgm,            \
                                      __VA_ARGS__))                            \
                        return 1;                                              \
        } while (0)



/**
 * @brief Initialize group member struct for testing.
 *
 * va-args is a NULL-terminated list of (const char *) topics.
 *
 * Use rd_kafka_group_member_clear() to free fields.
 */
static void
ut_init_member(rd_kafka_group_member_t *rkgm, const char *member_id, ...) {
        va_list ap;
        const char *topic;

        memset(rkgm, 0, sizeof(*rkgm));

        rkgm->rkgm_member_id         = rd_kafkap_str_new(member_id, -1);
        rkgm->rkgm_group_instance_id = rd_kafkap_str_new(member_id, -1);
        rd_list_init(&rkgm->rkgm_eligible, 0, NULL);

        rkgm->rkgm_subscription = rd_kafka_topic_partition_list_new(4);

        va_start(ap, member_id);
        while ((topic = va_arg(ap, const char *)))
                rd_kafka_topic_partition_list_add(rkgm->rkgm_subscription,
                                                  topic, RD_KAFKA_PARTITION_UA);
        va_end(ap);

        rkgm->rkgm_assignment =
            rd_kafka_topic_partition_list_new(rkgm->rkgm_subscription->size);
}



static int ut_testOneConsumerNoTopic(rd_kafka_t *rk,
                                     const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata = rd_kafka_metadata_new_topic_mock(NULL, 0);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], NULL);
        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testOneConsumerNonexistentTopic(rd_kafka_t *rk,
                                              const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 0);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], NULL);
        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}



static int ut_testOneConsumerOneTopic(rd_kafka_t *rk,
                                      const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 3);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        RD_UT_ASSERT(members[0].rkgm_assignment->cnt == 3,
                     "expected assignment of 3 partitions, got %d partition(s)",
                     members[0].rkgm_assignment->cnt);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 1, "topic1", 2,
                         NULL);
        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testOnlyAssignsPartitionsFromSubscribedTopics(
    rd_kafka_t *rk,
    const rd_kafka_assignor_t *rkas) {

        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 3, "topic2", 3);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 1, "topic1", 2,
                         NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testOneConsumerMultipleTopics(rd_kafka_t *rk,
                                            const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 1, "topic2", 2);
        ut_init_member(&members[0], "consumer1", "topic1", "topic2", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic2", 0, "topic2", 1,
                         NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}

static int
ut_testTwoConsumersOneTopicOnePartition(rd_kafka_t *rk,
                                        const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 1);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, NULL);
        verifyAssignment(&members[1], NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testTwoConsumersOneTopicTwoPartitions(rd_kafka_t *rk,
                                         const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 2);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, NULL);
        verifyAssignment(&members[1], "topic1", 1, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testMultipleConsumersMixedTopicSubscriptions(
    rd_kafka_t *rk,
    const rd_kafka_assignor_t *rkas) {

        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[3];

        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 3, "topic2", 2);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", "topic2", NULL);
        ut_init_member(&members[2], "consumer3", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 2, NULL);
        verifyAssignment(&members[1], "topic2", 0, "topic2", 1, NULL);
        verifyAssignment(&members[2], "topic1", 1, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_group_member_clear(&members[2]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testTwoConsumersTwoTopicsSixPartitions(rd_kafka_t *rk,
                                          const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];

        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 3, "topic2", 3);
        ut_init_member(&members[0], "consumer1", "topic1", "topic2", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", "topic2", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 2, "topic2", 1,
                         NULL);
        verifyAssignment(&members[1], "topic1", 1, "topic2", 0, "topic2", 2,
                         NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testAddRemoveConsumerOneTopic(rd_kafka_t *rk,
                                            const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 3);
        ut_init_member(&members[0], "consumer1", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members, 1,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 1, "topic1", 2,
                         NULL);

        verifyValidityAndBalance(members, 1, metadata);
        isFullyBalanced(members, 1);

        /* Add consumer2 */
        ut_init_member(&members[1], "consumer2", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 1, "topic1", 2, NULL);
        verifyAssignment(&members[1], "topic1", 0, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));
        // FIXME: isSticky();


        /* Remove consumer1 */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, &members[1], 1,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[1], "topic1", 0, "topic1", 1, "topic1", 2,
                         NULL);

        verifyValidityAndBalance(&members[1], 1, metadata);
        isFullyBalanced(&members[1], 1);
        // FIXME: isSticky();

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}

/**
 * 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
 */
static int
ut_testPoorRoundRobinAssignmentScenario(rd_kafka_t *rk,
                                        const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[4];

        metadata = rd_kafka_metadata_new_topic_mockv(
            5, "topic1", 2, "topic2", 1, "topic3", 2, "topic4", 1, "topic5", 2);

        ut_init_member(&members[0], "consumer1", "topic1", "topic2", "topic3",
                       "topic4", "topic5", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", "topic3", "topic5",
                       NULL);
        ut_init_member(&members[2], "consumer3", "topic1", "topic3", "topic5",
                       NULL);
        ut_init_member(&members[3], "consumer4", "topic1", "topic2", "topic3",
                       "topic4", "topic5", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic2", 0, "topic3", 0, NULL);
        verifyAssignment(&members[1], "topic1", 0, "topic3", 1, NULL);
        verifyAssignment(&members[2], "topic1", 1, "topic5", 0, NULL);
        verifyAssignment(&members[3], "topic4", 0, "topic5", 1, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_group_member_clear(&members[2]);
        rd_kafka_group_member_clear(&members[3]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}



static int ut_testAddRemoveTopicTwoConsumers(rd_kafka_t *rk,
                                             const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 3);
        ut_init_member(&members[0], "consumer1", "topic1", "topic2", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", "topic2", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 2, NULL);
        verifyAssignment(&members[1], "topic1", 1, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        /*
         * Add topic2
         */
        RD_UT_SAY("Adding topic2");
        rd_kafka_metadata_destroy(metadata);
        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 3, "topic2", 3);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic1", 0, "topic1", 2, "topic2", 1,
                         NULL);
        verifyAssignment(&members[1], "topic1", 1, "topic2", 2, "topic2", 0,
                         NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));
        // FIXME: isSticky();


        /*
         * Remove topic1
         */
        RD_UT_SAY("Removing topic1");
        rd_kafka_metadata_destroy(metadata);
        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic2", 3);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyAssignment(&members[0], "topic2", 1, NULL);
        verifyAssignment(&members[1], "topic2", 0, "topic2", 2, NULL);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));
        // FIXME: isSticky();

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_group_member_clear(&members[1]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testReassignmentAfterOneConsumerLeaves(rd_kafka_t *rk,
                                          const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[19];
        int member_cnt = RD_ARRAYSIZE(members);
        rd_kafka_metadata_topic_t mt[19];
        int topic_cnt = RD_ARRAYSIZE(mt);
        int i;

        for (i = 0; i < topic_cnt; i++) {
                char topic[10];
                rd_snprintf(topic, sizeof(topic), "topic%d", i + 1);
                rd_strdupa(&mt[i].topic, topic);
                mt[i].partition_cnt = i + 1;
        }

        metadata = rd_kafka_metadata_new_topic_mock(mt, topic_cnt);


        for (i = 1; i <= member_cnt; i++) {
                char name[20];
                rd_kafka_topic_partition_list_t *subscription =
                    rd_kafka_topic_partition_list_new(i);
                int j;
                for (j = 1; j <= i; j++) {
                        char topic[16];
                        rd_snprintf(topic, sizeof(topic), "topic%d", j);
                        rd_kafka_topic_partition_list_add(
                            subscription, topic, RD_KAFKA_PARTITION_UA);
                }
                rd_snprintf(name, sizeof(name), "consumer%d", i);
                ut_init_member(&members[i - 1], name, NULL);
                rd_kafka_topic_partition_list_destroy(
                    members[i - 1].rkgm_subscription);
                members[i - 1].rkgm_subscription = subscription;
        }

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);


        /*
         * Remove consumer10.
         */
        rd_kafka_group_member_clear(&members[9]);
        memmove(&members[9], &members[10],
                sizeof(*members) * (member_cnt - 10));
        member_cnt--;

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);
        // FIXME: isSticky();

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testReassignmentAfterOneConsumerAdded(rd_kafka_t *rk,
                                         const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[9];
        int member_cnt = RD_ARRAYSIZE(members);
        int i;

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 20);

        for (i = 1; i <= member_cnt; i++) {
                char name[20];
                rd_kafka_topic_partition_list_t *subscription =
                    rd_kafka_topic_partition_list_new(1);
                rd_kafka_topic_partition_list_add(subscription, "topic1",
                                                  RD_KAFKA_PARTITION_UA);
                rd_snprintf(name, sizeof(name), "consumer%d", i);
                ut_init_member(&members[i - 1], name, NULL);
                rd_kafka_topic_partition_list_destroy(
                    members[i - 1].rkgm_subscription);
                members[i - 1].rkgm_subscription = subscription;
        }

        member_cnt--; /* Skip one consumer */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);


        /*
         * Add consumer.
         */
        member_cnt++;

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);
        // FIXME: isSticky();

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testSameSubscriptions(rd_kafka_t *rk,
                                    const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[9];
        int member_cnt = RD_ARRAYSIZE(members);
        rd_kafka_metadata_topic_t mt[15];
        int topic_cnt = RD_ARRAYSIZE(mt);
        rd_kafka_topic_partition_list_t *subscription =
            rd_kafka_topic_partition_list_new(topic_cnt);
        int i;

        for (i = 0; i < topic_cnt; i++) {
                char topic[10];
                rd_snprintf(topic, sizeof(topic), "topic%d", i + 1);
                rd_strdupa(&mt[i].topic, topic);
                mt[i].partition_cnt = i + 1;
                rd_kafka_topic_partition_list_add(subscription, topic,
                                                  RD_KAFKA_PARTITION_UA);
        }

        metadata = rd_kafka_metadata_new_topic_mock(mt, topic_cnt);

        for (i = 1; i <= member_cnt; i++) {
                char name[16];
                rd_snprintf(name, sizeof(name), "consumer%d", i);
                ut_init_member(&members[i - 1], name, NULL);
                rd_kafka_topic_partition_list_destroy(
                    members[i - 1].rkgm_subscription);
                members[i - 1].rkgm_subscription =
                    rd_kafka_topic_partition_list_copy(subscription);
        }

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);

        /*
         * Remove consumer5
         */
        rd_kafka_group_member_clear(&members[5]);
        memmove(&members[5], &members[6], sizeof(*members) * (member_cnt - 6));
        member_cnt--;

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);
        // FIXME: isSticky();

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);
        rd_kafka_topic_partition_list_destroy(subscription);

        RD_UT_PASS();
}


static int ut_testLargeAssignmentWithMultipleConsumersLeaving(
    rd_kafka_t *rk,
    const rd_kafka_assignor_t *rkas) {

        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[200];
        int member_cnt = RD_ARRAYSIZE(members);
        rd_kafka_metadata_topic_t mt[40];
        int topic_cnt = RD_ARRAYSIZE(mt);
        int i;

        for (i = 0; i < topic_cnt; i++) {
                char topic[10];
                rd_snprintf(topic, sizeof(topic), "topic%d", i + 1);
                rd_strdupa(&mt[i].topic, topic);
                mt[i].partition_cnt = i + 1;
        }

        metadata = rd_kafka_metadata_new_topic_mock(mt, topic_cnt);

        for (i = 0; i < member_cnt; i++) {
                /* Java tests use a random set, this is more deterministic. */
                int sub_cnt = ((i + 1) * 17) % topic_cnt;
                rd_kafka_topic_partition_list_t *subscription =
                    rd_kafka_topic_partition_list_new(sub_cnt);
                char name[16];
                int j;

                /* Subscribe to a subset of topics */
                for (j = 0; j < sub_cnt; j++)
                        rd_kafka_topic_partition_list_add(
                            subscription, metadata->topics[j].topic,
                            RD_KAFKA_PARTITION_UA);

                rd_snprintf(name, sizeof(name), "consumer%d", i + 1);
                ut_init_member(&members[i], name, NULL);
                rd_kafka_topic_partition_list_destroy(
                    members[i].rkgm_subscription);
                members[i].rkgm_subscription = subscription;
        }

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);

        /*
         * Remove every 4th consumer (~50)
         */
        for (i = member_cnt - 1; i >= 0; i -= 4) {
                rd_kafka_group_member_clear(&members[i]);
                memmove(&members[i], &members[i + 1],
                        sizeof(*members) * (member_cnt - (i + 1)));
                member_cnt--;
        }

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);
        // FIXME: isSticky();

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testNewSubscription(rd_kafka_t *rk,
                                  const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[3];
        int member_cnt = RD_ARRAYSIZE(members);
        int i;

        metadata = rd_kafka_metadata_new_topic_mockv(
            5, "topic1", 1, "topic2", 2, "topic3", 3, "topic4", 4, "topic5", 5);

        for (i = 0; i < member_cnt; i++) {
                char name[16];
                int j;

                rd_snprintf(name, sizeof(name), "consumer%d", i);
                ut_init_member(&members[i], name, NULL);

                rd_kafka_topic_partition_list_destroy(
                    members[i].rkgm_subscription);
                members[i].rkgm_subscription =
                    rd_kafka_topic_partition_list_new(5);

                for (j = metadata->topic_cnt - (1 + i); j >= 0; j--)
                        rd_kafka_topic_partition_list_add(
                            members[i].rkgm_subscription,
                            metadata->topics[j].topic, RD_KAFKA_PARTITION_UA);
        }

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        /*
         * Add topic1 to consumer1's subscription
         */
        RD_UT_SAY("Adding topic1 to consumer1");
        rd_kafka_topic_partition_list_add(members[0].rkgm_subscription,
                                          "topic1", RD_KAFKA_PARTITION_UA);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));
        // FIXME: isSticky();

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testMoveExistingAssignments(rd_kafka_t *rk,
                                          const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[4];
        int member_cnt                                  = RD_ARRAYSIZE(members);
        rd_kafka_topic_partition_list_t *assignments[4] = RD_ZERO_INIT;
        int i;
        int fails = 0;

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 3);

        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", NULL);
        ut_init_member(&members[2], "consumer3", "topic1", NULL);
        ut_init_member(&members[3], "consumer4", "topic1", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, member_cnt, metadata);

        for (i = 0; i < member_cnt; i++) {
                if (members[i].rkgm_assignment->cnt > 1) {
                        RD_UT_WARN("%s assigned %d partitions, expected <= 1",
                                   members[i].rkgm_member_id->str,
                                   members[i].rkgm_assignment->cnt);
                        fails++;
                } else if (members[i].rkgm_assignment->cnt == 1) {
                        assignments[i] = rd_kafka_topic_partition_list_copy(
                            members[i].rkgm_assignment);
                }
        }

        /*
         * Remove potential group leader consumer1
         */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, &members[1],
                                    member_cnt - 1, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(&members[1], member_cnt - 1, metadata);
        // FIXME: isSticky()

        for (i = 1; i < member_cnt; i++) {
                if (members[i].rkgm_assignment->cnt != 1) {
                        RD_UT_WARN("%s assigned %d partitions, expected 1",
                                   members[i].rkgm_member_id->str,
                                   members[i].rkgm_assignment->cnt);
                        fails++;
                } else if (assignments[i] &&
                           !rd_kafka_topic_partition_list_find(
                               assignments[i],
                               members[i].rkgm_assignment->elems[0].topic,
                               members[i]
                                   .rkgm_assignment->elems[0]
                                   .partition)) {
                        RD_UT_WARN(
                            "Stickiness was not honored for %s, "
                            "%s [%" PRId32 "] not in previous assignment",
                            members[i].rkgm_member_id->str,
                            members[i].rkgm_assignment->elems[0].topic,
                            members[i].rkgm_assignment->elems[0].partition);
                        fails++;
                }
        }

        RD_UT_ASSERT(!fails, "See previous errors");


        for (i = 0; i < member_cnt; i++) {
                rd_kafka_group_member_clear(&members[i]);
                if (assignments[i])
                        rd_kafka_topic_partition_list_destroy(assignments[i]);
        }
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}



static int ut_testStickiness(rd_kafka_t *rk, const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[3];
        int member_cnt = RD_ARRAYSIZE(members);
        int i;

        metadata = rd_kafka_metadata_new_topic_mockv(
            6, "topic1", 1, "topic2", 1, "topic3", 1, "topic4", 1, "topic5", 1,
            "topic6", 1);

        ut_init_member(&members[0], "consumer1", "topic1", "topic2", NULL);
        rd_kafka_topic_partition_list_destroy(members[0].rkgm_assignment);
        members[0].rkgm_assignment = rd_kafka_topic_partition_list_new(1);
        rd_kafka_topic_partition_list_add(members[0].rkgm_assignment, "topic1",
                                          0);

        ut_init_member(&members[1], "consumer2", "topic1", "topic2", "topic3",
                       "topic4", NULL);
        rd_kafka_topic_partition_list_destroy(members[1].rkgm_assignment);
        members[1].rkgm_assignment = rd_kafka_topic_partition_list_new(2);
        rd_kafka_topic_partition_list_add(members[1].rkgm_assignment, "topic2",
                                          0);
        rd_kafka_topic_partition_list_add(members[1].rkgm_assignment, "topic3",
                                          0);

        ut_init_member(&members[2], "consumer3", "topic4", "topic5", "topic6",
                       NULL);
        rd_kafka_topic_partition_list_destroy(members[2].rkgm_assignment);
        members[2].rkgm_assignment = rd_kafka_topic_partition_list_new(3);
        rd_kafka_topic_partition_list_add(members[2].rkgm_assignment, "topic4",
                                          0);
        rd_kafka_topic_partition_list_add(members[2].rkgm_assignment, "topic5",
                                          0);
        rd_kafka_topic_partition_list_add(members[2].rkgm_assignment, "topic6",
                                          0);


        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);


        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


/**
 * @brief Verify stickiness across three rebalances.
 */
static int ut_testStickiness2(rd_kafka_t *rk, const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[3];
        int member_cnt = RD_ARRAYSIZE(members);
        int i;

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 6);

        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        ut_init_member(&members[1], "consumer2", "topic1", NULL);
        ut_init_member(&members[2], "consumer3", "topic1", NULL);

        /* Just consumer1 */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members, 1,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, 1, metadata);
        isFullyBalanced(members, 1);
        verifyAssignment(&members[0], "topic1", 0, "topic1", 1, "topic1", 2,
                         "topic1", 3, "topic1", 4, "topic1", 5, NULL);

        /* consumer1 and consumer2 */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members, 2,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, 2, metadata);
        isFullyBalanced(members, 2);
        verifyAssignment(&members[0], "topic1", 3, "topic1", 4, "topic1", 5,
                         NULL);
        verifyAssignment(&members[1], "topic1", 0, "topic1", 1, "topic1", 2,
                         NULL);

        /* Run it twice, should be stable. */
        for (i = 0; i < 2; i++) {
                /* consumer1, consumer2, and consumer3 */
                err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata,
                                            members, 3, errstr, sizeof(errstr));
                RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

                verifyValidityAndBalance(members, 3, metadata);
                isFullyBalanced(members, 3);
                verifyAssignment(&members[0], "topic1", 4, "topic1", 5, NULL);
                verifyAssignment(&members[1], "topic1", 1, "topic1", 2, NULL);
                verifyAssignment(&members[2], "topic1", 0, "topic1", 3, NULL);
        }

        /* Remove consumer1 */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, &members[1], 2,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(&members[1], 2, metadata);
        isFullyBalanced(&members[1], 2);
        verifyAssignment(&members[1], "topic1", 1, "topic1", 2, "topic1", 5,
                         NULL);
        verifyAssignment(&members[2], "topic1", 0, "topic1", 3, "topic1", 4,
                         NULL);

        /* Remove consumer2 */
        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, &members[2], 1,
                                    errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(&members[2], 1, metadata);
        isFullyBalanced(&members[2], 1);
        verifyAssignment(&members[2], "topic1", 0, "topic1", 1, "topic1", 2,
                         "topic1", 3, "topic1", 4, "topic1", 5, NULL);

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testAssignmentUpdatedForDeletedTopic(rd_kafka_t *rk,
                                        const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata =
            rd_kafka_metadata_new_topic_mockv(2, "topic1", 1, "topic3", 100);
        ut_init_member(&members[0], "consumer1", "topic1", "topic2", "topic3",
                       NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        RD_UT_ASSERT(members[0].rkgm_assignment->cnt == 1 + 100,
                     "Expected %d assigned partitions, not %d", 1 + 100,
                     members[0].rkgm_assignment->cnt);

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int ut_testNoExceptionThrownWhenOnlySubscribedTopicDeleted(
    rd_kafka_t *rk,
    const rd_kafka_assignor_t *rkas) {

        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[1];

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 3);

        ut_init_member(&members[0], "consumer1", "topic", NULL);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        /*
         * Remove topic
         */
        rd_kafka_metadata_destroy(metadata);
        metadata = rd_kafka_metadata_new_topic_mock(NULL, 0);

        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    RD_ARRAYSIZE(members), errstr,
                                    sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));

        rd_kafka_group_member_clear(&members[0]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}


static int
ut_testConflictingPreviousAssignments(rd_kafka_t *rk,
                                      const rd_kafka_assignor_t *rkas) {
        rd_kafka_resp_err_t err;
        char errstr[512];
        rd_kafka_metadata_t *metadata;
        rd_kafka_group_member_t members[2];
        int member_cnt = RD_ARRAYSIZE(members);
        int i;

        // FIXME: removed from Java test suite, and fails for us, why, why?
        RD_UT_PASS();

        metadata = rd_kafka_metadata_new_topic_mockv(1, "topic1", 2);

        /* Both consumer and consumer2 have both partitions assigned */
        ut_init_member(&members[0], "consumer1", "topic1", NULL);
        rd_kafka_topic_partition_list_destroy(members[0].rkgm_assignment);
        members[0].rkgm_assignment = rd_kafka_topic_partition_list_new(2);
        rd_kafka_topic_partition_list_add(members[0].rkgm_assignment, "topic1",
                                          0);
        rd_kafka_topic_partition_list_add(members[0].rkgm_assignment, "topic1",
                                          1);

        ut_init_member(&members[1], "consumer2", "topic1", NULL);
        rd_kafka_topic_partition_list_destroy(members[1].rkgm_assignment);
        members[1].rkgm_assignment = rd_kafka_topic_partition_list_new(2);
        rd_kafka_topic_partition_list_add(members[1].rkgm_assignment, "topic1",
                                          0);
        rd_kafka_topic_partition_list_add(members[1].rkgm_assignment, "topic1",
                                          1);


        err = rd_kafka_assignor_run(rk->rk_cgrp, rkas, metadata, members,
                                    member_cnt, errstr, sizeof(errstr));
        RD_UT_ASSERT(!err, "assignor run failed: %s", errstr);

        RD_UT_ASSERT(members[0].rkgm_assignment->cnt == 1 &&
                         members[1].rkgm_assignment->cnt == 1,
                     "Expected consumers to have 1 partition each, "
                     "not %d and %d",
                     members[0].rkgm_assignment->cnt,
                     members[1].rkgm_assignment->cnt);
        RD_UT_ASSERT(members[0].rkgm_assignment->elems[0].partition !=
                         members[1].rkgm_assignment->elems[0].partition,
                     "Expected consumers to have different partitions "
                     "assigned, not same partition %" PRId32,
                     members[0].rkgm_assignment->elems[0].partition);

        verifyValidityAndBalance(members, RD_ARRAYSIZE(members), metadata);
        isFullyBalanced(members, RD_ARRAYSIZE(members));
        /* FIXME: isSticky() */

        for (i = 0; i < member_cnt; i++)
                rd_kafka_group_member_clear(&members[i]);
        rd_kafka_metadata_destroy(metadata);

        RD_UT_PASS();
}

/* testReassignmentWithRandomSubscriptionsAndChanges is not ported
 * from Java since random tests don't provide meaningful test coverage. */


static int rd_kafka_sticky_assignor_unittest(void) {
        rd_kafka_conf_t *conf;
        rd_kafka_t *rk;
        int fails = 0;
        char errstr[256];
        rd_kafka_assignor_t *rkas;
        static int (*tests[])(rd_kafka_t *, const rd_kafka_assignor_t *) = {
            ut_testOneConsumerNoTopic,
            ut_testOneConsumerNonexistentTopic,
            ut_testOneConsumerOneTopic,
            ut_testOnlyAssignsPartitionsFromSubscribedTopics,
            ut_testOneConsumerMultipleTopics,
            ut_testTwoConsumersOneTopicOnePartition,
            ut_testTwoConsumersOneTopicTwoPartitions,
            ut_testMultipleConsumersMixedTopicSubscriptions,
            ut_testTwoConsumersTwoTopicsSixPartitions,
            ut_testAddRemoveConsumerOneTopic,
            ut_testPoorRoundRobinAssignmentScenario,
            ut_testAddRemoveTopicTwoConsumers,
            ut_testReassignmentAfterOneConsumerLeaves,
            ut_testReassignmentAfterOneConsumerAdded,
            ut_testSameSubscriptions,
            ut_testLargeAssignmentWithMultipleConsumersLeaving,
            ut_testNewSubscription,
            ut_testMoveExistingAssignments,
            ut_testStickiness,
            ut_testStickiness2,
            ut_testAssignmentUpdatedForDeletedTopic,
            ut_testNoExceptionThrownWhenOnlySubscribedTopicDeleted,
            ut_testConflictingPreviousAssignments,
            NULL,
        };
        int i;


        conf = rd_kafka_conf_new();
        if (rd_kafka_conf_set(conf, "group.id", "test", errstr,
                              sizeof(errstr)) ||
            rd_kafka_conf_set(conf, "partition.assignment.strategy",
                              "cooperative-sticky", errstr, sizeof(errstr)))
                RD_UT_FAIL("sticky assignor conf failed: %s", errstr);

        rd_kafka_conf_set(conf, "debug", rd_getenv("TEST_DEBUG", NULL), NULL,
                          0);

        rk = rd_kafka_new(RD_KAFKA_CONSUMER, conf, errstr, sizeof(errstr));
        RD_UT_ASSERT(rk, "sticky assignor client instantiation failed: %s",
                     errstr);

        rkas = rd_kafka_assignor_find(rk, "cooperative-sticky");
        RD_UT_ASSERT(rkas, "sticky assignor not found");

        for (i = 0; tests[i]; i++) {
                rd_ts_t ts = rd_clock();
                int r;

                RD_UT_SAY("[ Test #%d ]", i);
                r = tests[i](rk, rkas);
                RD_UT_SAY("[ Test #%d ran for %.3fms ]", i,
                          (double)(rd_clock() - ts) / 1000.0);

                RD_UT_ASSERT(!r, "^ failed");

                fails += r;
        }

        rd_kafka_destroy(rk);

        return fails;
}


/**
 * @brief Initialzie and add sticky assignor.
 */
rd_kafka_resp_err_t rd_kafka_sticky_assignor_init(rd_kafka_t *rk) {
        return rd_kafka_assignor_add(rk, "consumer", "cooperative-sticky",
                                     RD_KAFKA_REBALANCE_PROTOCOL_COOPERATIVE,
                                     rd_kafka_sticky_assignor_assign_cb,
                                     rd_kafka_sticky_assignor_get_metadata,
                                     rd_kafka_sticky_assignor_on_assignment_cb,
                                     rd_kafka_sticky_assignor_state_destroy,
                                     rd_kafka_sticky_assignor_unittest, NULL);
}