KAFKA-19960: Close pending tasks on shutdown. by Nikita-Shupletsov · Pull Request #21365 · apache/kafka

Nikita-Shupletsov · 2026-01-27T19:52:06Z

This PR fixes a bug when KS doesn't close stores if the shutdown was
triggered during rebalance where an active tasks gets converted to a
standby one and put into pendingTasksToInit

Added logic to close pending tasks to init.
Made standby task closure similar to the one for active tasks.
Added a separate method for getting standby tasks from task registry.
Added an integration test that reproduces the issue.

Added logic to close pending tasks to init. Made standby task closure similar to the one for active tasks. Added a separate method for getting standby tasks from task registry. Added an integration test that reproduces the issue.

mjsax

Made a first pass.

Can you update the PR description adding context on what the bug exactly is, and when we hit it? It seems to be related to not closing "pending tasks", but might be good to give some more context.

.../src/test/java/org/apache/kafka/streams/integration/RebalanceTaskClosureIntegrationTest.java

...ntegration-tests/src/test/java/org/apache/kafka/streams/integration/KafkaStreamsWrapper.java

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java

…ms/integration/RebalanceTaskClosureIntegrationTest.java Co-authored-by: Matthias J. Sax <mjsax@apache.org>

Minor refactoring.

mjsax

Some more follow ups -- also need to get input from @lucasbru...

mjsax · 2026-02-03T03:35:11Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

+            } else {
+                standbyTasks.add(pendingTask);
+            }
+        }


Thanks for updating the PR description. It say "shutdown during rebalance when active task become standby tasks" but seems it goes either way, and the PR is actually fixing both direction?

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java

lucasbru

Good find, thanks for identifying the problem.

An alternative way to fix it would be, I think, to consider pendingTasksToInit as a subset of Tasks.allTasks. Then conceptually, it may be simpler to say "close all tasks". But we'd have to check carefully what other places Tasks.allTasks and related methods are used, and if we are changing one of those places. In that case, we'd still need the "remove" fix you did.

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java

lucasbru · 2026-02-03T14:15:39Z

An alternative way to fix it would be, I think, to consider pendingTasksToInit as a subset of Tasks.allTasks. Then conceptually, it may be simpler to say "close all tasks". But we'd have to check carefully what other places Tasks.allTasks and related methods are used, and if we are changing one of those places. In that case, we'd still need the "remove" fix you did.

On second thought, that would probably create more problems than it solves. But I wonder if we should rename "allTasks" to "allInitializedTasks" then?

Nikita-Shupletsov · 2026-02-03T17:22:17Z

Using drainPendingTasksToInit sounds like a good idea to me. Not sure if there are any complications though, but it would seem cleaner code wise.

the problem with drain is that we always try to remove the task we are closing from the task registry: https://github.com/Nikita-Shupletsov/kafka/blob/8ee99ee82df2cf89fbb769d26c66395fd3a63761/streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java#L1405

if we drain the pending tasks, there will be no way for us to confirm it ever existed there. so we either need to relax the check in the remove method or have a separate branch for closing pending to init tasks

mjsax · 2026-02-03T20:21:47Z

maybe we would need to add a new member to TasksRegistry taking "pending tasks" which got closed?

Yes, that is what I mean by this -- after we drain the pending-tasks-to-init, we would add them into some new "pending-task-to-close" member inside TaskRegistry and update removeTask so check ownership based on this new member variable.

It's to some extend "cosmetics" as we only move all "pending task" from one collection into a different one, but it preserves the invariant we have, and makes the code cleaner.

Nikita-Shupletsov · 2026-02-03T21:14:31Z

But I wonder if we should rename "allTasks" to "allInitializedTasks" then?
@lucasbru would you mind if I did that in a followup PR?

lucasbru · 2026-02-04T10:09:25Z

But I wonder if we should rename "allTasks" to "allInitializedTasks" then?
@lucasbru would you mind if I did that in a followup PR?

No problem, it's a minor thing.

mjsax

Few follow up comments. -- Still working through the integration test, but might be good to share what I have right away.

mjsax · 2026-02-04T19:54:11Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

+        final Set<Task> standbyTasks = new TreeSet<>(Comparator.comparing(Task::id));
+        standbyTasks.addAll(tasks.standbyTasks());
+
+        Set<Task> pendingActiveTasks = tasks.drainPendingActiveTasksToInit();


Seems this should be final -- wondering why the build did not fail with a checktsyle error? Can we check this?

mjsax · 2026-02-04T19:54:21Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

+        Set<Task> pendingActiveTasks = tasks.drainPendingActiveTasksToInit();
+        activeTasks.addAll(pendingActiveTasks);
+        tasks.addPendingTasksToClose(pendingActiveTasks);
+        Set<Task> pendingStandbyTasks = tasks.drainPendingStandbyTasksToInit();


Same. Should be final

mjsax · 2026-02-04T19:54:46Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java

