Search in sources :

Example 6 with SubscriptionInfo

use of org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo in project kafka by apache.

the class StreamPartitionAssignor method assign.

/*
     * This assigns tasks to consumer clients in the following steps.
     *
     * 0. check all repartition source topics and use internal topic manager to make sure
     *    they have been created with the right number of partitions.
     *
     * 1. using user customized partition grouper to generate tasks along with their
     *    assigned partitions; also make sure that the task's corresponding changelog topics
     *    have been created with the right number of partitions.
     *
     * 2. using TaskAssignor to assign tasks to consumer clients.
     *    - Assign a task to a client which was running it previously.
     *      If there is no such client, assign a task to a client which has its valid local state.
     *    - A client may have more than one stream threads.
     *      The assignor tries to assign tasks to a client proportionally to the number of threads.
     *    - We try not to assign the same set of tasks to two different clients
     *    We do the assignment in one-pass. The result may not satisfy above all.
     *
     * 3. within each client, tasks are assigned to consumer clients in round-robin manner.
     */
@Override
public Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions) {
    // construct the client metadata from the decoded subscription info
    Map<UUID, ClientMetadata> clientsMetadata = new HashMap<>();
    for (Map.Entry<String, Subscription> entry : subscriptions.entrySet()) {
        String consumerId = entry.getKey();
        Subscription subscription = entry.getValue();
        SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData());
        // create the new client metadata if necessary
        ClientMetadata clientMetadata = clientsMetadata.get(info.processId);
        if (clientMetadata == null) {
            clientMetadata = new ClientMetadata(info.userEndPoint);
            clientsMetadata.put(info.processId, clientMetadata);
        }
        // add the consumer to the client
        clientMetadata.addConsumer(consumerId, info);
    }
    log.info("stream-thread [{}] Constructed client metadata {} from the member subscriptions.", streamThread.getName(), clientsMetadata);
    // ---------------- Step Zero ---------------- //
    // parse the topology to determine the repartition source topics,
    // making sure they are created with the number of partitions as
    // the maximum of the depending sub-topologies source topics' number of partitions
    Map<Integer, TopologyBuilder.TopicsInfo> topicGroups = streamThread.builder.topicGroups();
    Map<String, InternalTopicMetadata> repartitionTopicMetadata = new HashMap<>();
    for (TopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
        for (InternalTopicConfig topic : topicsInfo.repartitionSourceTopics.values()) {
            repartitionTopicMetadata.put(topic.name(), new InternalTopicMetadata(topic));
        }
    }
    boolean numPartitionsNeeded;
    do {
        numPartitionsNeeded = false;
        for (TopologyBuilder.TopicsInfo topicsInfo : topicGroups.values()) {
            for (String topicName : topicsInfo.repartitionSourceTopics.keySet()) {
                int numPartitions = repartitionTopicMetadata.get(topicName).numPartitions;
                // try set the number of partitions for this repartition topic if it is not set yet
                if (numPartitions == UNKNOWN) {
                    for (TopologyBuilder.TopicsInfo otherTopicsInfo : topicGroups.values()) {
                        Set<String> otherSinkTopics = otherTopicsInfo.sinkTopics;
                        if (otherSinkTopics.contains(topicName)) {
                            // use the maximum of all its source topic partitions as the number of partitions
                            for (String sourceTopicName : otherTopicsInfo.sourceTopics) {
                                Integer numPartitionsCandidate;
                                // map().join().join(map())
                                if (repartitionTopicMetadata.containsKey(sourceTopicName)) {
                                    numPartitionsCandidate = repartitionTopicMetadata.get(sourceTopicName).numPartitions;
                                } else {
                                    numPartitionsCandidate = metadata.partitionCountForTopic(sourceTopicName);
                                    if (numPartitionsCandidate == null) {
                                        repartitionTopicMetadata.get(topicName).numPartitions = NOT_AVAILABLE;
                                    }
                                }
                                if (numPartitionsCandidate != null && numPartitionsCandidate > numPartitions) {
                                    numPartitions = numPartitionsCandidate;
                                }
                            }
                        }
                    }
                    // another iteration is needed
                    if (numPartitions == UNKNOWN)
                        numPartitionsNeeded = true;
                    else
                        repartitionTopicMetadata.get(topicName).numPartitions = numPartitions;
                }
            }
        }
    } while (numPartitionsNeeded);
    // augment the metadata with the newly computed number of partitions for all the
    // repartition source topics
    Map<TopicPartition, PartitionInfo> allRepartitionTopicPartitions = new HashMap<>();
    for (Map.Entry<String, InternalTopicMetadata> entry : repartitionTopicMetadata.entrySet()) {
        String topic = entry.getKey();
        Integer numPartitions = entry.getValue().numPartitions;
        for (int partition = 0; partition < numPartitions; partition++) {
            allRepartitionTopicPartitions.put(new TopicPartition(topic, partition), new PartitionInfo(topic, partition, null, new Node[0], new Node[0]));
        }
    }
    // ensure the co-partitioning topics within the group have the same number of partitions,
    // and enforce the number of partitions for those repartition topics to be the same if they
    // are co-partitioned as well.
    ensureCopartitioning(streamThread.builder.copartitionGroups(), repartitionTopicMetadata, metadata);
    // make sure the repartition source topics exist with the right number of partitions,
    // create these topics if necessary
    prepareTopic(repartitionTopicMetadata);
    metadataWithInternalTopics = metadata.withPartitions(allRepartitionTopicPartitions);
    log.debug("stream-thread [{}] Created repartition topics {} from the parsed topology.", streamThread.getName(), allRepartitionTopicPartitions.values());
    // ---------------- Step One ---------------- //
    // get the tasks as partition groups from the partition grouper
    Set<String> allSourceTopics = new HashSet<>();
    Map<Integer, Set<String>> sourceTopicsByGroup = new HashMap<>();
    for (Map.Entry<Integer, TopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
        allSourceTopics.addAll(entry.getValue().sourceTopics);
        sourceTopicsByGroup.put(entry.getKey(), entry.getValue().sourceTopics);
    }
    Map<TaskId, Set<TopicPartition>> partitionsForTask = streamThread.partitionGrouper.partitionGroups(sourceTopicsByGroup, metadataWithInternalTopics);
    // check if all partitions are assigned, and there are no duplicates of partitions in multiple tasks
    Set<TopicPartition> allAssignedPartitions = new HashSet<>();
    Map<Integer, Set<TaskId>> tasksByTopicGroup = new HashMap<>();
    for (Map.Entry<TaskId, Set<TopicPartition>> entry : partitionsForTask.entrySet()) {
        Set<TopicPartition> partitions = entry.getValue();
        for (TopicPartition partition : partitions) {
            if (allAssignedPartitions.contains(partition)) {
                log.warn("stream-thread [{}] Partition {} is assigned to more than one tasks: {}", streamThread.getName(), partition, partitionsForTask);
            }
        }
        allAssignedPartitions.addAll(partitions);
        TaskId id = entry.getKey();
        Set<TaskId> ids = tasksByTopicGroup.get(id.topicGroupId);
        if (ids == null) {
            ids = new HashSet<>();
            tasksByTopicGroup.put(id.topicGroupId, ids);
        }
        ids.add(id);
    }
    for (String topic : allSourceTopics) {
        List<PartitionInfo> partitionInfoList = metadataWithInternalTopics.partitionsForTopic(topic);
        if (!partitionInfoList.isEmpty()) {
            for (PartitionInfo partitionInfo : partitionInfoList) {
                TopicPartition partition = new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
                if (!allAssignedPartitions.contains(partition)) {
                    log.warn("stream-thread [{}] Partition {} is not assigned to any tasks: {}", streamThread.getName(), partition, partitionsForTask);
                }
            }
        } else {
            log.warn("stream-thread [{}] No partitions found for topic {}", streamThread.getName(), topic);
        }
    }
    // add tasks to state change log topic subscribers
    Map<String, InternalTopicMetadata> changelogTopicMetadata = new HashMap<>();
    for (Map.Entry<Integer, TopologyBuilder.TopicsInfo> entry : topicGroups.entrySet()) {
        final int topicGroupId = entry.getKey();
        final Map<String, InternalTopicConfig> stateChangelogTopics = entry.getValue().stateChangelogTopics;
        for (InternalTopicConfig topicConfig : stateChangelogTopics.values()) {
            // the expected number of partitions is the max value of TaskId.partition + 1
            int numPartitions = UNKNOWN;
            if (tasksByTopicGroup.get(topicGroupId) != null) {
                for (TaskId task : tasksByTopicGroup.get(topicGroupId)) {
                    if (numPartitions < task.partition + 1)
                        numPartitions = task.partition + 1;
                }
                InternalTopicMetadata topicMetadata = new InternalTopicMetadata(topicConfig);
                topicMetadata.numPartitions = numPartitions;
                changelogTopicMetadata.put(topicConfig.name(), topicMetadata);
            } else {
                log.debug("stream-thread [{}] No tasks found for topic group {}", streamThread.getName(), topicGroupId);
            }
        }
    }
    prepareTopic(changelogTopicMetadata);
    log.debug("stream-thread [{}] Created state changelog topics {} from the parsed topology.", streamThread.getName(), changelogTopicMetadata);
    // ---------------- Step Two ---------------- //
    // assign tasks to clients
    Map<UUID, ClientState> states = new HashMap<>();
    for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
        states.put(entry.getKey(), entry.getValue().state);
    }
    log.debug("stream-thread [{}] Assigning tasks {} to clients {} with number of replicas {}", streamThread.getName(), partitionsForTask.keySet(), states, numStandbyReplicas);
    final StickyTaskAssignor<UUID> taskAssignor = new StickyTaskAssignor<>(states, partitionsForTask.keySet());
    taskAssignor.assign(numStandbyReplicas);
    log.info("stream-thread [{}] Assigned tasks to clients as {}.", streamThread.getName(), states);
    // ---------------- Step Three ---------------- //
    // construct the global partition assignment per host map
    partitionsByHostState = new HashMap<>();
    for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
        HostInfo hostInfo = entry.getValue().hostInfo;
        if (hostInfo != null) {
            final Set<TopicPartition> topicPartitions = new HashSet<>();
            final ClientState state = entry.getValue().state;
            for (final TaskId id : state.activeTasks()) {
                topicPartitions.addAll(partitionsForTask.get(id));
            }
            partitionsByHostState.put(hostInfo, topicPartitions);
        }
    }
    // within the client, distribute tasks to its owned consumers
    Map<String, Assignment> assignment = new HashMap<>();
    for (Map.Entry<UUID, ClientMetadata> entry : clientsMetadata.entrySet()) {
        final Set<String> consumers = entry.getValue().consumers;
        final ClientState state = entry.getValue().state;
        final ArrayList<TaskId> taskIds = new ArrayList<>(state.assignedTaskCount());
        final int numActiveTasks = state.activeTaskCount();
        taskIds.addAll(state.activeTasks());
        taskIds.addAll(state.standbyTasks());
        final int numConsumers = consumers.size();
        int i = 0;
        for (String consumer : consumers) {
            Map<TaskId, Set<TopicPartition>> standby = new HashMap<>();
            ArrayList<AssignedPartition> assignedPartitions = new ArrayList<>();
            final int numTaskIds = taskIds.size();
            for (int j = i; j < numTaskIds; j += numConsumers) {
                TaskId taskId = taskIds.get(j);
                if (j < numActiveTasks) {
                    for (TopicPartition partition : partitionsForTask.get(taskId)) {
                        assignedPartitions.add(new AssignedPartition(taskId, partition));
                    }
                } else {
                    Set<TopicPartition> standbyPartitions = standby.get(taskId);
                    if (standbyPartitions == null) {
                        standbyPartitions = new HashSet<>();
                        standby.put(taskId, standbyPartitions);
                    }
                    standbyPartitions.addAll(partitionsForTask.get(taskId));
                }
            }
            Collections.sort(assignedPartitions);
            List<TaskId> active = new ArrayList<>();
            List<TopicPartition> activePartitions = new ArrayList<>();
            for (AssignedPartition partition : assignedPartitions) {
                active.add(partition.taskId);
                activePartitions.add(partition.partition);
            }
            // finally, encode the assignment before sending back to coordinator
            assignment.put(consumer, new Assignment(activePartitions, new AssignmentInfo(active, standby, partitionsByHostState).encode()));
            i++;
        }
    }
    return assignment;
}
Also used : ClientState(org.apache.kafka.streams.processor.internals.assignment.ClientState) HashMap(java.util.HashMap) Node(org.apache.kafka.common.Node) ArrayList(java.util.ArrayList) SubscriptionInfo(org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo) StickyTaskAssignor(org.apache.kafka.streams.processor.internals.assignment.StickyTaskAssignor) AssignmentInfo(org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo) PartitionInfo(org.apache.kafka.common.PartitionInfo) UUID(java.util.UUID) HashSet(java.util.HashSet) TopicPartition(org.apache.kafka.common.TopicPartition) HashMap(java.util.HashMap) Map(java.util.Map) HashSet(java.util.HashSet) Set(java.util.Set) TaskId(org.apache.kafka.streams.processor.TaskId) TopologyBuilder(org.apache.kafka.streams.processor.TopologyBuilder) HostInfo(org.apache.kafka.streams.state.HostInfo)

