lianetm commented on code in PR #20511: URL: https://github.com/apache/kafka/pull/20511#discussion_r2337642213
########## clients/src/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java: ########## @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals; + +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.WakeupException; +import org.apache.kafka.common.utils.LogContext; + +import org.slf4j.Logger; + +import java.util.Optional; +import java.util.Set; + +/** + * This class encapsulates the invocation of the callback methods defined in the {@link StreamsRebalanceListener} + * interface. When streams group task assignment changes, these methods are invoked. This class wraps those + * callback calls with some logging and error handling. + */ +public class StreamsRebalanceListenerInvoker { + + private final Logger log; + + private final StreamsRebalanceData streamsRebalanceData; + private Optional<StreamsRebalanceListener> listener; + + StreamsRebalanceListenerInvoker(LogContext logContext, StreamsRebalanceData streamsRebalanceData) { + this.log = logContext.logger(getClass()); + this.listener = Optional.empty(); + this.streamsRebalanceData = streamsRebalanceData; + } + + public void setRebalanceListener(StreamsRebalanceListener streamsRebalanceListener) { + this.listener = Optional.ofNullable(streamsRebalanceListener); + } + + public Exception invokeAllTasksRevoked() { + if (listener.isPresent()) { + return invokeTasksRevoked(streamsRebalanceData.reconciledAssignment().activeTasks()); + } + + return null; + } + + public Exception invokeTasksAssigned(final StreamsRebalanceData.Assignment assignment) { + if (listener.isPresent()) { + log.info("Adding newly assigned tasks: {}", assignment); Review Comment: This log seems a bit confusing here on the callback invocations. Shouldn't it be more along the lines of "trigger task assigned callbacks" (maybe is me misreading, missing KS details. Ex. on the consumer side the assigned callback runs when the partitions were already assigned/added). If this applies, let's align also the logs for revoked/lost ########## clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java: ########## @@ -278,44 +269,28 @@ private void processStreamsOnAllTasksLostCallbackNeededEvent(final StreamsOnAllT private StreamsOnTasksRevokedCallbackCompletedEvent invokeOnTasksRevokedCallback(final Set<StreamsRebalanceData.TaskId> activeTasksToRevoke, final CompletableFuture<Void> future) { - final Optional<Exception> exceptionFromCallback = streamsRebalanceListener().onTasksRevoked(activeTasksToRevoke); + final Optional<Exception> exceptionFromCallback = Optional.ofNullable(streamsRebalanceListenerInvoker().invokeTasksRevoked(activeTasksToRevoke)); final Optional<KafkaException> error = exceptionFromCallback.map(e -> ConsumerUtils.maybeWrapAsKafkaException(e, "Task revocation callback throws an error")); Review Comment: could we maybe simplify by pushing the optionals and maybeWrap logic to the completedEvent constructors? Keeping here simply ``` Exception error = invokeTasksRevoked(...) return new StreamsOnTasksRevokedCallbackCompletedEvent(future, error) // and then in the constructor is where magic happens to shape that error, have a custom msg, etc. // Applied consistently for the 3 StreamsOnTasksXXXCallbackCompletedEvent ``` Just trying to detangle all this optionals a bit ########## clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java: ########## @@ -1964,7 +1957,11 @@ public void subscribe(Collection<String> topics, ConsumerRebalanceListener liste public void subscribe(Collection<String> topics, StreamsRebalanceListener streamsRebalanceListener) { subscribeInternal(topics, Optional.empty()); - backgroundEventProcessor.setStreamsRebalanceListener(streamsRebalanceListener); + if (streamsRebalanceListenerInvoker.isPresent()) { + streamsRebalanceListenerInvoker.get().setRebalanceListener(streamsRebalanceListener); + } else { + throw new IllegalStateException("Consumer was not created to be used with Streams rebalance protocol events"); + } Review Comment: I would expect is probably best/clearer to do this setup/validation before any actual subscription action (`subscribeInternal`)? ########## clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java: ########## @@ -2210,6 +2210,75 @@ private void markOffsetsReadyForCommitEvent() { }).when(applicationEventHandler).add(ArgumentMatchers.isA(CommitEvent.class)); } + @Test + public void testCloseInvokesStreamsRebalanceListenerOnTasksRevokedWhenMemberEpochPositive() { + final String groupId = "streamsGroup"; + final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()); + + try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) { + consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData); + StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class); + when(mockStreamsListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + consumer.subscribe(singletonList("topic"), mockStreamsListener); + final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers); + final int memberEpoch = 42; + final String memberId = "memberId"; + groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId); + + consumer.close(CloseOptions.timeout(Duration.ZERO)); + + verify(mockStreamsListener).onTasksRevoked(any()); + } + } + + @Test + public void testCloseInvokesStreamsRebalanceListenerOnAllTasksLostWhenMemberEpochZeroOrNegative() { + // Test that close() calls streamsRebalanceListener.invokeAllTasksLost() when memberEpoch <= 0 Review Comment: this comment seems to state the same as the func name really, needed? ########## clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java: ########## @@ -0,0 +1,337 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals; + +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.WakeupException; +import org.apache.kafka.common.utils.LogContext; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; + +import java.util.Optional; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.STRICT_STUBS) +public class StreamsRebalanceListenerInvokerTest { + + @Mock + private StreamsRebalanceListener mockListener; + + @Mock + private StreamsRebalanceData streamsRebalanceData; + + private StreamsRebalanceListenerInvoker invoker; + private final LogContext logContext = new LogContext(); + + @BeforeEach + public void setup() { + invoker = new StreamsRebalanceListenerInvoker(logContext, streamsRebalanceData); + } + + @Test + public void testConstructorInitializesWithEmptyListener() { + // When invoker is constructed, it should have no listener set initially + // This is verified by testing that invoke methods return null when no listener is present + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListener() { + // Test setting a listener + invoker.setRebalanceListener(mockListener); + + // Verify listener is set by checking that methods no longer return null immediately + // (we'll mock the dependencies needed for actual invocation) + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(mockListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Should now invoke the listener instead of returning null immediately + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); // No exception thrown by mock listener + verify(mockListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testSetRebalanceListenerWithNull() { + // Test setting listener to null + invoker.setRebalanceListener(null); + + // Should behave as if no listener is set + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListenerOverwritesExisting() { + StreamsRebalanceListener firstListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + StreamsRebalanceListener secondListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(secondListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Set first listener + invoker.setRebalanceListener(firstListener); + + // Overwrite with second listener + invoker.setRebalanceListener(secondListener); + + // Should use second listener + invoker.invokeAllTasksRevoked(); + verify(firstListener, never()).onTasksRevoked(any()); + verify(secondListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testInvokeAllTasksRevokedWithNoListener() { + // When no listener is set, should return null + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); + } + + @Test + public void testInvokeAllTasksRevokedWithListener() { + invoker.setRebalanceListener(mockListener); + + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(mockListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + Exception result = invoker.invokeAllTasksRevoked(); + + assertNull(result); + verify(mockListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } Review Comment: This seems to be the same as `testSetRebalanceListener`? (I do like this name better though) ########## clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java: ########## @@ -2210,6 +2210,75 @@ private void markOffsetsReadyForCommitEvent() { }).when(applicationEventHandler).add(ArgumentMatchers.isA(CommitEvent.class)); } + @Test + public void testCloseInvokesStreamsRebalanceListenerOnTasksRevokedWhenMemberEpochPositive() { + final String groupId = "streamsGroup"; + final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()); + + try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) { + consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData); + StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class); + when(mockStreamsListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + consumer.subscribe(singletonList("topic"), mockStreamsListener); + final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers); + final int memberEpoch = 42; + final String memberId = "memberId"; + groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId); + + consumer.close(CloseOptions.timeout(Duration.ZERO)); + + verify(mockStreamsListener).onTasksRevoked(any()); + } + } + + @Test + public void testCloseInvokesStreamsRebalanceListenerOnAllTasksLostWhenMemberEpochZeroOrNegative() { + // Test that close() calls streamsRebalanceListener.invokeAllTasksLost() when memberEpoch <= 0 + final String groupId = "streamsGroup"; + final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()); + + try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) { + consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData); + StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class); + when(mockStreamsListener.onAllTasksLost()).thenReturn(Optional.empty()); + consumer.subscribe(singletonList("topic"), mockStreamsListener); + final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers); + final int memberEpoch = 0; + final String memberId = "memberId"; + groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId); + + consumer.close(CloseOptions.timeout(Duration.ZERO)); + + verify(mockStreamsListener).onAllTasksLost(); + } + } + + @Test + public void testCloseWrapsStreamsRebalanceListenerException() { + final String groupId = "streamsGroup"; + final StreamsRebalanceData streamsRebalanceData = new StreamsRebalanceData(UUID.randomUUID(), Optional.empty(), Map.of(), Map.of()); + + try (final MockedStatic<RequestManagers> requestManagers = mockStatic(RequestManagers.class)) { + consumer = newConsumerWithStreamRebalanceData(requiredConsumerConfigAndGroupId(groupId), streamsRebalanceData); + StreamsRebalanceListener mockStreamsListener = mock(StreamsRebalanceListener.class); + RuntimeException testException = new RuntimeException("Test streams listener exception"); + doThrow(testException).when(mockStreamsListener).onTasksRevoked(any()); + consumer.subscribe(singletonList("topic"), mockStreamsListener); + final MemberStateListener groupMetadataUpdateListener = captureGroupMetadataUpdateListener(requestManagers); + final int memberEpoch = 1; + final String memberId = "memberId"; + groupMetadataUpdateListener.onMemberEpochUpdated(Optional.of(memberEpoch), memberId); + + KafkaException thrownException = assertThrows(KafkaException.class, + () -> consumer.close(CloseOptions.timeout(Duration.ZERO))); + + assertNotNull(thrownException.getCause()); + assertTrue(thrownException.getCause() instanceof RuntimeException); Review Comment: ```suggestion assertInstanceOf(RuntimeException.class, thrownException.getCause()) ``` ########## clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java: ########## @@ -0,0 +1,337 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals; + +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.WakeupException; +import org.apache.kafka.common.utils.LogContext; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; + +import java.util.Optional; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.STRICT_STUBS) +public class StreamsRebalanceListenerInvokerTest { + + @Mock + private StreamsRebalanceListener mockListener; + + @Mock + private StreamsRebalanceData streamsRebalanceData; + + private StreamsRebalanceListenerInvoker invoker; + private final LogContext logContext = new LogContext(); + + @BeforeEach + public void setup() { + invoker = new StreamsRebalanceListenerInvoker(logContext, streamsRebalanceData); + } + + @Test + public void testConstructorInitializesWithEmptyListener() { + // When invoker is constructed, it should have no listener set initially + // This is verified by testing that invoke methods return null when no listener is present + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListener() { + // Test setting a listener + invoker.setRebalanceListener(mockListener); + + // Verify listener is set by checking that methods no longer return null immediately + // (we'll mock the dependencies needed for actual invocation) + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(mockListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Should now invoke the listener instead of returning null immediately + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); // No exception thrown by mock listener + verify(mockListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testSetRebalanceListenerWithNull() { + // Test setting listener to null + invoker.setRebalanceListener(null); + + // Should behave as if no listener is set + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListenerOverwritesExisting() { + StreamsRebalanceListener firstListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + StreamsRebalanceListener secondListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(secondListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Set first listener + invoker.setRebalanceListener(firstListener); + + // Overwrite with second listener + invoker.setRebalanceListener(secondListener); + + // Should use second listener + invoker.invokeAllTasksRevoked(); + verify(firstListener, never()).onTasksRevoked(any()); + verify(secondListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testInvokeAllTasksRevokedWithNoListener() { + // When no listener is set, should return null + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); + } + + @Test + public void testInvokeAllTasksRevokedWithListener() { + invoker.setRebalanceListener(mockListener); + + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(mockListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + Exception result = invoker.invokeAllTasksRevoked(); + + assertNull(result); + verify(mockListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testInvokeTasksAssignedWithNoListener() { + StreamsRebalanceData.Assignment assignment = createMockAssignment(); + + Exception result = invoker.invokeTasksAssigned(assignment); + + assertNull(result); + verify(mockListener, never()).onTasksAssigned(any()); + } + + @Test + public void testInvokeTasksAssignedWithListener() { + invoker.setRebalanceListener(mockListener); + StreamsRebalanceData.Assignment assignment = createMockAssignment(); + when(mockListener.onTasksAssigned(assignment)).thenReturn(Optional.empty()); + + Exception result = invoker.invokeTasksAssigned(assignment); + + assertNull(result); + verify(mockListener).onTasksAssigned(eq(assignment)); + } + + @Test + public void testInvokeTasksAssignedWithWakeupException() { + invoker.setRebalanceListener(mockListener); + StreamsRebalanceData.Assignment assignment = createMockAssignment(); + WakeupException wakeupException = new WakeupException(); + doThrow(wakeupException).when(mockListener).onTasksAssigned(assignment); + + WakeupException thrownException = assertThrows(WakeupException.class, + () -> invoker.invokeTasksAssigned(assignment)); + + assertEquals(wakeupException, thrownException); + verify(mockListener).onTasksAssigned(eq(assignment)); + } + + @Test + public void testInvokeTasksAssignedWithInterruptException() { + invoker.setRebalanceListener(mockListener); + StreamsRebalanceData.Assignment assignment = createMockAssignment(); + InterruptException interruptException = new InterruptException("Test interrupt"); + doThrow(interruptException).when(mockListener).onTasksAssigned(assignment); + + InterruptException thrownException = assertThrows(InterruptException.class, + () -> invoker.invokeTasksAssigned(assignment)); + + assertEquals(interruptException, thrownException); + verify(mockListener).onTasksAssigned(eq(assignment)); + } + + @Test + public void testInvokeTasksAssignedWithOtherException() { + invoker.setRebalanceListener(mockListener); + StreamsRebalanceData.Assignment assignment = createMockAssignment(); + RuntimeException runtimeException = new RuntimeException("Test exception"); + doThrow(runtimeException).when(mockListener).onTasksAssigned(assignment); + + Exception result = invoker.invokeTasksAssigned(assignment); + + assertEquals(runtimeException, result); + verify(mockListener).onTasksAssigned(eq(assignment)); + } + + @Test + public void testInvokeTasksRevokedWithNoListener() { + Set<StreamsRebalanceData.TaskId> tasks = createMockTasks(); + + Exception result = invoker.invokeTasksRevoked(tasks); + + assertNull(result); + verify(mockListener, never()).onTasksRevoked(any()); + } + + @Test + public void testInvokeTasksRevokedWithListener() { + invoker.setRebalanceListener(mockListener); + Set<StreamsRebalanceData.TaskId> tasks = createMockTasks(); + when(mockListener.onTasksRevoked(tasks)).thenReturn(Optional.empty()); + + Exception result = invoker.invokeTasksRevoked(tasks); + + assertNull(result); + verify(mockListener).onTasksRevoked(eq(tasks)); + } + + @Test + public void testInvokeTasksRevokedWithWakeupException() { + invoker.setRebalanceListener(mockListener); + Set<StreamsRebalanceData.TaskId> tasks = createMockTasks(); + WakeupException wakeupException = new WakeupException(); + doThrow(wakeupException).when(mockListener).onTasksRevoked(tasks); + + WakeupException thrownException = assertThrows(WakeupException.class, + () -> invoker.invokeTasksRevoked(tasks)); + + assertEquals(wakeupException, thrownException); + verify(mockListener).onTasksRevoked(eq(tasks)); + } + + @Test + public void testInvokeTasksRevokedWithInterruptException() { + invoker.setRebalanceListener(mockListener); + Set<StreamsRebalanceData.TaskId> tasks = createMockTasks(); + InterruptException interruptException = new InterruptException("Test interrupt"); + doThrow(interruptException).when(mockListener).onTasksRevoked(tasks); + + InterruptException thrownException = assertThrows(InterruptException.class, + () -> invoker.invokeTasksRevoked(tasks)); + + assertEquals(interruptException, thrownException); + verify(mockListener).onTasksRevoked(eq(tasks)); + } + + @Test + public void testInvokeTasksRevokedWithOtherException() { + invoker.setRebalanceListener(mockListener); + Set<StreamsRebalanceData.TaskId> tasks = createMockTasks(); + RuntimeException runtimeException = new RuntimeException("Test exception"); + doThrow(runtimeException).when(mockListener).onTasksRevoked(tasks); + + Exception result = invoker.invokeTasksRevoked(tasks); + + assertEquals(runtimeException, result); + verify(mockListener).onTasksRevoked(eq(tasks)); + } + + @Test + public void testInvokeAllTasksLostWithNoListener() { Review Comment: isn't this testing the same as `testConstructorInitializesWithEmptyListener`? ########## clients/src/test/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java: ########## @@ -0,0 +1,337 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.clients.consumer.internals; + +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.WakeupException; +import org.apache.kafka.common.utils.LogContext; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; + +import java.util.Optional; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.STRICT_STUBS) +public class StreamsRebalanceListenerInvokerTest { + + @Mock + private StreamsRebalanceListener mockListener; + + @Mock + private StreamsRebalanceData streamsRebalanceData; + + private StreamsRebalanceListenerInvoker invoker; + private final LogContext logContext = new LogContext(); + + @BeforeEach + public void setup() { + invoker = new StreamsRebalanceListenerInvoker(logContext, streamsRebalanceData); + } + + @Test + public void testConstructorInitializesWithEmptyListener() { + // When invoker is constructed, it should have no listener set initially + // This is verified by testing that invoke methods return null when no listener is present + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListener() { + // Test setting a listener + invoker.setRebalanceListener(mockListener); + + // Verify listener is set by checking that methods no longer return null immediately + // (we'll mock the dependencies needed for actual invocation) + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(mockListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Should now invoke the listener instead of returning null immediately + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); // No exception thrown by mock listener + verify(mockListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testSetRebalanceListenerWithNull() { + // Test setting listener to null + invoker.setRebalanceListener(null); + + // Should behave as if no listener is set + assertNull(invoker.invokeAllTasksRevoked()); + assertNull(invoker.invokeAllTasksLost()); + } + + @Test + public void testSetRebalanceListenerOverwritesExisting() { + StreamsRebalanceListener firstListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + StreamsRebalanceListener secondListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class); + + StreamsRebalanceData.Assignment mockAssignment = createMockAssignment(); + when(streamsRebalanceData.reconciledAssignment()).thenReturn(mockAssignment); + when(secondListener.onTasksRevoked(any())).thenReturn(Optional.empty()); + + // Set first listener + invoker.setRebalanceListener(firstListener); + + // Overwrite with second listener + invoker.setRebalanceListener(secondListener); + + // Should use second listener + invoker.invokeAllTasksRevoked(); + verify(firstListener, never()).onTasksRevoked(any()); + verify(secondListener).onTasksRevoked(eq(mockAssignment.activeTasks())); + } + + @Test + public void testInvokeAllTasksRevokedWithNoListener() { + // When no listener is set, should return null + Exception result = invoker.invokeAllTasksRevoked(); + assertNull(result); + } Review Comment: isn't this testing the same as `testConstructorInitializesWithEmptyListener`? ########## clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java: ########## @@ -187,23 +187,14 @@ public class AsyncKafkaConsumer<K, V> implements ConsumerDelegate<K, V> { */ private class BackgroundEventProcessor implements EventProcessor<BackgroundEvent> { - private Optional<StreamsRebalanceListener> streamsRebalanceListener = Optional.empty(); - private final Optional<StreamsRebalanceData> streamsRebalanceData; + private final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker; Review Comment: I would expect it's optional because it will only exist if the AsycKafkaConsumer is being used with the streams rebalance protocol. But why do we need to declare it here in the internal class if it already exist in the parent `AsyncKafkaConsumer`? https://github.com/apache/kafka/blob/45a410b8a8fc3cc95c92ae8da9748990e9ddf1b8/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java#L345 (can we just use that one? as we do with the `rebalanceListenerInvoker`) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
