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 2d8559e9ad7f..6e343a5c4523 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.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.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 extends Abstract .required(true) .build(); - private volatile IJMSConnectionFactoryProvider connectionFactoryProvider; - private volatile BlockingQueue workerPool; + private final AtomicReference> 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 @@ protected Collection customValidate(ValidationContext validati @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 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 lifecycle = workerLifecycle.get(); + if (lifecycle == null) { + return; + } + + final JmsWorkerLifecycle.Generation 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 void onTrigger(ProcessContext context, ProcessSession session) throws Pro } } + if (worker == null) { + return; + } + try { rendezvousWithJms(context, session, worker); } finally { @@ -245,37 +255,31 @@ public void onTrigger(ProcessContext context, ProcessSession session) throws Pro //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 void setup(final ProcessContext context) { 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 lifecycle = workerLifecycle.getAndSet(null); + if (lifecycle != null) { + lifecycle.closeCycle(); + } } @OnStopped public void close() { - T worker; - while ((worker = workerPool.poll()) != null) { - worker.shutdown(); + final JmsWorkerLifecycle lifecycle = workerLifecycle.getAndSet(null); + if (lifecycle != null) { + lifecycle.closeCycle(); } } @@ -323,12 +329,12 @@ public void close() { /** * 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 162b852d8d61..811f68c3c1d2 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.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 @@ public JMSWorker(CachingConnectionFactory connectionFactory, JmsTemplate jmsTemp } 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 000000000000..64084d6da38b --- /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 { + + private final IJMSConnectionFactoryProvider connectionFactoryProvider; + private final int maxPoolSize; + private final ComponentLog logger; + + private boolean cycleOpen = true; + private Generation 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 captureGeneration() { + return currentGeneration; + } + + synchronized T pollIdleWorker(final Generation generation) { + return isCurrentOpenGeneration(generation) ? generation.idleWorkers.poll() : null; + } + + synchronized boolean canCreateWorker(final Generation generation) { + return isCurrentOpenGeneration(generation); + } + + boolean registerWorker(final Generation 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 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 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 retireCurrentGeneration() { + return currentGeneration.retireAndSnapshot(); + } + + private synchronized List closeCurrentCycle() { + if (!cycleOpen) { + return List.of(); + } + + cycleOpen = false; + return currentGeneration.retireAndSnapshot(); + } + + private synchronized boolean isCurrentOpenGeneration(final Generation generation) { + return cycleOpen && currentGeneration == generation && !generation.retired; + } + + private void closeWorkers(final List 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 { + + private final BlockingQueue idleWorkers; + private final Set allWorkers = new HashSet<>(); + + private boolean retired; + + private Generation(final int maxPoolSize) { + idleWorkers = new LinkedBlockingQueue<>(maxPoolSize); + } + + private List retireAndSnapshot() { + if (retired) { + return List.of(); + } + + retired = true; + idleWorkers.clear(); + + final List 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 000000000000..d26e42a2fdf5 --- /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 lifecycle = new JmsWorkerLifecycle<>(provider, 2, logger); + final JmsWorkerLifecycle.Generation 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 lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object())); + final JmsWorkerLifecycle.Generation generation = lifecycle.captureGeneration(); + final CountDownLatch workerConstructed = new CountDownLatch(1); + final CountDownLatch retireRequested = new CountDownLatch(1); + final AtomicReference workerReference = new AtomicReference<>(); + final ExecutorService executorService = Executors.newSingleThreadExecutor(); + + try { + final Future 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 lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object())); + final JmsWorkerLifecycle.Generation 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 lifecycle = new JmsWorkerLifecycle<>(provider, 1, new MockComponentLog("processor", new Object())); + final JmsWorkerLifecycle.Generation 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 oldLifecycle = new JmsWorkerLifecycle<>(new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("old", new Object())); + final JmsWorkerLifecycle replacementLifecycle = new JmsWorkerLifecycle<>(new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("replacement", new Object())); + final JmsWorkerLifecycle.Generation oldGeneration = oldLifecycle.captureGeneration(); + final JmsWorkerLifecycle.Generation 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 lifecycle = new JmsWorkerLifecycle<>(provider, 2, new MockComponentLog("processor", new Object())); + final JmsWorkerLifecycle.Generation 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 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 lifecycle = new JmsWorkerLifecycle<>(provider, 3, logger); + final JmsWorkerLifecycle.Generation 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 lifecycle = new JmsWorkerLifecycle<>( + new RecordingConnectionFactoryProvider(), 1, new MockComponentLog("processor", new Object())); + final JmsWorkerLifecycle.Generation 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 fe31c2f36d5d..3d6a79d9b630 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.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.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.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.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 void validateConnectionClosedOnPrimaryOnlyNodeChange() throws Exception { } } + @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 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 @@ private static TestRunner createNonSharedDurableConsumer(ActiveMQConnectionFacto 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 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 000000000000..cbb6c9197965 --- /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> 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 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 firstFuture = executorService.submit(() -> { + ready.countDown(); + assertTrue(start.await(5, TimeUnit.SECONDS)); + worker.shutdown(); + return null; + }); + final Future 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 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 525c99dd83b8..40cc78551605 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 @@ 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.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.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 void validateNIFI7563UsingMultipleThreads() throws Exception { } } + @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 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 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 @@ private void assertProvenanceEvent(String expectedDetails) { 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();