From ea5ca8486126969fc8c19d63787746ce1a3e3e4e Mon Sep 17 00:00:00 2001 From: Lucas Brutschy Date: Mon, 16 Oct 2023 14:41:02 +0200 Subject: [PATCH 1/3] FIX: allow removing a suspended task from task registry. When we get a suspended task re-assigned in the eager rebalance protocol, we have to add the task back to the state updater so that it has a chance to catch up with its change log. This was prevented by a check in `Tasks`, which disallows removing SUSPENDED tasks from the task registry. I couldn't find a reason why this must be an invariant of the task registry, so this weakens the check. --- .../streams/processor/internals/Tasks.java | 5 +++-- .../processor/internals/TaskManagerTest.java | 20 +++++++++++++++++++ .../processor/internals/TasksTest.java | 17 ++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java index 894c2587578b3..a28fb22766d6b 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java @@ -20,6 +20,7 @@ import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.streams.processor.TaskId; import org.apache.kafka.streams.processor.internals.PendingUpdateAction.Action; +import org.apache.kafka.streams.processor.internals.Task.State; import org.slf4j.Logger; import java.util.Collection; @@ -240,8 +241,8 @@ public synchronized void addTask(final Task task) { public synchronized void removeTask(final Task taskToRemove) { final TaskId taskId = taskToRemove.id(); - if (taskToRemove.state() != Task.State.CLOSED) { - throw new IllegalStateException("Attempted to remove a task that is not closed: " + taskId); + if (taskToRemove.state() != Task.State.CLOSED && taskToRemove.state() != State.SUSPENDED) { + throw new IllegalStateException("Attempted to remove a task that is not closed or suspended: " + taskId); } if (taskToRemove.isActive()) { 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 a876d8ab4e11d..8ea9efa748111 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 @@ -434,6 +434,26 @@ public void shouldKeepReAssignedActiveTaskInStateUpdater() { Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap()); } + @Test + public void shouldMoveReAssignedSuspendedActiveTaskToStateUpdater() { + final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) + .inState(State.SUSPENDED) + .withInputPartitions(taskId03Partitions).build(); + final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); + final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); + when(tasks.allTasks()).thenReturn(mkSet(reassignedActiveTask)); + + taskManager.handleAssignment( + mkMap(mkEntry(reassignedActiveTask.id(), reassignedActiveTask.inputPartitions())), + Collections.emptyMap() + ); + + Mockito.verify(tasks).removeTask(reassignedActiveTask); + Mockito.verify(stateUpdater).add(reassignedActiveTask); + Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap()); + Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap()); + } + @Test public void shouldRemoveReAssignedRevokedActiveTaskInStateUpdaterFromPendingTaskToSuspend() { final StreamTask reAssignedRevokedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java index 23aee1937c0a3..70cbd4ae7fc66 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java @@ -19,6 +19,7 @@ import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.utils.LogContext; import org.apache.kafka.streams.processor.TaskId; +import org.apache.kafka.streams.processor.internals.Task.State; import org.junit.jupiter.api.Test; import java.util.Collections; @@ -31,10 +32,12 @@ import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.standbyTask; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statefulTask; import static org.apache.kafka.test.StreamsTestUtils.TaskBuilder.statelessTask; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class TasksTest { @@ -52,6 +55,20 @@ public class TasksTest { private final Tasks tasks = new Tasks(new LogContext()); + @Test + public void shouldCheckStateWhenClosedTasks() { + final StreamTask closedTask = statefulTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.CLOSED).build(); + final StandbyTask suspendedTask = standbyTask(TASK_0_1, mkSet(TOPIC_PARTITION_A_1)).inState(State.SUSPENDED).build(); + final StreamTask runningTask = statelessTask(TASK_1_0).inState(State.RUNNING).build(); + + tasks.addActiveTasks(mkSet(closedTask, runningTask)); + tasks.addStandbyTasks(Collections.singletonList(suspendedTask)); + + assertDoesNotThrow(() -> tasks.removeTask(closedTask)); + assertDoesNotThrow(() -> tasks.removeTask(suspendedTask)); + assertThrows(IllegalStateException.class, () -> tasks.removeTask(runningTask)); + } + @Test public void shouldKeepAddedTasks() { final StreamTask statefulTask = statefulTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).build(); From de85f63a8de858d5868219299b324f5813897d86 Mon Sep 17 00:00:00 2001 From: Lucas Brutschy Date: Mon, 16 Oct 2023 15:31:29 +0200 Subject: [PATCH 2/3] capitalization fix and code fix --- .../processor/internals/TaskManager.java | 12 ++++------ .../processor/internals/TaskManagerTest.java | 22 +++++++++---------- 2 files changed, 15 insertions(+), 19 deletions(-) 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 cf6f6b6326f56..99a786c580b55 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 @@ -494,7 +494,7 @@ private void handleRunningAndSuspendedTasks(final Map inputPartitions) { if (tasks.updateActiveTaskInputPartitions(task, inputPartitions)) { task.updateInputPartitions(inputPartitions, topologyMetadata.nodeToSourceTopics(task.id())); } if (task.state() == State.SUSPENDED) { + tasks.removeTask(task); task.resume(); - moveTaskFromTasksRegistryToStateUpdater(task); + stateUpdater.add(task); } } - private void moveTaskFromTasksRegistryToStateUpdater(final Task task) { - tasks.removeTask(task); - stateUpdater.add(task); - } - private void handleTasksInStateUpdater(final Map> activeTasksToCreate, final Map> standbyTasksToCreate) { for (final Task task : stateUpdater.getTasks()) { 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 8ea9efa748111..d37ba3a8d9e5f 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 @@ -417,7 +417,7 @@ public void shouldUpdateInputPartitionOfActiveTaskInStateUpdater() { } @Test - public void shouldKeepReAssignedActiveTaskInStateUpdater() { + public void shouldKeepReassignedActiveTaskInStateUpdater() { final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); @@ -435,7 +435,7 @@ public void shouldKeepReAssignedActiveTaskInStateUpdater() { } @Test - public void shouldMoveReAssignedSuspendedActiveTaskToStateUpdater() { + public void shouldMoveReassignedSuspendedActiveTaskToStateUpdater() { final StreamTask reassignedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.SUSPENDED) .withInputPartitions(taskId03Partitions).build(); @@ -455,21 +455,21 @@ public void shouldMoveReAssignedSuspendedActiveTaskToStateUpdater() { } @Test - public void shouldRemoveReAssignedRevokedActiveTaskInStateUpdaterFromPendingTaskToSuspend() { - final StreamTask reAssignedRevokedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) + public void shouldRemoveReassignedRevokedActiveTaskInStateUpdaterFromPendingTaskToSuspend() { + final StreamTask reassignedRevokedActiveTask = statefulTask(taskId03, taskId03ChangelogPartitions) .inState(State.RESTORING) .withInputPartitions(taskId03Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - when(stateUpdater.getTasks()).thenReturn(mkSet(reAssignedRevokedActiveTask)); + when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedRevokedActiveTask)); taskManager.handleAssignment( - mkMap(mkEntry(reAssignedRevokedActiveTask.id(), reAssignedRevokedActiveTask.inputPartitions())), + mkMap(mkEntry(reassignedRevokedActiveTask.id(), reassignedRevokedActiveTask.inputPartitions())), Collections.emptyMap() ); Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap()); - Mockito.verify(tasks).removePendingActiveTaskToSuspend(reAssignedRevokedActiveTask.id()); + Mockito.verify(tasks).removePendingActiveTaskToSuspend(reassignedRevokedActiveTask.id()); Mockito.verify(standbyTaskCreator).createTasks(Collections.emptyMap()); } @@ -496,17 +496,17 @@ public void shouldNeverUpdateInputPartitionsOfStandbyTaskInStateUpdater() { } @Test - public void shouldKeepReAssignedStandbyTaskInStateUpdater() { - final StandbyTask reAssignedStandbyTask = standbyTask(taskId02, taskId02ChangelogPartitions) + public void shouldKeepReassignedStandbyTaskInStateUpdater() { + final StandbyTask reassignedStandbyTask = standbyTask(taskId02, taskId02ChangelogPartitions) .inState(State.RUNNING) .withInputPartitions(taskId02Partitions).build(); final TasksRegistry tasks = Mockito.mock(TasksRegistry.class); final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks, true); - when(stateUpdater.getTasks()).thenReturn(mkSet(reAssignedStandbyTask)); + when(stateUpdater.getTasks()).thenReturn(mkSet(reassignedStandbyTask)); taskManager.handleAssignment( Collections.emptyMap(), - mkMap(mkEntry(reAssignedStandbyTask.id(), reAssignedStandbyTask.inputPartitions())) + mkMap(mkEntry(reassignedStandbyTask.id(), reassignedStandbyTask.inputPartitions())) ); Mockito.verify(activeTaskCreator).createTasks(consumer, Collections.emptyMap()); From 22643c99b779ae4aa25c34aef4a89019da94e19a Mon Sep 17 00:00:00 2001 From: Lucas Brutschy Date: Tue, 17 Oct 2023 11:21:08 +0200 Subject: [PATCH 3/3] fix name --- .../org/apache/kafka/streams/processor/internals/TasksTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java index 70cbd4ae7fc66..c65756d41f431 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java @@ -56,7 +56,7 @@ public class TasksTest { private final Tasks tasks = new Tasks(new LogContext()); @Test - public void shouldCheckStateWhenClosedTasks() { + public void shouldCheckStateWhenRemoveTask() { final StreamTask closedTask = statefulTask(TASK_0_0, mkSet(TOPIC_PARTITION_A_0)).inState(State.CLOSED).build(); final StandbyTask suspendedTask = standbyTask(TASK_0_1, mkSet(TOPIC_PARTITION_A_1)).inState(State.SUSPENDED).build(); final StreamTask runningTask = statelessTask(TASK_1_0).inState(State.RUNNING).build();