From d5a7d23c2078bd26717673cdeb6e0eba1c75bc65 Mon Sep 17 00:00:00 2001
From: Emmanuel Hugonnet <ehugonne@redhat.com>
Date: Mon, 12 Aug 2024 15:48:11 +0200
Subject: [PATCH] [WFLY-19567]: Removing ConcurrentHashMap.computeIfAbsent
 usage

* The blocking code is out of the compute part.
* Using a ReentrantLock
* Still need to work on Artemis pinned thread code.

Signed-off-by: Emmanuel Hugonnet <ehugonne@redhat.com>
---
 .../injection/AbstractJMSContext.java         | 24 ++++++++++++++-----
 .../injection/AbstractJMSContextTestCase.java |  4 +---
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/messaging-activemq/injection/src/main/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContext.java b/messaging-activemq/injection/src/main/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContext.java
index 60dd9e1f6657..7e2ad1d48d0f 100644
--- a/messaging-activemq/injection/src/main/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContext.java
+++ b/messaging-activemq/injection/src/main/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContext.java
@@ -2,17 +2,17 @@
  * Copyright The WildFly Authors
  * SPDX-License-Identifier: Apache-2.0
  */
-
 package org.wildfly.extension.messaging.activemq.deployment.injection;
 
 import static org.wildfly.extension.messaging.activemq._private.MessagingLogger.ROOT_LOGGER;
 
 import java.io.Serializable;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 import jakarta.jms.ConnectionFactory;
 import jakarta.jms.JMSContext;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Abstract class for managing JMS Contexts.
@@ -23,12 +23,24 @@
  */
 public abstract class AbstractJMSContext implements Serializable {
 
-    private final Map<String, JMSContext> contexts = new ConcurrentHashMap<>();
+    private final transient ReentrantLock lock = new ReentrantLock();
+    private final transient Map<String, JMSContext> contexts = new ConcurrentHashMap<>();
 
     JMSContext getContext(String injectionPointId, JMSInfo info, ConnectionFactory connectionFactory) {
-        return contexts.computeIfAbsent(injectionPointId, key -> {
-            return createContext(info, connectionFactory);
-        });
+        JMSContext context = contexts.get(injectionPointId);
+        if (context == null) {
+            lock.lock();
+            try {
+                context = contexts.get(injectionPointId);
+                if (context == null) {
+                    context = createContext(info, connectionFactory);
+                    contexts.put(injectionPointId, context);
+                }
+            } finally {
+                lock.unlock();
+            }
+        }
+        return context;
     }
 
     private JMSContext createContext(JMSInfo info, ConnectionFactory connectionFactory) {
diff --git a/messaging-activemq/injection/src/test/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContextTestCase.java b/messaging-activemq/injection/src/test/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContextTestCase.java
index 41d80548128c..a94922295427 100644
--- a/messaging-activemq/injection/src/test/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContextTestCase.java
+++ b/messaging-activemq/injection/src/test/java/org/wildfly/extension/messaging/activemq/deployment/injection/AbstractJMSContextTestCase.java
@@ -72,9 +72,7 @@ public JMSContext call() throws Exception {
         }
 
         // Execute all tasks
-        for (Callable<JMSContext> task : tasks) {
-            executor.submit(task);
-        }
+        executor.invokeAll(tasks);
 
         // Shutdown the executor
         executor.shutdown();