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

Example 41 with Subscription

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

the class RangeAssignorTest method testStaticMemberRangeAssignmentPersistentAfterMemberIdChanges.

@Test
public void testStaticMemberRangeAssignmentPersistentAfterMemberIdChanges() {
    Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(5, 5);
    Map<String, Subscription> consumers = new HashMap<>();
    for (MemberInfo m : staticMemberInfos) {
        Subscription subscription = new Subscription(topics(topic1, topic2), null, Collections.emptyList());
        subscription.setGroupInstanceId(m.groupInstanceId);
        consumers.put(m.memberId, subscription);
    }
    Map<String, List<TopicPartition>> expectedInstanceAssignment = new HashMap<>();
    expectedInstanceAssignment.put(instance1, partitions(tp(topic1, 0), tp(topic1, 1), tp(topic2, 0), tp(topic2, 1)));
    expectedInstanceAssignment.put(instance2, partitions(tp(topic1, 2), tp(topic1, 3), tp(topic2, 2), tp(topic2, 3)));
    expectedInstanceAssignment.put(instance3, partitions(tp(topic1, 4), tp(topic2, 4)));
    Map<String, List<TopicPartition>> staticAssignment = checkStaticAssignment(assignor, partitionsPerTopic, consumers);
    assertEquals(expectedInstanceAssignment, staticAssignment);
    // 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);
}

Example 42 with Subscription

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

the class RangeAssignorTest method testTwoConsumersOneTopicTwoPartitions.

@Test
public void testTwoConsumersOneTopicTwoPartitions() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic1, 2);
    Map<String, Subscription> consumers = new HashMap<>();
    consumers.put(consumer1, new Subscription(topics(topic1)));
    consumers.put(consumer2, 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)), assignment.get(consumer2));
}

Example 43 with Subscription

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

the class StickyAssignorTest method testSchemaBackwardCompatibility.

@Test
public void testSchemaBackwardCompatibility() {
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 3);
    subscriptions.put(consumer1, new Subscription(topics(topic)));
    subscriptions.put(consumer2, new Subscription(topics(topic)));
    subscriptions.put(consumer3, new Subscription(topics(topic)));
    TopicPartition tp0 = new TopicPartition(topic, 0);
    TopicPartition tp1 = new TopicPartition(topic, 1);
    TopicPartition tp2 = new TopicPartition(topic, 2);
    List<TopicPartition> c1partitions0 = partitions(tp0, tp2);
    List<TopicPartition> c2partitions0 = partitions(tp1);
    subscriptions.put(consumer1, buildSubscriptionWithGeneration(topics(topic), c1partitions0, 1));
    subscriptions.put(consumer2, buildSubscriptionWithOldSchema(topics(topic), c2partitions0));
    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() == 1 && c2partitions.size() == 1 && c3partitions.size() == 1);
    assertTrue(c1partitions0.containsAll(c1partitions));
    assertTrue(c2partitions0.containsAll(c2partitions));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}

Example 44 with Subscription

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

the class StickyAssignorTest method testAssignmentWithMultipleGenerations1.