Example 7 with SubscriptionInfo

use of org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo in project kafka by apache.

the class StreamPartitionAssignor method subscription.

@Override
public Subscription subscription(Set<String> topics) {
    // Adds the following information to subscription
    // 1. Client UUID (a unique id assigned to an instance of KafkaStreams)
    // 2. Task ids of previously running tasks
    // 3. Task ids of valid local states on the client's state directory.
    final Set<TaskId> previousActiveTasks = streamThread.prevActiveTasks();
    Set<TaskId> standbyTasks = streamThread.cachedTasks();
    standbyTasks.removeAll(previousActiveTasks);
    SubscriptionInfo data = new SubscriptionInfo(streamThread.processId, previousActiveTasks, standbyTasks, this.userEndPoint);
    if (streamThread.builder.sourceTopicPattern() != null) {
        SubscriptionUpdates subscriptionUpdates = new SubscriptionUpdates();
        log.debug("stream-thread [{}] found {} topics possibly matching regex", streamThread.getName(), topics);
        // update the topic groups with the returned subscription set for regex pattern subscriptions
        subscriptionUpdates.updateTopics(topics);
        streamThread.builder.updateSubscriptions(subscriptionUpdates, streamThread.getName());
    }
    return new Subscription(new ArrayList<>(topics), data.encode());
}
Also used : TaskId(org.apache.kafka.streams.processor.TaskId) SubscriptionInfo(org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo)

