Repository: neo4j-contrib/neo4j-jdbc
Pull request #261 (2021-03-25 08:26:00)

Make the BoltDriverCache thread safe.
Added a test for the thread safety of the BoltDriverCache and then fixed the failing test.

michael-simons: Re assertions: Without overlaps, 0 collisions on the next assert doesn't mean anything.
conker84: @michael-simons thank you so much! @fbiville I'm merging this as it's green

Patch:
diff --git a/neo4j-jdbc-bolt/src/main/java/org/neo4j/jdbc/bolt/cache/BoltDriverCache.java b/neo4j-jdbc-bolt/src/main/java/org/neo4j/jdbc/bolt/cache/BoltDriverCache.java
index 8bac5be0..1558b467 100644
--- a/neo4j-jdbc-bolt/src/main/java/org/neo4j/jdbc/bolt/cache/BoltDriverCache.java
+++ b/neo4j-jdbc-bolt/src/main/java/org/neo4j/jdbc/bolt/cache/BoltDriverCache.java
@@ -25,6 +25,7 @@ import org.neo4j.driver.Driver;
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
@@ -37,8 +38,8 @@ import java.util.function.Function;
  */
 public class BoltDriverCache {
 
-    private Map<BoltDriverCacheKey, Driver> cache;
-    private Function<BoltDriverCacheKey, Driver> builder;
+    private final Map<BoltDriverCacheKey, Driver> cache;
+    private final Function<BoltDriverCacheKey, Driver> builder;
 
     /**
      * Setup the cache for the specific building function
@@ -56,16 +57,10 @@ public class BoltDriverCache {
      * @param authToken
      * @param info
      * @return
-     * @throws URISyntaxException
      */
-    public Driver getDriver(List<URI> routingUris, Config config, AuthToken authToken, Properties info) throws URISyntaxException {
-        BoltDriverCacheKey key = new BoltDriverCacheKey(routingUris, config, authToken, info);
-        Driver driver = cache.get(key);
-        if(null == driver){
-            driver = new BoltDriverCached(builder.apply(key), this, key);
-            cache.put(key, driver);
-        }
-        return driver;
+    public Driver getDriver(List<URI> routingUris, Config config, AuthToken authToken, Properties info) {
+        return cache.computeIfAbsent(new BoltDriverCacheKey(routingUris, config, authToken, info),
+            key -> new BoltDriverCached(builder.apply(key), this, key));
     }
 
     public Driver removeFromCache(BoltDriverCacheKey key){
@@ -77,7 +72,7 @@ public class BoltDriverCache {
      * @return
      */
     public Map<BoltDriverCacheKey, Driver> getCache() {
-        return cache;
+        return Collections.unmodifiableMap(cache);
     }
 
     // visible for testing


Test patch:
diff --git a/neo4j-jdbc-bolt/src/test/java/org/neo4j/jdbc/bolt/cache/BoltDriverCacheTest.java b/neo4j-jdbc-bolt/src/test/java/org/neo4j/jdbc/bolt/cache/BoltDriverCacheTest.java
index 222f6d04..f8891b50 100644
--- a/neo4j-jdbc-bolt/src/test/java/org/neo4j/jdbc/bolt/cache/BoltDriverCacheTest.java
+++ b/neo4j-jdbc-bolt/src/test/java/org/neo4j/jdbc/bolt/cache/BoltDriverCacheTest.java
@@ -7,10 +7,23 @@ import org.neo4j.driver.internal.InternalBookmark;
 
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.Callable;
+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.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 
 import static org.junit.Assert.assertEquals;
@@ -20,6 +33,43 @@ public class BoltDriverCacheTest {
 
     Function<BoltDriverCacheKey, Driver> builder = (params -> Mockito.mock(Driver.class));
 
+    @Test
+    public void cacheShouldBeThreadSafe() throws ExecutionException, InterruptedException {
+        List<URI> url = Arrays.asList(URI.create("bolt://localhost"));
+        Config.ConfigBuilder configBuilder = Config.builder();
+        AuthToken authToken = AuthTokens.basic("neo4j", "password");
+
+        BoltDriverCache cache = new BoltDriverCache(builder);
+        assertEquals(0, cache.getCache().size());
+
+        int numThreads = Runtime.getRuntime().availableProcessors();
+        ExecutorService executor = Executors.newWorkStealingPool();
+
+        AtomicBoolean running = new AtomicBoolean();
+        AtomicInteger overlaps = new AtomicInteger();
+
+        Collection<Callable<Driver>> getDriver = new ArrayList<>();
+        for (int t = 0; t < numThreads; ++t) {
+            getDriver.add(() -> {
+                if (!running.compareAndSet(false, true)) {
+                    overlaps.incrementAndGet();
+                }
+                Driver d = cache.getDriver(url, configBuilder.build(), authToken, new Properties());
+                running.compareAndSet(true, false);
+                return d;
+            });
+        }
+
+        Map<Driver, Integer> drivers = new IdentityHashMap<>();
+        for (Future<Driver> getDriverFuture : executor.invokeAll(getDriver)) {
+            drivers.put(getDriverFuture.get(), 1);
+        }
+
+        assertTrue(overlaps.get() > 0); // This might get flaky...
+        assertEquals(1, drivers.size());
+        assertEquals(1, cache.getCache().size());
+    }
+
     @Test
     public void shouldBeSameInstance() throws URISyntaxException {
         List<URI> url = Arrays.asList(URI.create("bolt://localhost"));