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

Example 61 with Subscription

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

the class AbstractStickyAssignorTest method testLargeAssignmentAndGroupWithNonEqualSubscription.

@Timeout(60)
@Test
public void testLargeAssignmentAndGroupWithNonEqualSubscription() {
    // 1 million partitions!
    int topicCount = 500;
    int partitionCount = 2_000;
    int consumerCount = 2_000;
    List<String> topics = new ArrayList<>();
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 0; i < topicCount; i++) {
        String topicName = getTopicName(i, topicCount);
        topics.add(topicName);
        partitionsPerTopic.put(topicName, partitionCount);
    }
    for (int i = 0; i < consumerCount; i++) {
        if (i == consumerCount - 1) {
            subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics.subList(0, 1)));
        } else {
            subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
        }
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    for (int i = 1; i < consumerCount; i++) {
        String consumer = getConsumerName(i, consumerCount);
        if (i == consumerCount - 1) {
            subscriptions.put(consumer, buildSubscription(topics.subList(0, 1), assignment.get(consumer)));
        } else {
            subscriptions.put(consumer, buildSubscription(topics, assignment.get(consumer)));
        }
    }
    assignor.assign(partitionsPerTopic, subscriptions);
}

Example 62 with Subscription

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

the class AbstractStickyAssignorTest method testLargeAssignmentAndGroupWithUniformSubscription.

@Timeout(30)
@Test
public void testLargeAssignmentAndGroupWithUniformSubscription() {
    // 1 million partitions!
    int topicCount = 500;
    int partitionCount = 2_000;
    int consumerCount = 2_000;
    List<String> topics = new ArrayList<>();
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    for (int i = 0; i < topicCount; i++) {
        String topicName = getTopicName(i, topicCount);
        topics.add(topicName);
        partitionsPerTopic.put(topicName, partitionCount);
    }
    for (int i = 0; i < consumerCount; i++) {
        subscriptions.put(getConsumerName(i, consumerCount), new Subscription(topics));
    }
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    for (int i = 1; i < consumerCount; i++) {
        String consumer = getConsumerName(i, consumerCount);
        subscriptions.put(consumer, buildSubscription(topics, assignment.get(consumer)));
    }
    assignor.assign(partitionsPerTopic, subscriptions);
}

Example 63 with Subscription

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

the class AbstractStickyAssignorTest method testStickiness.

@Test
public void testStickiness() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic1, 3);
    subscriptions.put(consumer1, new Subscription(topics(topic1)));
    subscriptions.put(consumer2, new Subscription(topics(topic1)));
    subscriptions.put(consumer3, new Subscription(topics(topic1)));
    subscriptions.put(consumer4, new Subscription(topics(topic1)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    Map<String, TopicPartition> partitionsAssigned = new HashMap<>();
    Set<Map.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(size <= 1, "Consumer " + consumer + " is assigned more topic partitions than expected.");
        if (size == 1)
            partitionsAssigned.put(consumer, topicPartitions.get(0));
    }
    // removing the potential group leader
    subscriptions.remove(consumer1);
    subscriptions.put(consumer2, buildSubscription(topics(topic1), assignment.get(consumer2)));
    subscriptions.put(consumer3, buildSubscription(topics(topic1), assignment.get(consumer3)));
    subscriptions.put(consumer4, buildSubscription(topics(topic1), assignment.get(consumer4)));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assignments = assignment.entrySet();
    for (Map.Entry<String, List<TopicPartition>> entry : assignments) {
        String consumer = entry.getKey();
        List<TopicPartition> topicPartitions = entry.getValue();
        assertEquals(1, topicPartitions.size(), "Consumer " + consumer + " is assigned more topic partitions than expected.");
        assertTrue((!partitionsAssigned.containsKey(consumer)) || (assignment.get(consumer).contains(partitionsAssigned.get(consumer))), "Stickiness was not honored for consumer " + consumer);
    }
}

Example 64 with Subscription

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

the class AbstractStickyAssignorTest method testAddRemoveConsumerOneTopic.

@Test
public void testAddRemoveConsumerOneTopic() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 3);
    subscriptions.put(consumer1, new Subscription(topics(topic)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(partitions(tp(topic, 0), tp(topic, 1), tp(topic, 2)), assignment.get(consumer1));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    subscriptions.put(consumer1, buildSubscription(topics(topic), assignment.get(consumer1)));
    subscriptions.put(consumer2, buildSubscription(topics(topic), Collections.emptyList()));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(Collections.singletonMap(tp(topic, 2), consumer2), assignor.partitionsTransferringOwnership);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertEquals(partitions(tp(topic, 0), tp(topic, 1)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic, 2)), assignment.get(consumer2));
    assertTrue(isFullyBalanced(assignment));
    subscriptions.remove(consumer1);
    subscriptions.put(consumer2, buildSubscription(topics(topic), assignment.get(consumer2)));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    assertEquals(new HashSet<>(partitions(tp(topic, 2), tp(topic, 1), tp(topic, 0))), new HashSet<>(assignment.get(consumer2)));
    assertTrue(assignor.partitionsTransferringOwnership.isEmpty());
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}

Example 65 with Subscription

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

the class RoundRobinAssignorTest method testTwoConsumersOneTopicTwoPartitions.

@Test
public void testTwoConsumersOneTopicTwoPartitions() {
    String consumer1 = "consumer1";
    String consumer2 = "consumer2";
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 2);
    Map<String, Subscription> consumers = new HashMap<>();
    consumers.put(consumer1, new Subscription(topics(topic)));
    consumers.put(consumer2, new Subscription(topics(topic)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    assertEquals(partitions(tp(topic, 0)), assignment.get(consumer1));
    assertEquals(partitions(tp(topic, 1)), assignment.get(consumer2));
}