Examples with AssignmentInfo org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo used on opensource projects

Example 36 with AssignmentInfo

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

the class StreamsPartitionAssignor method onAssignment.

@Override
public void onAssignment(final Assignment assignment, final ConsumerGroupMetadata metadata) {
    final List<TopicPartition> partitions = new ArrayList<>(assignment.partitions());
    partitions.sort(PARTITION_COMPARATOR);
    final AssignmentInfo info = AssignmentInfo.decode(assignment.userData());
    if (info.errCode() != AssignorError.NONE.code()) {
        // set flag to shutdown streams app
        assignmentErrorCode.set(info.errCode());
        return;
    }
    /*
         * latestCommonlySupportedVersion belongs to [usedSubscriptionMetadataVersion, LATEST_SUPPORTED_VERSION]
         * receivedAssignmentMetadataVersion belongs to [EARLIEST_PROBEABLE_VERSION, usedSubscriptionMetadataVersion]
         *
         * usedSubscriptionMetadataVersion will be downgraded to receivedAssignmentMetadataVersion during a rolling
         * bounce upgrade with version probing.
         *
         * usedSubscriptionMetadataVersion will be upgraded to latestCommonlySupportedVersion when all members have
         * been bounced and it is safe to use the latest version.
         */
    final int receivedAssignmentMetadataVersion = info.version();
    final int latestCommonlySupportedVersion = info.commonlySupportedVersion();
    validateMetadataVersions(receivedAssignmentMetadataVersion, latestCommonlySupportedVersion);
    // version 1 field
    final Map<TaskId, Set<TopicPartition>> activeTasks;
    // version 2 fields
    final Map<TopicPartition, PartitionInfo> topicToPartitionInfo;
    final Map<HostInfo, Set<TopicPartition>> partitionsByHost;
    final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost;
    final long encodedNextScheduledRebalanceMs;
    switch(receivedAssignmentMetadataVersion) {
        case 1:
            validateActiveTaskEncoding(partitions, info, logPrefix);
            activeTasks = getActiveTasks(partitions, info);
            partitionsByHost = Collections.emptyMap();
            standbyPartitionsByHost = Collections.emptyMap();
            topicToPartitionInfo = Collections.emptyMap();
            encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
            break;
        case 2:
        case 3:
        case 4:
        case 5:
            validateActiveTaskEncoding(partitions, info, logPrefix);
            activeTasks = getActiveTasks(partitions, info);
            partitionsByHost = info.partitionsByHost();
            standbyPartitionsByHost = Collections.emptyMap();
            topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
            encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
            break;
        case 6:
            validateActiveTaskEncoding(partitions, info, logPrefix);
            activeTasks = getActiveTasks(partitions, info);
            partitionsByHost = info.partitionsByHost();
            standbyPartitionsByHost = info.standbyPartitionByHost();
            topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
            encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
            break;
        case 7:
        case 8:
        case 9:
        case 10:
            validateActiveTaskEncoding(partitions, info, logPrefix);
            activeTasks = getActiveTasks(partitions, info);
            partitionsByHost = info.partitionsByHost();
            standbyPartitionsByHost = info.standbyPartitionByHost();
            topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
            encodedNextScheduledRebalanceMs = info.nextRebalanceMs();
            break;
        default:
            throw new IllegalStateException("This code should never be reached." + " Please file a bug report at https://issues.apache.org/jira/projects/KAFKA/");
    }
    maybeScheduleFollowupRebalance(encodedNextScheduledRebalanceMs, receivedAssignmentMetadataVersion, latestCommonlySupportedVersion, partitionsByHost.keySet());
    final Cluster fakeCluster = Cluster.empty().withPartitions(topicToPartitionInfo);
    streamsMetadataState.onChange(partitionsByHost, standbyPartitionsByHost, fakeCluster);
    // we do not capture any exceptions but just let the exception thrown from consumer.poll directly
    // since when stream thread captures it, either we close all tasks as dirty or we close thread
    taskManager.handleAssignment(activeTasks, info.standbyTasks());
}

