Examples with Subscription org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription used on opensource projects

Example 16 with Subscription

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription in project kafka by apache.

the class AbstractStickyAssignorTest method testAddRemoveTopicTwoConsumers.

@Test
public void testAddRemoveTopicTwoConsumers() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 3);
    subscriptions.put(consumer1, new Subscription(topics(topic)));
    subscriptions.put(consumer2, new Subscription(topics(topic)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    // verify balance
    assertTrue(isFullyBalanced(assignment));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    // verify stickiness
    List<TopicPartition> consumer1Assignment1 = assignment.get(consumer1);
    List<TopicPartition> consumer2Assignment1 = assignment.get(consumer2);
    assertTrue((consumer1Assignment1.size() == 1 && consumer2Assignment1.size() == 2) || (consumer1Assignment1.size() == 2 && consumer2Assignment1.size() == 1));
    partitionsPerTopic.put(topic2, 3);
    subscriptions.put(consumer1, buildSubscription(topics(topic, topic2), assignment.get(consumer1)));
    subscriptions.put(consumer2, buildSubscription(topics(topic, topic2), assignment.get(consumer2)));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    // verify balance
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    // verify stickiness
    List<TopicPartition> consumer1assignment = assignment.get(consumer1);
    List<TopicPartition> consumer2assignment = assignment.get(consumer2);
    assertTrue(consumer1assignment.size() == 3 && consumer2assignment.size() == 3);
    assertTrue(consumer1assignment.containsAll(consumer1Assignment1));
    assertTrue(consumer2assignment.containsAll(consumer2Assignment1));
    partitionsPerTopic.remove(topic);
    subscriptions.put(consumer1, buildSubscription(topics(topic2), assignment.get(consumer1)));
    subscriptions.put(consumer2, buildSubscription(topics(topic2), assignment.get(consumer2)));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    // verify balance
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    // verify stickiness
    List<TopicPartition> consumer1Assignment3 = assignment.get(consumer1);
    List<TopicPartition> consumer2Assignment3 = assignment.get(consumer2);
    assertTrue((consumer1Assignment3.size() == 1 && consumer2Assignment3.size() == 2) || (consumer1Assignment3.size() == 2 && consumer2Assignment3.size() == 1));
    assertTrue(consumer1assignment.containsAll(consumer1Assignment3));
    assertTrue(consumer2assignment.containsAll(consumer2Assignment3));
}

Example 17 with Subscription

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription in project kafka by apache.

the class AbstractStickyAssignorTest method testPoorRoundRobinAssignmentScenario.

/**
 * This unit test performs sticky assignment for a scenario that round robin assignor handles poorly.
 * Topics (partitions per topic): topic1 (2), topic2 (1), topic3 (2), topic4 (1), topic5 (2)
 * Subscriptions:
 *  - consumer1: topic1, topic2, topic3, topic4, topic5
 *  - consumer2: topic1, topic3, topic5
 *  - consumer3: topic1, topic3, topic5
 *  - consumer4: topic1, topic2, topic3, topic4, topic5
 * Round Robin Assignment Result:
 *  - consumer1: topic1-0, topic3-0, topic5-0
 *  - consumer2: topic1-1, topic3-1, topic5-1
 *  - consumer3:
 *  - consumer4: topic2-0, topic4-0
 * Sticky Assignment Result:
 *  - consumer1: topic2-0, topic3-0
 *  - consumer2: topic1-0, topic3-1
 *  - consumer3: topic1-1, topic5-0
 *  - consumer4: topic4-0, topic5-1
 */
@Test
public void testPoorRoundRobinAssignmentScenario() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 1; i <= 5; i++) partitionsPerTopic.put(String.format("topic%d", i), (i % 2) + 1);
    subscriptions.put("consumer1", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
    subscriptions.put("consumer2", new Subscription(topics("topic1", "topic3", "topic5")));
    subscriptions.put("consumer3", new Subscription(topics("topic1", "topic3", "topic5")));
    subscriptions.put("consumer4", new Subscription(topics("topic1", "topic2", "topic3", "topic4", "topic5")));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
}

Example 18 with Subscription

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription in project kafka by apache.

the class AbstractStickyAssignorTest method testAddRemoveTwoConsumersTwoTopics.

