fix: Correct ActiveMQ client-side virtual thread support

activemq-client/src/main/java/org/apache/activemq/ActiveMQConnectionFactory.java

@@ -43,6 +43,7 @@
 import org.apache.activemq.transport.Transport;
 import org.apache.activemq.transport.TransportFactory;
 import org.apache.activemq.transport.TransportListener;
+import org.apache.activemq.transport.failover.FailoverTransport;
 import org.apache.activemq.util.*;
 import org.apache.activemq.util.URISupport.CompositeData;
 import org.slf4j.Logger;
@@ -392,6 +393,17 @@ protected ActiveMQConnection createActiveMQConnection(String userName, String pa
 
             configureConnection(connection);
 
+            // The failover reconnect task is created on first use, so the connection level
+            // virtual thread preference has to reach the transport before it starts. A value
+            // set on the failover URI itself takes precedence and is left alone, and a factory
+            // left at false says nothing, so the system property default still applies.
+            if (isUseVirtualThreadTaskRunner()) {
+                FailoverTransport failover = transport.narrow(FailoverTransport.class);
+                if (failover != null && failover.getUseVirtualThreadTaskRunner() == null) {
+                    failover.setUseVirtualThreadTaskRunner(true);
+                }
+            }
+
             transport.start();
 
             if (clientID != null) {
@@ -906,6 +918,8 @@ public void populateProperties(Properties props) {
         props.setProperty("transactedIndividualAck", Boolean.toString(isTransactedIndividualAck()));
         props.setProperty("nonBlockingRedelivery", Boolean.toString(isNonBlockingRedelivery()));
         props.setProperty("maxThreadPoolSize", Integer.toString(getMaxThreadPoolSize()));
+        props.setProperty("useDedicatedTaskRunner", Boolean.toString(isUseDedicatedTaskRunner()));
+        props.setProperty("useVirtualThreadTaskRunner", Boolean.toString(isUseVirtualThreadTaskRunner()));
         props.setProperty("nestedMapAndListEnabled", Boolean.toString(isNestedMapAndListEnabled()));
         props.setProperty("consumerFailoverRedeliveryWaitPeriod", Long.toString(getConsumerFailoverRedeliveryWaitPeriod()));
         props.setProperty("rmIdFromConnectionId", Boolean.toString(isRmIdFromConnectionId()));

activemq-client/src/main/java/org/apache/activemq/thread/TaskRunnerFactory.java

@@ -58,6 +58,9 @@ public class TaskRunnerFactory implements Executor {
     private final AtomicBoolean initDone = new AtomicBoolean(false);
     private int maxThreadPoolSize = getDefaultMaximumPoolSize();
     private RejectedExecutionHandler rejectedTaskHandler = null;
+    /** system property that turns on virtual thread task runners for every factory that is not told otherwise */
+    public static final String USE_VIRTUAL_THREAD_TASK_RUNNER_PROPERTY = "org.apache.activemq.UseVirtualThreadTaskRunner";
+
     private ClassLoader threadClassLoader;
 
     // [AMQ-9394] Virtual Thread support
@@ -106,7 +109,7 @@ public void init() {
             synchronized(this) {
                 //need to recheck if initDone is true under the lock
                 if (!initDone.get()) {
-                    if (virtualThreadTaskRunner || "true".equalsIgnoreCase(System.getProperty("org.apache.activemq.UseVirtualThreadTaskRunner")) ) {
+                    if (virtualThreadTaskRunner || isVirtualThreadTaskRunnerDefault()) {
                         executorRef.compareAndSet(null, createVirtualThreadExecutor());
                     } else if (dedicatedTaskRunner || "true".equalsIgnoreCase(System.getProperty("org.apache.activemq.UseDedicatedTaskRunner"))) {
                         executorRef.set(null);
@@ -239,11 +242,14 @@ public void uncaughtException(final Thread t, final Throwable e) {
         return rc;
     }
 
+    // The JDK 21 classes ship in this artifact under META-INF/versions/21, so they are
+    // loaded through this class's own loader. threadClassLoader is the context loader
+    // for spawned threads and is unset on the client, where it would mean the bootstrap loader.
     protected ExecutorService createVirtualThreadExecutor() {
         assertJDK21VirtualThreadSupport();
 
         try {
-            Class<?> virtualThreadExecutorClass = Class.forName("org.apache.activemq.thread.VirtualThreadExecutor", false, threadClassLoader);
+            Class<?> virtualThreadExecutorClass = Class.forName("org.apache.activemq.thread.VirtualThreadExecutor", false, TaskRunnerFactory.class.getClassLoader());
             Method method = virtualThreadExecutorClass.getMethod("createVirtualThreadExecutorService", String.class, AtomicLong.class, Logger.class);
             Object result = method.invoke(null, name, id, LOG);
             if(!ExecutorService.class.isAssignableFrom(result.getClass())) {
@@ -261,7 +267,7 @@ protected TaskRunner createVirtualThreadTaskRunner(Executor executor, Task task,
         assertJDK21VirtualThreadSupport();
 
         try {
-            Class<?> virtualThreadTaskRunnerClass = Class.forName("org.apache.activemq.thread.VirtualThreadTaskRunner", false, threadClassLoader);
+            Class<?> virtualThreadTaskRunnerClass = Class.forName("org.apache.activemq.thread.VirtualThreadTaskRunner", false, TaskRunnerFactory.class.getClassLoader());
             Constructor<?> constructor = virtualThreadTaskRunnerClass.getConstructor(Executor.class, Task.class, Integer.TYPE);
             Object result = constructor.newInstance(executor, task, maxIterations);
             if(!TaskRunner.class.isAssignableFrom(result.getClass())) {
@@ -274,6 +280,11 @@ protected TaskRunner createVirtualThreadTaskRunner(Executor executor, Task task,
         }
     }
 
+    /** the process wide default from {@link #USE_VIRTUAL_THREAD_TASK_RUNNER_PROPERTY} */
+    public static boolean isVirtualThreadTaskRunnerDefault() {
+        return "true".equalsIgnoreCase(System.getProperty(USE_VIRTUAL_THREAD_TASK_RUNNER_PROPERTY));
+    }
+
     private void assertJDK21VirtualThreadSupport() {
         if(!(Runtime.version().feature() >= 21)) {
             LOG.error("Virtual Thread support requires JDK 21 or higher");

activemq-client/src/main/java/org/apache/activemq/transport/failover/FailoverTransport.java

@@ -92,8 +92,11 @@ public class FailoverTransport implements CompositeTransport {
     private URI connectedTransportURI;
     private URI failedConnectTransportURI;
     private final AtomicReference<Transport> connectedTransport = new AtomicReference<Transport>();
-    private final TaskRunnerFactory reconnectTaskFactory;
-    private final TaskRunner reconnectTask;
+    private final Object reconnectTaskMutex = new Object();
+    private TaskRunnerFactory reconnectTaskFactory;
+    private volatile TaskRunner reconnectTask;
+    // null until set by a failover URI option; the connection factory fills it in only then
+    private Boolean useVirtualThreadTaskRunner;
     private volatile boolean started;
     private long initialReconnectDelay = DEFAULT_INITIAL_RECONNECT_DELAY;
     private long maxReconnectDelay = 1000 * 30;
@@ -140,10 +143,23 @@ public FailoverTransport() {
     public FailoverTransport(SslContext sslContext) {
         this.sslContext = sslContext;
         stateTracker.setTrackTransactions(true);
-        // Setup a task that is used to reconnect the a connection async.
-        reconnectTaskFactory = new TaskRunnerFactory();
-        reconnectTaskFactory.init();
-        reconnectTask = reconnectTaskFactory.createTaskRunner(new Task() {
+    }
+
+    /**
+     * The reconnect task is built on first use rather than in the constructor so
+     * that URI options and the owning connection factory can decide whether it
+     * runs on a virtual thread before any thread is created.
+     */
+    private TaskRunner getReconnectTask() {
+        TaskRunner task = reconnectTask;
+        if (task == null) {
+            synchronized (reconnectTaskMutex) {
+                task = reconnectTask;
+                if (task == null) {
+                    reconnectTaskFactory = new TaskRunnerFactory("ActiveMQ Task", Thread.NORM_PRIORITY, true, 1000,
+                            false, isUseVirtualThreadTaskRunner());
+                    reconnectTaskFactory.init();
+                    task = reconnectTaskFactory.createTaskRunner(new Task() {
             @Override
             public boolean iterate() {
                 boolean result = false;
@@ -162,10 +178,7 @@ public boolean iterate() {
                     if (priorityBackup && !connectedToPriority) {
                         try {
                             doDelay();
-                            if (reconnectTask == null) {
-                                return true;
-                            }
-                            reconnectTask.wakeup();
+                            getReconnectTask().wakeup();
                         } catch (InterruptedException e) {
                             LOG.debug("Reconnect task has been interrupted.", e);
                         }
@@ -174,10 +187,7 @@ public boolean iterate() {
                     // build backups on the next iteration
                     buildBackup = true;
                     try {
-                        if (reconnectTask == null) {
-                            return true;
-                        }
-                        reconnectTask.wakeup();
+                        getReconnectTask().wakeup();
                     } catch (InterruptedException e) {
                         LOG.debug("Reconnect task has been interrupted.", e);
                     }
@@ -186,6 +196,11 @@ public boolean iterate() {
             }
 
         }, "ActiveMQ Failover Worker: " + System.identityHashCode(this));
+                    reconnectTask = task;
+                }
+            }
+        }
+        return task;
     }
 
     private void processCommand(Object incoming) {
@@ -298,7 +313,7 @@ public final void handleTransportFailure(Transport failed, IOException e) throws
                         transportListener.transportInterupted();
                     }
 
-                    reconnectTask.wakeup();
+                    getReconnectTask().wakeup();
                 } else if (!isDisposed()) {
                     propagateFailureToExceptionListener(e);
                 }
@@ -408,8 +423,20 @@ public void stop() throws Exception {
                 sleepMutex.notifyAll();
             }
         } finally {
-            reconnectTask.shutdown();
-            reconnectTaskFactory.shutdownNow();
+            // Take the references under the mutex but shut down outside it: shutdown()
+            // waits for a running iteration, and that iteration may call getReconnectTask()
+            TaskRunner taskToStop;
+            TaskRunnerFactory factoryToStop;
+            synchronized (reconnectTaskMutex) {
+                taskToStop = reconnectTask;
+                factoryToStop = reconnectTaskFactory;
+            }
+            if (taskToStop != null) {
+                taskToStop.shutdown();
+            }
+            if (factoryToStop != null) {
+                factoryToStop.shutdownNow();
+            }
         }
 
         synchronized(backupMutex) {
@@ -796,7 +823,7 @@ public void reconnect(boolean rebalance) {
                 }
                 LOG.debug("Waking up reconnect task");
                 try {
-                    reconnectTask.wakeup();
+                    getReconnectTask().wakeup();
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
                 }
@@ -1449,6 +1476,31 @@ private URI addExtraQueryOptions(URI uri) {
         return uri;
     }
 
+    /**
+     * The effective setting: an explicit value from the URI or the connection
+     * factory, otherwise the process wide default from the
+     * {@code org.apache.activemq.UseVirtualThreadTaskRunner} system property.
+     */
+    public boolean isUseVirtualThreadTaskRunner() {
+        return useVirtualThreadTaskRunner != null ? useVirtualThreadTaskRunner : TaskRunnerFactory.isVirtualThreadTaskRunnerDefault();
+    }
+
+    /** the raw setting, null when neither a URI option nor the connection factory has set it */
+    public Boolean getUseVirtualThreadTaskRunner() {
+        return useVirtualThreadTaskRunner;
+    }
+
+    /**
+     * Run the reconnect task on a virtual thread (JDK 21+). A value given as a
+     * failover URI option is final; the owning connection factory turns this on
+     * when its own useVirtualThreadTaskRunner is set and the URI said nothing;
+     * left unset, the system property default applies. Takes effect only before
+     * the transport starts.
+     */
+    public void setUseVirtualThreadTaskRunner(Boolean useVirtualThreadTaskRunner) {
+        this.useVirtualThreadTaskRunner = useVirtualThreadTaskRunner;
+    }
+
     public void setNestedExtraQueryOptions(String nestedExtraQueryOptions) {
         this.nestedExtraQueryOptions = nestedExtraQueryOptions;
     }

activemq-client/src/main/java21/org/apache/activemq/thread/VirtualThreadTaskRunner.java

@@ -109,7 +109,13 @@ public void shutdown(long timeout) throws InterruptedException {
             // waiting for iterating to finish
             if (runningThread != Thread.currentThread()) {
                 if (iterating) {
-                    runnableCondition.await(timeout, TimeUnit.MILLISECONDS);
+                    // Same contract as Object.wait(timeout) in the other runners:
+                    // zero means wait until the running iteration signals, not return at once
+                    if (timeout > 0) {
+                        runnableCondition.await(timeout, TimeUnit.MILLISECONDS);
+                    } else {
+                        runnableCondition.await();
+                    }
                 }
             }
         } finally {

activemq-unit-tests/src/test/java/org/apache/activemq/ActiveMQConnectionFactoryTest.java

@@ -215,6 +215,27 @@ public void testFactorySerializable() throws Exception {
         assertEquals(cf.getClientID(), clientID);
     }
 
+    public void testTaskRunnerSettingsSurviveSerialization() throws Exception {
+        // Serialization and JNDI both go through populateProperties/buildFromProperties,
+        // so a setting missing from populateProperties is silently reset on the way back
+        var cf = new ActiveMQConnectionFactory();
+        cf.setUseDedicatedTaskRunner(true);
+        cf.setUseVirtualThreadTaskRunner(true);
+        cf.setMaxThreadPoolSize(7);
+
+        var bytesOut = new ByteArrayOutputStream();
+        try (var objectsOut = new ObjectOutputStream(bytesOut)) {
+            objectsOut.writeObject(cf);
+        }
+        try (var objectsIn = new ObjectInputStream(new ByteArrayInputStream(bytesOut.toByteArray()))) {
+            cf = (ActiveMQConnectionFactory) objectsIn.readObject();
+        }
+
+        assertTrue("useDedicatedTaskRunner lost", cf.isUseDedicatedTaskRunner());
+        assertTrue("useVirtualThreadTaskRunner lost", cf.isUseVirtualThreadTaskRunner());
+        assertEquals(7, cf.getMaxThreadPoolSize());
+    }
+
     public void testSetExceptionListener() throws Exception {
         ActiveMQConnectionFactory cf = new ActiveMQConnectionFactory("vm://localhost?broker.persistent=false");
         connection = (ActiveMQConnection)cf.createConnection();