Example 8 with SubscriptionInfo

use of org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo in project kafka by apache.

the class StreamPartitionAssignorTest method testAssignWithNewTasks.

@Test
public void testAssignWithNewTasks() throws Exception {
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addSource("source3", "topic3");
    builder.addProcessor("processor", new MockProcessorSupplier(), "source1", "source2", "source3");
    List<String> topics = Utils.mkList("topic1", "topic2", "topic3");
    Set<TaskId> allTasks = Utils.mkSet(task0, task1, task2, task3);
    // assuming that previous tasks do not have topic3
    final Set<TaskId> prevTasks10 = Utils.mkSet(task0);
    final Set<TaskId> prevTasks11 = Utils.mkSet(task1);
    final Set<TaskId> prevTasks20 = Utils.mkSet(task2);
    UUID uuid1 = UUID.randomUUID();
    UUID uuid2 = UUID.randomUUID();
    String client1 = "client1";
    StreamThread thread10 = new StreamThread(builder, config, mockClientSupplier, "test", client1, uuid1, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0);
    partitionAssignor.configure(config.getConsumerConfigs(thread10, "test", client1));
    partitionAssignor.setInternalTopicManager(new MockInternalTopicManager(thread10.config, mockClientSupplier.restoreConsumer));
    Map<String, PartitionAssignor.Subscription> subscriptions = new HashMap<>();
    subscriptions.put("consumer10", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid1, prevTasks10, Collections.<TaskId>emptySet(), userEndPoint).encode()));
    subscriptions.put("consumer11", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid1, prevTasks11, Collections.<TaskId>emptySet(), userEndPoint).encode()));
    subscriptions.put("consumer20", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid2, prevTasks20, Collections.<TaskId>emptySet(), userEndPoint).encode()));
    Map<String, PartitionAssignor.Assignment> assignments = partitionAssignor.assign(metadata, subscriptions);
    // 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
    Set<TaskId> allActiveTasks = new HashSet<>();
    Set<TopicPartition> allPartitions = new HashSet<>();
    AssignmentInfo info;
    info = AssignmentInfo.decode(assignments.get("consumer10").userData());
    allActiveTasks.addAll(info.activeTasks);
    allPartitions.addAll(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(Utils.mkSet(t1p0, t1p1, t1p2, t2p0, t2p1, t2p2, t3p0, t3p1, t3p2, t3p3), allPartitions);
}
Also used : TaskId(org.apache.kafka.streams.processor.TaskId) HashMap(java.util.HashMap) MockInternalTopicManager(org.apache.kafka.test.MockInternalTopicManager) SubscriptionInfo(org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo) AssignmentInfo(org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo) Metrics(org.apache.kafka.common.metrics.Metrics) MockProcessorSupplier(org.apache.kafka.test.MockProcessorSupplier) TopicPartition(org.apache.kafka.common.TopicPartition) PartitionAssignor(org.apache.kafka.clients.consumer.internals.PartitionAssignor) UUID(java.util.UUID) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 9 with SubscriptionInfo

