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

Example 66 with Subscription

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

the class RoundRobinAssignorTest method testStaticMemberRoundRobinAssignmentPersistent.

@Test
public void testStaticMemberRoundRobinAssignmentPersistent() {
    // Have 3 static members instance1, instance2, instance3 to be persistent
    // across generations. Their assignment shall be the same.
    String consumer1 = "consumer1";
    String instance1 = "instance1";
    String consumer2 = "consumer2";
    String instance2 = "instance2";
    String consumer3 = "consumer3";
    String instance3 = "instance3";
    List<MemberInfo> staticMemberInfos = new ArrayList<>();
    staticMemberInfos.add(new MemberInfo(consumer1, Optional.of(instance1)));
    staticMemberInfos.add(new MemberInfo(consumer2, Optional.of(instance2)));
    staticMemberInfos.add(new MemberInfo(consumer3, Optional.of(instance3)));
    // Consumer 4 is a dynamic member.
    String consumer4 = "consumer4";
    Map<String, Integer> partitionsPerTopic = setupPartitionsPerTopicWithTwoTopics(3, 3);
    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);
    }
    consumers.put(consumer4, new Subscription(topics(topic1, topic2)));
    Map<String, List<TopicPartition>> expectedAssignment = new HashMap<>();
    expectedAssignment.put(consumer1, partitions(tp(topic1, 0), tp(topic2, 1)));
    expectedAssignment.put(consumer2, partitions(tp(topic1, 1), tp(topic2, 2)));
    expectedAssignment.put(consumer3, partitions(tp(topic1, 2)));
    expectedAssignment.put(consumer4, partitions(tp(topic2, 0)));
    Map<String, List<TopicPartition>> assignment = assignor.assign(partitionsPerTopic, consumers);
    assertEquals(expectedAssignment, assignment);
    // Replace dynamic member 4 with a new dynamic member 5.
    consumers.remove(consumer4);
    String consumer5 = "consumer5";
    consumers.put(consumer5, new Subscription(topics(topic1, topic2)));
    expectedAssignment.remove(consumer4);
    expectedAssignment.put(consumer5, partitions(tp(topic2, 0)));
    assignment = assignor.assign(partitionsPerTopic, consumers);
    assertEquals(expectedAssignment, assignment);
}

Example 67 with Subscription

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

the class StreamsAssignmentScaleTest method completeLargeAssignment.

private void completeLargeAssignment(final int numPartitions, final int numClients, final int numThreadsPerClient, final int numStandbys, final Class<? extends TaskAssignor> taskAssignor) {
    final List<String> topic = singletonList("topic");
    final Map<TopicPartition, Long> changelogEndOffsets = new HashMap<>();
    for (int p = 0; p < numPartitions; ++p) {
        changelogEndOffsets.put(new TopicPartition(APPLICATION_ID + "-store-changelog", p), 100_000L);
    }
    final List<PartitionInfo> partitionInfos = new ArrayList<>();
    for (int p = 0; p < numPartitions; ++p) {
        partitionInfos.add(new PartitionInfo("topic", p, Node.noNode(), new Node[0], new Node[0]));
    }
    final Cluster clusterMetadata = new Cluster("cluster", Collections.singletonList(Node.noNode()), partitionInfos, emptySet(), emptySet());
    final Map<String, Object> configMap = new HashMap<>();
    configMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
    configMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:8080");
    final InternalTopologyBuilder builder = new InternalTopologyBuilder();
    builder.addSource(null, "source", null, null, null, "topic");
    builder.addProcessor("processor", new MockApiProcessorSupplier<>(), "source");
    builder.addStateStore(new MockKeyValueStoreBuilder("store", false), "processor");
    final TopologyMetadata topologyMetadata = new TopologyMetadata(builder, new StreamsConfig(configMap));
    topologyMetadata.buildAndRewriteTopology();
    final Consumer<byte[], byte[]> mainConsumer = EasyMock.createNiceMock(Consumer.class);
    final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
    expect(taskManager.topologyMetadata()).andStubReturn(topologyMetadata);
    expect(mainConsumer.committed(new HashSet<>())).andStubReturn(Collections.emptyMap());
    final AdminClient adminClient = createMockAdminClientForAssignor(changelogEndOffsets);
    final ReferenceContainer referenceContainer = new ReferenceContainer();
    referenceContainer.mainConsumer = mainConsumer;
    referenceContainer.adminClient = adminClient;
    referenceContainer.taskManager = taskManager;
    referenceContainer.streamsMetadataState = EasyMock.createNiceMock(StreamsMetadataState.class);
    referenceContainer.time = new MockTime();
    configMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
    configMap.put(InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS, taskAssignor.getName());
    configMap.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, numStandbys);
    final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(new MockTime(), new StreamsConfig(configMap), new MockClientSupplier().restoreConsumer, false);
    EasyMock.replay(taskManager, adminClient, mainConsumer);
    final StreamsPartitionAssignor partitionAssignor = new StreamsPartitionAssignor();
    partitionAssignor.configure(configMap);
    partitionAssignor.setInternalTopicManager(mockInternalTopicManager);
    final Map<String, Subscription> subscriptions = new HashMap<>();
    for (int client = 0; client < numClients; ++client) {
        for (int i = 0; i < numThreadsPerClient; ++i) {
            subscriptions.put(getConsumerName(i, client), new Subscription(topic, getInfo(uuidForInt(client), EMPTY_TASKS, EMPTY_TASKS).encode()));
        }
    }
    final long firstAssignmentStartMs = System.currentTimeMillis();
    final Map<String, Assignment> firstAssignments = partitionAssignor.assign(clusterMetadata, new GroupSubscription(subscriptions)).groupAssignment();
    final long firstAssignmentEndMs = System.currentTimeMillis();
    final long firstAssignmentDuration = firstAssignmentEndMs - firstAssignmentStartMs;
    if (firstAssignmentDuration > MAX_ASSIGNMENT_DURATION) {
        throw new AssertionError("The first assignment took too long to complete at " + firstAssignmentDuration + "ms.");
    } else {
        log.info("First assignment took {}ms.", firstAssignmentDuration);
    }
    // Use the assignment to generate the subscriptions' prev task data for the next rebalance
    for (int client = 0; client < numClients; ++client) {
        for (int i = 0; i < numThreadsPerClient; ++i) {
            final String consumer = getConsumerName(i, client);
            final Assignment assignment = firstAssignments.get(consumer);
            final AssignmentInfo info = AssignmentInfo.decode(assignment.userData());
            subscriptions.put(consumer, new Subscription(topic, getInfo(uuidForInt(client), new HashSet<>(info.activeTasks()), info.standbyTasks().keySet()).encode(), assignment.partitions()));
        }
    }
    final long secondAssignmentStartMs = System.currentTimeMillis();
    final Map<String, Assignment> secondAssignments = partitionAssignor.assign(clusterMetadata, new GroupSubscription(subscriptions)).groupAssignment();
    final long secondAssignmentEndMs = System.currentTimeMillis();
    final long secondAssignmentDuration = secondAssignmentEndMs - secondAssignmentStartMs;
    if (secondAssignmentDuration > MAX_ASSIGNMENT_DURATION) {
        throw new AssertionError("The second assignment took too long to complete at " + secondAssignmentDuration + "ms.");
    } else {
        log.info("Second assignment took {}ms.", secondAssignmentDuration);
    }
    assertThat(secondAssignments.size(), is(numClients * numThreadsPerClient));
}

