GH-2375 - Make LiteralReplacement-Cache threadsafe.

Fixes #2375.

Author: michael-simons

src/main/java/org/springframework/data/neo4j/repository/query/Neo4jSpelSupport.java

@@ -17,6 +17,7 @@
 
 import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.concurrent.locks.StampedLock;
 
 import org.apiguardian.api.API;
 import org.springframework.data.domain.Pageable;
@@ -111,12 +112,39 @@ protected boolean removeEldestEntry(Map.Entry<String, LiteralReplacement> eldest
 					}
 				};
 
+		private static final StampedLock LOCK = new StampedLock();
+
 		static LiteralReplacement withTargetAndValue(LiteralReplacement.Target target, @Nullable String value) {
 
 			String valueUsed = value == null ? "" : value;
-			StringBuilder key = new StringBuilder(target.name()).append("_").append(valueUsed);
-
-			return INSTANCES.computeIfAbsent(key.toString(), k -> new StringBasedLiteralReplacement(target, valueUsed));
+			String key = new StringBuilder(target.name()).append("_").append(valueUsed).toString();
+
+			long stamp = LOCK.tryOptimisticRead();
+			if (LOCK.validate(stamp) && INSTANCES.containsKey(key)) {
+				return INSTANCES.get(key);
+			}
+			try {
+				stamp = LOCK.readLock();
+				LiteralReplacement replacement = null;
+				while (replacement == null) {
+					if (INSTANCES.containsKey(key)) {
+						replacement = INSTANCES.get(key);
+					} else {
+						long writeStamp = LOCK.tryConvertToWriteLock(stamp);
+						if (LOCK.validate(writeStamp)) {
+							replacement = new StringBasedLiteralReplacement(target, valueUsed);
+							stamp = writeStamp;
+							INSTANCES.put(key, replacement);
+						} else {
+							LOCK.unlockRead(stamp);
+							stamp = LOCK.writeLock();
+						}
+					}
+				}
+				return replacement;
+			} finally {
+				LOCK.unlock(stamp);
+			}
 		}
 
 		private final Target target;

src/test/java/org/springframework/data/neo4j/repository/query/Neo4jSpelSupportTest.java

@@ -17,11 +17,26 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
+import static org.assertj.core.api.Assumptions.assumeThat;
+
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.IdentityHashMap;
+import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.jupiter.api.Test;
 import org.springframework.data.domain.PageRequest;
 import org.springframework.data.domain.Sort;
 import org.springframework.data.neo4j.repository.query.Neo4jSpelSupport.LiteralReplacement;
+import org.springframework.util.ReflectionUtils;
 
 /**
  * @author Michael J. Simons
@@ -65,16 +80,76 @@ void orderByShouldWork() {
 				.withMessageMatching(".+is not a valid order criteria.");
 	}
 
+	private Map<?, ?> getCacheInstance() throws ClassNotFoundException, IllegalAccessException {
+		Class<?> type = Class.forName(
+				"org.springframework.data.neo4j.repository.query.Neo4jSpelSupport$StringBasedLiteralReplacement");
+		Field cacheField = ReflectionUtils.findField(type, "INSTANCES");
+		cacheField.setAccessible(true);
+		return (Map<?, ?>) cacheField.get(null);
+	}
+
+	private void flushLiteralCache() {
+		try {
+			Map<?, ?> cache = getCacheInstance();
+			cache.clear();
+		} catch (Exception e) {
+			throw new RuntimeException(e);
+		}
+	}
+
+	private int getCacheSize() {
+		try {
+			Map<?, ?> cache = getCacheInstance();
+			return cache.size();
+		} catch (Exception e) {
+			throw new RuntimeException(e);
+		}
+	}
+
 	@Test // DATAGRAPH-1454
 	void cacheShouldWork() {
 
-		// Make sure we flush this before...
-		for (int i = 0; i < 16; ++i) {
-			LiteralReplacement literalReplacement = Neo4jSpelSupport.literal("y" + i);
-		}
+		flushLiteralCache();
 
 		LiteralReplacement literalReplacement1 = Neo4jSpelSupport.literal("x");
 		LiteralReplacement literalReplacement2 = Neo4jSpelSupport.literal("x");
 		assertThat(literalReplacement1).isSameAs(literalReplacement2);
 	}
+
+	@Test // GH-2375
+	void cacheShouldBeThreadSafe() throws ExecutionException, InterruptedException {
+
+		flushLiteralCache();
+
+		int numThreads = Runtime.getRuntime().availableProcessors();
+		ExecutorService executor = Executors.newWorkStealingPool();
+
+		AtomicBoolean running = new AtomicBoolean();
+		AtomicInteger overlaps = new AtomicInteger();
+
+		Collection<Callable<LiteralReplacement>> getReplacementCalls = new ArrayList<>();
+		for (int t = 0; t < numThreads; ++t) {
+			getReplacementCalls.add(() -> {
+				if (!running.compareAndSet(false, true)) {
+					overlaps.incrementAndGet();
+				}
+				Thread.sleep(100); // Make the chances of overlapping a bit bigger
+				LiteralReplacement d = Neo4jSpelSupport.literal("x");
+				running.compareAndSet(true, false);
+				return d;
+			});
+		}
+
+		Map<LiteralReplacement, Integer> replacements = new IdentityHashMap<>();
+		for (Future<LiteralReplacement> getDriverFuture : executor.invokeAll(getReplacementCalls)) {
+			replacements.put(getDriverFuture.get(), 1);
+		}
+		executor.shutdown();
+
+		// Assume things actually had been concurrent
+		assumeThat(overlaps.get()).isGreaterThan(0);
+
+		assertThat(getCacheSize()).isEqualTo(1);
+		assertThat(replacements).hasSize(1);
+	}
 }