You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@nifi.apache.org by jo...@apache.org on 2026/09/12 22:33:40 UTC
(nifi) branch main updated: NIFI-16287 - Ensure orderly shutdown of active and idle JMS workers This closes #11619
This is an automated email from the ASF dual-hosted git repository.
joewitt pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new 7298695681e NIFI-16287 - Ensure orderly shutdown of active and idle JMS workers This closes #11619
7298695681e is described below
commit 7298695681efcecf4d99d9d028b7d1fb98f471f7
Author: Pierre Villard <pi...@gmail.com>
AuthorDate: Thu Sep 3 15:05:52 2026 +0200
NIFI-16287 - Ensure orderly shutdown of active and idle JMS workers
This closes #11619
---
.../nifi/jms/processors/AbstractJMSProcessor.java | 92 +++---
.../org/apache/nifi/jms/processors/JMSWorker.java | 6 +-
.../nifi/jms/processors/JmsWorkerLifecycle.java | 188 ++++++++++++
.../jms/processors/AbstractJMSProcessorTest.java | 326 +++++++++++++++++++++
.../apache/nifi/jms/processors/ConsumeJMSIT.java | 247 ++++++++++++++++
.../apache/nifi/jms/processors/JMSWorkerTest.java | 173 +++++++++++
.../apache/nifi/jms/processors/PublishJMSIT.java | 259 ++++++++++++++++
7 files changed, 1247 insertions(+), 44 deletions(-)
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/AbstractJMSProcessor.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/AbstractJMSProcessor.java
index 2d8559e9ad7..6e343a5c452 100644
--- a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/AbstractJMSProcessor.java
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/AbstractJMSProcessor.java
@@ -17,7 +17,6 @@
package org.apache.nifi.jms.processors;
import jakarta.jms.ConnectionFactory;
-import jakarta.jms.Message;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.annotation.lifecycle.OnStopped;
import org.apache.nifi.annotation.lifecycle.OnUnscheduled;
@@ -31,7 +30,6 @@ import org.apache.nifi.expression.ExpressionLanguageScope;
import org.apache.nifi.jms.cf.IJMSConnectionFactoryProvider;
import org.apache.nifi.jms.cf.JMSConnectionFactoryHandler;
import org.apache.nifi.jms.cf.JMSConnectionFactoryProperties;
-import org.apache.nifi.jms.cf.JMSConnectionFactoryProvider;
import org.apache.nifi.jms.cf.JMSConnectionFactoryProviderDefinition;
import org.apache.nifi.jms.cf.JndiJmsConnectionFactoryHandler;
import org.apache.nifi.jms.cf.JndiJmsConnectionFactoryProperties;
@@ -55,10 +53,8 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
/**
@@ -170,10 +166,8 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
.required(true)
.build();
- private volatile IJMSConnectionFactoryProvider connectionFactoryProvider;
- private volatile BlockingQueue<T> workerPool;
+ private final AtomicReference<JmsWorkerLifecycle<T>> workerLifecycle = new AtomicReference<>();
private volatile boolean runOnPrimary;
- private final AtomicBoolean shutdownWorkers = new AtomicBoolean(false);
private final AtomicInteger clientIdCounter = new AtomicInteger(1);
protected static String getClientId(ProcessContext context) {
@@ -216,20 +210,32 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
@OnPrimaryNodeStateChange
public void onPrimaryNodeChange(final PrimaryNodeState newState) {
- if (isScheduled() && runOnPrimary && newState.equals(PrimaryNodeState.PRIMARY_NODE_REVOKED)) {
- shutdownWorkers.set(true);
- close();
- } else {
- shutdownWorkers.set(false);
+ final JmsWorkerLifecycle<T> lifecycle = workerLifecycle.get();
+ if (isScheduled() && runOnPrimary && lifecycle != null) {
+ if (newState.equals(PrimaryNodeState.PRIMARY_NODE_REVOKED)) {
+ lifecycle.retireGeneration();
+ } else if (newState.equals(PrimaryNodeState.ELECTED_PRIMARY_NODE)) {
+ lifecycle.activateFreshGeneration();
+ }
}
}
@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
- T worker = workerPool.poll();
- if (worker == null) {
+ final JmsWorkerLifecycle<T> lifecycle = workerLifecycle.get();
+ if (lifecycle == null) {
+ return;
+ }
+
+ final JmsWorkerLifecycle.Generation<T> generation = lifecycle.captureGeneration();
+ T worker = lifecycle.pollIdleWorker(generation);
+ if (worker == null && lifecycle.canCreateWorker(generation)) {
try {
- worker = buildTargetResource(context);
+ worker = buildTargetResource(context, lifecycle.getConnectionFactoryProvider());
+ if (!lifecycle.registerWorker(generation, worker)) {
+ worker.shutdown();
+ worker = null;
+ }
} catch (Exception e) {
getLogger().error("Failed to initialize JMS Connection Factory", e);
context.yield();
@@ -237,6 +243,10 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
}
}
+ if (worker == null) {
+ return;
+ }
+
try {
rendezvousWithJms(context, session, worker);
} finally {
@@ -245,37 +255,31 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
//if worker is not valid anymore, don't put it back into a pool, try to rebuild it first, or discard.
//this will be helpful in a situation, when JNDI has changed, or JMS server is not available
//and reconnection is required.
- if (worker == null || !worker.isValid()) {
+ if (!worker.isValid()) {
getLogger().debug("Worker is invalid. Will try re-create... ");
try {
- if (worker != null) {
- worker.shutdown();
- }
// Safe to cast. Method #buildTargetResource(ProcessContext context) sets only CachingConnectionFactory
CachingConnectionFactory currentCF = (CachingConnectionFactory) worker.jmsTemplate.getConnectionFactory();
- connectionFactoryProvider.resetConnectionFactory(currentCF.getTargetConnectionFactory());
- worker = buildTargetResource(context);
+ if (lifecycle.handleInvalidWorker(generation, worker, currentCF.getTargetConnectionFactory())) {
+ final T replacementWorker = buildTargetResource(context, lifecycle.getConnectionFactoryProvider());
+ if (lifecycle.registerWorker(generation, replacementWorker)) {
+ lifecycle.releaseWorker(generation, replacementWorker);
+ } else {
+ replacementWorker.shutdown();
+ }
+ }
} catch (Exception e) {
- getLogger().error("Failed to rebuild: {}", connectionFactoryProvider);
- worker = null;
- }
- }
- if (worker != null) {
- worker.jmsTemplate.setExplicitQosEnabled(false);
- worker.jmsTemplate.setDeliveryMode(Message.DEFAULT_DELIVERY_MODE);
- worker.jmsTemplate.setTimeToLive(Message.DEFAULT_TIME_TO_LIVE);
- worker.jmsTemplate.setPriority(Message.DEFAULT_PRIORITY);
- if (!shutdownWorkers.get()) {
- workerPool.offer(worker);
- } else {
- worker.shutdown();
+ getLogger().error("Failed to rebuild: {}", lifecycle.getConnectionFactoryProvider());
}
+ } else {
+ lifecycle.releaseWorker(generation, worker);
}
}
}
@OnScheduled
public void setup(final ProcessContext context) {
+ final IJMSConnectionFactoryProvider connectionFactoryProvider;
if (context.getProperty(CF_SERVICE).isSet()) {
connectionFactoryProvider = context.getProperty(CF_SERVICE).asControllerService(JMSConnectionFactoryProviderDefinition.class);
} else if (context.getProperty(JndiJmsConnectionFactoryProperties.JNDI_CONNECTION_FACTORY_NAME).isSet()) {
@@ -286,21 +290,23 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
throw new ProcessException("No Connection Factory configured.");
}
- workerPool = new LinkedBlockingQueue<>(context.getMaxConcurrentTasks());
+ workerLifecycle.set(new JmsWorkerLifecycle<>(connectionFactoryProvider, context.getMaxConcurrentTasks(), getLogger()));
runOnPrimary = context.getExecutionNode().equals(ExecutionNode.PRIMARY);
- shutdownWorkers.set(false);
}
@OnUnscheduled
public void shutdownConnectionFactoryProvider(final ProcessContext context) {
- connectionFactoryProvider = null;
+ final JmsWorkerLifecycle<T> lifecycle = workerLifecycle.getAndSet(null);
+ if (lifecycle != null) {
+ lifecycle.closeCycle();
+ }
}
@OnStopped
public void close() {
- T worker;
- while ((worker = workerPool.poll()) != null) {
- worker.shutdown();
+ final JmsWorkerLifecycle<T> lifecycle = workerLifecycle.getAndSet(null);
+ if (lifecycle != null) {
+ lifecycle.closeCycle();
}
}
@@ -323,12 +329,12 @@ public abstract class AbstractJMSProcessor<T extends JMSWorker> extends Abstract
/**
* This method essentially performs initialization of this Processor by
* obtaining an instance of the {@link ConnectionFactory} from the
- * {@link JMSConnectionFactoryProvider} (ControllerService) and performing a
+ * {@link JMSConnectionFactoryProviderDefinition} (ControllerService) and performing a
* series of {@link ConnectionFactory} adaptations which eventually results
* in an instance of the {@link CachingConnectionFactory} used to construct
* {@link JmsTemplate} used by this Processor.
*/
- private T buildTargetResource(ProcessContext context) {
+ private T buildTargetResource(final ProcessContext context, final IJMSConnectionFactoryProvider connectionFactoryProvider) {
final ConnectionFactory connectionFactory = connectionFactoryProvider.getConnectionFactory();
final UserCredentialsConnectionFactoryAdapter cfCredentialsAdapter = new UserCredentialsConnectionFactoryAdapter();
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JMSWorker.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JMSWorker.java
index 162b852d8d6..811f68c3c1d 100644
--- a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JMSWorker.java
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JMSWorker.java
@@ -22,6 +22,7 @@ import org.springframework.jms.connection.CachingConnectionFactory;
import org.springframework.jms.core.JmsTemplate;
import java.nio.channels.Channel;
+import java.util.concurrent.atomic.AtomicBoolean;
/**
* Base class for implementing publishing and consuming JMS workers.
@@ -34,6 +35,7 @@ abstract class JMSWorker {
protected final JmsTemplate jmsTemplate;
protected final ComponentLog processLog;
private final CachingConnectionFactory connectionFactory;
+ private final AtomicBoolean shutdown = new AtomicBoolean(false);
private boolean isValid = true;
/**
@@ -51,7 +53,9 @@ abstract class JMSWorker {
}
public void shutdown() {
- connectionFactory.destroy();
+ if (shutdown.compareAndSet(false, true)) {
+ connectionFactory.destroy();
+ }
}
@Override
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JmsWorkerLifecycle.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JmsWorkerLifecycle.java
new file mode 100644
index 00000000000..64084d6da38
--- /dev/null
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/main/java/org/apache/nifi/jms/processors/JmsWorkerLifecycle.java
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.jms.processors;
+
+import jakarta.jms.ConnectionFactory;
+import org.apache.nifi.jms.cf.IJMSConnectionFactoryProvider;
+import org.apache.nifi.logging.ComponentLog;
+import org.springframework.jms.core.JmsTemplate;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+
+final class JmsWorkerLifecycle<T extends JMSWorker> {
+
+ private final IJMSConnectionFactoryProvider connectionFactoryProvider;
+ private final int maxPoolSize;
+ private final ComponentLog logger;
+
+ private boolean cycleOpen = true;
+ private Generation<T> currentGeneration;
+
+ JmsWorkerLifecycle(final IJMSConnectionFactoryProvider connectionFactoryProvider, final int maxPoolSize, final ComponentLog logger) {
+ this.connectionFactoryProvider = connectionFactoryProvider;
+ this.maxPoolSize = maxPoolSize;
+ this.logger = logger;
+ this.currentGeneration = new Generation<>(maxPoolSize);
+ }
+
+ synchronized Generation<T> captureGeneration() {
+ return currentGeneration;
+ }
+
+ synchronized T pollIdleWorker(final Generation<T> generation) {
+ return isCurrentOpenGeneration(generation) ? generation.idleWorkers.poll() : null;
+ }
+
+ synchronized boolean canCreateWorker(final Generation<T> generation) {
+ return isCurrentOpenGeneration(generation);
+ }
+
+ boolean registerWorker(final Generation<T> generation, final T worker) {
+ boolean accepted;
+ synchronized (this) {
+ generation.allWorkers.add(worker);
+ accepted = isCurrentOpenGeneration(generation);
+ if (!accepted) {
+ generation.allWorkers.remove(worker);
+ }
+ }
+ return accepted;
+ }
+
+ void releaseWorker(final Generation<T> generation, final T worker) {
+ if (worker == null) {
+ return;
+ }
+
+ boolean closeWorker = true;
+ synchronized (this) {
+ if (worker.isValid() && isCurrentOpenGeneration(generation)) {
+ resetWorker(worker.jmsTemplate);
+ if (generation.idleWorkers.offer(worker)) {
+ closeWorker = false;
+ } else {
+ generation.allWorkers.remove(worker);
+ }
+ } else {
+ generation.allWorkers.remove(worker);
+ }
+ }
+
+ if (closeWorker) {
+ worker.shutdown();
+ }
+ }
+
+ boolean handleInvalidWorker(final Generation<T> generation, final T worker, final ConnectionFactory cachedConnectionFactory) {
+ synchronized (this) {
+ generation.allWorkers.remove(worker);
+ }
+
+ worker.shutdown();
+
+ synchronized (this) {
+ if (!isCurrentOpenGeneration(generation)) {
+ return false;
+ }
+
+ connectionFactoryProvider.resetConnectionFactory(cachedConnectionFactory);
+ return true;
+ }
+ }
+
+ void retireGeneration() {
+ closeWorkers(retireCurrentGeneration());
+ }
+
+ synchronized void activateFreshGeneration() {
+ if (cycleOpen && currentGeneration.retired) {
+ currentGeneration = new Generation<>(maxPoolSize);
+ }
+ }
+
+ void closeCycle() {
+ closeWorkers(closeCurrentCycle());
+ }
+
+ IJMSConnectionFactoryProvider getConnectionFactoryProvider() {
+ return connectionFactoryProvider;
+ }
+
+ private synchronized List<T> retireCurrentGeneration() {
+ return currentGeneration.retireAndSnapshot();
+ }
+
+ private synchronized List<T> closeCurrentCycle() {
+ if (!cycleOpen) {
+ return List.of();
+ }
+
+ cycleOpen = false;
+ return currentGeneration.retireAndSnapshot();
+ }
+
+ private synchronized boolean isCurrentOpenGeneration(final Generation<T> generation) {
+ return cycleOpen && currentGeneration == generation && !generation.retired;
+ }
+
+ private void closeWorkers(final List<T> workers) {
+ for (final T worker : workers) {
+ try {
+ worker.shutdown();
+ } catch (final Exception e) {
+ logger.error("Failed to close JMS worker {}", worker, e);
+ }
+ }
+ }
+
+ private void resetWorker(final JmsTemplate jmsTemplate) {
+ jmsTemplate.setExplicitQosEnabled(false);
+ jmsTemplate.setDeliveryMode(jakarta.jms.Message.DEFAULT_DELIVERY_MODE);
+ jmsTemplate.setTimeToLive(jakarta.jms.Message.DEFAULT_TIME_TO_LIVE);
+ jmsTemplate.setPriority(jakarta.jms.Message.DEFAULT_PRIORITY);
+ }
+
+ static final class Generation<T extends JMSWorker> {
+
+ private final BlockingQueue<T> idleWorkers;
+ private final Set<T> allWorkers = new HashSet<>();
+
+ private boolean retired;
+
+ private Generation(final int maxPoolSize) {
+ idleWorkers = new LinkedBlockingQueue<>(maxPoolSize);
+ }
+
+ private List<T> retireAndSnapshot() {
+ if (retired) {
+ return List.of();
+ }
+
+ retired = true;
+ idleWorkers.clear();
+
+ final List<T> workers = new ArrayList<>(allWorkers);
+ allWorkers.clear();
+ return workers;
+ }
+ }
+}
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/AbstractJMSProcessorTest.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/AbstractJMSProcessorTest.java
new file mode 100644
index 00000000000..d26e42a2fdf
--- /dev/null
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/AbstractJMSProcessorTest.java
@@ -0,0 +1,326 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.jms.processors;
+
+import jakarta.jms.ConnectionFactory;
+import org.apache.nifi.jms.cf.IJMSConnectionFactoryProvider;
+import org.apache.nifi.util.LogMessage;
+import org.apache.nifi.util.MockComponentLog;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.util.Arrays;
+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 java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+public class AbstractJMSProcessorTest {
+
+ @Test
+ @Timeout(value = 10)
+ public void earlyUnscheduleShouldCloseIdleAndActiveWorkers() {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final MockComponentLog logger = new MockComponentLog("processor", new Object());
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 2, logger);
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+
+ final TestWorker idleWorker = new TestWorker("idle");
+ final TestWorker activeWorker = new TestWorker("active");
+
+ assertTrue(lifecycle.registerWorker(generation, idleWorker));
+ lifecycle.releaseWorker(generation, idleWorker);
+ assertTrue(lifecycle.registerWorker(generation, activeWorker));
+
+ lifecycle.closeCycle();
+
+ assertEquals(1, idleWorker.getDestroyCalls());
+ assertEquals(1, activeWorker.getDestroyCalls());
+ assertNull(lifecycle.pollIdleWorker(generation));
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void retiredGenerationShouldCloseConstructedWorkerBeforeJmsUse() throws Exception {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+ final CountDownLatch workerConstructed = new CountDownLatch(1);
+ final CountDownLatch retireRequested = new CountDownLatch(1);
+ final AtomicReference<TestWorker> workerReference = new AtomicReference<>();
+ final ExecutorService executorService = Executors.newSingleThreadExecutor();
+
+ try {
+ final Future<Boolean> workerUsedFuture = executorService.submit(() -> {
+ assertTrue(lifecycle.canCreateWorker(generation));
+
+ final TestWorker worker = new TestWorker("constructing");
+ workerReference.set(worker);
+ workerConstructed.countDown();
+
+ assertTrue(retireRequested.await(5, TimeUnit.SECONDS));
+ final boolean accepted = lifecycle.registerWorker(generation, worker);
+ if (accepted) {
+ worker.markUsed();
+ } else {
+ worker.shutdown();
+ }
+
+ return worker.wasUsed();
+ });
+
+ assertTrue(workerConstructed.await(5, TimeUnit.SECONDS));
+ lifecycle.retireGeneration();
+ retireRequested.countDown();
+
+ assertFalse(workerUsedFuture.get(5, TimeUnit.SECONDS));
+ } finally {
+ executorService.shutdownNow();
+ executorService.awaitTermination(5, TimeUnit.SECONDS);
+ }
+
+ final TestWorker worker = workerReference.get();
+ assertNotNull(worker);
+ assertEquals(1, worker.getDestroyCalls());
+ assertFalse(worker.wasUsed());
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void returnedWorkerShouldBeClosedInsteadOfRePooledAfterRetirement() {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+ final TestWorker worker = new TestWorker("returned");
+
+ assertTrue(lifecycle.registerWorker(generation, worker));
+
+ lifecycle.retireGeneration();
+ lifecycle.releaseWorker(generation, worker);
+ lifecycle.activateFreshGeneration();
+
+ assertEquals(1, worker.getDestroyCalls());
+ assertNull(lifecycle.pollIdleWorker(generation));
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void retiredGenerationShouldNotResetOrRebuildInvalidWorker() {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+ final TestWorker invalidWorker = new TestWorker("invalid");
+
+ assertTrue(lifecycle.registerWorker(generation, invalidWorker));
+ invalidWorker.setValid(false);
+ lifecycle.retireGeneration();
+
+ final boolean rebuildAllowed = lifecycle.handleInvalidWorker(generation, invalidWorker, provider.getConnectionFactory());
+
+ assertFalse(rebuildAllowed);
+ assertEquals(0, provider.getResetCalls());
+ assertEquals(1, invalidWorker.getDestroyCalls());
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void oldCycleCleanupShouldNotMutateReplacementCycle() {
+ final JmsWorkerLifecycle<TestWorker> oldLifecycle = new JmsWorkerLifecycle<>(new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("old", new Object()));
+ final JmsWorkerLifecycle<TestWorker> replacementLifecycle = new JmsWorkerLifecycle<>(new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("replacement", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> oldGeneration = oldLifecycle.captureGeneration();
+ final JmsWorkerLifecycle.Generation<TestWorker> replacementGeneration = replacementLifecycle.captureGeneration();
+ final TestWorker oldWorker = new TestWorker("old-worker");
+ final TestWorker oldActiveWorker = new TestWorker("old-active-worker");
+ final TestWorker replacementWorker = new TestWorker("replacement-worker");
+
+ assertTrue(oldLifecycle.registerWorker(oldGeneration, oldWorker));
+ assertTrue(oldLifecycle.registerWorker(oldGeneration, oldActiveWorker));
+ assertTrue(replacementLifecycle.registerWorker(replacementGeneration, replacementWorker));
+ replacementLifecycle.releaseWorker(replacementGeneration, replacementWorker);
+
+ oldLifecycle.closeCycle();
+ oldLifecycle.releaseWorker(oldGeneration, oldActiveWorker);
+
+ assertEquals(1, oldWorker.getDestroyCalls());
+ assertEquals(1, oldActiveWorker.getDestroyCalls());
+ assertEquals(0, replacementWorker.getDestroyCalls());
+ assertSame(replacementWorker, replacementLifecycle.pollIdleWorker(replacementLifecycle.captureGeneration()));
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void primaryRevocationShouldRetireGenerationAndElectionShouldUseFreshGenerationOnly() {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 2, new MockComponentLog("processor", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> revokedGeneration = lifecycle.captureGeneration();
+ final TestWorker idleWorker = new TestWorker("idle");
+ final TestWorker activeWorker = new TestWorker("active");
+
+ assertTrue(lifecycle.registerWorker(revokedGeneration, idleWorker));
+ lifecycle.releaseWorker(revokedGeneration, idleWorker);
+ assertTrue(lifecycle.registerWorker(revokedGeneration, activeWorker));
+
+ lifecycle.retireGeneration();
+ lifecycle.activateFreshGeneration();
+
+ final JmsWorkerLifecycle.Generation<TestWorker> freshGeneration = lifecycle.captureGeneration();
+ final TestWorker freshWorker = new TestWorker("fresh");
+ assertTrue(lifecycle.registerWorker(freshGeneration, freshWorker));
+ lifecycle.releaseWorker(freshGeneration, freshWorker);
+ lifecycle.releaseWorker(revokedGeneration, activeWorker);
+
+ assertEquals(1, idleWorker.getDestroyCalls());
+ assertEquals(1, activeWorker.getDestroyCalls());
+ assertSame(freshWorker, lifecycle.pollIdleWorker(freshGeneration));
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void bulkCloseShouldContinueAfterShutdownFailureAndLogWorker() {
+ final RecordingConnectionFactoryProvider provider = new RecordingConnectionFactoryProvider();
+ final MockComponentLog logger = new MockComponentLog("processor", new Object());
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(provider, 3, logger);
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+ final TestWorker healthyWorker = new TestWorker("healthy");
+ final TestWorker failingWorker = new TestWorker("failing", true);
+ final TestWorker trailingWorker = new TestWorker("trailing");
+
+ assertTrue(lifecycle.registerWorker(generation, healthyWorker));
+ assertTrue(lifecycle.registerWorker(generation, failingWorker));
+ assertTrue(lifecycle.registerWorker(generation, trailingWorker));
+
+ lifecycle.closeCycle();
+
+ assertEquals(1, healthyWorker.getDestroyCalls());
+ assertEquals(1, failingWorker.getDestroyCalls());
+ assertEquals(1, trailingWorker.getDestroyCalls());
+ assertTrue(logger.getErrorMessages().stream().anyMatch(message -> containsWorkerReference(message, failingWorker)));
+ }
+
+ @Test
+ public void closeCycleShouldBeIdempotent() {
+ final JmsWorkerLifecycle<TestWorker> lifecycle = new JmsWorkerLifecycle<>(
+ new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("processor", new Object()));
+ final JmsWorkerLifecycle.Generation<TestWorker> generation = lifecycle.captureGeneration();
+ final TestWorker worker = new TestWorker("worker");
+ assertTrue(lifecycle.registerWorker(generation, worker));
+
+ lifecycle.closeCycle();
+ lifecycle.closeCycle();
+
+ assertEquals(1, worker.getDestroyCalls());
+ }
+
+ private boolean containsWorkerReference(final LogMessage message, final TestWorker worker) {
+ final Object[] arguments = message.getArgs();
+ return arguments != null && Arrays.stream(arguments).anyMatch(argument -> argument == worker);
+ }
+
+ private static final class RecordingConnectionFactoryProvider implements IJMSConnectionFactoryProvider {
+
+ private final ConnectionFactory connectionFactory = mock(ConnectionFactory.class);
+ private final AtomicInteger resetCalls = new AtomicInteger();
+
+ @Override
+ public ConnectionFactory getConnectionFactory() {
+ return connectionFactory;
+ }
+
+ @Override
+ public void resetConnectionFactory(final ConnectionFactory cachedFactory) {
+ resetCalls.incrementAndGet();
+ }
+
+ private int getResetCalls() {
+ return resetCalls.get();
+ }
+
+ }
+
+ private static final class TestWorker extends JMSWorker {
+
+ private final DestroyTrackingCachingConnectionFactory cachingConnectionFactory;
+ private final AtomicBoolean used = new AtomicBoolean();
+
+ private TestWorker(final String destinationName) {
+ this(destinationName, false);
+ }
+
+ private TestWorker(final String destinationName, final boolean failDestroy) {
+ this(new DestroyTrackingCachingConnectionFactory(mock(ConnectionFactory.class), failDestroy), destinationName);
+ }
+
+ private TestWorker(final DestroyTrackingCachingConnectionFactory connectionFactory, final String destinationName) {
+ super(connectionFactory, createJmsTemplate(connectionFactory, destinationName), new MockComponentLog(destinationName, destinationName));
+ this.cachingConnectionFactory = connectionFactory;
+ }
+
+ private static org.springframework.jms.core.JmsTemplate createJmsTemplate(final DestroyTrackingCachingConnectionFactory connectionFactory, final String destinationName) {
+ final org.springframework.jms.core.JmsTemplate jmsTemplate = new org.springframework.jms.core.JmsTemplate();
+ jmsTemplate.setConnectionFactory(connectionFactory);
+ jmsTemplate.setDefaultDestinationName(destinationName);
+ return jmsTemplate;
+ }
+
+ private void markUsed() {
+ used.set(true);
+ }
+
+ private boolean wasUsed() {
+ return used.get();
+ }
+
+ private int getDestroyCalls() {
+ return cachingConnectionFactory.getDestroyCalls();
+ }
+ }
+
+ private static final class DestroyTrackingCachingConnectionFactory extends org.springframework.jms.connection.CachingConnectionFactory {
+
+ private final AtomicInteger destroyCalls = new AtomicInteger();
+ private final boolean failFirstDestroy;
+
+ private DestroyTrackingCachingConnectionFactory(final ConnectionFactory targetConnectionFactory, final boolean failFirstDestroy) {
+ super(targetConnectionFactory);
+ this.failFirstDestroy = failFirstDestroy;
+ }
+
+ @Override
+ public void destroy() {
+ final int destroyCall = destroyCalls.incrementAndGet();
+ if (failFirstDestroy && destroyCall == 1) {
+ throw new RuntimeException("destroy failed");
+ }
+ }
+
+ private int getDestroyCalls() {
+ return destroyCalls.get();
+ }
+ }
+}
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/ConsumeJMSIT.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/ConsumeJMSIT.java
index fe31c2f36d5..3d6a79d9b63 100644
--- a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/ConsumeJMSIT.java
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/ConsumeJMSIT.java
@@ -26,6 +26,7 @@ import jakarta.jms.ConnectionFactory;
import jakarta.jms.JMSException;
import jakarta.jms.MapMessage;
import jakarta.jms.Message;
+import jakarta.jms.MessageConsumer;
import jakarta.jms.MessageProducer;
import jakarta.jms.ObjectMessage;
import jakarta.jms.Session;
@@ -51,8 +52,11 @@ import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.io.OutputStreamCallback;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.scheduling.ExecutionNode;
+import org.apache.nifi.state.MockStateManager;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.MockProcessContext;
+import org.apache.nifi.util.MockProcessSession;
+import org.apache.nifi.util.SharedSessionState;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.jupiter.api.Test;
@@ -63,13 +67,24 @@ import org.springframework.jms.core.MessageCreator;
import org.springframework.jms.support.JmsHeaders;
import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
import java.net.URI;
import java.net.UnknownHostException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
+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 java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
import javax.net.SocketFactory;
import static java.util.Arrays.asList;
@@ -78,7 +93,9 @@ import static org.apache.nifi.jms.processors.helpers.JMSTestUtil.createJsonRecor
import static org.apache.nifi.jms.processors.helpers.JMSTestUtil.createJsonRecordSetWriterService;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+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.Mockito.doThrow;
@@ -690,6 +707,97 @@ public class ConsumeJMSIT {
}
}
+ @Test
+ @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS)
+ public void activeReceiveShouldBeInterruptedOnUnscheduledWithoutFlowFileOrReplacementConnection() throws Exception {
+ final BrokerService broker = new BrokerService();
+ final ExecutorService executorService = Executors.newSingleThreadExecutor();
+ try {
+ broker.setPersistent(false);
+ broker.setBrokerName("blocked-receive-broker");
+ broker.start();
+
+ final ActiveMQConnectionFactory innerConnectionFactory = new ActiveMQConnectionFactory("vm://blocked-receive-broker");
+ final BlockingReceiveConnectionFactory controlledConnectionFactory = new BlockingReceiveConnectionFactory(innerConnectionFactory);
+ final CountingConsumeJMS processor = new CountingConsumeJMS();
+ final TestRunner runner = initializeTestRunner(processor, controlledConnectionFactory.getConnectionFactory(), "blocked-receive-destination");
+ runner.setProperty(ConsumeJMS.TIMEOUT, "30 sec");
+
+ final ProcessContext processContext = runner.getProcessContext();
+ processor.onSchedule(processContext);
+ processor.setup(processContext);
+
+ final MockProcessSession processSession = new MockProcessSession(new SharedSessionState(processor, new AtomicLong(0L)), processor);
+ final Future<Throwable> future = executorService.submit(() -> {
+ try {
+ processor.onTrigger(processContext, processSession);
+ return null;
+ } catch (final Throwable throwable) {
+ return throwable;
+ }
+ });
+
+ assertTrue(controlledConnectionFactory.awaitReceiveEntered(5, TimeUnit.SECONDS));
+
+ processor.shutdownConnectionFactoryProvider(processContext);
+
+ final Throwable thrown = future.get(5, TimeUnit.SECONDS);
+ assertTrue(thrown == null || thrown instanceof ProcessException || thrown instanceof JMSException
+ || thrown instanceof org.springframework.jms.IllegalStateException,
+ "Unexpected exception type returned from blocked receive shutdown path: " + thrown);
+ assertTrue(controlledConnectionFactory.awaitConnectionClosed(5, TimeUnit.SECONDS));
+ assertEquals(0, broker.getCurrentConnections());
+ assertEquals(1, controlledConnectionFactory.getOpenedConnections());
+ assertEquals(1, processor.getBuildCount());
+ assertTrue(processSession.getFlowFilesForRelationship(ConsumeJMS.REL_SUCCESS).isEmpty());
+ assertTrue(processSession.getFlowFilesForRelationship(ConsumeJMS.REL_PARSE_FAILURE).isEmpty());
+ } finally {
+ executorService.shutdownNow();
+ broker.stop();
+ }
+ }
+
+ @Test
+ @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS)
+ public void delayedCommitCallbackShouldFailAcknowledgeAfterShutdownAndAllowRedelivery() throws Exception {
+ final BrokerService broker = new BrokerService();
+ try {
+ broker.setPersistent(false);
+ broker.setBrokerName("delayed-commit-broker");
+ broker.start();
+
+ final String destinationName = "delayed-commit-destination";
+ final ActiveMQConnectionFactory connectionFactory = new ActiveMQConnectionFactory("vm://delayed-commit-broker");
+ publishQueueMessage(connectionFactory, destinationName, "delayed-ack-message");
+
+ final ConsumeJMS processor = new ConsumeJMS();
+ final TestRunner runner = initializeTestRunner(processor, connectionFactory, destinationName);
+ final ProcessContext processContext = runner.getProcessContext();
+ processor.onSchedule(processContext);
+ processor.setup(processContext);
+
+ final DelayedCommitProcessSession processSession = new DelayedCommitProcessSession(new SharedSessionState(processor, new AtomicLong(0L)), processor);
+
+ processor.onTrigger(processContext, processSession);
+
+ assertEquals(1, processSession.getFlowFilesForRelationship(ConsumeJMS.REL_SUCCESS).size());
+ assertNotNull(processSession.getDelayedSuccessCallback());
+
+ processor.shutdownConnectionFactoryProvider(processContext);
+
+ final ProcessException exception = assertThrows(ProcessException.class, processSession::runDelayedSuccessCallback);
+ assertNotNull(exception.getCause());
+
+ final JmsTemplate jmsTemplate = new JmsTemplate(connectionFactory);
+ jmsTemplate.setReceiveTimeout(5000L);
+ final Message redelivered = jmsTemplate.receive(destinationName);
+ assertInstanceOf(TextMessage.class, redelivered);
+ assertEquals("delayed-ack-message", ((TextMessage) redelivered).getText());
+ } finally {
+ broker.stop();
+ }
+ }
+
private static ArrayNode createTestJsonInput() {
final ObjectMapper mapper = new ObjectMapper();
@@ -766,4 +874,143 @@ public class ConsumeJMSIT {
return c1Consumer;
}
+ private static void publishQueueMessage(final ConnectionFactory connectionFactory, final String destinationName, final String messageText) {
+ final JmsTemplate jmsTemplate = new JmsTemplate(connectionFactory);
+ jmsTemplate.send(destinationName, session -> session.createTextMessage(messageText));
+ }
+
+ private static Object invokeTarget(final Object target, final Method method, final Object[] args) throws Throwable {
+ try {
+ return method.invoke(target, args);
+ } catch (final InvocationTargetException exception) {
+ throw exception.getCause();
+ }
+ }
+
+ private static final class CountingConsumeJMS extends ConsumeJMS {
+
+ private final AtomicInteger buildCount = new AtomicInteger();
+
+ @Override
+ protected JMSConsumer finishBuildingJmsWorker(final CachingConnectionFactory connectionFactory, final JmsTemplate jmsTemplate,
+ final ProcessContext processContext) {
+ buildCount.incrementAndGet();
+ return super.finishBuildingJmsWorker(connectionFactory, jmsTemplate, processContext);
+ }
+
+ private int getBuildCount() {
+ return buildCount.get();
+ }
+ }
+
+ private static final class DelayedCommitProcessSession extends MockProcessSession {
+
+ private Runnable delayedSuccessCallback;
+
+ private DelayedCommitProcessSession(final SharedSessionState sharedState, final ConsumeJMS processor) {
+ super(sharedState, processor, new MockStateManager(processor));
+ }
+
+ @Override
+ public void commitAsync(final Runnable onSuccess, final Consumer<Throwable> onFailure) {
+ delayedSuccessCallback = onSuccess;
+ super.commitAsync(null, onFailure);
+ }
+
+ private Runnable getDelayedSuccessCallback() {
+ return delayedSuccessCallback;
+ }
+
+ private void runDelayedSuccessCallback() {
+ delayedSuccessCallback.run();
+ }
+ }
+
+ private static final class BlockingReceiveConnectionFactory {
+
+ private final ConnectionFactory connectionFactory;
+ private final CountDownLatch receiveEntered = new CountDownLatch(1);
+ private final CountDownLatch connectionClosed = new CountDownLatch(1);
+ private final AtomicInteger openedConnections = new AtomicInteger();
+
+ private BlockingReceiveConnectionFactory(final ConnectionFactory targetConnectionFactory) {
+ Objects.requireNonNull(targetConnectionFactory);
+ connectionFactory = (ConnectionFactory) Proxy.newProxyInstance(
+ ConnectionFactory.class.getClassLoader(),
+ new Class[] {ConnectionFactory.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetConnectionFactory, method, args);
+ if ("createConnection".equals(method.getName())) {
+ openedConnections.incrementAndGet();
+ return createConnectionProxy((Connection) result);
+ }
+ return result;
+ });
+ }
+
+ private ConnectionFactory getConnectionFactory() {
+ return connectionFactory;
+ }
+
+ private boolean awaitReceiveEntered(final long timeout, final TimeUnit timeUnit) throws InterruptedException {
+ return receiveEntered.await(timeout, timeUnit);
+ }
+
+ private boolean awaitConnectionClosed(final long timeout, final TimeUnit timeUnit) throws InterruptedException {
+ return connectionClosed.await(timeout, timeUnit);
+ }
+
+ private int getOpenedConnections() {
+ return openedConnections.get();
+ }
+
+ private Connection createConnectionProxy(final Connection targetConnection) {
+ return (Connection) Proxy.newProxyInstance(
+ Connection.class.getClassLoader(),
+ new Class[] {Connection.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetConnection, method, args);
+ if ("createSession".equals(method.getName())) {
+ return createSessionProxy((Session) result);
+ }
+ if ("close".equals(method.getName())) {
+ connectionClosed.countDown();
+ }
+ return result;
+ });
+ }
+
+ private Session createSessionProxy(final Session targetSession) {
+ return (Session) Proxy.newProxyInstance(
+ Session.class.getClassLoader(),
+ new Class[] {Session.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetSession, method, args);
+ if (isConsumerFactoryMethod(method.getName())) {
+ return createMessageConsumerProxy((MessageConsumer) result);
+ }
+ return result;
+ });
+ }
+
+ private MessageConsumer createMessageConsumerProxy(final MessageConsumer targetConsumer) {
+ return (MessageConsumer) Proxy.newProxyInstance(
+ MessageConsumer.class.getClassLoader(),
+ new Class[] {MessageConsumer.class},
+ (proxy, method, args) -> {
+ if ("receive".equals(method.getName())) {
+ receiveEntered.countDown();
+ }
+ return invokeTarget(targetConsumer, method, args);
+ });
+ }
+
+ private boolean isConsumerFactoryMethod(final String methodName) {
+ return "createConsumer".equals(methodName)
+ || "createDurableConsumer".equals(methodName)
+ || "createSharedConsumer".equals(methodName)
+ || "createSharedDurableConsumer".equals(methodName);
+ }
+ }
+
}
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/JMSWorkerTest.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/JMSWorkerTest.java
new file mode 100644
index 00000000000..cbb6c919796
--- /dev/null
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/JMSWorkerTest.java
@@ -0,0 +1,173 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.nifi.jms.processors;
+
+import jakarta.jms.ConnectionFactory;
+import org.apache.nifi.logging.ComponentLog;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.springframework.jms.connection.CachingConnectionFactory;
+import org.springframework.jms.core.JmsTemplate;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+
+public class JMSWorkerTest {
+
+ @Test
+ @Timeout(value = 10)
+ public void shutdownShouldDestroyConnectionFactoryOnceAcrossConcurrentCallers() throws Exception {
+ final DestroyTrackingCachingConnectionFactory cachingConnectionFactory = new DestroyTrackingCachingConnectionFactory(mock(ConnectionFactory.class), false);
+ final TestWorker worker = new TestWorker(cachingConnectionFactory);
+ final int callers = 8;
+ final CountDownLatch ready = new CountDownLatch(callers);
+ final CountDownLatch start = new CountDownLatch(1);
+ final ExecutorService executorService = Executors.newFixedThreadPool(callers);
+
+ try {
+ final List<Future<Void>> futures = new ArrayList<>();
+ for (int index = 0; index < callers; index++) {
+ futures.add(executorService.submit(() -> {
+ ready.countDown();
+ assertTrue(start.await(5, TimeUnit.SECONDS));
+ worker.shutdown();
+ return null;
+ }));
+ }
+
+ assertTrue(ready.await(5, TimeUnit.SECONDS));
+ start.countDown();
+
+ for (final Future<Void> future : futures) {
+ future.get(5, TimeUnit.SECONDS);
+ }
+ } finally {
+ executorService.shutdownNow();
+ executorService.awaitTermination(5, TimeUnit.SECONDS);
+ }
+
+ assertEquals(1, cachingConnectionFactory.getDestroyCalls());
+ }
+
+ @Test
+ @Timeout(value = 10)
+ public void shutdownShouldNotRetryDestroyAfterInitialFailure() throws Exception {
+ final DestroyTrackingCachingConnectionFactory cachingConnectionFactory = new DestroyTrackingCachingConnectionFactory(mock(ConnectionFactory.class), true);
+ final TestWorker worker = new TestWorker(cachingConnectionFactory);
+
+ final RuntimeException thrown = getRuntimeException(worker);
+
+ assertNotNull(thrown, "Expected one shutdown caller to receive the destroy failure");
+ assertEquals("destroy failed", thrown.getMessage());
+ assertEquals(1, cachingConnectionFactory.getDestroyCalls());
+ assertDoesNotThrow(worker::shutdown);
+ assertEquals(1, cachingConnectionFactory.getDestroyCalls());
+ }
+
+ private RuntimeException getRuntimeException(final TestWorker worker) throws Exception {
+ final ExecutorService executorService = Executors.newFixedThreadPool(2);
+ final CountDownLatch ready = new CountDownLatch(2);
+ final CountDownLatch start = new CountDownLatch(1);
+
+ try {
+ final Future<Void> firstFuture = executorService.submit(() -> {
+ ready.countDown();
+ assertTrue(start.await(5, TimeUnit.SECONDS));
+ worker.shutdown();
+ return null;
+ });
+ final Future<Void> secondFuture = executorService.submit(() -> {
+ ready.countDown();
+ assertTrue(start.await(5, TimeUnit.SECONDS));
+ worker.shutdown();
+ return null;
+ });
+
+ assertTrue(ready.await(5, TimeUnit.SECONDS));
+ start.countDown();
+
+ RuntimeException runtimeException = null;
+ runtimeException = mergeRuntimeException(runtimeException, firstFuture);
+ runtimeException = mergeRuntimeException(runtimeException, secondFuture);
+ return runtimeException;
+ } finally {
+ executorService.shutdownNow();
+ executorService.awaitTermination(5, TimeUnit.SECONDS);
+ }
+ }
+
+ private RuntimeException mergeRuntimeException(final RuntimeException current, final Future<Void> future) throws Exception {
+ try {
+ future.get(5, TimeUnit.SECONDS);
+ return current;
+ } catch (final ExecutionException executionException) {
+ assertInstanceOf(RuntimeException.class, executionException.getCause());
+ return current == null ? (RuntimeException) executionException.getCause() : current;
+ }
+ }
+
+ private static final class TestWorker extends JMSWorker {
+
+ private TestWorker(final DestroyTrackingCachingConnectionFactory connectionFactory) {
+ super(connectionFactory, createJmsTemplate(connectionFactory), mock(ComponentLog.class));
+ }
+
+ private static JmsTemplate createJmsTemplate(final CachingConnectionFactory connectionFactory) {
+ final JmsTemplate jmsTemplate = new JmsTemplate();
+ jmsTemplate.setConnectionFactory(connectionFactory);
+ jmsTemplate.setDefaultDestinationName("test-destination");
+ return jmsTemplate;
+ }
+ }
+
+ private static final class DestroyTrackingCachingConnectionFactory extends CachingConnectionFactory {
+
+ private final AtomicInteger destroyCalls = new AtomicInteger();
+ private final boolean failFirstDestroy;
+
+ private DestroyTrackingCachingConnectionFactory(final ConnectionFactory targetConnectionFactory, final boolean failFirstDestroy) {
+ super(targetConnectionFactory);
+ this.failFirstDestroy = failFirstDestroy;
+ }
+
+ @Override
+ public void destroy() {
+ final int destroyCall = destroyCalls.incrementAndGet();
+ if (failFirstDestroy && destroyCall == 1) {
+ throw new RuntimeException("destroy failed");
+ }
+ }
+
+ private int getDestroyCalls() {
+ return destroyCalls.get();
+ }
+ }
+}
diff --git a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/PublishJMSIT.java b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/PublishJMSIT.java
index 525c99dd83b..40cc7855160 100644
--- a/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/PublishJMSIT.java
+++ b/nifi-extension-bundles/nifi-jms-bundle/nifi-jms-processors/src/test/java/org/apache/nifi/jms/processors/PublishJMSIT.java
@@ -19,9 +19,12 @@ package org.apache.nifi.jms.processors;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import jakarta.jms.BytesMessage;
+import jakarta.jms.Connection;
import jakarta.jms.ConnectionFactory;
import jakarta.jms.Message;
+import jakarta.jms.MessageProducer;
import jakarta.jms.Queue;
+import jakarta.jms.Session;
import jakarta.jms.TextMessage;
import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.broker.BrokerService;
@@ -40,6 +43,8 @@ import org.apache.nifi.provenance.ProvenanceEventType;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.MockProcessContext;
+import org.apache.nifi.util.MockProcessSession;
+import org.apache.nifi.util.SharedSessionState;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
import org.junit.jupiter.api.AfterEach;
@@ -50,12 +55,21 @@ import org.springframework.jms.core.JmsTemplate;
import org.springframework.jms.support.JmsHeaders;
import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.net.URI;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
+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 java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import javax.net.SocketFactory;
@@ -550,6 +564,110 @@ public class PublishJMSIT {
}
}
+ @Test
+ @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS)
+ public void activeSendShouldRouteSuccessWhenSendCompletesBeforeShutdown() throws Exception {
+ final BrokerService broker = new BrokerService();
+ final ExecutorService executorService = Executors.newSingleThreadExecutor();
+ try {
+ broker.setPersistent(false);
+ broker.setBrokerName("publisher-success-broker");
+ broker.start();
+
+ final BlockingProducerConnectionFactory controlledConnectionFactory = new BlockingProducerConnectionFactory(
+ new ActiveMQConnectionFactory("vm://publisher-success-broker"));
+ final CountingPublishJMS processor = new CountingPublishJMS();
+ final TestRunner runner = initializeTestRunner(processor, controlledConnectionFactory.getConnectionFactory(), "publisher-success-destination");
+
+ final ProcessContext processContext = runner.getProcessContext();
+ processor.onScheduled(processContext);
+ processor.setup(processContext);
+
+ final MockProcessSession processSession = createEnqueuedSession(processor, "success-message");
+ final Future<Throwable> future = executorService.submit(() -> {
+ try {
+ processor.onTrigger(processContext, processSession);
+ return null;
+ } catch (final Throwable throwable) {
+ return throwable;
+ }
+ });
+
+ assertTrue(controlledConnectionFactory.awaitSendEntered(5, TimeUnit.SECONDS));
+ controlledConnectionFactory.allowSendToProceed();
+ assertTrue(controlledConnectionFactory.awaitSendCompleted(5, TimeUnit.SECONDS));
+ assertTrue(processor.awaitRendezvousCompleted(5, TimeUnit.SECONDS));
+
+ processor.shutdownConnectionFactoryProvider(processContext);
+ processor.allowRendezvousReturn();
+
+ assertNull(future.get(5, TimeUnit.SECONDS));
+ assertEquals(1, processSession.getFlowFilesForRelationship(REL_SUCCESS).size());
+ assertTrue(processSession.getFlowFilesForRelationship(REL_FAILURE).isEmpty());
+ assertEquals(1, controlledConnectionFactory.getOpenedConnections());
+ assertEquals(1, processor.getBuildCount());
+
+ final JmsTemplate verifyTemplate = new JmsTemplate(new ActiveMQConnectionFactory("vm://publisher-success-broker"));
+ verifyTemplate.setReceiveTimeout(5000L);
+ final Message message = verifyTemplate.receive("publisher-success-destination");
+ assertInstanceOf(BytesMessage.class, message);
+ } finally {
+ executorService.shutdownNow();
+ broker.stop();
+ }
+ }
+
+ @Test
+ @Timeout(value = 10000, unit = TimeUnit.MILLISECONDS)
+ public void activeSendShouldRouteFailureWithoutReplacementAfterShutdownInterruptsSend() throws Exception {
+ final BrokerService broker = new BrokerService();
+ final ExecutorService executorService = Executors.newSingleThreadExecutor();
+ try {
+ broker.setPersistent(false);
+ broker.setBrokerName("publisher-failure-broker");
+ broker.start();
+
+ final BlockingProducerConnectionFactory controlledConnectionFactory = new BlockingProducerConnectionFactory(
+ new ActiveMQConnectionFactory("vm://publisher-failure-broker"));
+ final CountingPublishJMS processor = new CountingPublishJMS();
+ final TestRunner runner = initializeTestRunner(processor, controlledConnectionFactory.getConnectionFactory(), "publisher-failure-destination");
+
+ final ProcessContext processContext = runner.getProcessContext();
+ processor.onScheduled(processContext);
+ processor.setup(processContext);
+
+ final MockProcessSession processSession = createEnqueuedSession(processor, "failure-message");
+ final Future<Throwable> future = executorService.submit(() -> {
+ try {
+ processor.onTrigger(processContext, processSession);
+ return null;
+ } catch (final Throwable throwable) {
+ return throwable;
+ }
+ });
+
+ assertTrue(controlledConnectionFactory.awaitSendEntered(5, TimeUnit.SECONDS));
+
+ processor.shutdownConnectionFactoryProvider(processContext);
+ controlledConnectionFactory.allowSendToProceed();
+ assertTrue(processor.awaitRendezvousCompleted(5, TimeUnit.SECONDS));
+ processor.allowRendezvousReturn();
+
+ assertNull(future.get(5, TimeUnit.SECONDS));
+ assertTrue(processSession.getFlowFilesForRelationship(REL_SUCCESS).isEmpty());
+ assertEquals(1, processSession.getFlowFilesForRelationship(REL_FAILURE).size());
+ assertEquals(1, controlledConnectionFactory.getOpenedConnections());
+ assertEquals(1, processor.getBuildCount());
+
+ final JmsTemplate verifyTemplate = new JmsTemplate(new ActiveMQConnectionFactory("vm://publisher-failure-broker"));
+ verifyTemplate.setReceiveTimeout(250L);
+ assertNull(verifyTemplate.receive("publisher-failure-destination"));
+ } finally {
+ executorService.shutdownNow();
+ broker.stop();
+ }
+ }
+
@Test
public void whenExceptionIsRaisedDuringConnectionFactoryInitializationTheProcessorShouldBeYielded() {
final String nonExistentClassName = "DummyInitialContextFactoryClass";
@@ -758,6 +876,147 @@ public class PublishJMSIT {
assertEquals(expectedDetails, event.getDetails());
}
+ private MockProcessSession createEnqueuedSession(final PublishJMS processor, final String content) {
+ final SharedSessionState sharedSessionState = new SharedSessionState(processor, new AtomicLong(0L));
+ final MockProcessSession processSession = new MockProcessSession(sharedSessionState, processor);
+ sharedSessionState.getFlowFileQueue().offer(processSession.createFlowFile(content.getBytes()));
+ return processSession;
+ }
+
+ private static Object invokeTarget(final Object target, final Method method, final Object[] args) throws Throwable {
+ try {
+ return method.invoke(target, args);
+ } catch (final InvocationTargetException exception) {
+ throw exception.getCause();
+ }
+ }
+
+ private static final class CountingPublishJMS extends PublishJMS {
+
+ private final AtomicInteger buildCount = new AtomicInteger();
+ private final CountDownLatch rendezvousCompleted = new CountDownLatch(1);
+ private final CountDownLatch allowRendezvousReturn = new CountDownLatch(1);
+
+ @Override
+ protected void rendezvousWithJms(final ProcessContext context, final ProcessSession processSession, final JMSPublisher publisher) throws ProcessException {
+ super.rendezvousWithJms(context, processSession, publisher);
+ rendezvousCompleted.countDown();
+ try {
+ assertTrue(allowRendezvousReturn.await(5, TimeUnit.SECONDS));
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new ProcessException(e);
+ }
+ }
+
+ @Override
+ protected JMSPublisher finishBuildingJmsWorker(final org.springframework.jms.connection.CachingConnectionFactory connectionFactory,
+ final JmsTemplate jmsTemplate, final ProcessContext processContext) {
+ buildCount.incrementAndGet();
+ return super.finishBuildingJmsWorker(connectionFactory, jmsTemplate, processContext);
+ }
+
+ private int getBuildCount() {
+ return buildCount.get();
+ }
+
+ private boolean awaitRendezvousCompleted(final long timeout, final TimeUnit timeUnit) throws InterruptedException {
+ return rendezvousCompleted.await(timeout, timeUnit);
+ }
+
+ private void allowRendezvousReturn() {
+ allowRendezvousReturn.countDown();
+ }
+ }
+
+ private static final class BlockingProducerConnectionFactory {
+
+ private final ConnectionFactory connectionFactory;
+ private final CountDownLatch sendEntered = new CountDownLatch(1);
+ private final CountDownLatch allowSend = new CountDownLatch(1);
+ private final CountDownLatch sendCompleted = new CountDownLatch(1);
+ private final AtomicInteger openedConnections = new AtomicInteger();
+
+ private BlockingProducerConnectionFactory(final ConnectionFactory targetConnectionFactory) {
+ Objects.requireNonNull(targetConnectionFactory);
+ connectionFactory = (ConnectionFactory) Proxy.newProxyInstance(
+ ConnectionFactory.class.getClassLoader(),
+ new Class[] {ConnectionFactory.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetConnectionFactory, method, args);
+ if ("createConnection".equals(method.getName())) {
+ openedConnections.incrementAndGet();
+ return createConnectionProxy((Connection) result);
+ }
+ return result;
+ });
+ }
+
+ private ConnectionFactory getConnectionFactory() {
+ return connectionFactory;
+ }
+
+ private boolean awaitSendEntered(final long timeout, final TimeUnit timeUnit) throws InterruptedException {
+ return sendEntered.await(timeout, timeUnit);
+ }
+
+ private boolean awaitSendCompleted(final long timeout, final TimeUnit timeUnit) throws InterruptedException {
+ return sendCompleted.await(timeout, timeUnit);
+ }
+
+ private void allowSendToProceed() {
+ allowSend.countDown();
+ }
+
+ private int getOpenedConnections() {
+ return openedConnections.get();
+ }
+
+ private Connection createConnectionProxy(final Connection targetConnection) {
+ return (Connection) Proxy.newProxyInstance(
+ Connection.class.getClassLoader(),
+ new Class[] {Connection.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetConnection, method, args);
+ if ("createSession".equals(method.getName())) {
+ return createSessionProxy((Session) result);
+ }
+ return result;
+ });
+ }
+
+ private Session createSessionProxy(final Session targetSession) {
+ return (Session) Proxy.newProxyInstance(
+ Session.class.getClassLoader(),
+ new Class[] {Session.class},
+ (proxy, method, args) -> {
+ final Object result = invokeTarget(targetSession, method, args);
+ if ("createProducer".equals(method.getName())) {
+ return createProducerProxy((MessageProducer) result);
+ }
+ return result;
+ });
+ }
+
+ private MessageProducer createProducerProxy(final MessageProducer targetProducer) {
+ return (MessageProducer) Proxy.newProxyInstance(
+ MessageProducer.class.getClassLoader(),
+ new Class[] {MessageProducer.class},
+ (proxy, method, args) -> {
+ if ("send".equals(method.getName())) {
+ sendEntered.countDown();
+ assertTrue(allowSend.await(5, TimeUnit.SECONDS));
+ try {
+ return invokeTarget(targetProducer, method, args);
+ } finally {
+ sendCompleted.countDown();
+ }
+ }
+ return invokeTarget(targetProducer, method, args);
+ });
+ }
+ }
+
private static ArrayNode createTestJsonInput() {
final ObjectMapper mapper = new ObjectMapper();