@Test
public void testAddRemoveTwoConsumersTwoTopics() {
    List<String> allTopics = topics(topic1, topic2);
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic1, 3);
    partitionsPerTopic.put(topic2, 4);
    subscriptions.put(consumer1, new Subscription(allTopics));
    subscriptions.put(consumer2, new Subscription(allTopics));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1), tp(topic2, 3)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    // add 2 consumers
    subscriptions.put(consumer1, buildSubscription(allTopics, assignment.get(consumer1)));
    subscriptions.put(consumer2, buildSubscription(allTopics, assignment.get(consumer2)));
    subscriptions.put(consumer3, buildSubscription(allTopics, Collections.emptyList()));
    subscriptions.put(consumer4, buildSubscription(allTopics, Collections.emptyList()));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    Map<TopicPartition, String> expectedPartitionsTransferringOwnership = new HashMap<>();
    expectedPartitionsTransferringOwnership.put(tp(topic2, 1), consumer3);
    expectedPartitionsTransferringOwnership.put(tp(topic2, 3), consumer3);
    expectedPartitionsTransferringOwnership.put(tp(topic2, 2), consumer4);
    assertEquals(expectedPartitionsTransferringOwnership, assignor.partitionsTransferringOwnership);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertEquals(partitions(tp(topic1, 0), tp(topic1, 2)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic1, 1), tp(topic2, 0)), assignment.get(consumer2));
    assertEquals(partitions(tp(topic2, 1), tp(topic2, 3)), assignment.get(consumer3));
    assertEquals(partitions(tp(topic2, 2)), assignment.get(consumer4));
    assertTrue(isFullyBalanced(assignment));
    // remove 2 consumers
    subscriptions.remove(consumer1);
    subscriptions.remove(consumer2);
    subscriptions.put(consumer3, buildSubscription(allTopics, assignment.get(consumer3)));
    subscriptions.put(consumer4, buildSubscription(allTopics, assignment.get(consumer4)));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(partitions(tp(topic2, 1), tp(topic2, 3), tp(topic1, 0), tp(topic2, 0)), assignment.get(consumer3));
    assertEquals(partitions(tp(topic2, 2), tp(topic1, 1), tp(topic1, 2)), assignment.get(consumer4));
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}

Example 19 with Subscription

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription in project kafka by apache.

the class AbstractStickyAssignorTest method testSameSubscriptions.

@Test
public void testSameSubscriptions() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 1; i < 15; i++) partitionsPerTopic.put(getTopicName(i, 15), i);
    for (int i = 1; i < 9; i++) {
        List<String> topics = new ArrayList<>();
        for (int j = 1; j <= partitionsPerTopic.size(); j++) topics.add(getTopicName(j, 15));
        subscriptions.put(getConsumerName(i, 9), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    for (int i = 1; i < 9; i++) {
        String consumer = getConsumerName(i, 9);
        subscriptions.put(consumer, buildSubscription(subscriptions.get(consumer).topics(), assignment.get(consumer)));
    }
    subscriptions.remove(getConsumerName(5, 9));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
}

Example 20 with Subscription

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Subscription in project kafka by apache.

the class AbstractStickyAssignorTest method testReassignmentWithRandomSubscriptionsAndChanges.

@Test
public void testReassignmentWithRandomSubscriptionsAndChanges() {
    final int minNumConsumers = 20;
    final int maxNumConsumers = 40;
    final int minNumTopics = 10;
    final int maxNumTopics = 20;
    for (int round = 1; round <= 100; ++round) {
        int numTopics = minNumTopics + new Random().nextInt(maxNumTopics - minNumTopics);
        ArrayList<String> topics = new ArrayList<>();
        Map<String, Integer> partitionsPerTopic = new HashMap<>();
        for (int i = 0; i < numTopics; ++i) {
            topics.add(getTopicName(i, maxNumTopics));
            partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1);
        }
        int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
        for (int i = 0; i < numConsumers; ++i) {
            List<String> sub = Utils.sorted(getRandomSublist(topics));
            subscriptions.put(getConsumerName(i, maxNumConsumers), new Subscription(sub));
        }
        assignor = createAssignor();
        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
        subscriptions.clear();
        for (int i = 0; i < numConsumers; ++i) {
            List<String> sub = Utils.sorted(getRandomSublist(topics));
            String consumer = getConsumerName(i, maxNumConsumers);
            subscriptions.put(consumer, buildSubscription(sub, assignment.get(consumer)));
        }
        assignment = assignor.assign(partitionsPerTopic, subscriptions);
        verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
        assertTrue(assignor.isSticky());
    }
}