use of org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo in project kafka by apache.

the class StreamPartitionAssignorTest method shouldNotLoopInfinitelyOnMissingMetadataAndShouldNotCreateRelatedTasks.

@Test
public void shouldNotLoopInfinitelyOnMissingMetadataAndShouldNotCreateRelatedTasks() {
    final String applicationId = "application-id";
    final KStreamBuilder builder = new KStreamBuilder();
    builder.setApplicationId(applicationId);
    KStream<Object, Object> stream1 = builder.stream("topic1").selectKey(new KeyValueMapper<Object, Object, Object>() {

        @Override
        public Object apply(Object key, Object value) {
            return null;
        }
    }).groupByKey().count("count").toStream().map(new KeyValueMapper<Object, Long, KeyValue<Object, Object>>() {

        @Override
        public KeyValue<Object, Object> apply(Object key, Long value) {
            return null;
        }
    });
    builder.stream("unknownTopic").selectKey(new KeyValueMapper<Object, Object, Object>() {

        @Override
        public Object apply(Object key, Object value) {
            return null;
        }
    }).join(stream1, new ValueJoiner() {

        @Override
        public Object apply(Object value1, Object value2) {
            return null;
        }
    }, JoinWindows.of(0));
    final UUID uuid = UUID.randomUUID();
    final String client = "client1";
    final StreamThread streamThread = new StreamThread(builder, config, mockClientSupplier, applicationId, client, uuid, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0);
    partitionAssignor.configure(config.getConsumerConfigs(streamThread, applicationId, client));
    final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(streamThread.config, mockClientSupplier.restoreConsumer);
    partitionAssignor.setInternalTopicManager(mockInternalTopicManager);
    final Map<String, PartitionAssignor.Subscription> subscriptions = new HashMap<>();
    final Set<TaskId> emptyTasks = Collections.emptySet();
    subscriptions.put(client, new PartitionAssignor.Subscription(Collections.singletonList("unknownTopic"), new SubscriptionInfo(uuid, emptyTasks, emptyTasks, userEndPoint).encode()));
    final Map<String, PartitionAssignor.Assignment> assignment = partitionAssignor.assign(metadata, subscriptions);
    final Map<String, Integer> expectedCreatedInternalTopics = new HashMap<>();
    expectedCreatedInternalTopics.put(applicationId + "-count-repartition", 3);
    expectedCreatedInternalTopics.put(applicationId + "-count-changelog", 3);
    assertThat(mockInternalTopicManager.readyTopics, equalTo(expectedCreatedInternalTopics));
    final List<TopicPartition> expectedAssignment = Arrays.asList(new TopicPartition("topic1", 0), new TopicPartition("topic1", 1), new TopicPartition("topic1", 2), new TopicPartition(applicationId + "-count-repartition", 0), new TopicPartition(applicationId + "-count-repartition", 1), new TopicPartition(applicationId + "-count-repartition", 2));
    assertThat(new HashSet(assignment.get(client).partitions()), equalTo(new HashSet(expectedAssignment)));
}
Also used : KeyValue(org.apache.kafka.streams.KeyValue) TaskId(org.apache.kafka.streams.processor.TaskId) HashMap(java.util.HashMap) MockInternalTopicManager(org.apache.kafka.test.MockInternalTopicManager) KeyValueMapper(org.apache.kafka.streams.kstream.KeyValueMapper) SubscriptionInfo(org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo) ValueJoiner(org.apache.kafka.streams.kstream.ValueJoiner) Metrics(org.apache.kafka.common.metrics.Metrics) PartitionAssignor(org.apache.kafka.clients.consumer.internals.PartitionAssignor) UUID(java.util.UUID) HashSet(java.util.HashSet) KStreamBuilder(org.apache.kafka.streams.kstream.KStreamBuilder) TopicPartition(org.apache.kafka.common.TopicPartition) Test(org.junit.Test)

