NIFI-16290 - Node offload should give processors a bounded grace period to stop cleanly before forced termination

nifi-docs/src/main/asciidoc/administration-guide.adoc

@@ -2920,7 +2920,7 @@ This cleanup mechanism takes into account only automatically created archived _f
 |`nifi.flow.configuration.archive.max.storage`*|The total data size allowed for the archived _flow.json_ files. NiFi will delete the oldest archive files until the total archived file size becomes less than this configuration value, if this property is specified. If no archive limitation is specified in _nifi.properties_, NiFi uses `500 MB` for this.
 |`nifi.flow.configuration.archive.max.count`*|The number of archive files allowed. NiFi will delete the oldest archive files so that only N latest archives can be kept, if this property is specified.
 |`nifi.flowcontroller.autoResumeState`|Indicates whether -upon restart- the components on the NiFi graph, including Processors, Controller Services, and other schedulable components, should return to their last state. When running in cluster, all nodes should have the same value. The default value is `true`.
-|`nifi.flowcontroller.graceful.shutdown.period`|Indicates the shutdown period. The default value is `10 secs`.
+|`nifi.flowcontroller.graceful.shutdown.period`|Indicates the amount of time to wait for components to stop during graceful shutdown and before forced processor termination during node offload. The default value is `10 secs`.
 |`nifi.flowcontroller.registry.sync.interval`|Specifies the default recurring interval at which NiFi synchronizes the flow configuration with Flow Registry Clients. The default value is `30 min`. This value is used for any Flow Registry Client that does not configure its own `Synchronization Interval` property; a Flow Registry Client that sets that property is synchronized at its own interval instead.
 |`nifi.flowservice.writedelay.interval`|When many changes are made to the _flow.json_, this property specifies how long to wait before writing out the changes, so as to batch the changes into a single write. The default value is `500 ms`.
 |`nifi.administrative.yield.duration`|If a component allows an unexpected exception to escape, it is considered a bug. As a result, the framework will pause (or administratively yield) the component for this amount of time. This is done so that the component does not use up massive amounts of system resources, since it is known to have problems in the existing state. The default value is `30 secs`.

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

@@ -83,14 +83,20 @@
 import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
 import java.util.Calendar;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Date;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -415,13 +421,7 @@ public ProtocolMessage handle(final ProtocolMessage request, final Set<String> n
                     return new ReconnectionResponseMessage();
                 }
                 case OFFLOAD_REQUEST: {
-                    final Thread t = new Thread(() -> {
-                        try {
-                            handleOffloadRequest((OffloadMessage) request);
-                        } catch (InterruptedException e) {
-                            throw new ProtocolException("Could not complete offload request", e);
-                        }
-                    }, "Offload FlowFiles from Node");
+                    final Thread t = new Thread(() -> handleOffloadRequest((OffloadMessage) request), "Offload FlowFiles from Node");
                     t.setDaemon(true);
                     t.start();
 
@@ -657,12 +657,13 @@ private void handleReconnectionRequest(final ReconnectionRequestMessage request)
         }
     }
 
