Threading Refactor P2: Task Manager Impl

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<Task> 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<Task> tasks(final Collection<TaskId> taskIds) {
         return tasks;
     }
 
-    Collection<Task> activeTasks() {
+    public Collection<Task> 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 {

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<StreamTask> 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<StreamTask> pause() {
+    @Override
+    public KafkaFuture<StreamTask> unassign() {
         final KafkaFutureImpl<StreamTask> future = new KafkaFutureImpl<>();
 
         if (taskExecutorThread != null) {

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<TaskId> lockedTasks = new ArrayList<>();
+    private final Map<TaskId, TaskExecutor> assignedTasks = new HashMap<>();
+
+    private final List<TaskExecutor> 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<Void> lockTasks(final Set<TaskId> taskIds) {
+        return returnWithTasksLocked(() -> {
+            lockedTasks.addAll(taskIds);
+
+            final KafkaFutureImpl<Void> result = new KafkaFutureImpl<>();
+            final Set<TaskId> 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<StreamTask> 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<Void> lockAllTasks() {
+        return returnWithTasksLocked(() ->
+            lockTasks(tasks.activeTasks().stream().map(Task::id).collect(Collectors.toSet()))
+        );
+    }
+
+    public void unlockTasks(final Set<TaskId> taskIds) {
+        executeWithTasksLocked(() -> lockedTasks.removeAll(taskIds));
+    }
+
+    public void unlockAllTasks() {
+        executeWithTasksLocked(() -> unlockTasks(tasks.activeTasks().stream().map(Task::id).collect(Collectors.toSet())));
+    }
+
+    public void add(final Set<StreamTask> 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<ReadOnlyTask> 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> T returnWithTasksLocked(final Supplier<T> action) {
+        tasksLock.lock();
+        try {
+            return action.get();
+        } finally {
+            tasksLock.unlock();
+        }
+    }
+}

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<StreamTask> pause();
+    KafkaFuture<StreamTask> unassign();
 }

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<Set<TaskId>> lockTasks(final Set<TaskId> taskIds);
+    KafkaFuture<Void> lockTasks(final Set<TaskId> 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<Void> lockAllTasks();
+
+    /**
+     * Unlock the tasks so that they can be assigned to executors
+     */
+    void unlockTasks(final Set<TaskId> 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<Set<TaskId>> lockAllTasks();
+    void unlockAllTasks();
 
     /**
      * Add a new active task to the task manager.