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

Example 6 with Subscription

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

the class RangeAssignorTest method testMultipleConsumersMixedTopics.

@Test
public void testMultipleConsumersMixedTopics() {
    Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 2);
    Map<String, Subscription> consumers = new HashMap<>();
    consumers.put(consumer1, new Subscription(topics(topic1)));
    consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
    consumers.put(consumer3, new Subscription(topics(topic1)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    assertAssignment(partitions(tp(topic1, 0)), assignment.get(consumer1));
    assertAssignment(partitions(tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)), assignment.get(consumer2));
    assertAssignment(partitions(tp(topic1, 2)), assignment.get(consumer3));
}

Example 7 with Subscription

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

the class StickyAssignorTest method testAssignmentWithConflictingPreviousGenerations.

@ParameterizedTest(name = "testAssignmentWithConflictingPreviousGenerations with isAllSubscriptionsEqual: {0}")
@ValueSource(booleans = { true, false })
public void testAssignmentWithConflictingPreviousGenerations(boolean isAllSubscriptionsEqual) {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 4);
    partitionsPerTopic.put(topic2, 4);
    partitionsPerTopic.put(topic3, 4);
    List<String> allTopics = topics(topic, topic2, topic3);
    List<String> consumer1SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic);
    List<String> consumer2SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic, topic2);
    subscriptions.put(consumer1, new Subscription(consumer1SubscribedTopics));
    subscriptions.put(consumer2, new Subscription(consumer2SubscribedTopics));
    subscriptions.put(consumer3, new Subscription(allTopics));
    TopicPartition tp0 = new TopicPartition(topic, 0);
    TopicPartition tp1 = new TopicPartition(topic, 1);
    TopicPartition tp2 = new TopicPartition(topic, 2);
    TopicPartition tp3 = new TopicPartition(topic, 3);
    TopicPartition t2p0 = new TopicPartition(topic2, 0);
    TopicPartition t2p1 = new TopicPartition(topic2, 1);
    TopicPartition t2p2 = new TopicPartition(topic2, 2);
    TopicPartition t2p3 = new TopicPartition(topic2, 3);
    TopicPartition t3p0 = new TopicPartition(topic3, 0);
    TopicPartition t3p1 = new TopicPartition(topic3, 1);
    TopicPartition t3p2 = new TopicPartition(topic3, 2);
    TopicPartition t3p3 = new TopicPartition(topic3, 3);
    List<TopicPartition> c1partitions0 = isAllSubscriptionsEqual ? partitions(tp0, tp1, tp2, t2p2, t2p3, t3p0) : partitions(tp0, tp1, tp2, tp3);
    List<TopicPartition> c2partitions0 = partitions(tp0, tp1, t2p0, t2p1, t2p2, t2p3);
    List<TopicPartition> c3partitions0 = partitions(tp2, tp3, t3p0, t3p1, t3p2, t3p3);
    subscriptions.put(consumer1, buildSubscriptionWithGeneration(consumer1SubscribedTopics, c1partitions0, 1));
    subscriptions.put(consumer2, buildSubscriptionWithGeneration(consumer2SubscribedTopics, c2partitions0, 2));
    subscriptions.put(consumer3, buildSubscriptionWithGeneration(allTopics, c3partitions0, 2));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> c1partitions = assignment.get(consumer1);
    List<TopicPartition> c2partitions = assignment.get(consumer2);
    List<TopicPartition> c3partitions = assignment.get(consumer3);
    assertTrue(c1partitions.size() == 4 && c2partitions.size() == 4 && c3partitions.size() == 4);
    assertTrue(c2partitions0.containsAll(c2partitions));
    assertTrue(c3partitions0.containsAll(c3partitions));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}

Example 8 with Subscription

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

the class RoundRobinAssignorTest method testMultipleConsumersMixedTopics.

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

Example 9 with Subscription

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

the class RoundRobinAssignorTest method testOneStaticConsumerAndOneDynamicConsumerTwoTopicsSixPartitions.

@Test
public void testOneStaticConsumerAndOneDynamicConsumerTwoTopicsSixPartitions() {
    // although consumer 2 has a higher rank than 1, consumer 1 will win the comparison
    // because it has instance id while consumer 2 doesn't.
    String consumer1 = "consumer-b";
    String instance1 = "instance1";
    String consumer2 = "consumer-a";
    Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
    Map<String, Subscription> consumers = new HashMap<>();
    Subscription consumer1Subscription = new Subscription(topics(topic1, topic2), null);
    consumer1Subscription.setGroupInstanceId(Optional.of(instance1));
    consumers.put(consumer1, consumer1Subscription);
    consumers.put(consumer2, new Subscription(topics(topic1, topic2)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    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));
}

Example 10 with Subscription

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

the class RoundRobinAssignorTest method testStaticMemberRoundRobinAssignmentPersistentAfterMemberIdChanges.

@Test
public void testStaticMemberRoundRobinAssignmentPersistentAfterMemberIdChanges() {
    String consumer1 = "consumer1";
    String instance1 = "instance1";
    String consumer2 = "consumer2";
    String instance2 = "instance2";
    String consumer3 = "consumer3";
    String instance3 = "instance3";
    Map<String, String> memberIdToInstanceId = new HashMap<>();
    memberIdToInstanceId.put(consumer1, instance1);
    memberIdToInstanceId.put(consumer2, instance2);
    memberIdToInstanceId.put(consumer3, instance3);
    Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 5);
    Map<String, List<TopicPartition>> expectedInstanceAssignment = new HashMap<>();
    expectedInstanceAssignment.put(instance1, partitions(tp(topic1, 0), tp(topic1, 3), tp(topic2, 1), tp(topic2, 4)));
    expectedInstanceAssignment.put(instance2, partitions(tp(topic1, 1), tp(topic1, 4), tp(topic2, 2)));
    expectedInstanceAssignment.put(instance3, partitions(tp(topic1, 2), tp(topic2, 0), tp(topic2, 3)));
    List<AbstractPartitionAssignor.MemberInfo> staticMemberInfos = new ArrayList<>();
    for (Map.Entry<String, String> entry : memberIdToInstanceId.entrySet()) {
        staticMemberInfos.add(new AbstractPartitionAssignor.MemberInfo(entry.getKey(), Optional.of(entry.getValue())));
    }
    Map<String, Subscription> consumers = new HashMap<>();
    for (MemberInfo m : staticMemberInfos) {
        Subscription subscription = new Subscription(topics(topic1, topic2), null);
        subscription.setGroupInstanceId(m.groupInstanceId);
        consumers.put(m.memberId, subscription);
    }
    Map<String, List<TopicPartition>> staticAssignment = checkStaticAssignment(assignor, partitionsPerTopic, consumers);
    assertEquals(expectedInstanceAssignment, staticAssignment);
    memberIdToInstanceId.clear();
    // Now switch the member.id fields for each member info, the assignment should
    // stay the same as last time.
    String consumer4 = "consumer4";
    String consumer5 = "consumer5";
    consumers.put(consumer4, consumers.get(consumer3));
    consumers.remove(consumer3);
    consumers.put(consumer5, consumers.get(consumer2));
    consumers.remove(consumer2);
    Map<String, List<TopicPartition>> newStaticAssignment = checkStaticAssignment(assignor, partitionsPerTopic, consumers);
    assertEquals(staticAssignment, newStaticAssignment);
}