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

Example 1 with ReferenceContainer

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

the class StreamsPartitionAssignor method configure.

/**
 * We need to have the PartitionAssignor and its StreamThread to be mutually accessible since the former needs
 * latter's cached metadata while sending subscriptions, and the latter needs former's returned assignment when
 * adding tasks.
 *
 * @throws KafkaException if the stream thread is not specified
 */
@Override
public void configure(final Map<String, ?> configs) {
    final AssignorConfiguration assignorConfiguration = new AssignorConfiguration(configs);
    logPrefix = assignorConfiguration.logPrefix();
    log = new LogContext(logPrefix).logger(getClass());
    usedSubscriptionMetadataVersion = assignorConfiguration.configuredMetadataVersion(usedSubscriptionMetadataVersion);
    final ReferenceContainer referenceContainer = assignorConfiguration.referenceContainer();
    mainConsumerSupplier = () -> Objects.requireNonNull(referenceContainer.mainConsumer, "Main consumer was not specified");
    adminClient = Objects.requireNonNull(referenceContainer.adminClient, "Admin client was not specified");
    taskManager = Objects.requireNonNull(referenceContainer.taskManager, "TaskManager was not specified");
    streamsMetadataState = Objects.requireNonNull(referenceContainer.streamsMetadataState, "StreamsMetadataState was not specified");
    assignmentErrorCode = referenceContainer.assignmentErrorCode;
    nextScheduledRebalanceMs = referenceContainer.nextScheduledRebalanceMs;
    nonFatalExceptionsToHandle = referenceContainer.nonFatalExceptionsToHandle;
    time = Objects.requireNonNull(referenceContainer.time, "Time was not specified");
    assignmentConfigs = assignorConfiguration.assignmentConfigs();
    partitionGrouper = new PartitionGrouper();
    userEndPoint = assignorConfiguration.userEndPoint();
    internalTopicManager = assignorConfiguration.internalTopicManager();
    copartitionedTopicsEnforcer = assignorConfiguration.copartitionedTopicsEnforcer();
    rebalanceProtocol = assignorConfiguration.rebalanceProtocol();
    taskAssignorSupplier = assignorConfiguration::taskAssignor;
    assignmentListener = assignorConfiguration.assignmentListener();
    uniqueField = 0;
}

Example 2 with ReferenceContainer

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

the class StreamThread method create.

public static StreamThread create(final TopologyMetadata topologyMetadata, final StreamsConfig config, final KafkaClientSupplier clientSupplier, final Admin adminClient, final UUID processId, final String clientId, final StreamsMetricsImpl streamsMetrics, final Time time, final StreamsMetadataState streamsMetadataState, final long cacheSizeBytes, final StateDirectory stateDirectory, final StateRestoreListener userStateRestoreListener, final int threadIdx, final Runnable shutdownErrorHook, final BiConsumer<Throwable, Boolean> streamsUncaughtExceptionHandler) {
    final String threadId = clientId + "-StreamThread-" + threadIdx;
    final String logPrefix = String.format("stream-thread [%s] ", threadId);
    final LogContext logContext = new LogContext(logPrefix);
    final Logger log = logContext.logger(StreamThread.class);
    final ReferenceContainer referenceContainer = new ReferenceContainer();
    referenceContainer.adminClient = adminClient;
    referenceContainer.streamsMetadataState = streamsMetadataState;
    referenceContainer.time = time;
    log.info("Creating restore consumer client");
    final Map<String, Object> restoreConsumerConfigs = config.getRestoreConsumerConfigs(getRestoreConsumerClientId(threadId));
    final Consumer<byte[], byte[]> restoreConsumer = clientSupplier.getRestoreConsumer(restoreConsumerConfigs);
    final StoreChangelogReader changelogReader = new StoreChangelogReader(time, config, logContext, adminClient, restoreConsumer, userStateRestoreListener);
    final ThreadCache cache = new ThreadCache(logContext, cacheSizeBytes, streamsMetrics);
    final ActiveTaskCreator activeTaskCreator = new ActiveTaskCreator(topologyMetadata, config, streamsMetrics, stateDirectory, changelogReader, cache, time, clientSupplier, threadId, processId, log);
    final StandbyTaskCreator standbyTaskCreator = new StandbyTaskCreator(topologyMetadata, config, streamsMetrics, stateDirectory, changelogReader, threadId, log);
    final TaskManager taskManager = new TaskManager(time, changelogReader, processId, logPrefix, streamsMetrics, activeTaskCreator, standbyTaskCreator, topologyMetadata, adminClient, stateDirectory);
    referenceContainer.taskManager = taskManager;
    log.info("Creating consumer client");
    final String applicationId = config.getString(StreamsConfig.APPLICATION_ID_CONFIG);
    final Map<String, Object> consumerConfigs = config.getMainConsumerConfigs(applicationId, getConsumerClientId(threadId), threadIdx);
    consumerConfigs.put(StreamsConfig.InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
    final String originalReset = (String) consumerConfigs.get(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG);
    // If there are any overrides, we never fall through to the consumer, but only handle offset management ourselves.
    if (topologyMetadata.hasOffsetResetOverrides()) {
        consumerConfigs.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "none");
    }
    final Consumer<byte[], byte[]> mainConsumer = clientSupplier.getConsumer(consumerConfigs);
    changelogReader.setMainConsumer(mainConsumer);
    taskManager.setMainConsumer(mainConsumer);
    referenceContainer.mainConsumer = mainConsumer;
    final StreamThread streamThread = new StreamThread(time, config, adminClient, mainConsumer, restoreConsumer, changelogReader, originalReset, taskManager, streamsMetrics, topologyMetadata, threadId, logContext, referenceContainer.assignmentErrorCode, referenceContainer.nextScheduledRebalanceMs, referenceContainer.nonFatalExceptionsToHandle, shutdownErrorHook, streamsUncaughtExceptionHandler, cache::resize);
    return streamThread.updateThreadMetadata(getSharedAdminClientId(clientId));
}

Example 3 with ReferenceContainer

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

the class HighAvailabilityStreamsPartitionAssignorTest method configProps.

private Map<String, Object> configProps() {
    final Map<String, Object> configurationMap = new HashMap<>();
    configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
    configurationMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, USER_END_POINT);
    referenceContainer = new ReferenceContainer();
    referenceContainer.mainConsumer = EasyMock.mock(Consumer.class);
    referenceContainer.adminClient = adminClient;
    referenceContainer.taskManager = taskManager;
    referenceContainer.streamsMetadataState = streamsMetadataState;
    referenceContainer.time = time;
    configurationMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
    return configurationMap;
}

Example 4 with ReferenceContainer

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