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

Code before:

                    @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");

        assertTrue(streams.close(Duration.ofSeconds(60)));
    }
}


Review comment (bug):
I thought `waitForCondition` would throw `AssertionError` if we don't meet the condition on time? (I guess it _could_ also throw an `InterruptedException` if something goes wrong, but that's would be an error during test execution).

So I guess we would actually need to catch `Error` if we need a "catch all".

Code after:
     */
    @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);
        topology.addProcessor("my-processor", new MockApiProcessorSupplier<>(), "ingest");
        topology.addStateStore(storeBuilder, "my-processor");

        streams = new KafkaStreams(topology, streamsConfiguration);
        streams.setStateListener((newState, oldState) -> {
            if (newState == KafkaStreams.State.PENDING_SHUTDOWN) {
                pendingShutdownLatch.countDown();
            }
        });
        streams.start();

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

        streams.removeStreamThread();