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

Example 26 with Subscription

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

the class StickyAssignorTest 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);
    Map<String, Subscription> subscriptions = new HashMap<>();
    for (int i = 1; i < 9; i++) {
        List<String> topics = new ArrayList<String>();
        for (int j = 1; j <= partitionsPerTopic.size(); j++) topics.add(getTopicName(j, 15));
        subscriptions.put(getConsumerName(i, 9), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    for (int i = 1; i < 9; i++) {
        String consumer = getConsumerName(i, 9);
        subscriptions.put(consumer, new Subscription(subscriptions.get(consumer).topics(), StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
    }
    subscriptions.remove(getConsumerName(5, 9));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
}

Example 27 with Subscription

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

the class StickyAssignorTest method testMultipleConsumersMixedTopicSubscriptions.

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

Example 28 with Subscription

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

the class StickyAssignorTest method testStickiness.

@Test
public void testStickiness() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put("topic01", 3);
    Map<String, Subscription> subscriptions = new HashMap<>();
    subscriptions.put("consumer01", new Subscription(topics("topic01")));
    subscriptions.put("consumer02", new Subscription(topics("topic01")));
    subscriptions.put("consumer03", new Subscription(topics("topic01")));
    subscriptions.put("consumer04", new Subscription(topics("topic01")));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    Map<String, TopicPartition> partitionsAssigned = new HashMap<>();
    Set<Entry<String, List<TopicPartition>>> assignments = assignment.entrySet();
    for (Map.Entry<String, List<TopicPartition>> entry : assignments) {
        String consumer = entry.getKey();
        List<TopicPartition> topicPartitions = entry.getValue();
        int size = topicPartitions.size();
        assertTrue("Consumer " + consumer + " is assigned more topic partitions than expected.", size <= 1);
        if (size == 1)
            partitionsAssigned.put(consumer, topicPartitions.get(0));
    }
    // removing the potential group leader
    subscriptions.remove("consumer01");
    subscriptions.put("consumer02", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer02"))));
    subscriptions.put("consumer03", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer03"))));
    subscriptions.put("consumer04", new Subscription(topics("topic01"), StickyAssignor.serializeTopicPartitionAssignment(assignment.get("consumer04"))));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    verifyValidityAndBalance(subscriptions, assignment);
    assertTrue(assignor.isSticky());
    assignments = assignment.entrySet();
    for (Map.Entry<String, List<TopicPartition>> entry : assignments) {
        String consumer = entry.getKey();
        List<TopicPartition> topicPartitions = entry.getValue();
        assertEquals("Consumer " + consumer + " is assigned more topic partitions than expected.", 1, topicPartitions.size());
        assertTrue("Stickiness was not honored for consumer " + consumer, (!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))));
    }
}

Example 29 with Subscription

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

the class StickyAssignorTest method testAddRemoveTopicTwoConsumers.

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

Example 30 with Subscription

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

the class StickyAssignorTest 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<>();
        for (int i = 0; i < numTopics; ++i) topics.add(getTopicName(i, maxNumTopics));
        Map<String, Integer> partitionsPerTopic = new HashMap<>();
        for (int i = 0; i < numTopics; ++i) partitionsPerTopic.put(getTopicName(i, maxNumTopics), i + 1);
        int numConsumers = minNumConsumers + new Random().nextInt(maxNumConsumers - minNumConsumers);
        Map<String, Subscription> subscriptions = new HashMap<>();
        for (int i = 0; i < numConsumers; ++i) {
            List<String> sub = Utils.sorted(getRandomSublist(topics));
            subscriptions.put(getConsumerName(i, maxNumConsumers), new Subscription(sub));
        }
        StickyAssignor assignor = new StickyAssignor();
        Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
        verifyValidityAndBalance(subscriptions, assignment);
        subscriptions.clear();
        for (int i = 0; i < numConsumers; ++i) {
            List<String> sub = Utils.sorted(getRandomSublist(topics));
            String consumer = getConsumerName(i, maxNumConsumers);
            subscriptions.put(consumer, new Subscription(sub, StickyAssignor.serializeTopicPartitionAssignment(assignment.get(consumer))));
        }
        assignment = assignor.assign(partitionsPerTopic, subscriptions);
        verifyValidityAndBalance(subscriptions, assignment);
        assertTrue(assignor.isSticky());
    }
}