root
/
kafka
mirror of https://github.com/apache/kafka.git
1
0
Fork 0
Code Issues Actions 2 Packages Projects Releases Wiki Activity

KAFKA-16107 [2/N]: Fix inverted args & more tests (#15285) 

Small fix for argument order in onPartitionsAssigned callback, and more tests for disabling/enabling partitions while callback runs.

Reviewers: David Jacot <djacot@confluent.io>
This commit is contained in:
Lianet Magrans 2024-02-01 11:02:24 -05:00 committed by GitHub
parent 7bc4afee11
commit 8da6508966
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 75 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
clients/src/main/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImpl.java
View File

@ -1130,8 +1130,8 @@ public class MembershipManagerImpl implements MembershipManager, ClusterResource
// Invoke user call back.
CompletableFuture<Void> result = invokeOnPartitionsAssignedCallback(addedPartitions);
result.whenComplete((error, callbackResult) -> {
if (error == null) {
result.whenComplete((__, exception) -> {
if (exception == null) {
// Enable newly added partitions to start fetching and updating positions for them.
subscriptions.enablePartitionsAwaitingCallback(addedPartitions);
} else {
@ -1141,7 +1141,7 @@ public class MembershipManagerImpl implements MembershipManager, ClusterResource
if (!addedPartitions.isEmpty()) {
log.warn("Leaving newly assigned partitions {} marked as non-fetchable and not " +
"requiring initializing positions after onPartitionsAssigned callback failed.",
addedPartitions, error);
addedPartitions, exception);
}
}
});

72
clients/src/test/java/org/apache/kafka/clients/consumer/internals/MembershipManagerImplTest.java
View File

@ -1383,6 +1383,61 @@ public class MembershipManagerImplTest {
assertEquals(0, listener.lostCount());
}
@Test
public void testAddedPartitionsTemporarilyDisabledAwaitingOnPartitionsAssignedCallback() {
MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
String topicName = "topic1";
ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
int partitionOwned = 0;
int partitionAdded = 1;
SortedSet<TopicPartition> assignedPartitions = topicPartitions(topicName, partitionOwned,
partitionAdded);
SortedSet<TopicPartition> addedPartitions = topicPartitions(topicName, partitionAdded);
mockPartitionOwnedAndNewPartitionAdded(topicName, partitionOwned, partitionAdded,
new CounterConsumerRebalanceListener(), membershipManager);
verify(subscriptionState).assignFromSubscribedAwaitingCallback(assignedPartitions, addedPartitions);
performCallback(
membershipManager,
invoker,
ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
addedPartitions,
true
);
verify(subscriptionState).enablePartitionsAwaitingCallback(addedPartitions);
}
@Test
public void testAddedPartitionsNotEnabledAfterFailedOnPartitionsAssignedCallback() {
MembershipManagerImpl membershipManager = createMembershipManagerJoiningGroup();
String topicName = "topic1";
ConsumerRebalanceListenerInvoker invoker = consumerRebalanceListenerInvoker();
int partitionOwned = 0;
int partitionAdded = 1;
SortedSet<TopicPartition> assignedPartitions = topicPartitions(topicName, partitionOwned,
partitionAdded);
SortedSet<TopicPartition> addedPartitions = topicPartitions(topicName, partitionAdded);
CounterConsumerRebalanceListener listener =
new CounterConsumerRebalanceListener(Optional.empty(),
Optional.of(new RuntimeException("onPartitionsAssigned failed!")),
Optional.empty());
mockPartitionOwnedAndNewPartitionAdded(topicName, partitionOwned, partitionAdded,
listener, membershipManager);
verify(subscriptionState).assignFromSubscribedAwaitingCallback(assignedPartitions, addedPartitions);
performCallback(
membershipManager,
invoker,
ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED,
addedPartitions,
true
);
verify(subscriptionState, never()).enablePartitionsAwaitingCallback(any());
}
@Test
public void testOnPartitionsLostNoError() {
MembershipManagerImpl membershipManager = createMemberInStableState();
@ -1401,6 +1456,23 @@ public class MembershipManagerImplTest {
testOnPartitionsLost(Optional.of(new KafkaException("Intentional error for test")));
}
private void mockPartitionOwnedAndNewPartitionAdded(String topicName,
int partitionOwned,
int partitionAdded,
CounterConsumerRebalanceListener listener,
MembershipManagerImpl membershipManager) {
Uuid topicId = Uuid.randomUuid();
TopicPartition owned = new TopicPartition(topicName, partitionOwned);
when(subscriptionState.assignedPartitions()).thenReturn(Collections.singleton(owned));
membershipManager.updateCurrentAssignment(Collections.singleton(new TopicIdPartition(topicId, owned)));
when(metadata.topicNames()).thenReturn(Collections.singletonMap(topicId, topicName));
when(subscriptionState.hasAutoAssignedPartitions()).thenReturn(true);
when(subscriptionState.rebalanceListener()).thenReturn(Optional.ofNullable(listener));
// Receive assignment adding a new partition
receiveAssignment(topicId, Arrays.asList(partitionOwned, partitionAdded), membershipManager);
}
private void testOnPartitionsLost(Optional<RuntimeException> lostError) {
// Step 1: set up mocks
MembershipManagerImpl membershipManager = createMemberInStableState();