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

Code before:
            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");
                        }
                    }
                };
            }
        };

        final TopologyWrapper topology = new TopologyWrapper();
        topology.addSource("ingest", INPUT_TOPIC_NAME);
        topology.addProcessor("my-processor", new MockApiProcessorSupplier<>(), "ingest");
        topology.addStateStore(storeBuilder, "my-processor");

        streams = new KafkaStreams(topology, streamsConfiguration);
        streams.setStateListener((newState, oldState) -> currentState.set(newState));
        streams.start();

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

        streams.removeStreamThread();

        TestUtils.waitForCondition(() -> numberOfStoreInits.get() == NUM_PARTITIONS * 1.5, "Streams never reinitialized the store enough times");


Review comment (bug):
Not sure if I understand this condition. Can you elaborate?

Also: if we fail here, how do we ensure that the test fails? This would be executed on the background `StreamsThread` not the actual "main" thread running the test.

Code after:
        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();
                    }

                    @Override
                    public void flush() {
                        // we want to throw the ProcessorStateException here only when the rebalance finished(we reassigned the 3 tasks from the removed thread to the existing thread)
                        // we use waitForCondition to wait until the current state is PENDING_SHUTDOWN to make sure the Stream Thread will not handle the exception and we can get to in TaskManager#shutdownStateUpdater
                        if (numberOfStoreInits.get() == 9) {
                            try {
                                pendingShutdownLatch.await();
                            } catch (final InterruptedException e) {
                                throw new RuntimeException(e);
                            }
                            throw new ProcessorStateException("flush");
                        }
                    }
                };
            }
        };

        final TopologyWrapper topology = new TopologyWrapper();
        topology.addSource("ingest", INPUT_TOPIC_NAME);