[kafka] branch trunk updated: KAFKA-14299: Fix incorrect pauses in separate state restoration (#12743)

[ca...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

cadonna pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git


The following commit(s) were added to refs/heads/trunk by this push:
     new cc582897bfb KAFKA-14299: Fix incorrect pauses in separate state restoration (#12743)
cc582897bfb is described below

commit cc582897bfb237572131369a598f7869220b43dc
Author: Lucas Brutschy <lu...@users.noreply.github.com>
AuthorDate: Tue Oct 18 12:16:44 2022 +0200

    KAFKA-14299: Fix incorrect pauses in separate state restoration (#12743)
    
    The original code path paused the main consumer for
    all tasks before entering the restoration section
    of the code, and then resumed all after restoration
    has finished.
    
    In the new state updater part of the code, tasks that
    do not require restoration skip the restoration completely.
    They remain with the TaskManger and are never transferred
    to the StateUpdater, and thus are never resumed.
    
    This change makes sure that tasks that remain with the
    TaskManager are not paused.
    
    Reviewers: Guozhang Wang <wa...@gmail.com>, Bruno Cadonna <ca...@apache.org>
---
 .../streams/processor/internals/TaskManager.java   | 13 +++++++++-
 .../processor/internals/TaskManagerTest.java       | 30 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
index 48afe51aeb9..da7da31a0e2 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java
@@ -175,7 +175,18 @@ public class TaskManager {
     void handleRebalanceComplete() {
         // we should pause consumer only within the listener since
         // before then the assignment has not been updated yet.
-        mainConsumer.pause(mainConsumer.assignment());
+        if (stateUpdater == null) {
+            mainConsumer.pause(mainConsumer.assignment());
+        } else {
+            // All tasks that are owned by the task manager are ready and do not need to be paused
+            final Set<TopicPartition> partitionsNotToPause = tasks.allTasks()
+                .stream()
+                .flatMap(task -> task.inputPartitions().stream())
+                .collect(Collectors.toSet());
+            final Set<TopicPartition> partitionsToPause = new HashSet<>(mainConsumer.assignment());
+            partitionsToPause.removeAll(partitionsNotToPause);
+            mainConsumer.pause(partitionsToPause);
+        }
 
         releaseLockedUnassignedTaskDirectories();
 
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index b1fa45c58d6..0568a83d500 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -1481,6 +1481,36 @@ public class TaskManagerTest {
         assertThat(taskManager.lockedTaskDirectories(), is(singleton(taskId01)));
     }
 
+    @Test
+    public void shouldPauseAllTopicsWithoutStateUpdaterOnRebalanceComplete() {
+        final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
+        expect(consumer.assignment()).andReturn(assigned);
+        consumer.pause(assigned);
+        replay(consumer);
+
+        taskManager.handleRebalanceComplete();
+
+        verify(consumer);
+    }
+
+    @Test
+    public void shouldNotPauseReadyTasksWithStateUpdaterOnRebalanceComplete() {
+        final StreamTask statefulTask0 = statefulTask(taskId00, taskId00ChangelogPartitions)
+            .inState(State.RUNNING)
+            .withInputPartitions(taskId00Partitions).build();
+        final TasksRegistry tasks = Mockito.mock(TasksRegistry.class);
+        final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true);
+        when(tasks.allTasks()).thenReturn(mkSet(statefulTask0));
+        final Set<TopicPartition> assigned = mkSet(t1p0, t1p1);
+        expect(consumer.assignment()).andReturn(assigned);
+        consumer.pause(mkSet(t1p1));
+        replay(consumer);
+
+        taskManager.handleRebalanceComplete();
+
+        verify(consumer);
+    }
+
     @Test
     public void shouldReleaseLockForUnassignedTasksAfterRebalance() throws Exception {
         expectLockObtainedFor(taskId00, taskId01, taskId02);