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

Example 1 with ConsumerPartitionAssignor

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

the class ConsumerCoordinator method onJoinComplete.

@Override
protected void onJoinComplete(int generation, String memberId, String assignmentStrategy, ByteBuffer assignmentBuffer) {
    log.debug("Executing onJoinComplete with generation {} and memberId {}", generation, memberId);
    // Only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
    if (!isLeader)
        assignmentSnapshot = null;
    ConsumerPartitionAssignor assignor = lookupAssignor(assignmentStrategy);
    if (assignor == null)
        throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
    // Give the assignor a chance to update internal state based on the received assignment
    groupMetadata = new ConsumerGroupMetadata(rebalanceConfig.groupId, generation, memberId, rebalanceConfig.groupInstanceId);
    SortedSet<TopicPartition> ownedPartitions = new TreeSet<>(COMPARATOR);
    ownedPartitions.addAll(subscriptions.assignedPartitions());
    // should at least encode the short version
    if (assignmentBuffer.remaining() < 2)
        throw new IllegalStateException("There are insufficient bytes available to read assignment from the sync-group response (" + "actual byte size " + assignmentBuffer.remaining() + ") , this is not expected; " + "it is possible that the leader's assign function is buggy and did not return any assignment for this member, " + "or because static member is configured and the protocol is buggy hence did not get the assignment for this member");
    Assignment assignment = ConsumerProtocol.deserializeAssignment(assignmentBuffer);
    SortedSet<TopicPartition> assignedPartitions = new TreeSet<>(COMPARATOR);
    assignedPartitions.addAll(assignment.partitions());
    if (!subscriptions.checkAssignmentMatchedSubscription(assignedPartitions)) {
        final String reason = String.format("received assignment %s does not match the current subscription %s; " + "it is likely that the subscription has changed since we joined the group, will re-join with current subscription", assignment.partitions(), subscriptions.prettyString());
        requestRejoin(reason);
        return;
    }
    final AtomicReference<Exception> firstException = new AtomicReference<>(null);
    SortedSet<TopicPartition> addedPartitions = new TreeSet<>(COMPARATOR);
    addedPartitions.addAll(assignedPartitions);
    addedPartitions.removeAll(ownedPartitions);
    if (protocol == RebalanceProtocol.COOPERATIVE) {
        SortedSet<TopicPartition> revokedPartitions = new TreeSet<>(COMPARATOR);
        revokedPartitions.addAll(ownedPartitions);
        revokedPartitions.removeAll(assignedPartitions);
        log.info("Updating assignment with\n" + "\tAssigned partitions:                       {}\n" + "\tCurrent owned partitions:                  {}\n" + "\tAdded partitions (assigned - owned):       {}\n" + "\tRevoked partitions (owned - assigned):     {}\n", assignedPartitions, ownedPartitions, addedPartitions, revokedPartitions);
        if (!revokedPartitions.isEmpty()) {
            // Revoke partitions that were previously owned but no longer assigned;
            // note that we should only change the assignment (or update the assignor's state)
            // AFTER we've triggered  the revoke callback
            firstException.compareAndSet(null, invokePartitionsRevoked(revokedPartitions));
            // If revoked any partitions, need to re-join the group afterwards
            final String reason = String.format("need to revoke partitions %s as indicated " + "by the current assignment and re-join", revokedPartitions);
            requestRejoin(reason);
        }
    }
    // The leader may have assigned partitions which match our subscription pattern, but which
    // were not explicitly requested, so we update the joined subscription here.
    maybeUpdateJoinedSubscription(assignedPartitions);
    // Catch any exception here to make sure we could complete the user callback.
    firstException.compareAndSet(null, invokeOnAssignment(assignor, assignment));
    // Reschedule the auto commit starting from now
    if (autoCommitEnabled)
        this.nextAutoCommitTimer.updateAndReset(autoCommitIntervalMs);
    subscriptions.assignFromSubscribed(assignedPartitions);
    // Add partitions that were not previously owned but are now assigned
    firstException.compareAndSet(null, invokePartitionsAssigned(addedPartitions));
    if (firstException.get() != null) {
        if (firstException.get() instanceof KafkaException) {
            throw (KafkaException) firstException.get();
        } else {
            throw new KafkaException("User rebalance callback throws an error", firstException.get());
        }
    }
}

Example 2 with ConsumerPartitionAssignor

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

the class ConsumerCoordinatorTest method testSelectRebalanceProtcol.

