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

Code before:

        taskManager.handleAssignment(assignment, emptyMap());

        verify(tasks).addPendingTasksToInit(singletonList(task00));

        // next, drain pending tasks, initialize them, and then add to stateupdater
        when(tasks.drainPendingTasksToInit()).thenReturn(Set.of(task00));

        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        verify(task00).initializeIfNeeded();
        verify(stateUpdater).add(task00);

        // last, drain the restored tasks from stateupdater and transition to running
        when(stateUpdater.restoresActiveTasks()).thenReturn(true);
        when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(Set.of(task00));

        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        verifyTransitionToRunningOfRestoredTask(Set.of(task00), tasks);
    }

    @Test
    public void shouldNotCompleteRestorationIfTasksCannotInitialize() {
        final StreamTask task00 = statefulTask(taskId00, taskId00ChangelogPartitions)
            .withInputPartitions(taskId00Partitions)
            .inState(State.CREATED)
            .build();

        final TasksRegistry tasks = mock(TasksRegistry.class);
        final TaskManager taskManager = setUpTaskManagerWithStateUpdater(ProcessingMode.AT_LEAST_ONCE, tasks);
        final Map<TaskId, Set<TopicPartition>> assignment = mkMap(
            mkEntry(taskId00, taskId00Partitions)
        );

        when(activeTaskCreator.createTasks(any(), eq(assignment)))
            .thenReturn(singletonList(task00));
        taskManager.handleAssignment(assignment, emptyMap());

        verify(tasks).addPendingTasksToInit(singletonList(task00));

        when(tasks.drainPendingTasksToInit()).thenReturn(Set.of(task00));
        final LockException lockException = new LockException("can't lock");
        doThrow(lockException).when(task00).initializeIfNeeded();
        when(tasks.hasPendingTasksToInit()).thenReturn(true);

        final boolean restorationComplete = taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        assertFalse(restorationComplete);
        verify(task00).initializeIfNeeded();
        verify(tasks, times(2)).addPendingTasksToInit(

Review comment (suggestion):
Wait. We should actually also handle TimeoutException in the production code. I suppose you removed TimeoutException because state updater doesn't handle it.

We should fix that, because `initializeIfNeeded` may throw a timeout exception.

If we get a TimeoutException inside addTasksToStateUpdater, we should 
```
            task.maybeInitTaskTimeoutOrThrow(now, timeoutException);
            tasks.addPendingTasksToInit(Collections.singleton(task));
            updateOrCreateBackoffRecord(task.id(), nowMs);
```
and otherwise, we should 
```
                task.clearTaskTimeout();
```
after the initialization.

And then we should bring back this behavior in the test

Code after:
        // first, we need to handle assignment -- creates tasks and adds to pending initialization
        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(singletonList(task00));

        taskManager.handleAssignment(assignment, emptyMap());

        verify(tasks).addPendingTasksToInit(singletonList(task00));

        // next, drain pending tasks, initialize them, and then add to stateupdater
        when(tasks.drainPendingTasksToInit()).thenReturn(Set.of(task00));

        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        verify(task00).initializeIfNeeded();
        verify(stateUpdater).add(task00);

        // last, drain the restored tasks from stateupdater and transition to running
        when(stateUpdater.restoresActiveTasks()).thenReturn(true);
        when(stateUpdater.drainRestoredActiveTasks(any(Duration.class))).thenReturn(Set.of(task00));

        taskManager.checkStateUpdater(time.milliseconds(), noOpResetter);

        verifyTransitionToRunningOfRestoredTask(Set.of(task00), tasks);
    }

    @Test
    public void shouldNotCompleteRestorationIfTasksCannotInitialize() {
        final Map<TaskId, Set<TopicPartition>> assignment = mkMap(
            mkEntry(taskId00, taskId00Partitions),
            mkEntry(taskId01, taskId01Partitions)
        );
        final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {
            @Override
            public void initializeIfNeeded() {
                throw new LockException("can't lock");
            }
        };
        final Task task01 = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) {
            @Override
            public void initializeIfNeeded() {
                throw new TimeoutException("timed out");
            }
        };

        when(activeTaskCreator.createTasks(any(), eq(assignment))).thenReturn(asList(task00, task01));

        taskManager.handleAssignment(assignment, emptyMap());

        assertThat(task00.state(), is(Task.State.CREATED));
        assertThat(task01.state(), is(Task.State.CREATED));

        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(false));