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

Example 11 with Assignment

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

the class StreamsPartitionAssignorTest method shouldReturnLowestAssignmentVersionForDifferentSubscriptionVersions.

private void shouldReturnLowestAssignmentVersionForDifferentSubscriptionVersions(final int smallestVersion, final int otherVersion) {
    subscriptions.put("consumer1", new Subscription(Collections.singletonList("topic1"), getInfoForOlderVersion(smallestVersion, UUID_1, EMPTY_TASKS, EMPTY_TASKS).encode()));
    subscriptions.put("consumer2", new Subscription(Collections.singletonList("topic1"), getInfoForOlderVersion(otherVersion, UUID_2, EMPTY_TASKS, EMPTY_TASKS).encode()));
    configureDefault();
    final Map<String, Assignment> assignment = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    assertThat(assignment.size(), equalTo(2));
    assertThat(AssignmentInfo.decode(assignment.get("consumer1").userData()).version(), equalTo(smallestVersion));
    assertThat(AssignmentInfo.decode(assignment.get("consumer2").userData()).version(), equalTo(smallestVersion));
}

Example 12 with Assignment

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

the class HighAvailabilityStreamsPartitionAssignorTest method shouldReturnAllActiveTasksToPreviousOwnerRegardlessOfBalanceAndTriggerRebalanceIfEndOffsetFetchFailsAndHighAvailabilityEnabled.

@Test
public void shouldReturnAllActiveTasksToPreviousOwnerRegardlessOfBalanceAndTriggerRebalanceIfEndOffsetFetchFailsAndHighAvailabilityEnabled() {
    final long rebalanceInterval = 5 * 60 * 1000L;
    builder.addSource(null, "source1", null, null, null, "topic1");
    builder.addProcessor("processor1", new MockApiProcessorSupplier<>(), "source1");
    builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1");
    final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2);
    createMockTaskManager(allTasks);
    adminClient = EasyMock.createMock(AdminClient.class);
    expect(adminClient.listOffsets(anyObject())).andThrow(new StreamsException("Should be handled"));
    configurePartitionAssignorWith(singletonMap(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG, rebalanceInterval));
    final String firstConsumer = "consumer1";
    final String newConsumer = "consumer2";
    subscriptions.put(firstConsumer, new Subscription(singletonList("source1"), getInfo(UUID_1, allTasks).encode()));
    subscriptions.put(newConsumer, new Subscription(singletonList("source1"), getInfo(UUID_2, EMPTY_TASKS).encode()));
    final Map<String, Assignment> assignments = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    final AssignmentInfo firstConsumerUserData = AssignmentInfo.decode(assignments.get(firstConsumer).userData());
    final List<TaskId> firstConsumerActiveTasks = firstConsumerUserData.activeTasks();
    final AssignmentInfo newConsumerUserData = AssignmentInfo.decode(assignments.get(newConsumer).userData());
    final List<TaskId> newConsumerActiveTasks = newConsumerUserData.activeTasks();
    // The tasks were returned to their prior owner
    final ArrayList<TaskId> sortedExpectedTasks = new ArrayList<>(allTasks);
    Collections.sort(sortedExpectedTasks);
    assertThat(firstConsumerActiveTasks, equalTo(sortedExpectedTasks));
    assertThat(newConsumerActiveTasks, empty());
    // There is a rebalance scheduled
    assertThat(time.milliseconds() + rebalanceInterval, anyOf(is(firstConsumerUserData.nextRebalanceMs()), is(newConsumerUserData.nextRebalanceMs())));
}

Example 13 with Assignment

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

the class HighAvailabilityStreamsPartitionAssignorTest method shouldScheduleProbingRebalanceOnThisClientIfWarmupTasksRequired.