+    }
+
+    @Override
+    public void addPendingTasksToClose(Collection<Task> tasks) {


Parameter should be final

mjsax · 2026-02-04T20:05:32Z

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

        // TODO: change type to `StreamTask`
        final Set<Task> activeTasks = new TreeSet<>(Comparator.comparing(Task::id));
        activeTasks.addAll(tasks.activeTasks());
+        final Set<Task> standbyTasks = new TreeSet<>(Comparator.comparing(Task::id));


nit: should this be StrandbyTask type (compare comment above about StreamTask type -- we would also just update the comment, and try to do a follow up PR to improve types later?

mjsax · 2026-02-04T20:22:53Z

streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java

        final StreamTask activeTask1 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId00Partitions).build();
-        final StreamTask activeTask2 = statefulTask(taskId01, taskId01ChangelogPartitions)


Why are we removing this task from the test? Is using a single task sufficient (and it's unclear why we did use two tasks to begin with)?

mjsax · 2026-02-04T20:28:27Z

streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java

+        final TaskManager taskManager = setUpTaskManager(ProcessingMode.AT_LEAST_ONCE, tasks);
+
+        final StandbyTask standbyTask00 = standbyTask(taskId00, taskId00ChangelogPartitions)
+                .inState(State.RUNNING)


Should the state be CREATED?

mjsax · 2026-02-04T20:28:30Z

streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java

+                .build();
+
+        final StreamTask activeTask01 = statefulTask(taskId01, taskId00ChangelogPartitions)
+                .inState(State.RUNNING)


mjsax · 2026-02-04T20:30:07Z

streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java

+        tasks.addPendingTasksToInit(Set.of(activeTask1, activeTask2, standbyTask1, standbyTask2));
+
+        final Set<Task> standbyTasksToInit = tasks.drainPendingStandbyTasksToInit();
+        assertEquals(2, standbyTasksToInit.size());


Suggested change

assertEquals(2, standbyTasksToInit.size());

assertEquals(2, standbyTasksToInit.size());

mjsax · 2026-02-04T20:33:36Z

streams/src/test/java/org/apache/kafka/streams/processor/internals/TasksTest.java

+        assertFalse(tasks.allTasks().contains(activeTask1));
+
+        tasks.addPendingTasksToClose(List.of(activeTask1));
+        assertTrue(tasks.pendingTasksToClose().contains(activeTask1));


Why do we add it back and verify? Could we not do the verification above, right after the added the task the first time?

mjsax · 2026-02-04T20:38:08Z

.../src/test/java/org/apache/kafka/streams/integration/RebalanceTaskClosureIntegrationTest.java

+     * The conditions that we need to meet:
+     * <p><ul>
+     * <li>There is a task with an open store in {@link org.apache.kafka.streams.processor.internals.TasksRegistry#pendingTasksToInit}</li>
+     * <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} return false


Suggested change

* <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} return false

* <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} returns false

KAFKA-19960: Close pending tasks on shutdown.

ecbe338

Added logic to close pending tasks to init. Made standby task closure similar to the one for active tasks. Added a separate method for getting standby tasks from task registry. Added an integration test that reproduces the issue.

github-actions bot added triage PRs from the community streams labels Jan 27, 2026

mjsax added ci-approved and removed triage PRs from the community labels Jan 29, 2026

mjsax reviewed Jan 30, 2026

View reviewed changes

Nikita-Shupletsov and others added 2 commits January 29, 2026 21:52

Update streams/integration-tests/src/test/java/org/apache/kafka/strea…

3527671

…ms/integration/RebalanceTaskClosureIntegrationTest.java Co-authored-by: Matthias J. Sax <mjsax@apache.org>

Added a comment on what the test actually tests.

8ee99ee

Minor refactoring.

mjsax reviewed Feb 3, 2026

View reviewed changes

lucasbru reviewed Feb 3, 2026

View reviewed changes

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java Outdated Show resolved Hide resolved

streams/src/main/java/org/apache/kafka/streams/processor/internals/Tasks.java Show resolved Hide resolved

Nikita-Shupletsov added 2 commits February 3, 2026 13:12

Merge branch 'apache:trunk' into KAFKA-19960

4f2bdf4

Added PendingTasksToClose.

8793347

Updated the comment to better explain the sequence of actions.

1d7586a

mjsax reviewed Feb 4, 2026

View reviewed changes

	assertEquals(2, standbyTasksToInit.size());

	assertEquals(2, standbyTasksToInit.size());

	* <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} return false
	* <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} returns false

Conversation

Nikita-Shupletsov commented Jan 27, 2026 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mjsax left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mjsax left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lucasbru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

lucasbru commented Feb 3, 2026

Uh oh!

Nikita-Shupletsov commented Feb 3, 2026

Uh oh!

mjsax commented Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Nikita-Shupletsov commented Feb 3, 2026

Uh oh!

lucasbru commented Feb 4, 2026

Uh oh!

mjsax left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Nikita-Shupletsov commented Jan 27, 2026 •

edited by github-actions bot

Loading

mjsax commented Feb 3, 2026 •

edited

Loading