KAFKA-10317: Global thread should honor shutdown signal during bootstrapping

streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManager.java

@@ -18,15 +18,19 @@
 
 import org.apache.kafka.streams.errors.StreamsException;
 
+import java.util.Optional;
 import java.util.Set;
 
 public interface GlobalStateManager extends StateManager {
 
     void setGlobalProcessorContext(final InternalProcessorContext<?, ?> processorContext);
 
     /**
+     * Bootstraps all global state stores. Returns the set of registered store names on success,
+     * or {@link Optional#empty()} if bootstrap was interrupted by a shutdown request.
+     *
      * @throws IllegalStateException If store gets registered after initialized is already finished
      * @throws StreamsException if the store's change log does not contain the partition
      */
-    Set<String> initialize();
+    Optional<Set<String>> initialize();
 }

streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateManagerImpl.java

@@ -23,6 +23,7 @@
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.common.utils.Time;
@@ -64,6 +65,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.BooleanSupplier;
 import java.util.function.Supplier;
 
 import static org.apache.kafka.streams.StreamsConfig.PROCESSING_EXCEPTION_HANDLER_CLASS_CONFIG;
@@ -123,14 +125,16 @@ private static class StateStoreMetadata {
     private DeserializationExceptionHandler deserializationExceptionHandler;
     private ProcessingExceptionHandler processingExceptionHandler;
     private Sensor droppedRecordsSensor;
+    private BooleanSupplier shouldStopBootstrappingSupplier;
 
     public GlobalStateManagerImpl(final LogContext logContext,
                                   final Time time,
                                   final ProcessorTopology topology,
                                   final Consumer<byte[], byte[]> globalConsumer,
                                   final StateDirectory stateDirectory,
                                   final StateRestoreListener stateRestoreListener,
-                                  final StreamsConfig config) {
+                                  final StreamsConfig config,
+                                  final BooleanSupplier shouldStopBootstrappingSupplier) {
         this.time = time;
         this.topology = topology;
         this.stateDirectory = stateDirectory;
@@ -147,6 +151,7 @@ public GlobalStateManagerImpl(final LogContext logContext,
         logPrefix = logContext.logPrefix();
         this.globalConsumer = globalConsumer;
         this.stateRestoreListener = stateRestoreListener;
+        this.shouldStopBootstrappingSupplier = shouldStopBootstrappingSupplier;
 
         final Map<String, Object> consumerProps = config.getGlobalConsumerConfigs("dummy");
         // need to add mandatory configs; otherwise `QuietConsumerConfig` throws
@@ -173,66 +178,78 @@ public void setGlobalProcessorContext(final InternalProcessorContext<?, ?> globa
     }
 
     @Override
-    public Set<String> initialize() {
+    public Optional<Set<String>> initialize() {
         droppedRecordsSensor = droppedRecordsSensor(
             Thread.currentThread().getName(),
             globalProcessorContext.taskId().toString(),
             globalProcessorContext.metrics()
         );
 
-        final Map<TopicPartition, StateStore> wrappedStores = new HashMap<>();
-        for (final StateStore stateStore : topology.globalStateStores()) {
-            final List<TopicPartition> storePartitions = topicPartitionsForStore(stateStore);
-            final StateStore maybeWrappedStore = LegacyCheckpointingStateStore.maybeWrapStore(
-                    stateStore, eosEnabled, new HashSet<>(storePartitions), stateDirectory, null, logPrefix);
-            try {
-                maybeWrappedStore.init(globalProcessorContext, maybeWrappedStore);
-            } catch (final ProcessorStateException e) {
-                if (eosEnabled) {
-                    log.warn("{}Detected unclean shutdown for global store {}. " +
-                            "Wiping global state directory.", logPrefix, stateStore.name(), e);
-                    try {
-                        Utils.delete(stateDirectory.globalStateDir().getAbsoluteFile());
-                    } catch (final IOException ioe) {
-                        e.addSuppressed(ioe);
+        try {
+            final Map<TopicPartition, StateStore> wrappedStores = new HashMap<>();
+            for (final StateStore stateStore : topology.globalStateStores()) {
+                final List<TopicPartition> storePartitions = topicPartitionsForStore(stateStore);
+                final StateStore maybeWrappedStore = LegacyCheckpointingStateStore.maybeWrapStore(
+                        stateStore, eosEnabled, new HashSet<>(storePartitions), stateDirectory, null, logPrefix);
+                try {
+                    maybeWrappedStore.init(globalProcessorContext, maybeWrappedStore);
+                } catch (final ProcessorStateException e) {
+                    if (eosEnabled) {
+                        log.warn("{}Detected unclean shutdown for global store {}. " +
+                                "Wiping global state directory.", logPrefix, stateStore.name(), e);
+                        try {
+                            Utils.delete(stateDirectory.globalStateDir().getAbsoluteFile());
+                        } catch (final IOException ioe) {
+                            e.addSuppressed(ioe);
+                        }
                     }
+                    throw e;
                 }
-                throw e;
-            }
 
-            for (final TopicPartition storePartition : storePartitions) {
-                wrappedStores.put(storePartition, maybeWrappedStore);
+                for (final TopicPartition storePartition : storePartitions) {
+                    wrappedStores.put(storePartition, maybeWrappedStore);
+                }
             }
-        }
 
-        // migrate offsets from legacy checkpoint file into the stores
-        LegacyCheckpointingStateStore.migrateLegacyOffsets(logPrefix, stateDirectory, null, wrappedStores);
-
-        for (final StateStoreMetadata metadata : storeMetadata.values()) {
-            // load the committed offsets from the store
-            final StateStore store = metadata.stateStore;
-            if (store.persistent()) {
-                for (final TopicPartition partition : metadata.changelogPartitions) {
-                    final Long offset = store.committedOffset(partition);
-                    if (offset != null) {
-                        currentOffsets.put(partition, offset);
+            // migrate offsets from legacy checkpoint file into the stores
+            LegacyCheckpointingStateStore.migrateLegacyOffsets(logPrefix, stateDirectory, null, wrappedStores);
+
+            for (final StateStoreMetadata metadata : storeMetadata.values()) {
+                if (shouldStopBootstrappingSupplier.getAsBoolean()) {
+                    log.info("Global store bootstrap interrupted by shutdown before starting {}", metadata.stateStore.name());
+                    return Optional.empty();
+                }
+                // load the committed offsets from the store
+                final StateStore store = metadata.stateStore;
+                if (store.persistent()) {
+                    for (final TopicPartition partition : metadata.changelogPartitions) {
+                        final Long offset = store.committedOffset(partition);
+                        if (offset != null) {
+                            currentOffsets.put(partition, offset);
+                        }
+                    }
+                }
+
+                // restore or reprocess each registered store using the now-populated currentOffsets
+                try {
+                    if (metadata.reprocessFactory.isPresent()) {
+                        reprocessState(metadata);
+                    } else {
+                        restoreState(metadata);
                     }
+                } finally {
+                    globalConsumer.unsubscribe();
                 }
             }
 
-            // restore or reprocess each registered store using the now-populated currentOffsets
-            try {
-                if (metadata.reprocessFactory.isPresent()) {
-                    reprocessState(metadata);
-                } else {
-                    restoreState(metadata);
-                }
-            } finally {
-                globalConsumer.unsubscribe();
+            return Optional.of(Collections.unmodifiableSet(globalStoreNames));
+        } catch (final WakeupException e) {
+            if (!shouldStopBootstrappingSupplier.getAsBoolean()) {
+                throw e;
             }
+            log.info("Global store bootstrap interrupted by shutdown");
+            return Optional.empty();
         }
-
-        return Collections.unmodifiableSet(globalStoreNames);
     }
 
     public StateStore globalStore(final String name) {

streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStateUpdateTask.java

@@ -29,8 +29,10 @@
 import org.slf4j.Logger;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 
 import static org.apache.kafka.streams.processor.internals.metrics.TaskMetrics.droppedRecordsSensor;
@@ -79,9 +81,14 @@ public GlobalStateUpdateTask(final LogContext logContext,
      */
     @Override
     public Map<TopicPartition, Long> initialize() {
-        final Set<String> storeNames = stateMgr.initialize();
+        final Optional<Set<String>> storeNames = stateMgr.initialize();
+        if (storeNames.isEmpty()) {
+            // bootstrap was interrupted by shutdown; skip topology/processor init to avoid
+            // opening user resources via Processor#init() during a shutdown.
+            return Collections.emptyMap();
+        }
         final Map<String, String> storeNameToTopic = topology.storeToChangelogTopic();
-        for (final String storeName : storeNames) {
+        for (final String storeName : storeNames.get()) {
             final String sourceTopic = storeNameToTopic.get(storeName);
             final SourceNode<?, ?> source = topology.source(sourceTopic);
             deserializers.put(

streams/src/main/java/org/apache/kafka/streams/processor/internals/GlobalStreamThread.java

@@ -382,7 +382,8 @@ private StateConsumer initialize() {
                 globalConsumer,
                 stateDirectory,
                 stateRestoreListener,
-                config
+                config,
+                this::inErrorState
             );
 
             final GlobalProcessorContextImpl globalProcessorContext = new GlobalProcessorContextImpl(
@@ -429,6 +430,11 @@ private StateConsumer initialize() {
                 );
             }
 
+            if (inErrorState()) {
+                closeStateConsumer(stateConsumer, false);
+                return null;
+            }
+
             setState(RUNNING);
             return stateConsumer;
         } catch (final StreamsException fatalException) {
@@ -475,8 +481,15 @@ public synchronized void start() {
     public void shutdown() {
         // one could call shutdown() multiple times, so ignore subsequent calls
         // if already shutting down or dead
-        setState(PENDING_SHUTDOWN);
+        final boolean wakeupBootstrap;
+        synchronized (stateLock) {
+            wakeupBootstrap = (state == State.CREATED);
+            setState(PENDING_SHUTDOWN);
+        }
         initializationLatch.countDown();
+        if (wakeupBootstrap) {
+            globalConsumer.wakeup();
+        }
     }
 
     public Map<MetricName, Metric> consumerMetrics() {