@Test
public void shouldScheduleProbingRebalanceOnThisClientIfWarmupTasksRequired() {
    final long rebalanceInterval = 5 * 60 * 1000L;
    builder.addSource(null, "source1", null, null, null, "topic1");
    builder.addProcessor("processor1", new MockApiProcessorSupplier<>(), "source1");
    builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1");
    final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2);
    createMockTaskManager(allTasks);
    createMockAdminClient(getTopicPartitionOffsetsMap(singletonList(APPLICATION_ID + "-store1-changelog"), singletonList(3)));
    configurePartitionAssignorWith(singletonMap(StreamsConfig.PROBING_REBALANCE_INTERVAL_MS_CONFIG, rebalanceInterval));
    final String firstConsumer = "consumer1";
    final String newConsumer = "consumer2";
    subscriptions.put(firstConsumer, new Subscription(singletonList("source1"), getInfo(UUID_1, allTasks).encode()));
    subscriptions.put(newConsumer, new Subscription(singletonList("source1"), getInfo(UUID_2, EMPTY_TASKS).encode()));
    final Map<String, Assignment> assignments = partitionAssignor.assign(metadata, new GroupSubscription(subscriptions)).groupAssignment();
    final List<TaskId> firstConsumerActiveTasks = AssignmentInfo.decode(assignments.get(firstConsumer).userData()).activeTasks();
    final List<TaskId> newConsumerActiveTasks = AssignmentInfo.decode(assignments.get(newConsumer).userData()).activeTasks();
    final ArrayList<TaskId> sortedExpectedTasks = new ArrayList<>(allTasks);
    Collections.sort(sortedExpectedTasks);
    assertThat(firstConsumerActiveTasks, equalTo(sortedExpectedTasks));
    assertThat(newConsumerActiveTasks, empty());
    assertThat(referenceContainer.assignmentErrorCode.get(), equalTo(AssignorError.NONE.code()));
    final long nextScheduledRebalanceOnThisClient = AssignmentInfo.decode(assignments.get(firstConsumer).userData()).nextRebalanceMs();
    final long nextScheduledRebalanceOnOtherClient = AssignmentInfo.decode(assignments.get(newConsumer).userData()).nextRebalanceMs();
    assertThat(nextScheduledRebalanceOnThisClient, equalTo(time.milliseconds() + rebalanceInterval));
    assertThat(nextScheduledRebalanceOnOtherClient, equalTo(Long.MAX_VALUE));
}

Example 14 with Assignment

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

the class ConsumerCoordinator method validateCooperativeAssignment.

/**
 * Used by COOPERATIVE rebalance protocol only.
 *
 * Validate the assignments returned by the assignor such that no owned partitions are going to
 * be reassigned to a different consumer directly: if the assignor wants to reassign an owned partition,
 * it must first remove it from the new assignment of the current owner so that it is not assigned to any
 * member, and then in the next rebalance it can finally reassign those partitions not owned by anyone to consumers.
 */
private void validateCooperativeAssignment(final Map<String, List<TopicPartition>> ownedPartitions, final Map<String, Assignment> assignments) {
    Set<TopicPartition> totalRevokedPartitions = new HashSet<>();
    SortedSet<TopicPartition> totalAddedPartitions = new TreeSet<>(COMPARATOR);
    for (final Map.Entry<String, Assignment> entry : assignments.entrySet()) {
        final Assignment assignment = entry.getValue();
        final Set<TopicPartition> addedPartitions = new HashSet<>(assignment.partitions());
        addedPartitions.removeAll(ownedPartitions.get(entry.getKey()));
        final Set<TopicPartition> revokedPartitions = new HashSet<>(ownedPartitions.get(entry.getKey()));
        revokedPartitions.removeAll(assignment.partitions());
        totalAddedPartitions.addAll(addedPartitions);
        totalRevokedPartitions.addAll(revokedPartitions);
    }
    // if there are overlap between revoked partitions and added partitions, it means some partitions
    // immediately gets re-assigned to another member while it is still claimed by some member
    totalAddedPartitions.retainAll(totalRevokedPartitions);
    if (!totalAddedPartitions.isEmpty()) {
        log.error("With the COOPERATIVE protocol, owned partitions cannot be " + "reassigned to other members; however the assignor has reassigned partitions {} which are still owned " + "by some members", totalAddedPartitions);
        throw new IllegalStateException("Assignor supporting the COOPERATIVE protocol violates its requirements");
    }
}

Example 15 with Assignment

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