Example 68 with Subscription

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

the class StreamsPartitionAssignorTest method shouldThrowTimeoutExceptionWhenCreatingRepartitionTopicsTimesOut.

@Test
public void shouldThrowTimeoutExceptionWhenCreatingRepartitionTopicsTimesOut() {
    final StreamsBuilder streamsBuilder = new StreamsBuilder();
    streamsBuilder.stream("topic1").repartition();
    final String client = "client1";
    builder = TopologyWrapper.getInternalTopologyBuilder(streamsBuilder.build());
    createDefaultMockTaskManager();
    EasyMock.replay(taskManager);
    partitionAssignor.configure(configProps());
    final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(time, new StreamsConfig(configProps()), mockClientSupplier.restoreConsumer, false) {

        @Override
        public Set<String> makeReady(final Map<String, InternalTopicConfig> topics) {
            throw new TimeoutException("KABOOM!");
        }
    };
    partitionAssignor.setInternalTopicManager(mockInternalTopicManager);
    subscriptions.put(client, new Subscription(singletonList("topic1"), defaultSubscriptionInfo.encode()));
    assertThrows(TimeoutException.class, () -> partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)));
}

Example 69 with Subscription

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

the class StreamsPartitionAssignorTest method shouldAddUserDefinedEndPointToSubscription.

@Test
public void shouldAddUserDefinedEndPointToSubscription() {
    builder.addSource(null, "source", null, null, null, "input");
    builder.addProcessor("processor", new MockApiProcessorSupplier<>(), "source");
    builder.addSink("sink", "output", null, null, null, "processor");
    createDefaultMockTaskManager();
    configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.APPLICATION_SERVER_CONFIG, USER_END_POINT));
    final Set<String> topics = mkSet("input");
    final ByteBuffer userData = partitionAssignor.subscriptionUserData(topics);
    final Subscription subscription = new Subscription(new ArrayList<>(topics), userData);
    final SubscriptionInfo subscriptionInfo = SubscriptionInfo.decode(subscription.userData());
    assertEquals("localhost:8080", subscriptionInfo.userEndPoint());
}

Example 70 with Subscription

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

the class StreamsPartitionAssignorTest method shouldSkipListOffsetsRequestForNewlyCreatedChangelogTopics.

@Test
public void shouldSkipListOffsetsRequestForNewlyCreatedChangelogTopics() {
    adminClient = EasyMock.createMock(AdminClient.class);
    final ListOffsetsResult result = EasyMock.createNiceMock(ListOffsetsResult.class);
    final KafkaFutureImpl<Map<TopicPartition, ListOffsetsResultInfo>> allFuture = new KafkaFutureImpl<>();
    allFuture.complete(emptyMap());
    expect(adminClient.listOffsets(emptyMap())).andStubReturn(result);
    expect(result.all()).andReturn(allFuture);
    builder.addSource(null, "source1", null, null, null, "topic1");
    builder.addProcessor("processor1", new MockApiProcessorSupplier<>(), "source1");
    builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1");
    subscriptions.put("consumer10", new Subscription(singletonList("topic1"), defaultSubscriptionInfo.encode()));
    EasyMock.replay(result);
    configureDefault();
    overwriteInternalTopicManagerWithMock(true);
    partitionAssignor.assign(metadata, new GroupSubscription(subscriptions));
    EasyMock.verify(adminClient);
}