@ParameterizedTest(name = "testAssignmentWithMultipleGenerations1 with isAllSubscriptionsEqual: {0}")
@ValueSource(booleans = { true, false })
public void testAssignmentWithMultipleGenerations1(boolean isAllSubscriptionsEqual) {
    List<String> allTopics = topics(topic, topic2);
    List<String> consumer2SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic);
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 6);
    partitionsPerTopic.put(topic2, 6);
    subscriptions.put(consumer1, new Subscription(allTopics));
    subscriptions.put(consumer2, new Subscription(consumer2SubscribedTopics));
    subscriptions.put(consumer3, new Subscription(allTopics));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r1partitions1 = assignment.get(consumer1);
    List<TopicPartition> r1partitions2 = assignment.get(consumer2);
    List<TopicPartition> r1partitions3 = assignment.get(consumer3);
    assertTrue(r1partitions1.size() == 4 && r1partitions2.size() == 4 && r1partitions3.size() == 4);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    subscriptions.put(consumer1, buildSubscription(allTopics, r1partitions1));
    subscriptions.put(consumer2, buildSubscription(consumer2SubscribedTopics, r1partitions2));
    subscriptions.remove(consumer3);
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r2partitions1 = assignment.get(consumer1);
    List<TopicPartition> r2partitions2 = assignment.get(consumer2);
    assertTrue(r2partitions1.size() == 6 && r2partitions2.size() == 6);
    if (isAllSubscriptionsEqual) {
        // only true in all subscription equal case
        assertTrue(r2partitions1.containsAll(r1partitions1));
    }
    assertTrue(r2partitions2.containsAll(r1partitions2));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    assertFalse(Collections.disjoint(r2partitions2, r1partitions3));
    subscriptions.remove(consumer1);
    subscriptions.put(consumer2, buildSubscriptionWithGeneration(consumer2SubscribedTopics, r2partitions2, 2));
    subscriptions.put(consumer3, buildSubscriptionWithGeneration(allTopics, r1partitions3, 1));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r3partitions2 = assignment.get(consumer2);
    List<TopicPartition> r3partitions3 = assignment.get(consumer3);
    assertTrue(r3partitions2.size() == 6 && r3partitions3.size() == 6);
    assertTrue(Collections.disjoint(r3partitions2, r3partitions3));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}

Example 45 with Subscription

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

the class StickyAssignorTest method testAssignmentWithMultipleGenerations2.

@ParameterizedTest(name = "testAssignmentWithMultipleGenerations2 with isAllSubscriptionsEqual: {0}")
@ValueSource(booleans = { true, false })
public void testAssignmentWithMultipleGenerations2(boolean isAllSubscriptionsEqual) {
    List<String> allTopics = topics(topic, topic2, topic3);
    List<String> consumer1SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic);
    List<String> consumer3SubscribedTopics = isAllSubscriptionsEqual ? allTopics : topics(topic, topic2);
    Map<String, Integer> partitionsPerTopic = new HashMap<>();
    partitionsPerTopic.put(topic, 4);
    partitionsPerTopic.put(topic2, 4);
    partitionsPerTopic.put(topic3, 4);
    subscriptions.put(consumer1, new Subscription(consumer1SubscribedTopics));
    subscriptions.put(consumer2, new Subscription(allTopics));
    subscriptions.put(consumer3, new Subscription(consumer3SubscribedTopics));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r1partitions1 = assignment.get(consumer1);
    List<TopicPartition> r1partitions2 = assignment.get(consumer2);
    List<TopicPartition> r1partitions3 = assignment.get(consumer3);
    assertTrue(r1partitions1.size() == 4 && r1partitions2.size() == 4 && r1partitions3.size() == 4);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    subscriptions.remove(consumer1);
    subscriptions.put(consumer2, buildSubscriptionWithGeneration(allTopics, r1partitions2, 1));
    subscriptions.remove(consumer3);
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r2partitions2 = assignment.get(consumer2);
    assertEquals(12, r2partitions2.size());
    assertTrue(r2partitions2.containsAll(r1partitions2));
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
    subscriptions.put(consumer1, buildSubscriptionWithGeneration(consumer1SubscribedTopics, r1partitions1, 1));
    subscriptions.put(consumer2, buildSubscriptionWithGeneration(allTopics, r2partitions2, 2));
    subscriptions.put(consumer3, buildSubscriptionWithGeneration(consumer3SubscribedTopics, r1partitions3, 1));
    assignment = assignor.assign(partitionsPerTopic, subscriptions);
    List<TopicPartition> r3partitions1 = assignment.get(consumer1);
    List<TopicPartition> r3partitions2 = assignment.get(consumer2);
    List<TopicPartition> r3partitions3 = assignment.get(consumer3);
    assertTrue(r3partitions1.size() == 4 && r3partitions2.size() == 4 && r3partitions3.size() == 4);
    verifyValidityAndBalance(subscriptions, assignment, partitionsPerTopic);
    assertTrue(isFullyBalanced(assignment));
}