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

Code before:
        final CountDownLatch pendingShutdownLatch = new CountDownLatch(1);
        // Count how many times we initialize and close stores
        final AtomicInteger initCount = new AtomicInteger();
        final AtomicInteger closeCount = new AtomicInteger();
        final StoreBuilder<KeyValueStore<Bytes, byte[]>> storeBuilder = new AbstractStoreBuilder<>("testStateStore", Serdes.Integer(), Serdes.ByteArray(), new MockTime()) {

            @Override
            public KeyValueStore<Bytes, byte[]> build() {
                return new TestRocksDBStore(name, recycleLatch, pendingShutdownLatch, initCount, closeCount);
            }
        };

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

        streams1 = new KafkaStreamsWrapper(topology, props("1"));
        streams1.setStreamThreadStateListener((t, newState, oldState) -> {
            if (newState == StreamThread.State.PENDING_SHUTDOWN) {
                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 to disable retries in the consumer client. if there are no messages, it retries the whole sequence of actions,
        // including the rebalance data. which we don't want, because we just staged the right condition
        IntegrationTestUtils.produceKeyValuesSynchronously(INPUT_TOPIC_NAME, List.of(new KeyValue<>(1L, "key")),
                TestUtils.producerConfig(cluster.bootstrapServers(), LongSerializer.class, StringSerializer.class, new Properties()), cluster.time);
        // Now we can close both apps. The StreamThreadStateListener will unblock the clearCache call, letting the rebalance finish.
        // We don't want it to happen any sooner, because we want the stream thread to stop before it gets to moving messages from task registry to state updater.
        streams1.close(CloseOptions.groupMembershipOperation(CloseOptions.GroupMembershipOperation.LEAVE_GROUP));
        streams2.close(CloseOptions.groupMembershipOperation(CloseOptions.GroupMembershipOperation.LEAVE_GROUP));

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

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


Review comment (suggestion):
Might be good to add a test description at the top, explaining what we are exactly testing.
```
Starting two KS app
Blocking the first KS app during the rebalance (which is started when the second apps starts) because <explanation>
```

and so forth. I am not sure I understand the test yet, and I am sure if somebody come back to it at some point in the future they might struggle, too.

Code after:
     * <li>There is a task with an open store in {@link org.apache.kafka.streams.processor.internals.TasksRegistry#pendingTasksToInit}</li>
     * <li>StreamThread gets into PENDING_SHUTDOWN state, so that {@link StreamThread#isStartingRunningOrPartitionAssigned} returns false
     * before we call {@link StreamThread#checkStateUpdater} that would move the task to the StateUpdater </li>
     * </ul><p>
     * If all conditions are met, {@code TaskManager} needs to correctly close the open store during shutdown.
     * <p>
     * In order to have a task with an open store in the pending task list we first need to have an active task that gets converted
     * to a standby one during rebalance(see {@link org.apache.kafka.streams.processor.internals.TaskManager#closeAndRecycleTasks}).
     * Second, we need to avoid the second rebalance, to avoid that the pending tasks is closed during such a rebalance, ie, before we enter the shutdown phase.
     * <p>
     * For that this test:
     * <p><ul>
     * <li>starts a KS app and waits for it to fully start</li>
     * <li>starts another KS app which will trigger reassignment</li>
     * <li>waits for {@link CachedStateStore#clearCache} to be called(it's called during task recycle) and locks on it</li>
     * <li>sends a message with wrong types to crash the stream thread (this avoids a second rebalance, and enters shutdown directly)</li>
     * <li>shutdowns the first KS app</li>
     * <li>releases the lock</li>
     * </ul><p>
     * At this point {@link org.apache.kafka.streams.processor.internals.TaskManager#shutdown} will be called,
     * and we will have a pending task to init with an open store(because tasks keep their stores open during recycle).
     * <p>
     * This test verifies that the open store is closed during shutdown.
     */
    @Test
    public void shouldClosePendingTasksToInitAfterRebalance() throws Exception {
        final CountDownLatch recycleLatch = new CountDownLatch(1);
        final CountDownLatch pendingShutdownLatch = new CountDownLatch(1);
        // Count how many times we initialize and close stores
        final AtomicInteger initCount = new AtomicInteger();
        final AtomicInteger closeCount = new AtomicInteger();
        final StoreBuilder<KeyValueStore<Bytes, byte[]>> storeBuilder = new AbstractStoreBuilder<>("testStateStore", Serdes.Integer(), Serdes.ByteArray(), new MockTime()) {

            @Override
            public KeyValueStore<Bytes, byte[]> build() {
                return new TestRocksDBStore(name, recycleLatch, pendingShutdownLatch, initCount, closeCount);
            }
        };

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

        streams1 = new KafkaStreamsWrapper(topology, props("1"));
        streams1.setStreamThreadStateListener((t, newState, oldState) -> {
            if (newState == StreamThread.State.PENDING_SHUTDOWN) {
                pendingShutdownLatch.countDown();
            }
        });
        streams1.start();