From e1c84725008c6cc00682ca18d1a6b9a6cf271043 Mon Sep 17 00:00:00 2001 From: Wojciech Kochman Date: Tue, 11 Aug 2026 15:05:41 +0200 Subject: [PATCH] NIFI-16183 Add active-polling stopConnector(Duration) to connector mock framework The mock server previously stopped a Connector with a single fixed 10-second blocking wait, which could time out on slower shutdown paths (for example a failed table still draining) even though the Connector node retries its stop internally and eventually reaches STOPPED. - ConnectorTestRunner: add a default stopConnector(Duration) overload - StandardConnectorMockServer: poll the Connector state until STOPPED within the given timeout instead of a single fixed wait - StandardConnectorTestRunner: pass the timeout through to the mock server - CreateConnectorIT: exercise the new overload end-to-end --- .../connector/server/ConnectorTestRunner.java | 14 +++++++ .../server/StandardConnectorMockServer.java | 38 +++++++++++++++---- .../connectors/tests/CreateConnectorIT.java | 21 ++++++++++ .../StandardConnectorTestRunner.java | 6 +++ 4 files changed, 72 insertions(+), 7 deletions(-) diff --git a/nifi-connector-mock-bundle/nifi-connector-mock-api/src/main/java/org/apache/nifi/mock/connector/server/ConnectorTestRunner.java b/nifi-connector-mock-bundle/nifi-connector-mock-api/src/main/java/org/apache/nifi/mock/connector/server/ConnectorTestRunner.java index 449ca64c5de9..42f4c1821ee9 100644 --- a/nifi-connector-mock-bundle/nifi-connector-mock-api/src/main/java/org/apache/nifi/mock/connector/server/ConnectorTestRunner.java +++ b/nifi-connector-mock-bundle/nifi-connector-mock-api/src/main/java/org/apache/nifi/mock/connector/server/ConnectorTestRunner.java @@ -32,6 +32,7 @@ import java.time.Duration; import java.util.List; import java.util.Map; +import java.util.concurrent.TimeoutException; public interface ConnectorTestRunner extends Closeable { @@ -144,6 +145,19 @@ public interface ConnectorTestRunner extends Closeable { */ void stopConnector(); + /** + * Stops the Connector, waiting up to the given timeout for it to reach a stopped state. Implementations + * should actively poll the Connector's state until it is stopped or the timeout elapses, which is more + * tolerant of a flow that takes a while to quiesce (for example a failed table still draining) than the + * default stop budget. + * + * @param timeout the maximum duration to wait for the Connector to stop + * @throws TimeoutException if the timeout elapses before the Connector stops + */ + default void stopConnector(final Duration timeout) throws TimeoutException { + stopConnector(); + } + /** * Blocks until the Connector has received at least one FlowFile, or until the specified timeout elapses. * diff --git a/nifi-connector-mock-bundle/nifi-connector-mock-server/src/main/java/org/apache/nifi/mock/connector/server/StandardConnectorMockServer.java b/nifi-connector-mock-bundle/nifi-connector-mock-server/src/main/java/org/apache/nifi/mock/connector/server/StandardConnectorMockServer.java index 314445d5ad32..15e54d6d971f 100644 --- a/nifi-connector-mock-bundle/nifi-connector-mock-server/src/main/java/org/apache/nifi/mock/connector/server/StandardConnectorMockServer.java +++ b/nifi-connector-mock-bundle/nifi-connector-mock-server/src/main/java/org/apache/nifi/mock/connector/server/StandardConnectorMockServer.java @@ -102,7 +102,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; import java.util.jar.JarFile; import java.util.stream.Stream; @@ -113,6 +113,8 @@ public class StandardConnectorMockServer implements ConnectorMockServer { private static final String NAR_DEPENDENCIES_PATH = "NAR-INF/bundled-dependencies"; private static final String CONNECTOR_WAR_MANIFEST_PATH = "META-INF/nifi-connector"; private static final String WAR_EXTENSION = ".war"; + private static final Duration DEFAULT_STOP_TIMEOUT = Duration.ofSeconds(60); + private static final long STOP_POLL_INTERVAL_MILLIS = 250L; private Bundle systemBundle; private Set bundles; @@ -362,13 +364,35 @@ public void startConnector() { @Override public void stopConnector() { + // The no-arg convenience method keeps its unchecked contract: a timeout at the default budget is not + // something a caller is expected to recover from, so wrap the checked TimeoutException. try { - connectorNode.stop(flowEngine).get(10, TimeUnit.SECONDS); - } catch (final InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted while waiting for connector to stop", e); - } catch (final Exception e) { - throw new RuntimeException("Failed to stop Connector", e); + stopConnector(DEFAULT_STOP_TIMEOUT); + } catch (final TimeoutException e) { + throw new RuntimeException(e.getMessage(), e); + } + } + + @Override + public void stopConnector(final Duration maxWaitTime) throws TimeoutException { + // Initiate the asynchronous stop, then actively poll the Connector's state until it reports STOPPED. + // The node flips its state to STOPPED at the same point it completes the stop future, and it retries a + // failed component stop internally (every 10 seconds), so polling the state rides through those retries + // up to maxWaitTime instead of being capped by a single fixed blocking wait. + connectorNode.stop(flowEngine); + + final long expirationTime = System.currentTimeMillis() + maxWaitTime.toMillis(); + while (connectorNode.getCurrentState() != ConnectorState.STOPPED) { + if (System.currentTimeMillis() > expirationTime) { + throw new TimeoutException("Timed out waiting for the Connector to stop after " + maxWaitTime); + } + + try { + Thread.sleep(STOP_POLL_INTERVAL_MILLIS); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while waiting for the Connector to stop", e); + } } } diff --git a/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java b/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java index 8d15b2c6c71c..1b8c49e49f40 100644 --- a/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java +++ b/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java @@ -28,10 +28,12 @@ import java.io.File; import java.io.IOException; +import java.time.Duration; import java.util.List; import java.util.Optional; import java.util.Set; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -65,6 +67,25 @@ public void testCreateStartAndStopGenerateAndUpdateConnector() throws IOExceptio } } + @Test + public void testStopConnectorWithTimeoutStopsRunningConnector() throws IOException { + try (final ConnectorTestRunner testRunner = new StandardConnectorTestRunner.Builder() + .connectorClassName("org.apache.nifi.mock.connectors.GenerateAndLog") + .narLibraryDirectory(new File("target/libDir")) + .build()) { + + testRunner.startConnector(); + + // Exercises the timeout-aware overload: it initiates the asynchronous stop and actively polls the + // Connector's state until it reaches STOPPED, returning as soon as it does rather than after a single + // fixed blocking wait. Because start is asynchronous, a stop issued immediately afterwards may have to + // ride through the node's internal stop retries (every 10 seconds) before the state settles, so the + // budget is generous enough to stay deterministic on a slow CI runner; the poll still returns the + // instant the Connector reports STOPPED. + assertDoesNotThrow(() -> testRunner.stopConnector(Duration.ofSeconds(120))); + } + } + @Test public void testConnectorWithMissingBundleFailsValidate() throws IOException { diff --git a/nifi-connector-mock-bundle/nifi-connector-mock/src/main/java/org/apache/nifi/mock/connector/StandardConnectorTestRunner.java b/nifi-connector-mock-bundle/nifi-connector-mock/src/main/java/org/apache/nifi/mock/connector/StandardConnectorTestRunner.java index 7153f8d868b3..186867d7e802 100644 --- a/nifi-connector-mock-bundle/nifi-connector-mock/src/main/java/org/apache/nifi/mock/connector/StandardConnectorTestRunner.java +++ b/nifi-connector-mock-bundle/nifi-connector-mock/src/main/java/org/apache/nifi/mock/connector/StandardConnectorTestRunner.java @@ -50,6 +50,7 @@ import java.util.Map; import java.util.Properties; import java.util.Set; +import java.util.concurrent.TimeoutException; public class StandardConnectorTestRunner implements ConnectorTestRunner, Closeable { private final File narLibraryDirectory; @@ -198,6 +199,11 @@ public void stopConnector() { mockServer.stopConnector(); } + @Override + public void stopConnector(final Duration timeout) throws TimeoutException { + mockServer.stopConnector(timeout); + } + @Override public void waitForDataIngested(final Duration maxWaitTime) { mockServer.waitForDataIngested(maxWaitTime);