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

Example 21 with Assignment

use of org.apache.kafka.clients.consumer.ConsumerPartitionAssignor.Assignment 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 22 with Assignment

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

the class StreamsPartitionAssignorTest method testAssignWithStandbyReplicasAndLoggingDisabled.

@Test
public void testAssignWithStandbyReplicasAndLoggingDisabled() {
    builder.addSource(null, "source1", null, null, null, "topic1", "topic2");
    builder.addProcessor("processor", new MockApiProcessorSupplier<>(), "source1");
    builder.addStateStore(new MockKeyValueStoreBuilder("store1", false).withLoggingDisabled(), "processor");
    final List<String> topics = asList("topic1", "topic2");
    createMockTaskManager(mkSet(TASK_0_0), emptySet());
    configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1));
    subscriptions.put("consumer10", new Subscription(topics, getInfo(UUID_1, mkSet(TASK_0_0), emptySet()).encode()));
    subscriptions.put("consumer20", new Subscription(topics, getInfo(UUID_2, mkSet(TASK_0_2), emptySet()).encode()));
    final Map<String, Assignment> assignments = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    final AssignmentInfo info10 = checkAssignment(allTopics, assignments.get("consumer10"));
    assertTrue(info10.standbyTasks().isEmpty());
    final AssignmentInfo info20 = checkAssignment(allTopics, assignments.get("consumer20"));
    assertTrue(info20.standbyTasks().isEmpty());
}

Example 23 with Assignment

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

the class StreamsPartitionAssignorTest method testAssignWithNewTasks.

@Test
public void testAssignWithNewTasks() {
    builder.addSource(null, "source1", null, null, null, "topic1");
    builder.addSource(null, "source2", null, null, null, "topic2");
    builder.addSource(null, "source3", null, null, null, "topic3");
    builder.addProcessor("processor", new MockApiProcessorSupplier<>(), "source1", "source2", "source3");
    final List<String> topics = asList("topic1", "topic2", "topic3");
    final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3);
    // assuming that previous tasks do not have topic3
    final Set<TaskId> prevTasks10 = mkSet(TASK_0_0);
    final Set<TaskId> prevTasks11 = mkSet(TASK_0_1);
    final Set<TaskId> prevTasks20 = mkSet(TASK_0_2);
    createMockTaskManager(prevTasks10, EMPTY_TASKS);
    configureDefaultPartitionAssignor();
    subscriptions.put("consumer10", new Subscription(topics, getInfo(UUID_1, prevTasks10, EMPTY_TASKS).encode()));
    subscriptions.put("consumer11", new Subscription(topics, getInfo(UUID_1, prevTasks11, EMPTY_TASKS).encode()));
    subscriptions.put("consumer20", new Subscription(topics, getInfo(UUID_2, prevTasks20, EMPTY_TASKS).encode()));
    final Map<String, Assignment> assignments = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    // check assigned partitions: since there is no previous task for topic 3 it will be assigned randomly so we cannot check exact match
    // also note that previously assigned partitions / tasks may not stay on the previous host since we may assign the new task first and
    // then later ones will be re-assigned to other hosts due to load balancing
    AssignmentInfo info = AssignmentInfo.decode(assignments.get("consumer10").userData());
    final Set<TaskId> allActiveTasks = new HashSet<>(info.activeTasks());
    final Set<TopicPartition> allPartitions = new HashSet<>(assignments.get("consumer10").partitions());
    info = AssignmentInfo.decode(assignments.get("consumer11").userData());
    allActiveTasks.addAll(info.activeTasks());
    allPartitions.addAll(assignments.get("consumer11").partitions());
    info = AssignmentInfo.decode(assignments.get("consumer20").userData());
    allActiveTasks.addAll(info.activeTasks());
    allPartitions.addAll(assignments.get("consumer20").partitions());
    assertEquals(allTasks, allActiveTasks);
    assertEquals(mkSet(t1p0, t1p1, t1p2, t2p0, t2p1, t2p2, t3p0, t3p1, t3p2, t3p3), allPartitions);
}

Example 24 with Assignment

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

the class StreamsPartitionAssignorTest method shouldNotAddStandbyTaskPartitionsToPartitionsForHost.

