Examples with State org.apache.kafka.streams.processor.internals.Task.State used on opensource projects

Example 1 with State

use of org.apache.kafka.streams.processor.internals.Task.State in project kafka by apache.

the class TaskManagerTest method shouldReviveCorruptTasks.

@Test
public void shouldReviveCorruptTasks() {
    final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class);
    stateManager.markChangelogAsCorrupted(taskId00Partitions);
    EasyMock.expectLastCall().once();
    replay(stateManager);
    final AtomicBoolean enforcedCheckpoint = new AtomicBoolean(false);
    final StateMachineTask task00 = new StateMachineTask(taskId00, taskId00Partitions, true, stateManager) {

        @Override
        public void postCommit(final boolean enforceCheckpoint) {
            if (enforceCheckpoint) {
                enforcedCheckpoint.set(true);
            }
            super.postCommit(enforceCheckpoint);
        }
    };
    // `handleAssignment`
    expectRestoreToBeCompleted(consumer, changeLogReader);
    expect(activeTaskCreator.createTasks(anyObject(), eq(taskId00Assignment))).andStubReturn(singletonList(task00));
    topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString());
    expectLastCall().anyTimes();
    expect(consumer.assignment()).andReturn(taskId00Partitions);
    replay(activeTaskCreator, topologyBuilder, consumer, changeLogReader);
    taskManager.handleAssignment(taskId00Assignment, emptyMap());
    assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), tp -> assertThat(tp, is(empty()))), is(true));
    assertThat(task00.state(), is(Task.State.RUNNING));
    task00.setChangelogOffsets(singletonMap(t1p0, 0L));
    taskManager.handleCorruption(singleton(taskId00));
    assertThat(task00.commitPrepared, is(true));
    assertThat(task00.state(), is(Task.State.CREATED));
    assertThat(task00.partitionsForOffsetReset, equalTo(taskId00Partitions));
    assertThat(enforcedCheckpoint.get(), is(true));
    assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
    assertThat(taskManager.standbyTaskMap(), Matchers.anEmptyMap());
    verify(stateManager);
    verify(consumer);
}

Example 2 with State

use of org.apache.kafka.streams.processor.internals.Task.State in project kafka by apache.

the class TaskManagerTest method handleAssignment.

private Map<TaskId, StateMachineTask> handleAssignment(final Map<TaskId, Set<TopicPartition>> runningActiveAssignment, final Map<TaskId, Set<TopicPartition>> standbyAssignment, final Map<TaskId, Set<TopicPartition>> restoringActiveAssignment) {
    final Set<Task> runningTasks = runningActiveAssignment.entrySet().stream().map(t -> new StateMachineTask(t.getKey(), t.getValue(), true)).collect(Collectors.toSet());
    final Set<Task> standbyTasks = standbyAssignment.entrySet().stream().map(t -> new StateMachineTask(t.getKey(), t.getValue(), false)).collect(Collectors.toSet());
    final Set<Task> restoringTasks = restoringActiveAssignment.entrySet().stream().map(t -> new StateMachineTask(t.getKey(), t.getValue(), true)).collect(Collectors.toSet());
    // give the restoring tasks some uncompleted changelog partitions so they'll stay in restoring
    restoringTasks.forEach(t -> ((StateMachineTask) t).setChangelogOffsets(singletonMap(new TopicPartition("changelog", 0), 0L)));
    // Initially assign only the active tasks we want to complete restoration
    final Map<TaskId, Set<TopicPartition>> allActiveTasksAssignment = new HashMap<>(runningActiveAssignment);
    allActiveTasksAssignment.putAll(restoringActiveAssignment);
    final Set<Task> allActiveTasks = new HashSet<>(runningTasks);
    allActiveTasks.addAll(restoringTasks);
    expect(standbyTaskCreator.createTasks(eq(standbyAssignment))).andStubReturn(standbyTasks);
    expect(activeTaskCreator.createTasks(anyObject(), eq(allActiveTasksAssignment))).andStubReturn(allActiveTasks);
    expectRestoreToBeCompleted(consumer, changeLogReader);
    replay(activeTaskCreator, standbyTaskCreator, consumer, changeLogReader);
    taskManager.handleAssignment(allActiveTasksAssignment, standbyAssignment);
    taskManager.tryToCompleteRestoration(time.milliseconds(), null);
    final Map<TaskId, StateMachineTask> allTasks = new HashMap<>();
    // Just make sure all tasks ended up in the expected state
    for (final Task task : runningTasks) {
        assertThat(task.state(), is(Task.State.RUNNING));
        allTasks.put(task.id(), (StateMachineTask) task);
    }
    for (final Task task : restoringTasks) {
        assertThat(task.state(), is(Task.State.RESTORING));
        allTasks.put(task.id(), (StateMachineTask) task);
    }
    for (final Task task : standbyTasks) {
        assertThat(task.state(), is(Task.State.RUNNING));
        allTasks.put(task.id(), (StateMachineTask) task);
    }
    return allTasks;
}