@Test
public void testSelectRebalanceProtcol() {
    List<ConsumerPartitionAssignor> assignors = new ArrayList<>();
    assignors.add(new MockPartitionAssignor(Collections.singletonList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER)));
    assignors.add(new MockPartitionAssignor(Collections.singletonList(COOPERATIVE)));
    // no commonly supported protocols
    assertThrows(IllegalArgumentException.class, () -> buildCoordinator(rebalanceConfig, new Metrics(), assignors, false, subscriptions));
    assignors.clear();
    assignors.add(new MockPartitionAssignor(Arrays.asList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER, COOPERATIVE)));
    assignors.add(new MockPartitionAssignor(Arrays.asList(ConsumerPartitionAssignor.RebalanceProtocol.EAGER, COOPERATIVE)));
    // select higher indexed (more advanced) protocols
    try (ConsumerCoordinator coordinator = buildCoordinator(rebalanceConfig, new Metrics(), assignors, false, subscriptions)) {
        assertEquals(COOPERATIVE, coordinator.getProtocol());
    }
}

Example 3 with ConsumerPartitionAssignor

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

the class ConsumerCoordinator method onLeaderElected.

@Override
protected Map<String, ByteBuffer> onLeaderElected(String leaderId, String assignmentStrategy, List<JoinGroupResponseData.JoinGroupResponseMember> allSubscriptions, boolean skipAssignment) {
    ConsumerPartitionAssignor assignor = lookupAssignor(assignmentStrategy);
    if (assignor == null)
        throw new IllegalStateException("Coordinator selected invalid assignment protocol: " + assignmentStrategy);
    String assignorName = assignor.name();
    Set<String> allSubscribedTopics = new HashSet<>();
    Map<String, Subscription> subscriptions = new HashMap<>();
    // collect all the owned partitions
    Map<String, List<TopicPartition>> ownedPartitions = new HashMap<>();
    for (JoinGroupResponseData.JoinGroupResponseMember memberSubscription : allSubscriptions) {
        Subscription subscription = ConsumerProtocol.deserializeSubscription(ByteBuffer.wrap(memberSubscription.metadata()));
        subscription.setGroupInstanceId(Optional.ofNullable(memberSubscription.groupInstanceId()));
        subscriptions.put(memberSubscription.memberId(), subscription);
        allSubscribedTopics.addAll(subscription.topics());
        ownedPartitions.put(memberSubscription.memberId(), subscription.ownedPartitions());
    }
    // the leader will begin watching for changes to any of the topics the group is interested in,
    // which ensures that all metadata changes will eventually be seen
    updateGroupSubscription(allSubscribedTopics);
    isLeader = true;
    if (skipAssignment) {
        log.info("Skipped assignment for returning static leader at generation {}. The static leader " + "will continue with its existing assignment.", generation().generationId);
        assignmentSnapshot = metadataSnapshot;
        return Collections.emptyMap();
    }
    log.debug("Performing assignment using strategy {} with subscriptions {}", assignorName, subscriptions);
    Map<String, Assignment> assignments = assignor.assign(metadata.fetch(), new GroupSubscription(subscriptions)).groupAssignment();
    // the "generation" of ownedPartition inside the assignor
    if (protocol == RebalanceProtocol.COOPERATIVE && !assignorName.equals(COOPERATIVE_STICKY_ASSIGNOR_NAME)) {
        validateCooperativeAssignment(ownedPartitions, assignments);
    }
    maybeUpdateGroupSubscription(assignorName, assignments, allSubscribedTopics);
    // metadataSnapshot could be updated when the subscription is updated therefore
    // we must take the assignment snapshot after.
    assignmentSnapshot = metadataSnapshot;
    log.info("Finished assignment for group at generation {}: {}", generation().generationId, assignments);
    Map<String, ByteBuffer> groupAssignment = new HashMap<>();
    for (Map.Entry<String, Assignment> assignmentEntry : assignments.entrySet()) {
        ByteBuffer buffer = ConsumerProtocol.serializeAssignment(assignmentEntry.getValue());
        groupAssignment.put(assignmentEntry.getKey(), buffer);
    }
    return groupAssignment;
}

Example 4 with ConsumerPartitionAssignor

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

the class ConsumerCoordinator method metadata.

@Override
protected JoinGroupRequestData.JoinGroupRequestProtocolCollection metadata() {
    log.debug("Joining group with current subscription: {}", subscriptions.subscription());
    this.joinedSubscription = subscriptions.subscription();
    JoinGroupRequestData.JoinGroupRequestProtocolCollection protocolSet = new JoinGroupRequestData.JoinGroupRequestProtocolCollection();
    List<String> topics = new ArrayList<>(joinedSubscription);
    for (ConsumerPartitionAssignor assignor : assignors) {
        Subscription subscription = new Subscription(topics, assignor.subscriptionUserData(joinedSubscription), subscriptions.assignedPartitionsList());
        ByteBuffer metadata = ConsumerProtocol.serializeSubscription(subscription);
        protocolSet.add(new JoinGroupRequestData.JoinGroupRequestProtocol().setName(assignor.name()).setMetadata(Utils.toArray(metadata)));
    }
    return protocolSet;
}

