Merge pull request #204 from emaccaull/fix/pool-lock-contention

Speed up concurrent pool creation

Author: rusher

src/main/java/org/mariadb/jdbc/pool/Pool.java

@@ -534,8 +534,10 @@ private void registerJmx() throws Exception {
     String jmxName = poolTag.replace(":", "_");
     ObjectName name = new ObjectName("org.mariadb.jdbc.pool:type=" + jmxName);
 
-    if (!mbs.isRegistered(name)) {
-      mbs.registerMBean(this, name);
+    synchronized (mbs) {
+      if (!mbs.isRegistered(name)) {
+        mbs.registerMBean(this, name);
+      }
     }
   }
 
@@ -544,8 +546,10 @@ private void unRegisterJmx() throws Exception {
     String jmxName = poolTag.replace(":", "_");
     ObjectName name = new ObjectName("org.mariadb.jdbc.pool:type=" + jmxName);
 
-    if (mbs.isRegistered(name)) {
-      mbs.unregisterMBean(name);
+    synchronized (mbs) {
+      if (mbs.isRegistered(name)) {
+        mbs.unregisterMBean(name);
+      }
     }
   }
 

src/main/java/org/mariadb/jdbc/pool/Pools.java

@@ -14,31 +14,53 @@
 public final class Pools {
 
   private static final AtomicInteger poolIndex = new AtomicInteger();
-  private static final Map<Configuration, Pool> poolMap = new ConcurrentHashMap<>();
+  private static final Map<Configuration, PoolHolder> poolMap = new ConcurrentHashMap<>();
   private static ScheduledThreadPoolExecutor poolExecutor = null;
 
+  static class PoolHolder {
+    private final Configuration conf;
+    private final int poolIndex;
+    private final ScheduledThreadPoolExecutor executor;
+    private Pool pool;
+
+    PoolHolder(Configuration conf, int poolIndex, ScheduledThreadPoolExecutor executor) {
+      this.conf = conf;
+      this.poolIndex = poolIndex;
+      this.executor = executor;
+    }
+
+    synchronized Pool getPool() {
+      if (pool == null) {
+        pool = new Pool(conf, poolIndex, executor);
+      }
+      return pool;
+    }
+  }
+
   /**
    * Get existing pool for a configuration. Create it if it doesn't exist.
    *
    * @param conf configuration parser
    * @return pool
    */
   public static Pool retrievePool(Configuration conf) {
-    if (!poolMap.containsKey(conf)) {
+    PoolHolder holder = poolMap.get(conf);
+    if (holder == null) {
       synchronized (poolMap) {
-        if (!poolMap.containsKey(conf)) {
+        holder = poolMap.get(conf);
+        if (holder == null) {
           if (poolExecutor == null) {
             poolExecutor =
                 new ScheduledThreadPoolExecutor(
                     1, new PoolThreadFactory("MariaDbPool-maxTimeoutIdle-checker"));
           }
-          Pool pool = new Pool(conf, poolIndex.incrementAndGet(), poolExecutor);
-          poolMap.put(conf, pool);
-          return pool;
+          holder = new PoolHolder(conf, poolIndex.incrementAndGet(), poolExecutor);
+          poolMap.put(conf, holder);
         }
       }
     }
-    return poolMap.get(conf);
+    // Don't initialize a pool while holding a lock on `poolMap`.
+    return holder.getPool();
   }
 
   /**
@@ -49,12 +71,9 @@ public static Pool retrievePool(Configuration conf) {
   public static void remove(Pool pool) {
     if (poolMap.containsKey(pool.getConf())) {
       synchronized (poolMap) {
-        if (poolMap.containsKey(pool.getConf())) {
-          poolMap.remove(pool.getConf());
-
-          if (poolMap.isEmpty()) {
-            shutdownExecutor();
-          }
+        PoolHolder previous = poolMap.remove(pool.getConf());
+        if (previous != null && poolMap.isEmpty()) {
+          shutdownExecutor();
         }
       }
     }
@@ -63,9 +82,9 @@ public static void remove(Pool pool) {
   /** Close all pools. */
   public static void close() {
     synchronized (poolMap) {
-      for (Pool pool : poolMap.values()) {
+      for (PoolHolder holder : poolMap.values()) {
         try {
-          pool.close();
+          holder.getPool().close();
         } catch (Exception exception) {
           // eat
         }
@@ -85,10 +104,12 @@ public static void close(String poolName) {
       return;
     }
     synchronized (poolMap) {
-      for (Pool pool : poolMap.values()) {
-        if (poolName.equals(pool.getConf().poolName())) {
+      for (PoolHolder holder : poolMap.values()) {
+        if (poolName.equals(holder.conf.poolName())) {
           try {
-            pool.close(); // Pool.close() calls Pools.remove(), which does the rest of the cleanup
+            holder
+                .getPool()
+                .close(); // Pool.close() calls Pools.remove(), which does the rest of the cleanup
           } catch (Exception exception) {
             // eat
           }

src/test/java/org/mariadb/jdbc/integration/PoolDataSourceTest.java

@@ -8,11 +8,18 @@
 import java.lang.management.ManagementFactory;
 import java.sql.*;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Locale;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import javax.management.MBeanInfo;
 import javax.management.MBeanServer;
 import javax.management.ObjectName;
@@ -23,6 +30,7 @@
 import org.junit.jupiter.api.Assumptions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.mariadb.jdbc.MariaDbPoolDataSource;
 import org.mariadb.jdbc.pool.PoolThreadFactory;
 import org.mariadb.jdbc.pool.Pools;
@@ -774,4 +782,40 @@ public void ensureConnectionClose() throws Exception {
     assertFalse(xac.getConnection().isClosed());
     xac.close();
   }
+
+  @Timeout(value = 5, unit = TimeUnit.SECONDS)
+  @Test
+  public void testConcurrentCreationForDifferentHosts() throws Exception {
+    CountDownLatch ready = new CountDownLatch(5);
+    CountDownLatch start = new CountDownLatch(1);
+    ExecutorService executor = Executors.newCachedThreadPool();
+    try {
+      // When many pools are created concurrently
+      List<Future<MariaDbPoolDataSource>> futures =
+          IntStream.rangeClosed(1, 5)
+              .mapToObj(
+                  hostIndex ->
+                      executor.submit(
+                          () -> {
+                            ready.countDown();
+                            start.await();
+                            MariaDbPoolDataSource ds = new MariaDbPoolDataSource();
+                            ds.setUrl(
+                                "jdbc:mariadb://myhost" + hostIndex + ":5500/db?someOption=val");
+                            return ds;
+                          }))
+              .collect(Collectors.toList());
+
+      ready.await();
+      start.countDown();
+
+      // Then they should all be created in a timely manner
+      for (Future<MariaDbPoolDataSource> future : futures) {
+        future.get().close();
+      }
+
+    } finally {
+      executor.shutdown();
+    }
+  }
 }