From 1f3faabeb856a507e0aab8addd1a5b97380e6b03 Mon Sep 17 00:00:00 2001 From: Pierre Villard Date: Thu, 3 Sep 2026 22:17:15 +0200 Subject: [PATCH] NIFI-16290 - Node offload should give processors a bounded grace period to stop cleanly before forced termination --- .../main/asciidoc/administration-guide.adoc | 2 +- .../nifi/controller/StandardFlowService.java | 119 +++++- .../controller/TestStandardFlowService.java | 403 ++++++++++++++++++ .../tests/system/clustering/OffloadIT.java | 17 +- 4 files changed, 516 insertions(+), 25 deletions(-) diff --git a/nifi-docs/src/main/asciidoc/administration-guide.adoc b/nifi-docs/src/main/asciidoc/administration-guide.adoc index 5836a0d8b943..b0cb7026e348 100644 --- a/nifi-docs/src/main/asciidoc/administration-guide.adoc +++ b/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`. diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowService.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowService.java index 91c1197b2712..25da34effb92 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/StandardFlowService.java +++ b/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 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 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 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 processors) { + final long startNanos = System.nanoTime(); + GracefulOffloadOutcome outcome = GracefulOffloadOutcome.COMPLETED; + Throwable cause = null; + + final CompletableFuture 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()); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/TestStandardFlowService.java b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/TestStandardFlowService.java index f19d6c2d6138..ad4fbab9652d 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/TestStandardFlowService.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/TestStandardFlowService.java @@ -30,15 +30,23 @@ import org.apache.nifi.components.state.Scope; import org.apache.nifi.components.state.StateManager; import org.apache.nifi.components.state.StateManagerProvider; +import org.apache.nifi.connectable.Connection; +import org.apache.nifi.controller.flow.FlowManager; +import org.apache.nifi.controller.queue.FlowFileQueue; import org.apache.nifi.controller.serialization.FlowSynchronizationException; +import org.apache.nifi.controller.status.ProcessGroupStatus; import org.apache.nifi.groups.BundleUpdateStrategy; +import org.apache.nifi.groups.ProcessGroup; +import org.apache.nifi.groups.RemoteProcessGroup; import org.apache.nifi.nar.ExtensionManager; import org.apache.nifi.nar.NarManager; +import org.apache.nifi.reporting.UserAwareEventAccess; import org.apache.nifi.state.MockStateMap; import org.apache.nifi.util.NiFiProperties; import org.apache.nifi.web.revision.RevisionManager; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.api.io.TempDir; import org.mockito.ArgumentCaptor; import org.mockito.InOrder; @@ -48,19 +56,32 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.same; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -176,6 +197,292 @@ public void testNodeConnectionStateConnectingBeforeFlowSynchronization() throws verify(clusteredController, never()).setClustered(anyBoolean(), any()); } + @Test + @Timeout(10) + public void testOffloadWaitsForHeldStopFutureBeforeTerminatingProcessors() throws Exception { + final TestContext context = createTestContext("10 secs"); + final CompletableFuture stopFuture = new CompletableFuture<>(); + final CountDownLatch stopRequested = new CountDownLatch(1); + final CountDownLatch terminationRequested = new CountDownLatch(1); + + when(context.rootGroup.stopProcessing()).thenAnswer(invocation -> { + stopRequested.countDown(); + return stopFuture; + }); + + final ProcessGroup owningGroup = mock(ProcessGroup.class); + final ProcessorNode processor = createProcessorNode(owningGroup, ScheduledState.STOPPED, ScheduledState.STOPPED); + when(context.rootGroup.findAllProcessors()).thenReturn(List.of(processor)); + doAnswer(invocation -> { + terminationRequested.countDown(); + return null; + }).when(owningGroup).terminateProcessor(same(processor)); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + final Future offloadFuture = executorService.submit(() -> context.flowService.offloadNode("held-stop-future")); + try { + assertTrue(stopRequested.await(2, TimeUnit.SECONDS), "Expected offload to request processor stop before waiting"); + assertFalse(terminationRequested.await(200, TimeUnit.MILLISECONDS), "Termination should not begin while the aggregate stop future is incomplete"); + + stopFuture.complete(null); + offloadFuture.get(5, TimeUnit.SECONDS); + } finally { + stopFuture.complete(null); + executorService.shutdownNow(); + } + + verify(owningGroup).terminateProcessor(processor); + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testOffloadTerminatesLogicallyStoppedProcessorAfterGracefulCompletionEvenWhenPhysicallyStopping() throws Exception { + final TestContext context = createTestContext("10 secs"); + final ProcessGroup owningGroup = mock(ProcessGroup.class); + final ProcessorNode processor = createProcessorNode(owningGroup, ScheduledState.STOPPED, ScheduledState.STOPPING); + when(context.rootGroup.findAllProcessors()).thenReturn(List.of(processor)); + + context.flowService.offloadNode("graceful-complete"); + + final InOrder inOrder = inOrder(context.rootGroup, owningGroup, context.clusterCoordinator); + inOrder.verify(context.rootGroup).stopProcessing(); + inOrder.verify(owningGroup).terminateProcessor(processor); + inOrder.verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testOffloadFallsBackToTerminationWhenGracefulStopTimesOut() throws Exception { + final TestContext context = createTestContext("0 secs"); + final ProcessGroup owningGroup = mock(ProcessGroup.class); + final ProcessorNode processor = createProcessorNode(owningGroup, ScheduledState.STOPPED, ScheduledState.RUNNING); + when(context.rootGroup.stopProcessing()).thenReturn(new CompletableFuture<>()); + when(context.rootGroup.findAllProcessors()).thenReturn(List.of(processor)); + + context.flowService.offloadNode("timeout"); + + verify(owningGroup).terminateProcessor(processor); + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testOffloadFallsBackToTerminationWhenGracefulStopCompletesExceptionally() throws Exception { + final TestContext context = createTestContext("10 secs"); + final ProcessGroup owningGroup = mock(ProcessGroup.class); + final ProcessorNode processor = createProcessorNode(owningGroup, ScheduledState.STOPPED, ScheduledState.RUNNING); + final CompletableFuture stopFuture = new CompletableFuture<>(); + stopFuture.completeExceptionally(new IllegalStateException("stop failed")); + when(context.rootGroup.stopProcessing()).thenReturn(stopFuture); + when(context.rootGroup.findAllProcessors()).thenReturn(List.of(processor)); + + context.flowService.offloadNode("exceptional"); + + verify(owningGroup).terminateProcessor(processor); + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testOffloadFallsBackToTerminationWhenGracefulStopInterruptedAndRestoresInterruptStatusAfterExit() throws Exception { + final TestContext context = createTestContext("10 secs"); + final ProcessGroup owningGroup = mock(ProcessGroup.class); + final ProcessorNode processor = createProcessorNode(owningGroup, ScheduledState.STOPPED, ScheduledState.RUNNING); + when(context.rootGroup.findAllProcessors()).thenReturn(List.of(processor)); + when(context.rootGroup.stopProcessing()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return new CompletableFuture<>(); + }); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + try { + final Future interruptedAfterReturn = executorService.submit(() -> { + context.flowService.offloadNode("interrupt-during-grace-wait"); + return Thread.currentThread().isInterrupted(); + }); + + assertTrue(interruptedAfterReturn.get(5, TimeUnit.SECONDS), "Offload thread should restore interrupt status only after the critical section exits"); + } finally { + executorService.shutdownNow(); + } + + verify(owningGroup).terminateProcessor(processor); + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + } + + @Test + @Timeout(10) + public void testOffloadContinuesWhenInterruptedDuringRemoteProcessGroupWaitAndRestoresInterruptStatus() throws Exception { + final TestContext context = createTestContext("10 secs"); + final RemoteProcessGroup remoteProcessGroup = mock(RemoteProcessGroup.class); + when(remoteProcessGroup.isTransmitting()).thenReturn(true); + when(remoteProcessGroup.stopTransmitting()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return new CompletableFuture<>(); + }); + when(remoteProcessGroup.getCommunicationsTimeout(TimeUnit.MILLISECONDS)).thenReturn(10_000); + when(context.rootGroup.findAllRemoteProcessGroups()).thenReturn(List.of(remoteProcessGroup)); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + try { + final Future interruptedAfterReturn = executorService.submit(() -> { + context.flowService.offloadNode("interrupt-during-remote-process-group-wait"); + return Thread.currentThread().isInterrupted(); + }); + + assertTrue(interruptedAfterReturn.get(5, TimeUnit.SECONDS)); + } finally { + executorService.shutdownNow(); + } + + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testOffloadContinuesQueueDrainingWhenInterruptedDuringQueueWaitAndReachesOffloaded() throws Exception { + final TestContext context = createTestContext("10 secs"); + final FlowFileQueue queue = mock(FlowFileQueue.class); + final Connection connection = mock(Connection.class); + when(connection.getFlowFileQueue()).thenReturn(queue); + when(context.flowManager.findAllConnections()).thenReturn(Set.of(connection)); + + final ProcessGroupStatus queuedStatus = queueStatusWithQueuedCount(1); + final ProcessGroupStatus drainedStatus = queueStatusWithQueuedCount(0); + when(context.eventAccess.getControllerStatus()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return queuedStatus; + }).thenReturn(drainedStatus); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + try { + final Future interruptedAfterReturn = executorService.submit(() -> { + context.flowService.offloadNode("interrupt-during-queue-wait"); + return Thread.currentThread().isInterrupted(); + }); + + assertTrue(interruptedAfterReturn.get(5, TimeUnit.SECONDS), "Queue-wait interruption should be restored after offload finishes"); + } finally { + executorService.shutdownNow(); + } + + verify(queue).offloadQueue(); + verify(queue).resetOffloadedQueue(); + verify(context.eventAccess, times(2)).getControllerStatus(); + verify(context.clusterCoordinator).finishNodeOffload(any(NodeIdentifier.class)); + verifyOffloadStatusTransitions(context.controller); + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopClassifiesNormalCompletion() { + final TestContext context = createUncheckedTestContext("10 secs"); + final ProcessorNode stoppedProcessor = createProcessorNode(mock(ProcessGroup.class), ScheduledState.STOPPED, ScheduledState.STOPPED); + + final StandardFlowService.GracefulOffloadResult result = context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of(stoppedProcessor)); + + assertEquals(StandardFlowService.GracefulOffloadOutcome.COMPLETED, result.outcome()); + assertEquals(1, result.processorCount()); + assertEquals(0, result.processorsNotFullyStopped()); + assertNull(result.cause()); + assertTrue(result.elapsedMillis() >= 0); + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopClassifiesTimeoutAndCountsProcessorsNotFullyStopped() { + final TestContext context = createUncheckedTestContext("0 secs"); + when(context.rootGroup.stopProcessing()).thenReturn(new CompletableFuture<>()); + + final ProcessorNode stoppingProcessor = createProcessorNode(mock(ProcessGroup.class), ScheduledState.STOPPED, ScheduledState.STOPPING); + final ProcessorNode runningProcessor = createProcessorNode(mock(ProcessGroup.class), ScheduledState.RUNNING, ScheduledState.RUNNING); + + final StandardFlowService.GracefulOffloadResult result = context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of(stoppingProcessor, runningProcessor)); + + assertEquals(StandardFlowService.GracefulOffloadOutcome.TIMED_OUT, result.outcome()); + assertEquals(2, result.processorCount()); + assertEquals(2, result.processorsNotFullyStopped()); + assertNull(result.cause()); + assertTrue(result.elapsedMillis() >= 0); + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopClassifiesExceptionalCompletionWithCause() { + final TestContext context = createUncheckedTestContext("10 secs"); + final IllegalStateException failure = new IllegalStateException("stop failed"); + final CompletableFuture stopFuture = new CompletableFuture<>(); + stopFuture.completeExceptionally(failure); + when(context.rootGroup.stopProcessing()).thenReturn(stopFuture); + + final StandardFlowService.GracefulOffloadResult result = context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of()); + + assertEquals(StandardFlowService.GracefulOffloadOutcome.EXCEPTIONAL, result.outcome()); + assertEquals(0, result.processorCount()); + assertEquals(0, result.processorsNotFullyStopped()); + assertEquals(failure, result.cause()); + assertTrue(result.elapsedMillis() >= 0); + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopClassifiesCancellationAsExceptionalCompletion() { + final TestContext context = createUncheckedTestContext("10 secs"); + final CompletableFuture stopFuture = new CompletableFuture<>(); + stopFuture.cancel(false); + when(context.rootGroup.stopProcessing()).thenReturn(stopFuture); + + final StandardFlowService.GracefulOffloadResult result = context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of()); + + assertEquals(StandardFlowService.GracefulOffloadOutcome.EXCEPTIONAL, result.outcome()); + assertTrue(result.cause() instanceof CancellationException); + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopClassifiesInterruption() throws Exception { + final TestContext context = createUncheckedTestContext("10 secs"); + when(context.rootGroup.stopProcessing()).thenAnswer(invocation -> { + Thread.currentThread().interrupt(); + return new CompletableFuture<>(); + }); + + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + try { + final Future resultFuture = executorService.submit( + () -> context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of())); + + final StandardFlowService.GracefulOffloadResult result = resultFuture.get(5, TimeUnit.SECONDS); + assertEquals(StandardFlowService.GracefulOffloadOutcome.INTERRUPTED, result.outcome()); + assertEquals(0, result.processorCount()); + assertEquals(0, result.processorsNotFullyStopped()); + assertNull(result.cause()); + assertTrue(result.elapsedMillis() >= 0); + } finally { + executorService.shutdownNow(); + } + } + + @Test + @Timeout(10) + public void testAwaitGracefulProcessorStopBoundsIncompleteAggregateFutureWithoutProcessors() { + final TestContext context = createUncheckedTestContext("0 secs"); + when(context.rootGroup.stopProcessing()).thenReturn(new CompletableFuture<>()); + + final StandardFlowService.GracefulOffloadResult result = context.flowService.awaitGracefulProcessorStop(context.rootGroup, List.of()); + + assertEquals(StandardFlowService.GracefulOffloadOutcome.TIMED_OUT, result.outcome()); + assertEquals(0, result.processorCount()); + assertEquals(0, result.processorsNotFullyStopped()); + assertNull(result.cause()); + } + private DisconnectMessage createDisconnectMessage(final String explanation) { final NodeIdentifier nodeIdentifier = createNodeIdentifier(); final DisconnectMessage message = new DisconnectMessage(); @@ -188,4 +495,100 @@ private NodeIdentifier createNodeIdentifier() { return new NodeIdentifier("node-1", "localhost", 8443, "localhost", 9090, "localhost", 6342, 10443, false); } + private void verifyOffloadStatusTransitions(final FlowController testController) { + final ArgumentCaptor statusCaptor = ArgumentCaptor.forClass(NodeConnectionStatus.class); + verify(testController, times(2)).setConnectionStatus(statusCaptor.capture()); + assertEquals(NodeConnectionState.OFFLOADING, statusCaptor.getAllValues().get(0).getState()); + assertEquals(NodeConnectionState.OFFLOADED, statusCaptor.getAllValues().get(1).getState()); + } + + private ProcessGroupStatus queueStatusWithQueuedCount(final int queuedCount) { + final ProcessGroupStatus status = mock(ProcessGroupStatus.class); + when(status.getQueuedCount()).thenReturn(queuedCount); + return status; + } + + private ProcessorNode createProcessorNode(final ProcessGroup processGroup, final ScheduledState logicalState, final ScheduledState physicalState) { + final ProcessorNode processorNode = mock(ProcessorNode.class); + when(processorNode.getProcessGroup()).thenReturn(processGroup); + when(processorNode.getScheduledState()).thenReturn(logicalState); + when(processorNode.getPhysicalScheduledState()).thenReturn(physicalState); + return processorNode; + } + + private TestContext createUncheckedTestContext(final String gracefulShutdownPeriod) { + try { + return createTestContext(gracefulShutdownPeriod); + } catch (final IOException e) { + throw new IllegalStateException(e); + } + } + + private TestContext createTestContext(final String gracefulShutdownPeriod) throws IOException { + final FlowController testController = mock(FlowController.class); + final ClusterCoordinator testClusterCoordinator = mock(ClusterCoordinator.class); + final NodeProtocolSenderListener senderListener = mock(NodeProtocolSenderListener.class); + final RevisionManager revisionManager = mock(RevisionManager.class); + final NarManager narManager = mock(NarManager.class); + final AssetSynchronizer parameterContextAssetSynchronizer = mock(AssetSynchronizer.class); + final AssetSynchronizer connectorAssetSynchronizer = mock(AssetSynchronizer.class); + final Authorizer authorizer = mock(Authorizer.class); + final FlowManager testFlowManager = mock(FlowManager.class); + final ProcessGroup testRootGroup = mock(ProcessGroup.class); + final UserAwareEventAccess eventAccess = mock(UserAwareEventAccess.class); + + final StateManagerProvider stateManagerProvider = mock(StateManagerProvider.class); + final StateManager stateManager = mock(StateManager.class); + when(stateManager.getState(any(Scope.class))).thenReturn(new MockStateMap(Collections.emptyMap(), 1)); + when(stateManagerProvider.getStateManager(anyString())).thenReturn(stateManager); + when(testController.getStateManagerProvider()).thenReturn(stateManagerProvider); + when(testController.getExtensionManager()).thenReturn(mock(ExtensionManager.class)); + when(testController.getFlowManager()).thenReturn(testFlowManager); + when(testController.getEventAccess()).thenReturn(eventAccess); + when(testFlowManager.getRootGroup()).thenReturn(testRootGroup); + when(testFlowManager.findAllConnections()).thenReturn(Collections.emptySet()); + when(testRootGroup.stopProcessing()).thenReturn(CompletableFuture.completedFuture(null)); + when(testRootGroup.findAllProcessors()).thenReturn(Collections.emptyList()); + when(testRootGroup.findAllRemoteProcessGroups()).thenReturn(Collections.emptyList()); + final ProcessGroupStatus drainedStatus = queueStatusWithQueuedCount(0); + when(eventAccess.getControllerStatus()).thenReturn(drainedStatus); + + final Path flowConfigFile = tempDir.resolve("flow-" + gracefulShutdownPeriod.replace(' ', '-') + ".json.gz"); + final NiFiProperties nifiProperties = NiFiProperties.createBasicNiFiProperties(null, Map.of( + NiFiProperties.FLOW_CONFIGURATION_FILE, flowConfigFile.toString(), + NiFiProperties.FLOW_CONTROLLER_GRACEFUL_SHUTDOWN_PERIOD, gracefulShutdownPeriod, + NiFiProperties.WEB_HTTPS_HOST, "localhost", + NiFiProperties.WEB_HTTPS_PORT, "8443", + NiFiProperties.CLUSTER_NODE_ADDRESS, "localhost", + NiFiProperties.CLUSTER_NODE_PROTOCOL_PORT, "9090", + NiFiProperties.LOAD_BALANCE_HOST, "localhost", + NiFiProperties.LOAD_BALANCE_PORT, "6342", + NiFiProperties.FLOW_CONFIGURATION_ARCHIVE_ENABLED, "false" + )); + + final StandardFlowService testFlowService = StandardFlowService.createClusteredInstance(testController, nifiProperties, senderListener, + testClusterCoordinator, revisionManager, narManager, parameterContextAssetSynchronizer, + connectorAssetSynchronizer, authorizer); + + return new TestContext(testController, testClusterCoordinator, testFlowManager, testRootGroup, eventAccess, testFlowService); + } + + private static final class TestContext { + private final FlowController controller; + private final ClusterCoordinator clusterCoordinator; + private final FlowManager flowManager; + private final ProcessGroup rootGroup; + private final UserAwareEventAccess eventAccess; + private final StandardFlowService flowService; + + private TestContext(final FlowController controller, final ClusterCoordinator clusterCoordinator, final FlowManager flowManager, + final ProcessGroup rootGroup, final UserAwareEventAccess eventAccess, final StandardFlowService flowService) { + this.controller = controller; + this.clusterCoordinator = clusterCoordinator; + this.flowManager = flowManager; + this.rootGroup = rootGroup; + this.eventAccess = eventAccess; + this.flowService = flowService; + } + } } diff --git a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/clustering/OffloadIT.java b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/clustering/OffloadIT.java index 99334a3a3e38..dfc07cc84d8d 100644 --- a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/clustering/OffloadIT.java +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/clustering/OffloadIT.java @@ -17,8 +17,10 @@ package org.apache.nifi.tests.system.clustering; +import org.apache.nifi.tests.system.InstanceConfiguration; import org.apache.nifi.tests.system.NiFiInstanceFactory; import org.apache.nifi.tests.system.NiFiSystemIT; +import org.apache.nifi.tests.system.SpawnedClusterNiFiInstanceFactory; import org.apache.nifi.toolkit.client.NiFiClientException; import org.apache.nifi.web.api.dto.NodeDTO; import org.apache.nifi.web.api.dto.ProcessorConfigDTO; @@ -36,16 +38,19 @@ public class OffloadIT extends NiFiSystemIT { private static final Logger logger = LoggerFactory.getLogger(OffloadIT.class); + private static final String GRACEFUL_SHUTDOWN_PERIOD = "1 sec"; @Override public NiFiInstanceFactory getInstanceFactory() { - return createTwoNodeInstanceFactory(); + return new SpawnedClusterNiFiInstanceFactory(createClusterNodeConfiguration(1), createClusterNodeConfiguration(2)); } @Test @Timeout(value = 10, unit = TimeUnit.MINUTES) // Test to ensure that node can be offloaded, reconnected, offloaded several times. This test typically takes only about 1-2 minutes - // but can occasionally take 5-6 minutes on Github Actions so we set the timeout to 10 minutes to allow for these occasions + // but can occasionally take 5-6 minutes on Github Actions so we set the timeout to 10 minutes to allow for these occasions. + // Runs against a short nifi.flowcontroller.graceful.shutdown.period (see createClusterNodeConfiguration) so each iteration also + // exercises the bounded graceful processor-stop wait added in front of the existing forced termination fallback. public void testOffload() throws InterruptedException, IOException, NiFiClientException { for (int i = 0; i < 5; i++) { logger.info("Running iteration {}", i); @@ -85,6 +90,14 @@ private void testIteration() throws NiFiClientException, IOException, Interrupte waitForAllNodesConnected(); } + private InstanceConfiguration createClusterNodeConfiguration(final int nodeIndex) { + return new InstanceConfiguration.Builder() + .bootstrapConfig("src/test/resources/conf/clustered/node" + nodeIndex + "/bootstrap.conf") + .instanceDirectory("target/node" + nodeIndex) + .overrideNifiProperties(Map.of("nifi.flowcontroller.graceful.shutdown.period", GRACEFUL_SHUTDOWN_PERIOD)) + .build(); + } + /** * Verifies that a node can complete offload after a Processor that retains Sessions across onTrigger * invocations is stopped and terminated. The HoldInput Processor extends AbstractSessionFactoryProcessor