KAFKA-10754: fix flaky tests by waiting kafka streams be in running state before assert

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

@@ -350,7 +350,7 @@ public void setStateListener(final KafkaStreams.StateListener listener) {
     }
 
     /**
-     * Set the handler invoked when a {@link StreamsConfig#NUM_STREAM_THREADS_CONFIG internal thread} abruptly
+     * Set the handler invoked when an internal {@link StreamsConfig#NUM_STREAM_THREADS_CONFIG stream thread} abruptly
      * terminates due to an uncaught exception.
      *
      * @param uncaughtExceptionHandler the uncaught exception handler for all internal threads; {@code null} deletes the current handler
@@ -379,13 +379,12 @@ public void setUncaughtExceptionHandler(final Thread.UncaughtExceptionHandler un
     }
 
     /**
-     * Set the handler invoked when an {@link StreamsConfig#NUM_STREAM_THREADS_CONFIG internal thread}
+     * Set the handler invoked when an internal {@link StreamsConfig#NUM_STREAM_THREADS_CONFIG stream thread}
      * throws an unexpected exception.
      * These might be exceptions indicating rare bugs in Kafka Streams, or they
-     * might be exceptions thrown by your code, for example a NullPointerException thrown from your processor
-     * logic.
+     * might be exceptions thrown by your code, for example a NullPointerException thrown from your processor logic.
      * The handler will execute on the thread that produced the exception.
-     * In order to get the thread that threw the exception, Thread.currentThread().
+     * In order to get the thread that threw the exception, use {@code Thread.currentThread()}.
      * <p>
      * Note, this handler must be threadsafe, since it will be shared among all threads, and invoked from any
      * thread that encounters such an exception.

streams/src/test/java/org/apache/kafka/streams/integration/StreamsUncaughtExceptionHandlerIntegrationTest.java

@@ -48,8 +48,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.apache.kafka.common.utils.Utils.mkEntry;
 import static org.apache.kafka.common.utils.Utils.mkMap;
@@ -68,6 +67,7 @@
 public class StreamsUncaughtExceptionHandlerIntegrationTest {
     @ClassRule
     public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(1);
+    public static final Duration DEFAULT_DURATION = Duration.ofSeconds(30);
 
     @Rule
     public TestName testName = new TestName();
@@ -110,83 +110,49 @@ public void teardown() throws IOException {
     }
 
     @Test
-    public void shouldShutdownThreadUsingOldHandler() throws Exception {
+    public void shouldShutdownThreadUsingOldHandler() throws InterruptedException {
         try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), properties)) {
-            final CountDownLatch latch = new CountDownLatch(1);
-            final AtomicBoolean flag = new AtomicBoolean(false);
-            kafkaStreams.setUncaughtExceptionHandler((t, e) -> flag.set(true));
+            final AtomicInteger counter = new AtomicInteger(0);
+            kafkaStreams.setUncaughtExceptionHandler((t, e) -> counter.incrementAndGet());
 
             StreamsTestUtils.startKafkaStreamsAndWaitForRunningState(kafkaStreams);
-
             produceMessages(0L, inputTopic, "A");
-            waitForApplicationState(Collections.singletonList(kafkaStreams), KafkaStreams.State.ERROR, Duration.ofSeconds(15));
 
-            TestUtils.waitForCondition(flag::get, "Handler was called");
+            // should call the UncaughtExceptionHandler in current thread
+            TestUtils.waitForCondition(() -> counter.get() == 1, "Handler was called 1st time");
+            // should call the UncaughtExceptionHandler after rebalancing to another thread
+            TestUtils.waitForCondition(() -> counter.get() == 2, DEFAULT_DURATION.toMillis(), "Handler was called 2nd time");
+            // the stream should now turn into ERROR state after 2 threads are dead
+            waitForApplicationState(Collections.singletonList(kafkaStreams), KafkaStreams.State.ERROR, DEFAULT_DURATION);
+
             assertThat(processorValueCollector.size(), equalTo(2));
         }
     }
 
     @Test
-    public void shouldShutdownClient() throws Exception {
+    public void shouldShutdownClient() throws InterruptedException {
         try (final KafkaStreams kafkaStreams = new KafkaStreams(builder.build(), properties)) {
-            final CountDownLatch latch = new CountDownLatch(1);
             kafkaStreams.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
 
             kafkaStreams.setUncaughtExceptionHandler(exception -> SHUTDOWN_CLIENT);
 
             StreamsTestUtils.startKafkaStreamsAndWaitForRunningState(kafkaStreams);
 
             produceMessages(0L, inputTopic, "A");
-            waitForApplicationState(Collections.singletonList(kafkaStreams), KafkaStreams.State.NOT_RUNNING, Duration.ofSeconds(15));
+            waitForApplicationState(Collections.singletonList(kafkaStreams), KafkaStreams.State.NOT_RUNNING, DEFAULT_DURATION);
 
             assertThat(processorValueCollector.size(), equalTo(1));
         }
     }
 
     @Test
-    public void shouldShutdownApplication() throws Exception {
-        final Topology topology = builder.build();
-
-        try (final KafkaStreams kafkaStreams = new KafkaStreams(topology, properties)) {
-            final KafkaStreams kafkaStreams1 = new KafkaStreams(topology, properties);
-            final CountDownLatch latch = new CountDownLatch(1);
-            kafkaStreams.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
-            kafkaStreams1.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
-            kafkaStreams.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
-            kafkaStreams1.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
-
-            kafkaStreams.start();
-            kafkaStreams1.start();
-
-            produceMessages(0L, inputTopic, "A");
-            waitForApplicationState(Arrays.asList(kafkaStreams, kafkaStreams1), KafkaStreams.State.ERROR, Duration.ofSeconds(30));
-
-            assertThat(processorValueCollector.size(), equalTo(1));
-        }
+    public void shouldShutdownMultipleThreadApplication() throws InterruptedException {
+        testShutdownApplication(2);
     }
 
     @Test
-    public void shouldShutdownSingleThreadApplication() throws Exception {
-        properties.setProperty(StreamsConfig.NUM_STREAM_THREADS_CONFIG, "1");
-
-        final Topology topology = builder.build();
-
-        try (final KafkaStreams kafkaStreams = new KafkaStreams(topology, properties)) {
-            final KafkaStreams kafkaStreams1 = new KafkaStreams(topology, properties);
-            final CountDownLatch latch = new CountDownLatch(1);
-            kafkaStreams.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
-            kafkaStreams1.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
-            kafkaStreams.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
-            kafkaStreams1.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
-
-            kafkaStreams.start();
-            kafkaStreams1.start();
-
-            produceMessages(0L, inputTopic, "A");
-            waitForApplicationState(Arrays.asList(kafkaStreams, kafkaStreams1), KafkaStreams.State.ERROR, Duration.ofSeconds(30));
-
-            assertThat(processorValueCollector.size(), equalTo(1));
-        }
+    public void shouldShutdownSingleThreadApplication() throws InterruptedException {
+        testShutdownApplication(1);
     }
 
     private void produceMessages(final long timestamp, final String streamOneInput, final String msg) {
@@ -202,7 +168,6 @@ private void produceMessages(final long timestamp, final String streamOneInput,
     }
 
     private static class ShutdownProcessor extends AbstractProcessor<String, String> {
-
         final List<String> valueList;
 
         ShutdownProcessor(final List<String> valueList) {
@@ -215,6 +180,28 @@ public void process(final String key, final String value) {
             throw new StreamsException(Thread.currentThread().getName());
         }
     }
+
+    private void testShutdownApplication(final int numThreads) throws InterruptedException {
+        properties.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, numThreads);
+
+        final Topology topology = builder.build();
+
+        try (final KafkaStreams kafkaStreams1 = new KafkaStreams(topology, properties);
+             final KafkaStreams kafkaStreams2 = new KafkaStreams(topology, properties)) {
+            kafkaStreams1.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
+            kafkaStreams2.setUncaughtExceptionHandler((t, e) -> fail("should not hit old handler"));
+            kafkaStreams1.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
+            kafkaStreams2.setUncaughtExceptionHandler(exception -> SHUTDOWN_APPLICATION);
+
+            StreamsTestUtils.startKafkaStreamsAndWaitForRunningState(kafkaStreams1);
+            StreamsTestUtils.startKafkaStreamsAndWaitForRunningState(kafkaStreams2);
+
+            produceMessages(0L, inputTopic, "A");
+            waitForApplicationState(Arrays.asList(kafkaStreams1, kafkaStreams2), KafkaStreams.State.ERROR, DEFAULT_DURATION);
+
+            assertThat(processorValueCollector.size(), equalTo(1));
+        }
+    }
 }
 
 

streams/src/test/java/org/apache/kafka/streams/integration/utils/IntegrationTestUtils.java

@@ -833,7 +833,7 @@ private static void waitUntilMetadataIsPropagated(final List<KafkaServer> server
      * {@link State#RUNNING} state at the same time. Note that states may change between the time
      * that this method returns and the calling function executes its next statement.<p>
      *
-     * When the application is already started use {@link #waitForApplicationState(List, State, Duration)}
+     * If the application is already started, use {@link #waitForApplicationState(List, State, Duration)}
      * to wait for instances to reach {@link State#RUNNING} state.
      *
      * @param streamsList the list of streams instances to run.
@@ -903,16 +903,19 @@ public static void startApplicationAndWaitUntilRunning(final List<KafkaStreams>
     }
 
     /**
-     * Waits for the given {@link KafkaStreams} instances to all be in a {@link State#RUNNING}
-     * state. Prefer {@link #startApplicationAndWaitUntilRunning(List, Duration)} when possible
+     * Waits for the given {@link KafkaStreams} instances to all be in a specific {@link State}.
+     * Prefer {@link #startApplicationAndWaitUntilRunning(List, Duration)} when possible
      * because this method uses polling, which can be more error prone and slightly slower.
      *
      * @param streamsList the list of streams instances to run.
-     * @param timeout the time to wait for the streams to all be in {@link State#RUNNING} state.
+     * @param state the expected state that all the streams to be in within timeout
+     * @param timeout the time to wait for the streams to all be in the specific state.
+     *
+     * @throws InterruptedException if the streams doesn't change to the expected state in time.
      */
     public static void waitForApplicationState(final List<KafkaStreams> streamsList,
                                                final State state,
-                                               final Duration timeout) throws Exception {
+                                               final Duration timeout) throws InterruptedException {
         retryOnExceptionWithTimeout(timeout.toMillis(), () -> {
             final Map<KafkaStreams, State> streamsToStates = streamsList
                 .stream()