Repository: apache/kafka
File: streams/src/main/java/org/apache/kafka/streams/processor/internals/StateDirectory.java (Java)

Code before:
                // because it's possible that the topology has changed since that data was written, and is now stateless
                // this therefore prevents us from creating unnecessary stores just because of some left-over state
                if (subTopology.hasStateWithChangelogs()) {
                    final Set<TopicPartition> inputPartitions = topologyMetadata.nodeToSourceTopics(id).values().stream()
                            .flatMap(Collection::stream)
                            .map(t -> new TopicPartition(t, id.partition()))
                            .collect(Collectors.toSet());
                    // Open a temporary state manager that will open the stores inside the subtopology
                    final ProcessorStateManager temporaryStateManager = ProcessorStateManager.createStartupTaskStateManager(
                        id,
                        eosEnabled,
                        logContext,
                        this,
                        subTopology.storeToChangelogTopic(),
                        inputPartitions
                    );

                    final StartupContext initContext = new StartupContext(id, config, temporaryStateManager, metricsImpl, cache);
                    try {
                        StateManagerUtil.registerStateStores(log, threadLogPrefix, subTopology, temporaryStateManager, this, initContext);
                    } catch (final TaskCorruptedException tce) {
                        log.warn("Failed to register startup state stores for task {}: {}", id, tce.getMessage());
                    } finally {
                        // Make sure the state manager writes the local checkpoint file before closing the stores
                        // This will be replaced in the future when removing the checkpoint file dependency.
                        temporaryStateManager.checkpoint();
                        temporaryStateManager.close();
                    }
                    tasksForLocalState.put(id, new StartupState(id, subTopology, temporaryStateManager));
                }
            }
        }
    }

    public boolean hasStartupTasks() {
        return !tasksForLocalState.isEmpty();
    }

    public boolean removeStartupState(final TaskId taskId) {
        final StartupState startupState = tasksForLocalState.remove(taskId);
        final boolean removed = startupState != null;
        if (removed) {
            lockedTasksToOwner.replace(taskId, Thread.currentThread());
        }
        return removed;
    }

    public void closeStartupTasks() {
        closeStartupTasks(t -> true);
    }


Review comment (suggestion):
Same as above `temporaryStateManager.checkpoint()` can throw a `ProcessorStateException` but we don't handle it here, meaning `close()` is not called.  So do we consider this a fatal exception and allow it to bubble up and kill the main thread? If not, I'm guessing we'll need to catch this exception and do something.  Either way let's add a couple of comments on the handling

Code after:
                final ProcessorTopology subTopology = topologyMetadata.buildSubtopology(task);

                // we still check if the task's sub-topology is stateful, even though we know its directory contains state,
                // because it's possible that the topology has changed since that data was written, and is now stateless
                // this therefore prevents us from creating unnecessary stores just because of some left-over state
                if (subTopology.hasStateWithChangelogs()) {
                    final Set<TopicPartition> inputPartitions = topologyMetadata.nodeToSourceTopics(task).values().stream()
                            .flatMap(Collection::stream)
                            .map(t -> new TopicPartition(t, task.partition()))
                            .collect(Collectors.toSet());
                    // Open a temporary state manager that will open the stores inside the subtopology
                    final ProcessorStateManager temporaryStateManager = ProcessorStateManager.createStartupTaskStateManager(
                        task,
                        eosEnabled,
                        logContext,
                        this,
                        subTopology.storeToChangelogTopic(),
                        inputPartitions
                    );

                    final StartupContext initContext = new StartupContext(task, config, temporaryStateManager, metricsImpl, dummyCache);
                    try {
                        // We only handle TaskCorruptedException at this point. Any other exception is considered fatal.
                        StateManagerUtil.registerStateStores(log, threadLogPrefix, subTopology, temporaryStateManager, this, initContext);
                        temporaryStateManager.checkpoint();
                    } catch (final TaskCorruptedException tce) {
                        // At this point, we only log a warning and continue with the startup store initialization.
                        // The task-corrupted exception will be handled in the first Task assignment phase.
                        log.warn("Failed to register startup state stores for task {}: {}", task, tce.getMessage());
                    } finally {
                        // Make sure the state manager writes the local checkpoint file before closing the stores
                        // This will be replaced in the future when removing the checkpoint file dependency.
                        temporaryStateManager.close();
                    }
                    tasksInLocalState.add(task);
                }
            }
        }
    }

    public boolean hasStartupTasks() {
        return !tasksInLocalState.isEmpty();
    }

    public synchronized boolean removeStartupState(final TaskId taskId) {
        final boolean removed = tasksInLocalState.remove(taskId);
        if (removed) {
            lockedTasksToOwner.put(taskId, Thread.currentThread());
        }
        return removed;
    }