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

Code before:
            streams1.close(Duration.ofSeconds(30));
        }
        if (streams2 != null) {
            streams2.close(Duration.ofSeconds(30));
        }
    }

    /**
     * The conditions that we need to meet:
     * <p><ul>
     * <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.
     * <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</li>
     * <li>shutdowns the first KS app and waits for the stream tread to get into PENDING_SHUTDOWN state</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);

Review comment (bug):
```suggestion
     * <li>sends a message with wrong types to crash the stream thread (this avoids a second rebalance, and enters shutdown directly)</li>
```

Code after:
            streams1.close(Duration.ofSeconds(30));
        }
        if (streams2 != null) {
            streams2.close(Duration.ofSeconds(30));
        }
    }

    /**
     * The conditions that we need to meet:
     * <p><ul>
     * <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);