KAFKA-9113: StreamThread unit test

clients/src/main/java/org/apache/kafka/clients/producer/MockProducer.java

@@ -67,7 +67,7 @@ public class MockProducer<K, V> implements Producer<K, V> {
     private boolean transactionCommitted;
     private boolean transactionAborted;
     private boolean producerFenced;
-    private boolean producerFencedOnClose;
+    private boolean producerFencedOnCommitTxn;
     private boolean sentOffsets;
     private long commitCount = 0L;
     private Map<MetricName, Metric> mockMetrics;
@@ -182,6 +182,10 @@ public void sendOffsetsToTransaction(Map<TopicPartition, OffsetAndMetadata> offs
 
     @Override
     public void commitTransaction() throws ProducerFencedException {
+        if (producerFencedOnCommitTxn) {
+            throw new ProducerFencedException("Producer is fenced");
+        }
+
         verifyProducerState();
         verifyTransactionsInitialized();
         verifyNoTransactionInFlight();
@@ -325,9 +329,6 @@ public void close() {
 
     @Override
     public void close(Duration timeout) {
-        if (producerFencedOnClose) {
-            throw new ProducerFencedException("MockProducer is fenced.");
-        }
         this.closed = true;
     }
 
@@ -341,10 +342,10 @@ public synchronized void fenceProducer() {
         this.producerFenced = true;
     }
 
-    public void fenceProducerOnClose() {
+    public void fenceProducerOnCommitTxn() {
         verifyProducerState();
         verifyTransactionsInitialized();
-        this.producerFencedOnClose = true;
+        this.producerFencedOnCommitTxn = true;
     }
 
     public boolean transactionInitialized() {

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

@@ -1226,14 +1226,15 @@ synchronized Collection<String> sourceTopicCollection() {
     }
 
     synchronized Pattern sourceTopicPattern() {
-        log.debug("Found pattern subscribed source topics, falling back to pattern subscription for the main consumer.");
-
         if (topicPattern == null) {
             final List<String> allSourceTopics = maybeDecorateInternalSourceTopics(sourceTopicNames);
             Collections.sort(allSourceTopics);
             topicPattern = buildPattern(allSourceTopics, nodeToSourcePatterns.values());
         }
 
+        log.debug("Found pattern subscribed source topics, falling back to pattern " +
+            "subscription for the main consumer: {}", topicPattern);
+
         return topicPattern;
     }
 
@@ -1870,19 +1871,7 @@ private boolean hasSubscriptionUpdates() {
         return !subscriptionUpdates.isEmpty();
     }
 
-    synchronized void updateSubscribedTopics(final Set<String> topics,
-                                             final String logPrefix) {
-        log.debug("{}found {} topics possibly matching subscription", logPrefix, topics);
-        subscriptionUpdates.clear();
-        subscriptionUpdates.addAll(topics);
-
-        log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
-            logPrefix, subscriptionUpdates);
-        setRegexMatchedTopicsToSourceNodes();
-        setRegexMatchedTopicToStateStore();
-    }
-
-    void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, final String logPrefix) {
+    synchronized void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, final String logPrefix) {
         if (sourceTopicPattern() != null) {
             final Set<String> assignedTopics = new HashSet<>();
             for (final TopicPartition topicPartition : partitions) {
@@ -1892,21 +1881,33 @@ void addSubscribedTopicsFromAssignment(final List<TopicPartition> partitions, fi
             final Collection<String> existingTopics = subscriptionUpdates();
             if (!existingTopics.containsAll(assignedTopics)) {
                 assignedTopics.addAll(existingTopics);
+                updateSubscribedTopics(assignedTopics, logPrefix);
+            }
+        }
+    }
 
-                log.debug("{}found {} topics possibly matching subscription", logPrefix, assignedTopics);
-                subscriptionUpdates.clear();
-                subscriptionUpdates.addAll(assignedTopics);
-
-                log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
-                    logPrefix, subscriptionUpdates);
-                setRegexMatchedTopicsToSourceNodes();
-                setRegexMatchedTopicToStateStore();
+    synchronized void addSubscribedTopicsFromMetadata(final Set<String> topics, final String logPrefix) {
+        if (sourceTopicPattern() != null) {
+            final Collection<String> existingTopics = subscriptionUpdates();
+            if (!existingTopics.equals(topics)) {
+                topics.addAll(existingTopics);
+                updateSubscribedTopics(topics, logPrefix);
             }
         }
     }
 
-    // following functions are for test only
+    private void updateSubscribedTopics(final Set<String> topics, final String logPrefix) {
+        log.debug("{}found {} topics possibly matching subscription", logPrefix, topics);
+        subscriptionUpdates.clear();
+        subscriptionUpdates.addAll(topics);
 
+        log.debug("{}updating builder with {} topic(s) with possible matching regex subscription(s)",
+            logPrefix, subscriptionUpdates);
+        setRegexMatchedTopicsToSourceNodes();
+        setRegexMatchedTopicToStateStore();
+    }
+
+    // following functions are for test only
     public synchronized Set<String> sourceTopicNames() {
         return sourceTopicNames;
     }

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

@@ -189,7 +189,8 @@ public StateStore getGlobalStore(final String name) {
         return globalStores.get(name);
     }
 
-    public void initializeStoreOffsetsFromCheckpoint() {
+    // package-private for test only
+    void initializeStoreOffsetsFromCheckpoint() {
         try {
             final Map<TopicPartition, Long> loadedCheckpoints = checkpointFile.read();
 

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

@@ -27,11 +27,6 @@
 
 public interface RecordCollector extends AutoCloseable {
 
-    /**
-     * Initialize the record collector
-     */
-    void initialize();
-
     <K, V> void send(final String topic,
                      final K key,
                      final V value,

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

@@ -73,6 +73,9 @@ public class RecordCollectorImpl implements RecordCollector {
     private Producer<byte[], byte[]> producer;
     private volatile KafkaException sendException;
 
+    /**
+     * @throws StreamsException fatal error that should cause the thread to die (from producer.initTxn)
+     */
     public RecordCollectorImpl(final TaskId taskId,
                                final StreamsConfig config,
                                final LogContext logContext,
@@ -93,16 +96,6 @@ public RecordCollectorImpl(final TaskId taskId,
 
         producer = producerSupplier.get(taskId);
 
-        // TODO K9113: this should be moved to task when it transits to running from created / restoring
-        // then even if there's a long time between the initialization and the first txn that is also fine.
-        initialize();
-    }
-
-    /**
-     * @throws StreamsException fatal error that should cause the thread to die
-     */
-    @Override
-    public void initialize() {
         maybeInitTxns();
     }
 
@@ -333,7 +326,9 @@ private void checkForException() {
     @Override
     public void flush() {
         log.debug("Flushing record collector");
+
         producer.flush();
+
         checkForException();
     }
 

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

@@ -38,33 +38,37 @@
  */
 public class StandbyTask extends AbstractTask implements Task {
     private final TaskId id;
+    private final Logger log;
+    private final String logPrefix;
     private final ProcessorTopology topology;
+    private final Set<TopicPartition> partitions;
     private final ProcessorStateManager stateMgr;
-    private final String logPrefix;
-    private final Logger log;
     private final StateDirectory stateDirectory;
     private final Sensor closeTaskSensor;
     private final InternalProcessorContext processorContext;
 
     /**
      * @param id             the ID of this task
+     * @param partitions     input topic partitions, used for thread metadata only
      * @param topology       the instance of {@link ProcessorTopology}
      * @param config         the {@link StreamsConfig} specified by the user
      * @param metrics        the {@link StreamsMetrics} created by the thread
+     * @param stateMgr       the {@link ProcessorStateManager} for this task
      * @param stateDirectory the {@link StateDirectory} created by the thread
      */
     StandbyTask(final TaskId id,
+                final Set<TopicPartition> partitions,
                 final ProcessorTopology topology,
                 final StreamsConfig config,
                 final StreamsMetricsImpl metrics,
                 final ProcessorStateManager stateMgr,
                 final StateDirectory stateDirectory) {
         this.id = id;
+        this.stateMgr = stateMgr;
         this.topology = topology;
+        this.partitions = partitions;
         this.stateDirectory = stateDirectory;
 
-        this.stateMgr = stateMgr;
-
         final String threadIdPrefix = String.format("stream-thread [%s] ", Thread.currentThread().getName());
         logPrefix = threadIdPrefix + String.format("%s [%s] ", "standby-task", id);
         final LogContext logContext = new LogContext(logPrefix);
@@ -194,7 +198,7 @@ public TaskId id() {
 
     @Override
     public Set<TopicPartition> inputPartitions() {
-        return Collections.emptySet();
+        return partitions;
     }
 
     @Override

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

@@ -643,7 +643,7 @@ private void addChangelogsToRestoreConsumer(final Set<TopicPartition> partitions
 
         // the current assignment should not contain any of the new partitions
         if (assignment.removeAll(partitions)) {
-            throw new IllegalStateException("The current assignment " + assignment + " " +
+            throw new IllegalStateException("The current assignment " + restoreConsumer.assignment() + " " +
                 "already contains some of the new partitions " + partitions);
         }
         assignment.addAll(partitions);

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

@@ -20,7 +20,6 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.InvalidOffsetException;
 import org.apache.kafka.clients.producer.Producer;
@@ -53,7 +52,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -472,6 +470,7 @@ StandbyTask createTask(final Consumer<byte[], byte[]> consumer,
 
                 return new StandbyTask(
                     taskId,
+                    partitions,
                     topology,
                     config,
                     streamsMetrics,
@@ -513,7 +512,6 @@ StandbyTask createTask(final Consumer<byte[], byte[]> consumer,
     private volatile State state = State.CREATED;
     private volatile ThreadMetadata threadMetadata;
     private StreamThread.StateListener stateListener;
-    private Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> standbyRecords;
 
     private final ChangelogReader changelogReader;
 
@@ -633,11 +631,9 @@ public StreamThread(final Time time,
                         final LogContext logContext,
                         final AtomicInteger assignmentErrorCode) {
         super(threadId);
-        this.adminClient = adminClient;
-
         this.stateLock = new Object();
-        this.standbyRecords = new HashMap<>();
 
+        this.adminClient = adminClient;
         this.streamsMetrics = streamsMetrics;
         this.commitSensor = ThreadMetrics.commitSensor(threadId, streamsMetrics);
         this.pollSensor = ThreadMetrics.pollSensor(threadId, streamsMetrics);
@@ -1111,12 +1107,6 @@ private void completeShutdown(final boolean cleanRun) {
         log.info("Shutdown complete");
     }
 
-    void clearStandbyRecords(final List<TopicPartition> partitions) {
-        for (final TopicPartition tp : partitions) {
-            standbyRecords.remove(tp);
-        }
-    }
-
     /**
      * Return information about the current {@link StreamThread}.
      *
@@ -1167,7 +1157,7 @@ private void updateThreadMetadata(final Map<TaskId, Task> activeTasks,
             standbyTasksMetadata);
     }
 
-    public Iterable<Task> activeTasks() {
+    public List<Task> activeTasks() {
         return taskManager.activeTaskIterable();
     }
 
@@ -1238,15 +1228,11 @@ TaskManager taskManager() {
         return taskManager;
     }
 
-    Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> standbyRecords() {
-        return standbyRecords;
-    }
-
     int currentNumIterations() {
         return numIterations;
     }
 
-    public StreamThread.StateListener stateListener() {
-        return stateListener;
+    Throwable rebalanceException() {
+        return rebalanceException;
     }
 }