@Test
public void shouldNotAddStandbyTaskPartitionsToPartitionsForHost() {
    final Map<String, Object> props = configProps();
    props.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
    props.put(StreamsConfig.APPLICATION_SERVER_CONFIG, USER_END_POINT);
    final StreamsBuilder streamsBuilder = new StreamsBuilder();
    streamsBuilder.stream("topic1").groupByKey().count();
    builder = TopologyWrapper.getInternalTopologyBuilder(streamsBuilder.build());
    topologyMetadata = new TopologyMetadata(builder, new StreamsConfig(props));
    createDefaultMockTaskManager();
    adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(singletonList(APPLICATION_ID + "-KSTREAM-AGGREGATE-STATE-STORE-0000000001-changelog"), singletonList(3)));
    configurePartitionAssignorWith(props);
    subscriptions.put("consumer1", new Subscription(Collections.singletonList("topic1"), getInfo(UUID_1, EMPTY_TASKS, EMPTY_TASKS, USER_END_POINT).encode()));
    subscriptions.put("consumer2", new Subscription(Collections.singletonList("topic1"), getInfo(UUID_2, EMPTY_TASKS, EMPTY_TASKS, OTHER_END_POINT).encode()));
    final Set<TopicPartition> allPartitions = mkSet(t1p0, t1p1, t1p2);
    final Map<String, Assignment> assign = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    final Assignment consumer1Assignment = assign.get("consumer1");
    final AssignmentInfo assignmentInfo = AssignmentInfo.decode(consumer1Assignment.userData());
    final Set<TopicPartition> consumer1ActivePartitions = assignmentInfo.partitionsByHost().get(new HostInfo("localhost", 8080));
    final Set<TopicPartition> consumer2ActivePartitions = assignmentInfo.partitionsByHost().get(new HostInfo("other", 9090));
    final Set<TopicPartition> consumer1StandbyPartitions = assignmentInfo.standbyPartitionByHost().get(new HostInfo("localhost", 8080));
    final Set<TopicPartition> consumer2StandbyPartitions = assignmentInfo.standbyPartitionByHost().get(new HostInfo("other", 9090));
    final HashSet<TopicPartition> allAssignedPartitions = new HashSet<>(consumer1ActivePartitions);
    allAssignedPartitions.addAll(consumer2ActivePartitions);
    assertThat(consumer1ActivePartitions, not(allPartitions));
    assertThat(consumer2ActivePartitions, not(allPartitions));
    assertThat(consumer1ActivePartitions, equalTo(consumer2StandbyPartitions));
    assertThat(consumer2ActivePartitions, equalTo(consumer1StandbyPartitions));
    assertThat(allAssignedPartitions, equalTo(allPartitions));
}

Example 25 with Assignment

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

the class StreamsPartitionAssignorTest method shouldTriggerImmediateRebalanceOnTasksRevoked.

@Test
public void shouldTriggerImmediateRebalanceOnTasksRevoked() {
    builder.addSource(null, "source1", null, null, null, "topic1");
    final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2);
    final List<TopicPartition> allPartitions = asList(t1p0, t1p1, t1p2);
    subscriptions.put(CONSUMER_1, new Subscription(Collections.singletonList("topic1"), getInfo(UUID_1, allTasks, EMPTY_TASKS).encode(), allPartitions));
    subscriptions.put(CONSUMER_2, new Subscription(Collections.singletonList("topic1"), getInfo(UUID_1, EMPTY_TASKS, allTasks).encode(), emptyList()));
    createMockTaskManager(allTasks, allTasks);
    configurePartitionAssignorWith(singletonMap(StreamsConfig.ACCEPTABLE_RECOVERY_LAG_CONFIG, 0L));
    final Map<String, Assignment> assignment = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    // Verify at least one partition was revoked
    assertThat(assignment.get(CONSUMER_1).partitions(), not(allPartitions));
    assertThat(assignment.get(CONSUMER_2).partitions(), equalTo(emptyList()));
    // Verify that stateless revoked tasks would not be assigned as standbys
    assertThat(AssignmentInfo.decode(assignment.get(CONSUMER_2).userData()).activeTasks(), equalTo(emptyList()));
    assertThat(AssignmentInfo.decode(assignment.get(CONSUMER_2).userData()).standbyTasks(), equalTo(emptyMap()));
    partitionAssignor.onAssignment(assignment.get(CONSUMER_2), null);
    assertThat(referenceContainer.nextScheduledRebalanceMs.get(), is(0L));
}