Example 5 with ConsumerPartitionAssignor

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

the class TaskAssignorIntegrationTest method shouldProperlyConfigureTheAssignor.

@SuppressWarnings("unchecked")
@Test
public void shouldProperlyConfigureTheAssignor() throws NoSuchFieldException, IllegalAccessException {
    // This test uses reflection to check and make sure that all the expected configurations really
    // make it all the way to configure the task assignor. There's no other use case for being able
    // to extract all these fields, so reflection is a good choice until we find that the maintenance
    // burden is too high.
    // 
    // Also note that this is an integration test because so many components have to come together to
    // ensure these configurations wind up where they belong, and any number of future code changes
    // could break this change.
    final String testId = safeUniqueTestName(getClass(), testName);
    final String appId = "appId_" + testId;
    final String inputTopic = "input" + testId;
    IntegrationTestUtils.cleanStateBeforeTest(CLUSTER, inputTopic);
    // Maybe I'm paranoid, but I don't want the compiler deciding that my lambdas are equal to the identity
    // function and defeating my identity check
    final AtomicInteger compilerDefeatingReference = new AtomicInteger(0);
    // the implementation doesn't matter, we're just going to verify the reference.
    final AssignmentListener configuredAssignmentListener = stable -> compilerDefeatingReference.incrementAndGet();
    final Properties properties = mkObjectProperties(mkMap(mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers()), mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, appId), mkEntry(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath()), mkEntry(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, "5"), mkEntry(StreamsConfig.ACCEPTABLE_RECOVERY_LAG_CONFIG, "6"), mkEntry(StreamsConfig.MAX_WARMUP_REPLICAS_CONFIG, "7"), mkEntry(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG, "480000"), mkEntry(StreamsConfig.InternalConfig.ASSIGNMENT_LISTENER, configuredAssignmentListener), mkEntry(StreamsConfig.InternalConfig.INTERNAL_TASK_ASSIGNOR_CLASS, MyTaskAssignor.class.getName())));
    final StreamsBuilder builder = new StreamsBuilder();
    builder.stream(inputTopic);
    try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), properties)) {
        kafkaStreams.start();
        final Field threads = KafkaStreams.class.getDeclaredField("threads");
        threads.setAccessible(true);
        final List<StreamThread> streamThreads = (List<StreamThread>) threads.get(kafkaStreams);
        final StreamThread streamThread = streamThreads.get(0);
        final Field mainConsumer = StreamThread.class.getDeclaredField("mainConsumer");
        mainConsumer.setAccessible(true);
        final KafkaConsumer<?, ?> consumer = (KafkaConsumer<?, ?>) mainConsumer.get(streamThread);
        final Field assignors = KafkaConsumer.class.getDeclaredField("assignors");
        assignors.setAccessible(true);
        final List<ConsumerPartitionAssignor> consumerPartitionAssignors = (List<ConsumerPartitionAssignor>) assignors.get(consumer);
        final StreamsPartitionAssignor streamsPartitionAssignor = (StreamsPartitionAssignor) consumerPartitionAssignors.get(0);
        final Field assignmentConfigs = StreamsPartitionAssignor.class.getDeclaredField("assignmentConfigs");
        assignmentConfigs.setAccessible(true);
        final AssignorConfiguration.AssignmentConfigs configs = (AssignorConfiguration.AssignmentConfigs) assignmentConfigs.get(streamsPartitionAssignor);
        final Field assignmentListenerField = StreamsPartitionAssignor.class.getDeclaredField("assignmentListener");
        assignmentListenerField.setAccessible(true);
        final AssignmentListener actualAssignmentListener = (AssignmentListener) assignmentListenerField.get(streamsPartitionAssignor);
        final Field taskAssignorSupplierField = StreamsPartitionAssignor.class.getDeclaredField("taskAssignorSupplier");
        taskAssignorSupplierField.setAccessible(true);
        final Supplier<TaskAssignor> taskAssignorSupplier = (Supplier<TaskAssignor>) taskAssignorSupplierField.get(streamsPartitionAssignor);
        final TaskAssignor taskAssignor = taskAssignorSupplier.get();
        assertThat(configs.numStandbyReplicas, is(5));
        assertThat(configs.acceptableRecoveryLag, is(6L));
        assertThat(configs.maxWarmupReplicas, is(7));
        assertThat(configs.probingRebalanceIntervalMs, is(480000L));
        assertThat(actualAssignmentListener, sameInstance(configuredAssignmentListener));
        assertThat(taskAssignor, instanceOf(MyTaskAssignor.class));
    }
}