Repository: apache/kafka
File: streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java (Java)

Code before:

        when(stateUpdater.tasks()).thenReturn(Set.of(migratedTask01, migratedTask02));

        // mock futures for removing tasks from StateUpdater
        final CompletableFuture<StateUpdater.RemovedTaskResult> future01 = new CompletableFuture<>();
        when(stateUpdater.remove(taskId01)).thenReturn(future01);
        future01.complete(new StateUpdater.RemovedTaskResult(migratedTask01));

        final CompletableFuture<StateUpdater.RemovedTaskResult> future02 = new CompletableFuture<>();
        when(stateUpdater.remove(taskId02)).thenReturn(future02);
        future02.complete(new StateUpdater.RemovedTaskResult(migratedTask02));

        final TaskMigratedException thrown = assertThrows(
            TaskMigratedException.class,
            () -> taskManager.handleAssignment(emptyMap(), emptyMap())
        );

        // iteration order here is non-deterministic due to hashset.
        // The last exception encountered wins, so we accept either
        // t1 or t2's exception message.
        assertThat(
            thrown.getMessage(),
            anyOf(
                equalTo("t1 close exception; it means all tasks belonging to this thread should be migrated."),
                equalTo("t2 close exception; it means all tasks belonging to this thread should be migrated.")
            )
        );

        verify(migratedTask01, times(2)).suspend();
        verify(migratedTask02, times(2)).suspend();
        verify(stateUpdater).remove(taskId01);
        verify(stateUpdater).remove(taskId02);
    }

    @Test
    public void shouldThrowRuntimeExceptionWhenEncounteredUnknownExceptionDuringTaskClose() {
        final StandbyTask migratedTask01 = standbyTask(taskId01, taskId01ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId01Partitions)
            .build();
        final StandbyTask migratedTask02 = standbyTask(taskId02, taskId02ChangelogPartitions)
            .inState(State.RUNNING)
            .withInputPartitions(taskId02Partitions)
            .build();

        doThrow(new TaskMigratedException("t1 close exception", new RuntimeException()))
            .when(migratedTask01).suspend();
        doThrow(new IllegalStateException("t2 illegal state exception", new RuntimeException()))
            .when(migratedTask02).suspend();

        final TasksRegistry tasks = mock(TasksRegistry.class);

Review comment (refactor):
The iteration order here with the new stateupdater code is non-deterministic, whereas earlier it was deterministic due to the use of `TreeSet`. I am not 100% sure if this was intentional and I think it may have been an oversight while refactoring?

Should we update the production code to use `TreeSet` instead? Please advise.

Code after:

        final Map<TopicPartition, OffsetAndMetadata> offsets = singletonMap(t1p0, new OffsetAndMetadata(0L, null));
        when(task00.prepareCommit(false)).thenReturn(offsets);

        final TasksRegistry tasks = mock(TasksRegistry.class);
        final TaskManager taskManager = setUpTaskManagerWithStateUpdater(ProcessingMode.AT_LEAST_ONCE, tasks);
        when(tasks.allTasks()).thenReturn(Set.of(task00));

        try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(TaskManager.class)) {
            appender.setClassLogger(TaskManager.class, Level.DEBUG);

            taskManager.handleRevocation(Set.of(t1p0, new TopicPartition("unknown", 0)));

            verify(task00).suspend();

            final List<String> messages = appender.getMessages();
            assertThat(
                messages,
                hasItem("taskManagerTestThe following revoked partitions [unknown-0] are missing " +
                    "from the current task partitions. It could potentially be due to race " +
                    "condition of consumer detecting the heartbeat failure, or the " +
                    "tasks have been cleaned up by the handleAssignment callback.")
            );
        }
    }

    @Test
    public void shouldThrowTaskMigratedWhenAllTaskCloseExceptionsAreTaskMigrated() {
        final StateMachineTask migratedTask01 = new StateMachineTask(taskId01, taskId01Partitions, false, stateManager) {
            @Override
            public void suspend() {
                super.suspend();
                throw new TaskMigratedException("t1 close exception", new RuntimeException());
            }
        };

        final StateMachineTask migratedTask02 = new StateMachineTask(taskId02, taskId02Partitions, false, stateManager) {
            @Override
            public void suspend() {
                super.suspend();
                throw new TaskMigratedException("t2 close exception", new RuntimeException());
            }
        };
        taskManager.addTask(migratedTask01);
        taskManager.addTask(migratedTask02);

        final TaskMigratedException thrown = assertThrows(
            TaskMigratedException.class,
            () -> taskManager.handleAssignment(emptyMap(), emptyMap())
        );
        // The task map orders tasks based on topic group id and partition, so here