Example 10 with SubscriptionInfo

use of org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo in project kafka by apache.

the class StreamPartitionAssignorTest method testAssignWithStates.

@Test
public void testAssignWithStates() throws Exception {
    String applicationId = "test";
    builder.setApplicationId(applicationId);
    builder.addSource("source1", "topic1");
    builder.addSource("source2", "topic2");
    builder.addProcessor("processor-1", new MockProcessorSupplier(), "source1");
    builder.addStateStore(new MockStateStoreSupplier("store1", false), "processor-1");
    builder.addProcessor("processor-2", new MockProcessorSupplier(), "source2");
    builder.addStateStore(new MockStateStoreSupplier("store2", false), "processor-2");
    builder.addStateStore(new MockStateStoreSupplier("store3", false), "processor-2");
    List<String> topics = Utils.mkList("topic1", "topic2");
    TaskId task00 = new TaskId(0, 0);
    TaskId task01 = new TaskId(0, 1);
    TaskId task02 = new TaskId(0, 2);
    TaskId task10 = new TaskId(1, 0);
    TaskId task11 = new TaskId(1, 1);
    TaskId task12 = new TaskId(1, 2);
    List<TaskId> tasks = Utils.mkList(task00, task01, task02, task10, task11, task12);
    UUID uuid1 = UUID.randomUUID();
    UUID uuid2 = UUID.randomUUID();
    String client1 = "client1";
    StreamThread thread10 = new StreamThread(builder, config, mockClientSupplier, applicationId, client1, uuid1, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0);
    partitionAssignor.configure(config.getConsumerConfigs(thread10, applicationId, client1));
    partitionAssignor.setInternalTopicManager(new MockInternalTopicManager(thread10.config, mockClientSupplier.restoreConsumer));
    Map<String, PartitionAssignor.Subscription> subscriptions = new HashMap<>();
    subscriptions.put("consumer10", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid1, Collections.<TaskId>emptySet(), Collections.<TaskId>emptySet(), userEndPoint).encode()));
    subscriptions.put("consumer11", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid1, Collections.<TaskId>emptySet(), Collections.<TaskId>emptySet(), userEndPoint).encode()));
    subscriptions.put("consumer20", new PartitionAssignor.Subscription(topics, new SubscriptionInfo(uuid2, Collections.<TaskId>emptySet(), Collections.<TaskId>emptySet(), userEndPoint).encode()));
    Map<String, PartitionAssignor.Assignment> assignments = partitionAssignor.assign(metadata, subscriptions);
    // check assigned partition size: since there is no previous task and there are two sub-topologies the assignment is random so we cannot check exact match
    assertEquals(2, assignments.get("consumer10").partitions().size());
    assertEquals(2, assignments.get("consumer11").partitions().size());
    assertEquals(2, assignments.get("consumer20").partitions().size());
    AssignmentInfo info10 = AssignmentInfo.decode(assignments.get("consumer10").userData());
    AssignmentInfo info11 = AssignmentInfo.decode(assignments.get("consumer11").userData());
    AssignmentInfo info20 = AssignmentInfo.decode(assignments.get("consumer20").userData());
    assertEquals(2, info10.activeTasks.size());
    assertEquals(2, info11.activeTasks.size());
    assertEquals(2, info20.activeTasks.size());
    Set<TaskId> allTasks = new HashSet<>();
    allTasks.addAll(info10.activeTasks);
    allTasks.addAll(info11.activeTasks);
    allTasks.addAll(info20.activeTasks);
    assertEquals(new HashSet<>(tasks), allTasks);
    // check tasks for state topics
    Map<Integer, TopologyBuilder.TopicsInfo> topicGroups = thread10.builder.topicGroups();
    assertEquals(Utils.mkSet(task00, task01, task02), tasksForState(applicationId, "store1", tasks, topicGroups));
    assertEquals(Utils.mkSet(task10, task11, task12), tasksForState(applicationId, "store2", tasks, topicGroups));
    assertEquals(Utils.mkSet(task10, task11, task12), tasksForState(applicationId, "store3", tasks, topicGroups));
}
Also used : TaskId(org.apache.kafka.streams.processor.TaskId) HashMap(java.util.HashMap) MockInternalTopicManager(org.apache.kafka.test.MockInternalTopicManager) SubscriptionInfo(org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo) AssignmentInfo(org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo) Metrics(org.apache.kafka.common.metrics.Metrics) MockStateStoreSupplier(org.apache.kafka.test.MockStateStoreSupplier) PartitionAssignor(org.apache.kafka.clients.consumer.internals.PartitionAssignor) UUID(java.util.UUID) HashSet(java.util.HashSet) MockProcessorSupplier(org.apache.kafka.test.MockProcessorSupplier) Test(org.junit.Test)

Aggregations

SubscriptionInfo (org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo)15 UUID (java.util.UUID)14 TaskId (org.apache.kafka.streams.processor.TaskId)14 PartitionAssignor (org.apache.kafka.clients.consumer.internals.PartitionAssignor)13 Metrics (org.apache.kafka.common.metrics.Metrics)13 Test (org.junit.Test)13 HashMap (java.util.HashMap)12 MockProcessorSupplier (org.apache.kafka.test.MockProcessorSupplier)11 HashSet (java.util.HashSet)10 MockInternalTopicManager (org.apache.kafka.test.MockInternalTopicManager)10 AssignmentInfo (org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo)9 Properties (java.util.Properties)5 TopicPartition (org.apache.kafka.common.TopicPartition)5 StreamsConfig (org.apache.kafka.streams.StreamsConfig)5 HostInfo (org.apache.kafka.streams.state.HostInfo)3 KStreamBuilder (org.apache.kafka.streams.kstream.KStreamBuilder)2 MockClientSupplier (org.apache.kafka.test.MockClientSupplier)2 MockStateStoreSupplier (org.apache.kafka.test.MockStateStoreSupplier)2 ArrayList (java.util.ArrayList)1 Map (java.util.Map)1