-    private void handleOffloadRequest(final OffloadMessage request) throws InterruptedException {
+    private void handleOffloadRequest(final OffloadMessage request) {
         logger.info("Received offload request message from cluster coordinator with explanation: {}", request.getExplanation());
-        offload(request.getExplanation());
+        offloadNode(request.getExplanation());
     }
 
-    private void offload(final String explanation) throws InterruptedException {
+    void offloadNode(final String explanation) {
+        boolean interrupted = false;
         writeLock.lock();
         try {
             logger.info("Offloading node due to {}", explanation);
@@ -671,26 +672,37 @@ private void offload(final String explanation) throws InterruptedException {
             controller.setConnectionStatus(new NodeConnectionStatus(nodeId, NodeConnectionState.OFFLOADING, OffloadCode.OFFLOADED, explanation));
 
             final FlowManager flowManager = controller.getFlowManager();
+            final ProcessGroup rootGroup = flowManager.getRootGroup();
+            final List<ProcessorNode> processors = List.copyOf(rootGroup.findAllProcessors());
 
             // request to stop all processors on node
-            flowManager.getRootGroup().stopProcessing();
+            final GracefulOffloadResult gracefulStopResult = awaitGracefulProcessorStop(rootGroup, processors);
+            interrupted = gracefulStopResult.outcome() == GracefulOffloadOutcome.INTERRUPTED;
+            logGracefulProcessorStop(gracefulStopResult);
 
             // terminate all processors
-            flowManager.getRootGroup().findAllProcessors()
+            processors
                     // filter stream, only stopped processors can be terminated
                     .stream().filter(pn -> pn.getScheduledState() == ScheduledState.STOPPED)
                     .forEach(pn -> pn.getProcessGroup().terminateProcessor(pn));
 
             // request to stop all remote process groups
-            flowManager.getRootGroup().findAllRemoteProcessGroups()
-                    .stream().filter(RemoteProcessGroup::isTransmitting)
-                    .forEach(rpg -> {
-                        try {
-                            rpg.stopTransmitting().get(rpg.getCommunicationsTimeout(TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS);
-                        } catch (final Exception e) {
-                            logger.warn("Encountered failure while waiting for {} to shutdown", rpg, e);
-                        }
-                    });
+            for (final RemoteProcessGroup remoteProcessGroup : rootGroup.findAllRemoteProcessGroups()) {
+                if (!remoteProcessGroup.isTransmitting()) {
+                    continue;
+                }
+
+                try {
+                    remoteProcessGroup.stopTransmitting().get(remoteProcessGroup.getCommunicationsTimeout(TimeUnit.MILLISECONDS), TimeUnit.MILLISECONDS);
+                } catch (final InterruptedException e) {
+                    if (!interrupted) {
+                        logger.warn("Interrupted while waiting for {} to shutdown; continuing node offload", remoteProcessGroup);
+                    }
+                    interrupted = true;
+                } catch (final ExecutionException | TimeoutException e) {
+                    logger.warn("Encountered failure while waiting for {} to shutdown", remoteProcessGroup, e);
+                }
+            }
 
             // offload all queues on node
             final Set<Connection> connections = flowManager.findAllConnections();
@@ -709,7 +721,14 @@ private void offload(final String explanation) throws InterruptedException {
                 }
 
                 logger.debug("Offloading queues on node {}, remaining queued count: {}", getNodeId(), controllerStatus.getQueuedCount());
-                Thread.sleep(1000);
+                try {
+                    Thread.sleep(1000);
+                } catch (final InterruptedException e) {
+                    if (!interrupted) {
+                        logger.warn("Interrupted while waiting for queued FlowFiles to offload; continuing until queues are drained");
+                    }
+                    interrupted = true;
+                }
             }
 
             // finish offload
@@ -724,9 +743,65 @@ private void offload(final String explanation) throws InterruptedException {
 
         } finally {
             writeLock.unlock();
+            if (interrupted) {
+                Thread.currentThread().interrupt();
+            }
         }
     }
 
+    GracefulOffloadResult awaitGracefulProcessorStop(final ProcessGroup rootGroup, final Collection<ProcessorNode> processors) {
+        final long startNanos = System.nanoTime();
+        GracefulOffloadOutcome outcome = GracefulOffloadOutcome.COMPLETED;
+        Throwable cause = null;
+
+        final CompletableFuture<Void> stopFuture = rootGroup.stopProcessing();
+        try {
+            stopFuture.get(gracefulShutdownSeconds, TimeUnit.SECONDS);
+        } catch (final TimeoutException e) {
+            outcome = GracefulOffloadOutcome.TIMED_OUT;
+        } catch (final InterruptedException e) {
+            outcome = GracefulOffloadOutcome.INTERRUPTED;
+        } catch (final ExecutionException e) {
+            outcome = GracefulOffloadOutcome.EXCEPTIONAL;
+            cause = e.getCause();
+        } catch (final CancellationException e) {
+            outcome = GracefulOffloadOutcome.EXCEPTIONAL;
+            cause = e;
+        }
+
+        final long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
+        final int notFullyStopped = (int) processors.stream()
+                .map(ProcessorNode::getPhysicalScheduledState)
+                .filter(state -> state != ScheduledState.STOPPED && state != ScheduledState.DISABLED)
+                .count();
+        return new GracefulOffloadResult(outcome, elapsedMillis, processors.size(), notFullyStopped, cause);
+    }
+
+    private void logGracefulProcessorStop(final GracefulOffloadResult result) {
+        if (result.outcome() == GracefulOffloadOutcome.COMPLETED) {
+            logger.info("Processor stop completed gracefully in {} ms within configured {} second window for {} processors; {} processors were not fully stopped",
+                    result.elapsedMillis(), gracefulShutdownSeconds, result.processorCount(), result.processorsNotFullyStopped());
+        } else {
+            logger.warn("Processor stop did not complete gracefully: outcome={}, elapsedMillis={}, configuredSeconds={}, processorCount={}, processorsNotFullyStopped={}",
+                    result.outcome(), result.elapsedMillis(), gracefulShutdownSeconds, result.processorCount(), result.processorsNotFullyStopped(), result.cause());
+        }
+    }
+
+    enum GracefulOffloadOutcome {
+        COMPLETED,
+        TIMED_OUT,
+        INTERRUPTED,
+        EXCEPTIONAL
+    }
+
+    record GracefulOffloadResult(
+            GracefulOffloadOutcome outcome,
+            long elapsedMillis,
+            int processorCount,
+            int processorsNotFullyStopped,
+            Throwable cause) {
+    }
+
     // Visible for testing
     void handleDisconnectionRequest(final DisconnectMessage request, final long expectedConnectionGeneration) {
         logger.info("Received disconnection request message from cluster coordinator with explanation: {}", request.getExplanation());