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

Example 11 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testReassignmentAfterOneConsumerAdded.

@Test
public void testReassignmentAfterOneConsumerAdded() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put("topic", 20);
    Map<String, Subscription> subscriptions = new HashMap<>();
    for (int i = 1; i < 10; i++) subscriptions.put(getConsumerName(i, 10), new Subscription(topics("topic")));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    // add a new consumer
    subscriptions.put(getConsumerName(10, 10), new Subscription(topics("topic")));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
}

Example 12 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testTwoConsumersTwoTopicsSixPartitions.

@Test
public void testTwoConsumersTwoTopicsSixPartitions() {
    String topic1 = "topic1";
    String topic2 = "topic2";
    String consumer1 = "consumer1";
    String consumer2 = "consumer2";
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic1, 3);
    partitionsPerTopic.put(topic2, 3);
    Map<String, Subscription> subscriptions = new HashMap<>();
    subscriptions.put(consumer1, new Subscription(topics(topic1, topic2)));
    subscriptions.put(consumer2, new Subscription(topics(topic1, topic2)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(partitions(tp(topic1, 0), tp(topic1, 2), tp(topic2, 1)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 2)), assignment.get(consumer2));
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(isFullyBalanced(assignment));
}

Example 13 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testMoveExistingAssignments.

@Test
public void testMoveExistingAssignments() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 1; i <= 6; i++) partitionsPerTopic.put(String.format("topic%02d", i), 1);
    Map<String, Subscription> subscriptions = new HashMap<>();
    subscriptions.put("consumer01", new Subscription(topics("topic01", "topic02"), StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic01", 0)))));
    subscriptions.put("consumer02", new Subscription(topics("topic01", "topic02", "topic03", "topic04"), StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic02", 0), tp("topic03", 0)))));
    subscriptions.put("consumer03", new Subscription(topics("topic02", "topic03", "topic04", "topic05", "topic06"), StickyAssignor.serializeTopicPartitionAssignment(partitions(tp("topic04", 0), tp("topic05", 0), tp("topic06", 0)))));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
}

Example 14 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest 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);
    Map<String, Subscription> subscriptions = new HashMap<>();
    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);
}

Example 15 with Subscription

use of org.apache.kafka.clients.consumer.internals.PartitionAssignor.Subscription in project apache-kafka-on-k8s by banzaicloud.

the class StickyAssignorTest method testLargeAssignmentWithMultipleConsumersLeaving.

@Test
public void testLargeAssignmentWithMultipleConsumersLeaving() {
    Random rand = new Random();
    int topicCount = 40;
    int consumerCount = 200;
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 0; i < topicCount; i++) partitionsPerTopic.put(getTopicName(i, topicCount), rand.nextInt(10) + 1);
    Map<String, Subscription> subscriptions = new HashMap<>();
    for (int i = 0; i < consumerCount; i++) {
        List<String> topics = new ArrayList<String>();
        for (int j = 0; j < rand.nextInt(20); j++) topics.add(getTopicName(rand.nextInt(topicCount), topicCount));
        subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    for (int i = 1; i < consumerCount; i++) {
        String consumer = getConsumerName(i, consumerCount);
        subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
    }
    for (int i = 0; i < 50; ++i) {
        String c = getConsumerName(rand.nextInt(consumerCount), consumerCount);
        subscriptions.remove(c);
    }
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
}