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

Code before:
        cluster.start();
        cluster.createTopic(INPUT_TOPIC_NAME, NUM_PARTITIONS, 1);
        streamsConfiguration = new Properties();
        final String safeTestName = safeUniqueTestName(testInfo);
        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, "app-" + safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
        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.IntegerSerde.class);
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
        streamsConfiguration.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, 2);

    }

    @AfterEach
    public void after() {
        cluster.stop();
        if (streams != null) {
            streams.close(Duration.ofSeconds(30));
        }
    }

    @Test
    public void correctlyHandleFlushErrorsDuringRebalance() throws InterruptedException {
        final AtomicInteger numberOfStoreInits = new AtomicInteger();
        final AtomicReference<KafkaStreams.State> currentState = new AtomicReference<>();

        final StoreBuilder<KeyValueStore<Object, Object>> storeBuilder = new AbstractStoreBuilder<>("testStateStore", Serdes.Integer(), Serdes.ByteArray(), new MockTime()) {

            @Override
            public KeyValueStore<Object, Object> build() {
                return new MockKeyValueStore(name, false) {

                    @Override
                    public void init(final StateStoreContext stateStoreContext, final StateStore root) {
                        super.init(stateStoreContext, root);
                        numberOfStoreInits.incrementAndGet();
                    }

                    @Override
                    public void flush() {
                        if (numberOfStoreInits.get() == NUM_PARTITIONS * 1.5) {
                            try {
                                TestUtils.waitForCondition(() -> currentState.get() == KafkaStreams.State.PENDING_SHUTDOWN, "Streams never reached PENDING_SHUTDOWN state");
                            } catch (final InterruptedException e) {
                                throw new RuntimeException(e);
                            }
                            throw new ProcessorStateException("flush");
                        }

Review comment (suggestion):
```suggestion
    public void correctlyHandleFlushErrorsDuringRebalance() throws Exception {
```

Code after:
    public void before(final TestInfo testInfo) throws InterruptedException, IOException {
        cluster.start();
        cluster.createTopic(INPUT_TOPIC_NAME, NUM_PARTITIONS, 1);
        streamsConfiguration = new Properties();
        final String safeTestName = safeUniqueTestName(testInfo);
        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, "app-" + safeTestName);
        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, cluster.bootstrapServers());
        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, TestUtils.tempDirectory().getPath());
        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.IntegerSerde.class);
        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.StringSerde.class);
        streamsConfiguration.put(StreamsConfig.NUM_STREAM_THREADS_CONFIG, 2);

    }

    @AfterEach
    public void after() {
        cluster.stop();
        if (streams != null) {
            streams.close(Duration.ofSeconds(30));
        }
    }

    /**
     * The conditions that we need to meet:
     * <p><ul>
     * <li>We have an unhandled task in {@link org.apache.kafka.streams.processor.internals.DefaultStateUpdater}</li>
     * <li>StreamThread is not running, so {@link org.apache.kafka.streams.processor.internals.TaskManager#handleExceptionsFromStateUpdater} is not called anymore</li>
     * <li>The task throws exception in {@link org.apache.kafka.streams.processor.internals.Task#maybeCheckpoint(boolean)} while being processed by {@code DefaultStateUpdater}</li>
     * <li>{@link org.apache.kafka.streams.processor.internals.TaskManager#shutdownStateUpdater} tries to clean up all tasks that are left in the {@code DefaultStateUpdater}</li>
     * </ul><p>
     * If all conditions are met, {@code TaskManager} needs to be able to handle the failed task from the {@code DefaultStateUpdater} correctly and not hang.
     */
    @Test
    public void correctlyHandleFlushErrorsDuringRebalance() throws Exception {
        final AtomicInteger numberOfStoreInits = new AtomicInteger();
        final CountDownLatch pendingShutdownLatch = new CountDownLatch(1);

        final StoreBuilder<KeyValueStore<Object, Object>> storeBuilder = new AbstractStoreBuilder<>("testStateStore", Serdes.Integer(), Serdes.ByteArray(), new MockTime()) {

            @Override
            public KeyValueStore<Object, Object> build() {
                return new MockKeyValueStore(name, false) {

                    @Override
                    public void init(final StateStoreContext stateStoreContext, final StateStore root) {
                        super.init(stateStoreContext, root);
                        numberOfStoreInits.incrementAndGet();
                    }