NIFI-14631 When stopping a Stateless Process Group, wait up to 10 seconds before interrupting (apache#9992)

Signed-off-by: David Handermann <exceptionfactory@apache.org>

Author: markap14

nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/tasks/StatelessFlowTask.java

@@ -56,6 +56,7 @@
 import org.apache.nifi.stateless.repository.StatelessProvenanceRepository;
 
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumSet;
@@ -71,6 +72,10 @@
 import java.util.stream.Collectors;
 
 public class StatelessFlowTask {
+    // The amount of time after a user stops a Stateless Group that we will wait before we disable the ProcessSession,
+    // throwing Exceptions on all ProcessSession access (effectively performing the same functionality as Terminating a processor).
+    private static final Duration GRACEFUL_STOP_DURATION = Duration.ofSeconds(10);
+
     private static final Set<ProvenanceEventType> eventTypesToKeepOnFailure = EnumSet.of(ProvenanceEventType.SEND, ProvenanceEventType.REMOTE_INVOCATION);
     private final StatelessGroupNode statelessGroupNode;
     private final StatelessDataflow flow;
@@ -88,6 +93,8 @@ public class StatelessFlowTask {
     private List<RepositoryRecord> outputRepositoryRecords;
     private List<ProvenanceEventRecord> cloneProvenanceEvents;
 
+    // State that is updated during invocation but do not need to be guarded by synchronized block
+    private volatile long shutdownInitiationTime = 0L;
 
     private StatelessFlowTask(final Builder builder) {
         this.statelessGroupNode = builder.statelessGroupNode;
@@ -147,12 +154,24 @@ private boolean isRunAsFastAsPossible(final ProcessorNode procNode) {
     }
 
     public void shutdown() {
-        this.flow.shutdown(false, true);
+        this.shutdownInitiationTime = System.nanoTime();
+        this.flow.shutdown(false, true, GRACEFUL_STOP_DURATION);
     }
 
     private boolean isAbort() {
         final ScheduledState desiredState = statelessGroupNode.getDesiredState();
-        return desiredState != ScheduledState.RUNNING && desiredState != ScheduledState.RUN_ONCE;
+        final boolean stopped = desiredState != ScheduledState.RUNNING && desiredState != ScheduledState.RUN_ONCE;
+        if (!stopped) {
+            return false;
+        }
+
+        final long shutdownTime = shutdownInitiationTime;
+        if (shutdownTime == 0) {
+            return false;
+        }
+
+        final long abortTime = shutdownTime + GRACEFUL_STOP_DURATION.toNanos();
+        return System.nanoTime() >= abortTime;
     }
 
 

nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/groups/StandardStatelessGroupNode.java

@@ -234,7 +234,7 @@ private void initialize(final ScheduledExecutorService executor, final Schedulin
             } catch (final Exception e) {
                 for (final StandardStatelessFlow flow : createdFlows) {
                     try {
-                        flow.shutdown(false, true);
+                        flow.shutdown(false, true, Duration.ofMillis(0));
                     } catch (final Exception ex) {
                         logger.error("Failed to shutdown Stateless Flow {}", flow, ex);
                     }
@@ -371,7 +371,7 @@ public boolean isEnableControllerServices() {
         try {
             dataflow.initialize(initializationContext);
         } catch (final Exception e) {
-            dataflow.shutdown(true, true);
+            dataflow.shutdown(true, true, Duration.ofMillis(0));
             throw e;
         }
 

nifi-stateless/nifi-stateless-api/src/main/java/org/apache/nifi/stateless/flow/StatelessDataflow.java

@@ -22,6 +22,7 @@
 import org.apache.nifi.reporting.BulletinRepository;
 
 import java.io.InputStream;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Set;
 
@@ -68,10 +69,16 @@ default DataflowTrigger trigger() {
     void initialize(StatelessDataflowInitializationContext initializationContext);
 
     default void shutdown() {
-        shutdown(true, false);
+        shutdown(true, false, Duration.ofMillis(0));
     }
 
-    void shutdown(boolean triggerComponentShutdown, boolean interruptProcessors);
+    /**
+     * Shuts down the dataflow, stopping all components and releasing all resources.
+     * @param triggerComponentShutdown whether or not to trigger the shutdown of components (e.g., invoking @OnShutdown methods)
+     * @param interruptProcessors whether or not to interrupt any processors and tasks that are running
+     * @param gracefulShutdownPeriod if interruptProcessors is true, this specifies the amount of time to wait for processors to finish before interrupting them.
+     */
+    void shutdown(boolean triggerComponentShutdown, boolean interruptProcessors, Duration gracefulShutdownPeriod);
 
     StatelessDataflowValidation performValidation();
 

nifi-stateless/nifi-stateless-bundle/nifi-stateless-engine/src/main/java/org/apache/nifi/stateless/flow/StandardStatelessFlow.java

@@ -87,6 +87,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -353,7 +354,7 @@ private void startReportingTask(final ReportingTaskNode taskNode) {
     }
 
     @Override
-    public void shutdown(final boolean triggerComponentShutdown, final boolean interruptProcessors) {
+    public void shutdown(final boolean triggerComponentShutdown, final boolean interruptProcessors, final Duration gracefulShutdownDuration) {
         if (shutdown) {
             return;
         }
@@ -362,18 +363,36 @@ public void shutdown(final boolean triggerComponentShutdown, final boolean inter
         logger.info("Shutting down dataflow {}", rootGroup.getName());
 
         if (runDataflowExecutor != null) {
-            if (interruptProcessors) {
-                runDataflowExecutor.shutdownNow();
-            } else {
-                runDataflowExecutor.shutdown();
-            }
+            runDataflowExecutor.shutdown();
         }
         if (backgroundTaskExecutor != null) {
             backgroundTaskExecutor.shutdown();
         }
 
         logger.info("Stopping all components");
-        rootGroup.stopComponents().join();
+        final CompletableFuture<Void> stopComponentsFuture = rootGroup.stopComponents();
+
+        // Wait for the graceful shutdown period for all processors to stop. If the processors do not stop within this time,
+        // then interrupt them.
+        if (runDataflowExecutor != null && interruptProcessors) {
+            if (gracefulShutdownDuration.isZero()) {
+                logger.info("Shutting down all components immediately without waiting for graceful shutdown period");
+                runDataflowExecutor.shutdownNow();
+            } else {
+                try {
+                    stopComponentsFuture.get(gracefulShutdownDuration.toMillis(), TimeUnit.MILLISECONDS);
+                } catch (final Exception e) {
+                    logger.warn("Stateless flow failed to stop all components gracefully after {} millis. Interrupting all running components.", gracefulShutdownDuration.toMillis(), e);
+                    if (e instanceof InterruptedException) {
+                        Thread.interrupted();
+                    }
+
+                    runDataflowExecutor.shutdownNow();
+                }
+            }
+        }
+
+        stopComponentsFuture.join();
         rootGroup.findAllRemoteProcessGroups().forEach(RemoteProcessGroup::shutdown);
 
         if (triggerComponentShutdown) {