KAFKA-14299: Fix pause and resume with state updater

streams/src/main/java/org/apache/kafka/streams/KafkaStreams.java

@@ -1751,6 +1751,7 @@ public void resume() {
         } else {
             topologyMetadata.resumeTopology(UNNAMED_TOPOLOGY);
         }
+        threads.forEach(StreamThread::signalResume);
     }
 
     /**

streams/src/main/java/org/apache/kafka/streams/processor/internals/DefaultStateUpdater.java

@@ -116,6 +116,7 @@ public void run() {
 
         private void runOnce() throws InterruptedException {
             performActionsOnTasks();
+            resumeTasks();
             restoreTasks();
             checkAllUpdatingTaskStates(time.milliseconds());
             waitIfAllChangelogsCompletelyRead();
@@ -140,6 +141,16 @@ private void performActionsOnTasks() {
             }
         }
 
+        private void resumeTasks() {
+            if (isTopologyResumed.compareAndSet(true, false)) {
+                for (final Task task : pausedTasks.values()) {
+                    if (!topologyMetadata.isPaused(task.id().topologyName())) {
+                        resumeTask(task);
+                    }
+                }
+            }
+        }
+
         private void restoreTasks() {
             try {
                 changelogReader.restore(updatingTasks);
@@ -229,7 +240,7 @@ private void waitIfAllChangelogsCompletelyRead() throws InterruptedException {
             if (isRunning.get() && changelogReader.allChangelogsCompleted()) {
                 tasksAndActionsLock.lock();
                 try {
-                    while (tasksAndActions.isEmpty()) {
+                    while (tasksAndActions.isEmpty() && !isTopologyResumed.get()) {
                         tasksAndActionsCondition.await();
                     }
                 } finally {
@@ -258,21 +269,39 @@ private List<TaskAndAction> getTasksAndActions() {
         }
 
         private void addTask(final Task task) {
+            final TaskId taskId = task.id();
+
+            Task existingTask = pausedTasks.get(taskId);
+            if (existingTask != null) {
+                throw new IllegalStateException(
+                    (existingTask.isActive() ? "Active" : "Standby") + " task " + taskId + " already exist in paused tasks, " +
+                        "should not try to add another " + (task.isActive() ? "active" : "standby") + " task with the same id. "
+                        + BUG_ERROR_MESSAGE);
+            }
+            existingTask = updatingTasks.get(taskId);
+            if (existingTask != null) {
+                throw new IllegalStateException(
+                    (existingTask.isActive() ? "Active" : "Standby") + " task " + taskId + " already exist in updating tasks, " +
+                        "should not try to add another " + (task.isActive() ? "active" : "standby") + " task with the same id. "
+                        + BUG_ERROR_MESSAGE);
+            }
+
             if (isStateless(task)) {
                 addToRestoredTasks((StreamTask) task);
-                log.info("Stateless active task " + task.id() + " was added to the restored tasks of the state updater");
+                log.info("Stateless active task " + taskId + " was added to the restored tasks of the state updater");
+            } else if (topologyMetadata.isPaused(taskId.topologyName())) {
+                pausedTasks.put(taskId, task);
+                changelogReader.register(task.changelogPartitions(), task.stateManager());
+                log.debug((task.isActive() ? "Active" : "Standby")
+                    + " task " + taskId + " was directly added to the paused tasks.");
             } else {
-                final Task existingTask = updatingTasks.putIfAbsent(task.id(), task);
-                if (existingTask != null) {
-                    throw new IllegalStateException((existingTask.isActive() ? "Active" : "Standby") + " task " + task.id() + " already exist, " +
-                        "should not try to add another " + (task.isActive() ? "active" : "standby") + " task with the same id. " + BUG_ERROR_MESSAGE);
-                }
+                updatingTasks.put(taskId, task);
                 changelogReader.register(task.changelogPartitions(), task.stateManager());
                 if (task.isActive()) {
-                    log.info("Stateful active task " + task.id() + " was added to the state updater");
+                    log.info("Stateful active task " + taskId + " was added to the state updater");
                     changelogReader.enforceRestoreActive();
                 } else {
-                    log.info("Standby task " + task.id() + " was added to the state updater");
+                    log.info("Standby task " + taskId + " was added to the state updater");
                     if (updatingTasks.size() == 1) {
                         changelogReader.transitToUpdateStandby();
                     }
@@ -388,12 +417,6 @@ private void checkAllUpdatingTaskStates(final long now) {
                     }
                 }
 
-                for (final Task task : pausedTasks.values()) {
-                    if (!topologyMetadata.isPaused(task.id().topologyName())) {
-                        resumeTask(task);
-                    }
-                }
-
                 lastCommitMs = now;
             }
         }
@@ -411,6 +434,7 @@ private void checkAllUpdatingTaskStates(final long now) {
     private final Condition restoredActiveTasksCondition = restoredActiveTasksLock.newCondition();
     private final BlockingQueue<ExceptionAndTasks> exceptionsAndFailedTasks = new LinkedBlockingQueue<>();
     private final BlockingQueue<Task> removedTasks = new LinkedBlockingQueue<>();
+    private final AtomicBoolean isTopologyResumed = new AtomicBoolean(false);
 
     private final long commitIntervalMs;
     private long lastCommitMs;
@@ -506,6 +530,17 @@ public void remove(final TaskId taskId) {
         }
     }
 
+    @Override
+    public void signalResume() {
+        tasksAndActionsLock.lock();
+        try {
+            isTopologyResumed.set(true);
+            tasksAndActionsCondition.signalAll();
+        } finally {
+            tasksAndActionsLock.unlock();
+        }
+    }
+
     @Override
     public Set<StreamTask> drainRestoredActiveTasks(final Duration timeout) {
         final long timeoutMs = timeout.toMillis();

streams/src/main/java/org/apache/kafka/streams/processor/internals/ReadOnlyTask.java

@@ -200,7 +200,7 @@ public StateStore getStore(final String name) {
 
     @Override
     public Map<TopicPartition, Long> changelogOffsets() {
-        throw new UnsupportedOperationException("This task is read-only");
+        return task.changelogOffsets();
     }
 
     @Override

streams/src/main/java/org/apache/kafka/streams/processor/internals/StateUpdater.java

@@ -94,6 +94,11 @@ public int hashCode() {
      */
     void remove(final TaskId taskId);
 
