ARTEMIS-6142 FollowUp to fix tests

artemis-core-client/src/main/java/org/apache/activemq/artemis/core/client/impl/ClientSessionFactoryImpl.java

@@ -1201,6 +1201,8 @@ protected Connection openTransportConnection(final Connector connector) {
             connector.close();
          } catch (Throwable t) {
          }
+      } else if (serverLocator.isConnected()) {
+         transportConnection.setConnected();
       }
 
       return transportConnection;

artemis-core-client/src/main/java/org/apache/activemq/artemis/core/client/impl/ServerLocatorImpl.java

@@ -131,6 +131,9 @@ private enum STATE {
     */
    private volatile boolean disableDiscoveryRetries = false;
 
+   // set when connect() is called, meaning this locator is used for clustering or topology discovery
+   private volatile boolean connected = false;
+
    // if the system should shutdown the pool when shutting down
    private transient boolean shutdownPool;
 
@@ -552,6 +555,7 @@ private ClientSessionFactoryInternal connect(final boolean skipWarnings) throws
       // if we used connect, we should control UDP reconnections at a different path.
       // and this belongs to a cluster connection, not client
       disableDiscoveryRetries = true;
+      connected = true;
       ClientSessionFactoryInternal returnFactory = null;
 
       synchronized (this) {
@@ -1382,6 +1386,11 @@ public boolean isClusterConnection() {
       return clusterConnection;
    }
 
+   @Override
+   public boolean isConnected() {
+      return connected;
+   }
+
    @Override
    public TransportConfiguration getClusterTransportConfiguration() {
       return clusterTransportConfiguration;

artemis-core-client/src/main/java/org/apache/activemq/artemis/core/client/impl/ServerLocatorInternal.java

@@ -99,4 +99,6 @@ default void notifyNodeDown(long uniqueEventID, String nodeID) {
    Pair<TransportConfiguration, TransportConfiguration> selectNextConnectorPair();
 
    long getNextRetryInterval(long retryInterval, double retryIntervalMultiplier, long maxRetryInterval);
+
+   boolean isConnected();
 }

artemis-core-client/src/main/java/org/apache/activemq/artemis/spi/core/remoting/Connection.java

@@ -134,6 +134,14 @@ default void disconnect() {
       close();
    }
 
+   // Marks this connection as being used for clustering or topology discovery
+   default void setConnected() {
+   }
+
+   default boolean isConnected() {
+      return false;
+   }
+
    /**
     * {@return the string representation of the remote address this connection is connected to}
     */

artemis-server/src/main/java/org/apache/activemq/artemis/core/remoting/impl/invm/InVMAcceptor.java

@@ -246,7 +246,7 @@ public void disconnect(final String connectionID, final boolean failed) {
       Connection conn = connections.get(connectionID);
 
       if (conn != null) {
-         if (failed) {
+         if (failed || conn.isConnected()) {
             conn.disconnect();
          } else {
             conn.close();

artemis-server/src/main/java/org/apache/activemq/artemis/core/remoting/impl/invm/InVMConnection.java

@@ -66,6 +66,9 @@ public class InVMConnection implements Connection {
 
    private RemotingConnection protocolConnection;
 
+   // set when this connection is used for clustering or topology discovery
+   private volatile boolean connected;
+
    private boolean bufferPoolingEnabled = TransportConstants.DEFAULT_BUFFER_POOLING;
 
    private boolean directDeliver = TransportConstants.DEFAULT_DIRECT_DELIVER;
@@ -138,6 +141,16 @@ public void setProtocolConnection(RemotingConnection connection) {
       this.protocolConnection = connection;
    }
 
+   @Override
+   public void setConnected() {
+      this.connected = true;
+   }
+
+   @Override
+   public boolean isConnected() {
+      return connected;
+   }
+
    @Override
    public void close() {
       internalClose(false);

artemis-server/src/main/java/org/apache/activemq/artemis/core/remoting/impl/invm/InVMConnector.java

@@ -222,7 +222,7 @@ public void disconnect(final String connectionID, final boolean failed) {
       Connection conn = connections.get(connectionID);
 
       if (conn != null) {
-         if (failed) {
+         if (failed || conn.isConnected()) {
             conn.disconnect();
          } else {
             conn.close();

tests/integration-tests/src/test/java/org/apache/activemq/artemis/tests/integration/jms/connection/InVMConnectionLeakStressTest.java

@@ -19,13 +19,16 @@
 import javax.jms.Connection;
 import javax.jms.Session;
 
+import java.lang.invoke.MethodHandles;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Callable;
+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 org.apache.activemq.artemis.api.core.TransportConfiguration;
 import org.apache.activemq.artemis.api.jms.ActiveMQJMSClient;
@@ -35,7 +38,10 @@
 import org.apache.activemq.artemis.tests.util.Wait;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
@@ -59,6 +65,8 @@
  */
 public class InVMConnectionLeakStressTest extends JMSTestBase {
 
+   private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
+
    private ActiveMQConnectionFactory floodCf;
 
    @Override
@@ -76,29 +84,31 @@ public void setUp() throws Exception {
    @Test
    public void testConcurrentGracefulCloseRemovesAllConnections() throws Exception {
       final int numConnections = 20_000;
-      final int threads = 100;
 
-      List<Callable<Void>> tasks = new ArrayList<>(numConnections);
+      AtomicInteger error = new AtomicInteger(0);
+      CountDownLatch latch = new CountDownLatch(numConnections);
+      ExecutorService executor = Executors.newFixedThreadPool(100);
+      runAfter(executor::shutdownNow);
+
       for (int i = 0; i < numConnections; i++) {
-         tasks.add(() -> {
-            Connection connection = floodCf.createConnection();
-            Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
-            session.createProducer(ActiveMQJMSClient.createQueue("stress-queue"));
-            // Graceful close
-            connection.close();
-            return null;
+         executor.execute(() -> {
+            try {
+               Connection connection = floodCf.createConnection();
+               Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
+               session.createProducer(ActiveMQJMSClient.createQueue("stress-queue"));
+               // Graceful close
+               connection.close();
+            } catch (Exception e) {
+               logger.warn(e.getMessage(), e);
+               error.incrementAndGet();
+            } finally {
+               latch.countDown();
+            }
          });
       }
 
-      ExecutorService executor = Executors.newFixedThreadPool(threads);
-      try {
-         for (Future<Void> future : executor.invokeAll(tasks)) {
-            future.get();
-         }
-      } finally {
-         executor.shutdown();
-         assertTrue(executor.awaitTermination(2, TimeUnit.MINUTES));
-      }
+      assertTrue(latch.await(10, TimeUnit.SECONDS));
+      assertEquals(0, error.get());
 
       // Every gracefully-closed connection must be removed from the server. InVM connection-ttl is -1 so the
       // failure-check reaper never removes them; if this never reaches 0 the connections have leaked.

tests/unit-tests/src/test/java/org/apache/activemq/artemis/tests/unit/core/remoting/impl/invm/InVMConnectionTest.java

@@ -29,14 +29,19 @@
 import org.apache.activemq.artemis.spi.core.protocol.ProtocolManager;
 import org.apache.activemq.artemis.spi.core.remoting.BaseConnectionLifeCycleListener;
 import org.apache.activemq.artemis.spi.core.remoting.Connection;
+import org.apache.activemq.artemis.tests.util.ActiveMQTestBase;
 import org.junit.jupiter.api.Test;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.CyclicBarrier;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
-public class InVMConnectionTest {
+public class InVMConnectionTest extends ActiveMQTestBase {
 
    @Test
    public void testIsTargetNode() throws Exception {
@@ -67,42 +72,47 @@ public void testIsTargetNode() throws Exception {
    @Test
    public void testConcurrentCloseFiresConnectionDestroyedExactlyOnce() throws Exception {
       final int threads = 16;
+
+      ExecutorService service = Executors.newFixedThreadPool(threads);
+      runAfter(service::shutdownNow);
+
+
       // Repeat several rounds to widen the window for catching the race.
       for (int round = 0; round < 50; round++) {
+         CountDownLatch done = new CountDownLatch(threads);
          final CountingLifeCycleListener listener = new CountingLifeCycleListener();
          final InVMConnection conn = new InVMConnection(0, null, listener, null);
 
          final CyclicBarrier barrier = new CyclicBarrier(threads);
-         final Thread[] workers = new Thread[threads];
          final AtomicInteger prematureReturns = new AtomicInteger();
 
          for (int i = 0; i < threads; i++) {
             final boolean disconnect = (i % 2 == 0);
-            workers[i] = new Thread(() -> {
+            service.execute(() -> {
                try {
                   // Line up all threads so they hit close()/disconnect() together.
                   barrier.await();
                } catch (Exception e) {
                   throw new RuntimeException(e);
                }
-               if (disconnect) {
-                  conn.disconnect();
-               } else {
-                  conn.close();
-               }
-               // by the time any close()/disconnect() call returns, connectionDestroyed must already have fired.
-               if (!listener.destroyFired) {
-                  prematureReturns.incrementAndGet();
+               try {
+                  if (disconnect) {
+                     conn.disconnect();
+                  } else {
+                     conn.close();
+                  }
+                  // by the time any close()/disconnect() call returns, connectionDestroyed must already have fired.
+                  if (!listener.destroyFired) {
+                     prematureReturns.incrementAndGet();
+                  }
+               } finally {
+                  done.countDown();
                }
             });
          }
 
-         for (Thread worker : workers) {
-            worker.start();
-         }
-         for (Thread worker : workers) {
-            worker.join();
-         }
+         assertTrue(done.await(10, TimeUnit.SECONDS));
+
 
          assertEquals(1, listener.destroyedCount.get(),
                       "connectionDestroyed must be fired exactly once per connection (round " + round + ")");