KAFKA-9113: StreamTask unit test

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

@@ -16,12 +16,7 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
-import static org.apache.kafka.streams.processor.internals.Task.State.CLOSED;
-import static org.apache.kafka.streams.processor.internals.Task.State.CLOSING;
 import static org.apache.kafka.streams.processor.internals.Task.State.CREATED;
-import static org.apache.kafka.streams.processor.internals.Task.State.RESTORING;
-import static org.apache.kafka.streams.processor.internals.Task.State.RUNNING;
-import static org.apache.kafka.streams.processor.internals.Task.State.SUSPENDED;
 
 public abstract class AbstractTask implements Task {
     private Task.State state = CREATED;
@@ -33,14 +28,11 @@ public final Task.State state() {
 
     final void transitionTo(final Task.State newState) {
         final State oldState = state();
-        if (oldState == CREATED && (newState == RESTORING || newState == CLOSING)) {
-        } else if (oldState == RESTORING && (newState == RUNNING || newState == SUSPENDED || newState == CLOSING)) {
-        } else if (oldState == RUNNING && (newState == SUSPENDED || newState == CLOSING)) {
-        } else if (oldState == SUSPENDED && (newState == RUNNING || newState == CLOSING)) {
-        } else if (oldState == CLOSING && (newState == CLOSING || newState == CLOSED)) {
+
+        if (oldState.isValidTransition(newState)) {
+            state = newState;
         } else {
             throw new IllegalStateException("Invalid transition from " + oldState + " to " + newState);
         }
-        state = newState;
     }
 }

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

@@ -207,11 +207,11 @@ boolean allPartitionsBuffered() {
     public void close() {
         clear();
         partitionQueues.clear();
+        streamTime = RecordQueue.UNKNOWN;
     }
 
     public void clear() {
         nonEmptyQueuesByTime.clear();
-        streamTime = RecordQueue.UNKNOWN;
         for (final RecordQueue queue : partitionQueues.values()) {
             queue.clear();
         }

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

@@ -435,8 +435,7 @@ public void checkpoint(final Map<TopicPartition, Long> writtenOffsets) {
             // store is logged, persistent, and has a valid current offset
             if (storeMetadata.changelogPartition != null &&
                 storeMetadata.stateStore.persistent() &&
-                storeMetadata.offset != null
-            ) {
+                storeMetadata.offset != null) {
                 checkpointingOffsets.put(storeMetadata.changelogPartition, storeMetadata.offset);
             }
         }

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

@@ -157,12 +157,11 @@ public long headRecordTimestamp() {
     }
 
     /**
-     * Clear the fifo queue of its elements, also clear the time tracker's kept stamped elements
+     * Clear the fifo queue of its elements
      */
     public void clear() {
         fifoQueue.clear();
         headRecord = null;
-        partitionTime = RecordQueue.UNKNOWN;
     }
 
     private void updateHead() {

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

@@ -79,58 +79,69 @@ public boolean isActive() {
         return false;
     }
 
+    /**
+     * @throws StreamsException fatal error, should close the thread
+     */
     @Override
     public void initializeIfNeeded() {
         if (state() == State.CREATED) {
-            initializeStateStores();
+            TaskUtils.registerStateStores(id, log, logPrefix, topology, stateMgr, stateDirectory, processorContext);
+
             // no topology needs initialized, we can transit to RUNNING
             // right after registered the stores
             transitionTo(State.RESTORING);
             transitionTo(State.RUNNING);
 
+            processorContext.initialize();
+
             log.debug("Initialized");
         }
     }
 
-    private void initializeStateStores() {
-        TaskUtils.registerStateStores(topology, stateDirectory, id, logPrefix, log, processorContext, stateMgr);
-
-        processorContext.initialize();
-    }
-
     @Override
-    public void completeInitializationAfterRestore() {
-        throw new IllegalStateException("StandbyTask " + id + " should already be in running state.");
+    public void completeRestoration() {
+        throw new IllegalStateException("Standby task " + id + " should never be completing restoration");
     }
 
     @Override
     public void suspend() {
-        log.debug("No-op suspend.");
+        log.trace("No-op suspend with state {}", state());
     }
 
     @Override
     public void resume() {
-        log.debug("No-op resume");
+        log.trace("No-op resume with state {}", state());
     }
 
     /**
-     * <pre>
-     * - flush store
-     * - checkpoint store
-     * </pre>
+     * 1. flush store
+     * 2. write checkpoint file
+     *
+     * @throws TaskMigratedException all the task has been migrated
+     * @throws StreamsException fatal error, should close the thread
      */
     @Override
     public void commit() {
-        if (state() == State.RUNNING) {
-            stateMgr.flush();
+        switch (state()) {
+            case RUNNING:
+                stateMgr.flush();
 
-            // since there's no written offsets we can checkpoint with empty map,
-            // and the state current offset would be used to checkpoint
-            stateMgr.checkpoint(Collections.emptyMap());
+                // since there's no written offsets we can checkpoint with empty map,
+                // and the state current offset would be used to checkpoint
+                stateMgr.checkpoint(Collections.emptyMap());
+
+                log.debug("Committed");
+                break;
+
+            case CLOSING:
+                // do nothing and also not throw
+                log.trace("Skip committing since task is closing");
+
+                break;
+
+            default:
+                throw new IllegalStateException("Illegal state " + state() + " while committing standby task " + id);
 
-            log.debug("Committed");
-        } else {
-            throw new IllegalStateException("Illegal state " + state() + " while committing standby task " + id);
         }
     }
 
@@ -145,37 +156,29 @@ public void closeDirty() {
     }
 
     /**
-     * 1. when unclean close, we do not need to commit;
-     * 2. when unclean close, we do not throw any exception;
+     * 1. commit if we are running and clean close;
+     * 2. close the state manager.
+     *
+     * @throws TaskMigratedException all the task has been migrated
+     * @throws StreamsException fatal error, should close the thread
      */
     private void close(final boolean clean) {
-        switch (state()) {
-            case CREATED:
-                transitionTo(State.CLOSING);
-                // the task is created and not initialized, do nothing
-                break;
-
-            case RUNNING:
-                // No resources are destroyed here, so we can do this before transit to CLOSING
-                if (clean) {
+        if (state() == State.CREATED) {
+            // the task is created and not initialized, do nothing
+            transitionTo(State.CLOSING);
+        } else {
+            if (state() == State.RUNNING) {
+                if (clean)
                     commit();
-                }
 
                 transitionTo(State.CLOSING);
+            }
 
-                try {
-                    TaskUtils.closeStateManager(log, logPrefix, stateMgr, stateDirectory, id);
-                } catch (final RuntimeException error) {
-                    if (clean) {
-                        throw error;
-                    } else {
-                        log.warn("Closing standby task " + id + " uncleanly throws an exception " + error);
-                    }
-                }
-                break;
-
-            default:
+            if (state() == State.CLOSING) {
+                TaskUtils.closeStateManager(id, log, logPrefix, clean, stateMgr, stateDirectory);
+            } else {
                 throw new IllegalStateException("Illegal state " + state() + " while closing standby task " + id);
+            }
         }
 
         closeTaskSensor.record();