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 9628b42d927c8..efa92e3484555 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 @@ -40,7 +40,7 @@ * i.e. all running active tasks are processed by the former and all restoring active tasks and standby tasks are * processed by the latter. */ -class Tasks { +public class Tasks { private final Logger log; // TODO: convert to Stream/StandbyTask when we remove TaskManager#StateMachineTask with mocks @@ -168,6 +168,13 @@ void addNewStandbyTasks(final Collection newTasks) { } } + public void removeActiveTask(final TaskId taskId) { + if (activeTasksPerId.remove(taskId) == null) { + throw new IllegalArgumentException("Attempted to remove an active task that is not owned: " + taskId); + } + removePartitionsForActiveTask(taskId); + } + void removeTask(final Task taskToRemove) { final TaskId taskId = taskToRemove.id(); @@ -255,7 +262,7 @@ private Task getTask(final TaskId taskId) { return null; } - Task task(final TaskId taskId) { + public Task task(final TaskId taskId) { final Task task = getTask(taskId); if (task != null) @@ -272,7 +279,7 @@ Collection tasks(final Collection taskIds) { return tasks; } - Collection activeTasks() { + public Collection activeTasks() { return Collections.unmodifiableCollection(activeTasksPerId.values()); } @@ -300,7 +307,7 @@ boolean owned(final TaskId taskId) { } // for testing only - void addTask(final Task task) { + public void addTask(final Task task) { if (task.isActive()) { activeTasksPerId.put(task.id(), task); } else { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskExecutor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskExecutor.java index b7d8c2091c33a..d5ab481e69a31 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskExecutor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskExecutor.java @@ -48,19 +48,19 @@ public void run() { private void runOnce(final long nowMs) { KafkaFutureImpl pauseFuture; if ((pauseFuture = pauseRequested.getAndSet(null)) != null) { - taskManager.unlockTask(currentTask, DefaultTaskExecutor.this); + taskManager.unassignTask(currentTask, DefaultTaskExecutor.this); pauseFuture.complete(currentTask); currentTask = null; } if (currentTask == null) { - currentTask = taskManager.lockNextTask(DefaultTaskExecutor.this); + currentTask = taskManager.assignNextTask(DefaultTaskExecutor.this); } if (currentTask.isProcessable(nowMs)) { currentTask.process(nowMs); } else { - taskManager.unlockTask(currentTask, DefaultTaskExecutor.this); + taskManager.unassignTask(currentTask, DefaultTaskExecutor.this); currentTask = null; } } @@ -79,6 +79,7 @@ public DefaultTaskExecutor(final TaskManager taskManager, this.taskManager = taskManager; } + @Override public void start() { if (taskExecutorThread == null) { taskExecutorThread = new TaskExecutorThread("task-executor"); @@ -87,6 +88,7 @@ public void start() { } } + @Override public void shutdown(final Duration timeout) { if (taskExecutorThread != null) { taskExecutorThread.isRunning.set(false); @@ -101,11 +103,13 @@ public void shutdown(final Duration timeout) { } } + @Override public ReadOnlyTask currentTask() { return new ReadOnlyTask(currentTask); } - public KafkaFuture pause() { + @Override + public KafkaFuture unassign() { final KafkaFutureImpl future = new KafkaFutureImpl<>(); if (taskExecutorThread != null) { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskManager.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskManager.java new file mode 100644 index 0000000000000..a8444ce379bb3 --- /dev/null +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/DefaultTaskManager.java @@ -0,0 +1,199 @@ +package org.apache.kafka.streams.processor.internals.tasks; + +import org.apache.kafka.common.KafkaFuture; +import org.apache.kafka.common.internals.KafkaFutureImpl; +import org.apache.kafka.common.utils.LogContext; +import org.apache.kafka.common.utils.Time; +import org.apache.kafka.streams.StreamsConfig; +import org.apache.kafka.streams.processor.TaskId; +import org.apache.kafka.streams.processor.internals.DefaultStateUpdater; +import org.apache.kafka.streams.processor.internals.ReadOnlyTask; +import org.apache.kafka.streams.processor.internals.StreamTask; +import org.apache.kafka.streams.processor.internals.Task; +import org.apache.kafka.streams.processor.internals.Tasks; +import org.slf4j.Logger; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentSkipListSet; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Supplier; +import java.util.stream.Collectors; + +/** + * An active task could only be in one of the following status: + * + * 1. It's assigned to one of the executors for processing. + * 2. It's locked by the stream thread for committing, removal, other manipulations etc. + * 3. Neither 1 or 2, i.e. it's not locked but also not assigned to any executors. It's possible + * that we do not have enough executors or because those tasks are not processible yet. + */ +public class DefaultTaskManager implements TaskManager { + + private final Time time; + private final Logger log; + private final Tasks tasks; + + private final Lock tasksLock = new ReentrantLock(); + private final List lockedTasks = new ArrayList<>(); + private final Map assignedTasks = new HashMap<>(); + + private final List taskExecutors; + + public DefaultTaskManager(final Time time, + final Tasks tasks, + final StreamsConfig config, + final String threadName) { + final String logPrefix = String.format("%s ", threadName); + final LogContext logContext = new LogContext(logPrefix); + this.log = logContext.logger(DefaultStateUpdater.class); + this.time = time; + this.tasks = tasks; + + final int numExecutors = config.getInt(StreamsConfig.NUM_STREAM_THREADS_CONFIG); + taskExecutors = new ArrayList<>(numExecutors); + for (int i = 1; i <= numExecutors; i++) { + taskExecutors.add(new DefaultTaskExecutor(this, time)); + } + } + + public StreamTask assignNextTask(final TaskExecutor executor) { + return returnWithTasksLocked(() -> { + if (!taskExecutors.contains(executor)) { + throw new IllegalArgumentException("The requested executor for getting next task to assign is unrecognized"); + } + + // the most naive scheduling algorithm for now: give the next unlocked, unassigned, and processible task + for (final Task task : tasks.activeTasks()) { + if (!assignedTasks.containsKey(task.id()) && + !lockedTasks.contains(task.id()) && + ((StreamTask) task).isProcessable(time.milliseconds())) { + assignedTasks.put(task.id(), executor); + return (StreamTask) task; + } + } + + return null; + }); + } + + public void unassignTask(final StreamTask task, final TaskExecutor executor) { + executeWithTasksLocked(() -> { + if (!taskExecutors.contains(executor)) { + throw new IllegalArgumentException("The requested executor for unassign task is unrecognized"); + } + + final TaskExecutor lockedExecutor = assignedTasks.get(task.id()); + if (lockedExecutor == null || lockedExecutor == executor) { + throw new IllegalArgumentException("Task " + task.id() + " is not locked by the executor"); + } + + assignedTasks.remove(task.id()); + }); + } + + public KafkaFuture lockTasks(final Set taskIds) { + return returnWithTasksLocked(() -> { + lockedTasks.addAll(taskIds); + + final KafkaFutureImpl result = new KafkaFutureImpl<>(); + final Set remainingTaskIds = new ConcurrentSkipListSet<>(taskIds); + + for (final TaskId taskId : taskIds) { + final Task task = tasks.task(taskId); + + if (task == null) { + throw new IllegalArgumentException("Trying to lock task " + taskId + " but it's not owned"); + } + + if (!task.isActive()) { + throw new IllegalArgumentException("The locking task " + taskId + " is not an active task"); + } + + if (assignedTasks.containsKey(taskId)) { + KafkaFuture future = assignedTasks.get(taskId).unassign(); + future.whenComplete((streamTask, throwable) -> { + if (throwable != null) { + result.completeExceptionally(throwable); + } else { + remainingTaskIds.remove(streamTask.id()); + if (remainingTaskIds.isEmpty()) { + result.complete(null); + } + } + }); + } else { + remainingTaskIds.remove(taskId); + } + } + + return result; + }); + } + + public KafkaFuture lockAllTasks() { + return returnWithTasksLocked(() -> + lockTasks(tasks.activeTasks().stream().map(Task::id).collect(Collectors.toSet())) + ); + } + + public void unlockTasks(final Set taskIds) { + executeWithTasksLocked(() -> lockedTasks.removeAll(taskIds)); + } + + public void unlockAllTasks() { + executeWithTasksLocked(() -> unlockTasks(tasks.activeTasks().stream().map(Task::id).collect(Collectors.toSet()))); + } + + public void add(final Set tasksToAdd) { + executeWithTasksLocked(() -> { + for (final StreamTask task : tasksToAdd) { + tasks.addTask(task); + } + }); + + log.info("Added tasks {} to the task manager to process", tasksToAdd); + } + + public void remove(final TaskId taskId) { + executeWithTasksLocked(() -> { + if (assignedTasks.containsKey(taskId)) { + throw new IllegalArgumentException("The task to remove is still assigned to executors"); + } + + if (!lockedTasks.contains(taskId)) { + throw new IllegalArgumentException("The task to remove is not locked yet"); + } + + tasks.removeActiveTask(taskId); + }); + + log.info("Removed task {} from the task manager", taskId); + } + + public Set getTasks() { + return returnWithTasksLocked(() -> tasks.activeTasks().stream().map(ReadOnlyTask::new).collect(Collectors.toSet())); + } + + private void executeWithTasksLocked(final Runnable action) { + tasksLock.lock(); + try { + action.run(); + } finally { + tasksLock.unlock(); + } + } + + private T returnWithTasksLocked(final Supplier action) { + tasksLock.lock(); + try { + return action.get(); + } finally { + tasksLock.unlock(); + } + } +} diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskExecutor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskExecutor.java index 9450d663770c5..795a64cd3a77f 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskExecutor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskExecutor.java @@ -24,18 +24,18 @@ public interface TaskExecutor { void shutdown(final Duration timeout); /** - * Get the current processing task. The task returned is read-only and cannot be modified. + * Get the current assigned processing task. The task returned is read-only and cannot be modified. * * @return the current processing task */ ReadOnlyTask currentTask(); /** - * Pause the current processing task from the task processor and unlock it in the state manager. + * Unassign the current processing task from the task processor and give it back to the state manager. * * The paused task must be flushed since it may be committed or closed by the task manager next. * * This method does not block, instead a future is returned. */ - KafkaFuture pause(); + KafkaFuture unassign(); } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskManager.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskManager.java index f3dd834e7a238..4d5ffbb4b4b36 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskManager.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/tasks/TaskManager.java @@ -9,46 +9,54 @@ public interface TaskManager { /** - * Lock the next processible active task for the requested executor. Once the task is locked by - * the requested task executor, it should not be locked to any other executors until it was + * Get the next processible active task for the requested executor. Once the task is assigned to + * the requested task executor, it should not be assigned to any other executors until it was * returned to the task manager. * * @param executor the requesting {@link TaskExecutor} */ - StreamTask lockNextTask(final TaskExecutor executor); + StreamTask assignNextTask(final TaskExecutor executor); /** - * Unlock the stream task so that it can be processed by other executors later + * Unassign the stream task so that it can be assigned to other executors later * or be removed from the task manager. The requested executor must have locked * the task already, otherwise an exception would be thrown. * * @param executor the requesting {@link TaskExecutor} */ - void unlockTask(final StreamTask task, final TaskExecutor executor); + void unassignTask(final StreamTask task, final TaskExecutor executor); /** - * Lock a set of active tasks from the task manager. + * Lock a set of active tasks from the task manager so that they will not be assigned to + * any {@link TaskExecutor}s anymore until they are unlocked. At the time this function + * is called, the requested tasks may already be locked by some {@link TaskExecutor}s, + * and in that case the task manager need to first unassign these tasks from the + * executors. * * This function is needed when we need to 1) commit these tasks, 2) remove these tasks. * - * The requested tasks may be locked by {@link TaskExecutor}s at the time, and in that case - * the task manager need to ask these executors to unlock the tasks first - * * This method does not block, instead a future is returned. */ - KafkaFuture> lockTasks(final Set taskIds); + KafkaFuture lockTasks(final Set taskIds); /** - * Lock all of the managed active tasks from the task manager. - * - * This function is needed when we need to 1) commit these tasks, 2) remove these tasks. + * Lock all of the managed active tasks from the task manager. Similar to {@link #lockTasks(Set)}. * - * The requested tasks may be locked by {@link TaskExecutor}s at the time, and in that case - * the task manager need to ask these executors to unlock the tasks first + * This method does not block, instead a future is returned. + */ + KafkaFuture lockAllTasks(); + + /** + * Unlock the tasks so that they can be assigned to executors + */ + void unlockTasks(final Set taskIds); + + /** + * Unlock all of the managed active tasks from the task manager. Similar to {@link #unlockTasks(Set)}. * * This method does not block, instead a future is returned. */ - KafkaFuture> lockAllTasks(); + void unlockAllTasks(); /** * Add a new active task to the task manager.