Example 37 with AssignmentInfo

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

the class StreamsPartitionAssignor method addClientAssignments.

/**
 * Adds the encoded assignment for each StreamThread consumer in the client to the overall assignment map
 * @return true if a followup rebalance will be required due to revoked tasks
 */
private boolean addClientAssignments(final Set<TaskId> statefulTasks, final Map<String, Assignment> assignment, final ClientMetadata clientMetadata, final Map<TaskId, Set<TopicPartition>> partitionsForTask, final Map<HostInfo, Set<TopicPartition>> partitionsByHostState, final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost, final Set<TopicPartition> allOwnedPartitions, final Map<String, List<TaskId>> activeTaskAssignments, final Map<String, List<TaskId>> standbyTaskAssignments, final int minUserMetadataVersion, final int minSupportedMetadataVersion, final boolean probingRebalanceNeeded) {
    boolean followupRebalanceRequiredForRevokedTasks = false;
    // We only want to encode a scheduled probing rebalance for a single member in this client
    boolean shouldEncodeProbingRebalance = probingRebalanceNeeded;
    // Loop through the consumers and build their assignment
    for (final String consumer : clientMetadata.consumers) {
        final List<TaskId> activeTasksForConsumer = activeTaskAssignments.get(consumer);
        // These will be filled in by populateActiveTaskAndPartitionsLists below
        final List<TopicPartition> activePartitionsList = new ArrayList<>();
        final List<TaskId> assignedActiveList = new ArrayList<>();
        final Set<TaskId> activeTasksRemovedPendingRevokation = populateActiveTaskAndPartitionsLists(activePartitionsList, assignedActiveList, consumer, clientMetadata.state, activeTasksForConsumer, partitionsForTask, allOwnedPartitions);
        final Map<TaskId, Set<TopicPartition>> standbyTaskMap = buildStandbyTaskMap(consumer, standbyTaskAssignments.get(consumer), activeTasksRemovedPendingRevokation, statefulTasks, partitionsForTask, clientMetadata.state);
        final AssignmentInfo info = new AssignmentInfo(minUserMetadataVersion, minSupportedMetadataVersion, assignedActiveList, standbyTaskMap, partitionsByHostState, standbyPartitionsByHost, AssignorError.NONE.code());
        if (!activeTasksRemovedPendingRevokation.isEmpty()) {
            // TODO: once KAFKA-10078 is resolved we can leave it to the client to trigger this rebalance
            log.info("Requesting followup rebalance be scheduled immediately by {} due to tasks changing ownership.", consumer);
            info.setNextRebalanceTime(0L);
            followupRebalanceRequiredForRevokedTasks = true;
            // Don't bother to schedule a probing rebalance if an immediate one is already scheduled
            shouldEncodeProbingRebalance = false;
        } else if (shouldEncodeProbingRebalance) {
            final long nextRebalanceTimeMs = time.milliseconds() + probingRebalanceIntervalMs();
            log.info("Requesting followup rebalance be scheduled by {} for {} ms to probe for caught-up replica tasks.", consumer, nextRebalanceTimeMs);
            info.setNextRebalanceTime(nextRebalanceTimeMs);
            shouldEncodeProbingRebalance = false;
        }
        assignment.put(consumer, new Assignment(activePartitionsList, info.encode()));
    }
    return followupRebalanceRequiredForRevokedTasks;
}

Example 38 with AssignmentInfo

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

the class StreamsPartitionAssignor method getActiveTasks.

// protected for upgrade test
protected static Map<TaskId, Set<TopicPartition>> getActiveTasks(final List<TopicPartition> partitions, final AssignmentInfo info) {
    final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
    for (int i = 0; i < partitions.size(); i++) {
        final TopicPartition partition = partitions.get(i);
        final TaskId id = info.activeTasks().get(i);
        activeTasks.computeIfAbsent(id, k1 -> new HashSet<>()).add(partition);
    }
    return activeTasks;
}

Example 39 with AssignmentInfo

use of org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo 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 40 with AssignmentInfo

use of org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo 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());
}