Repository: apache/kafka
File: streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RebalanceTaskClosureIntegrationTest.java (Java)

Code before:
                pendingShutdownLatch.countDown();
            }
        });
        streams1.start();

        TestUtils.waitForCondition(() -> streams1.state() == KafkaStreams.State.RUNNING, "Streams never reached RUNNING state");

        streams2 = new KafkaStreamsWrapper(topology, props("2"));
        streams2.start();

        TestUtils.waitForCondition(() -> streams2.state() == KafkaStreams.State.RUNNING, "Streams never reached RUNNING state");

        // starting the second KS app triggered a rebalance. Which in turn will recycle active tasks that need to become standby.
        // That's exactly what we are waiting for
        recycleLatch.await();

        // sending a message with wrong key and value types to trigger a stream thread failure and avoid the second rebalance
        // note that writing this message does not trigger the crash right away -- the thread is still blocked inside `poll()` waiting for the shutdown latch to unlock to complete the previous, still ongoing rebalance
        IntegrationTestUtils.produceKeyValuesSynchronously(INPUT_TOPIC_NAME, List.of(new KeyValue<>("key", 1L)),
                TestUtils.producerConfig(cluster.bootstrapServers(), StringSerializer.class, LongSerializer.class, new Properties()), cluster.time);
        // Now we can close both apps. The StreamThreadStateListener will unblock the clearCache call, letting the rebalance finish.
        // After the rebalance finished, the "poison pill" record gets picked up crashing the thread,
        // and starting the shutdown directly
        // We don't want to let the rebalance finish before we trigger the shutdown, because we want the stream thread to stop before it gets to moving pending tasks from task registry to state updater.
        streams1.close(new KafkaStreams.CloseOptions().leaveGroup(true));
        streams2.close(new KafkaStreams.CloseOptions().leaveGroup(true));

        assertEquals(initCount.get(), closeCount.get());
    }

    private Properties props(final String storePathSuffix) {
        final Properties streamsConfiguration = new Properties();

        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
        streamsConfiguration.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, 1000);
        streamsConfiguration.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 1000);
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath() + "/" + storePathSuffix);
        streamsConfiguration.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.LongSerde.class);
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
        streamsConfiguration.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);

        return streamsConfiguration;
    }

    private static class TestRocksDBStore extends RocksDBStore implements CachedStateStore<Bytes, byte[]> {

        private final CountDownLatch recycleLatch;
        private final CountDownLatch pendingShutdownLatch;

Review comment (suggestion):
This doesn't work as expected (that's why it was change in newer release). It only works for static members.

In 4.0, there is an internal config that we should use instead.

Code after:
                pendingShutdownLatch.countDown();
            }
        });
        streams1.start();

        TestUtils.waitForCondition(() -> streams1.state() == KafkaStreams.State.RUNNING, "Streams never reached RUNNING state");

        streams2 = new KafkaStreamsWrapper(topology, props("2"));
        streams2.start();

        TestUtils.waitForCondition(() -> streams2.state() == KafkaStreams.State.RUNNING, "Streams never reached RUNNING state");

        // starting the second KS app triggered a rebalance. Which in turn will recycle active tasks that need to become standby.
        // That's exactly what we are waiting for
        recycleLatch.await();

        // sending a message with wrong key and value types to trigger a stream thread failure and avoid the second rebalance
        // note that writing this message does not trigger the crash right away -- the thread is still blocked inside `poll()` waiting for the shutdown latch to unlock to complete the previous, still ongoing rebalance
        IntegrationTestUtils.produceKeyValuesSynchronously(INPUT_TOPIC_NAME, List.of(new KeyValue<>("key", 1L)),
                TestUtils.producerConfig(cluster.bootstrapServers(), StringSerializer.class, LongSerializer.class, new Properties()), cluster.time);
        // Now we can close both apps. The StreamThreadStateListener will unblock the clearCache call, letting the rebalance finish.
        // After the rebalance finished, the "poison pill" record gets picked up crashing the thread,
        // and starting the shutdown directly
        // We don't want to let the rebalance finish before we trigger the shutdown, because we want the stream thread to stop before it gets to moving pending tasks from task registry to state updater.
        streams1.close();
        streams2.close();

        assertEquals(initCount.get(), closeCount.get());
    }

    private Properties props(final String storePathSuffix) {
        final Properties streamsConfiguration = new Properties();

        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
        streamsConfiguration.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, 1000);
        streamsConfiguration.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 1000);
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath() + "/" + storePathSuffix);
        streamsConfiguration.put(StreamsConfig.STATESTORE_CACHE_MAX_BYTES_CONFIG, 0);
        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 100L);
        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.LongSerde.class);
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
        streamsConfiguration.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
        streamsConfiguration.put("internal.leave.group.on.close", true);

        return streamsConfiguration;
    }

    private static class TestRocksDBStore extends RocksDBStore implements CachedStateStore<Bytes, byte[]> {

        private final CountDownLatch recycleLatch;