+    /**
+     * Wakes up the state updater if it is currently dormant, to check if a paused task should be resumed.
+     */
+    void signalResume();
+
     /**
      * Drains the restored active tasks from the state updater.
      *

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java

@@ -1096,6 +1096,11 @@ public boolean isThreadAlive() {
         return isAlive();
     }
 
+    // Call method when a topology is resumed
+    public void signalResume() {
+        taskManager.signalResume();
+    }
+
     /**
      * Try to commit all active tasks owned by this thread.
      *
@@ -1113,7 +1118,7 @@ int maybeCommit() {
             }
 
             committed = taskManager.commit(
-                taskManager.allTasks()
+                taskManager.allOwnedTasks()
                     .values()
                     .stream()
                     .filter(t -> t.state() == Task.State.RUNNING || t.state() == Task.State.RESTORING)

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

@@ -1115,6 +1115,12 @@ private void removeLostActiveTasksFromStateUpdater() {
         }
     }
 
+    public void signalResume() {
+        if (stateUpdater != null) {
+            stateUpdater.signalResume();
+        }
+    }
+
     /**
      * Compute the offset total summed across all stores in a task. Includes offset sum for any tasks we own the
      * lock for, which includes assigned and unassigned tasks we locked in {@link #tryToLockAllNonEmptyTaskDirectories()}.
@@ -1532,15 +1538,38 @@ Set<TaskId> standbyTaskIds() {
     }
 
     Map<TaskId, Task> allTasks() {
+        // not bothering with an unmodifiable map, since the tasks themselves are mutable, but
+        // if any outside code modifies the map or the tasks, it would be a severe transgression.
+        if (stateUpdater != null) {
+            final Map<TaskId, Task> ret = stateUpdater.getTasks().stream().collect(Collectors.toMap(Task::id, x -> x));
+            ret.putAll(tasks.allTasksPerId());
+            return ret;
+        } else {
+            return tasks.allTasksPerId();
+        }
+    }
+
+    /**
+     * Returns tasks owned by the stream thread. With state updater disabled, these are all tasks. With
+     * state updater enabled, this does not return any tasks currently owned by the state updater.
+     * @return
+     */
+    Map<TaskId, Task> allOwnedTasks() {
         // not bothering with an unmodifiable map, since the tasks themselves are mutable, but
         // if any outside code modifies the map or the tasks, it would be a severe transgression.
         return tasks.allTasksPerId();
     }
 
     Set<Task> readOnlyAllTasks() {
-        // need to make sure the returned set is unmodifiable as it could be accessed
-        // by other thread than the StreamThread owning this task manager;
-        return Collections.unmodifiableSet(tasks.allTasks());
+        // not bothering with an unmodifiable map, since the tasks themselves are mutable, but
+        // if any outside code modifies the map or the tasks, it would be a severe transgression.
+        if (stateUpdater != null) {
+            final HashSet<Task> ret = new HashSet<>(stateUpdater.getTasks());
+            ret.addAll(tasks.allTasks());
+            return Collections.unmodifiableSet(ret);
+        } else {
+            return Collections.unmodifiableSet(tasks.allTasks());
+        }
     }
 
     Map<TaskId, Task> notPausedTasks() {

streams/src/main/java/org/apache/kafka/streams/processor/internals/namedtopology/KafkaStreamsNamedTopologyWrapper.java

@@ -40,6 +40,7 @@
 import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.internals.DefaultKafkaClientSupplier;
 import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder;
+import org.apache.kafka.streams.processor.internals.StreamThread;
 import org.apache.kafka.streams.processor.internals.Task;
 import org.apache.kafka.streams.processor.internals.TopologyMetadata;
 import org.slf4j.Logger;
@@ -276,6 +277,7 @@ public boolean isNamedTopologyPaused(final String topologyName) {
      */
     public void resumeNamedTopology(final String topologyName) {
         topologyMetadata.resumeTopology(topologyName);
+        threads.forEach(StreamThread::signalResume);
     }
 
     /**