From 461e7dbf3e5f90bb05b55bb621aa92da8ae4c569 Mon Sep 17 00:00:00 2001
From: Chan Young <lms062111@naver.com>
Date: Sat, 23 Mar 2024 20:59:22 +0900
Subject: [PATCH] Improve atomicity in DefaultRedisCacheWriter.doLock().

Closes #1686
---
 .../redis/cache/DefaultRedisCacheWriter.java  | 10 +++--
 .../cache/DefaultRedisCacheWriterTests.java   | 41 +++++++++++++++++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/src/main/java/org/springframework/data/redis/cache/DefaultRedisCacheWriter.java b/src/main/java/org/springframework/data/redis/cache/DefaultRedisCacheWriter.java
index 488442ca45..b91733f444 100644
--- a/src/main/java/org/springframework/data/redis/cache/DefaultRedisCacheWriter.java
+++ b/src/main/java/org/springframework/data/redis/cache/DefaultRedisCacheWriter.java
@@ -57,6 +57,7 @@
  * @author Mark Paluch
  * @author André Prata
  * @author John Blum
+ * @author ChanYoung Joung
  * @since 2.0
  */
 class DefaultRedisCacheWriter implements RedisCacheWriter {
@@ -319,12 +320,15 @@ void lock(String name) {
 	}
 
 	@Nullable
-	private Boolean doLock(String name, Object contextualKey, @Nullable Object contextualValue,
-			RedisConnection connection) {
+	protected Boolean doLock(String name, Object contextualKey, @Nullable Object contextualValue,
+							 RedisConnection connection) {
 
 		Expiration expiration = Expiration.from(this.lockTtl.getTimeToLive(contextualKey, contextualValue));
 
-		return connection.stringCommands().set(createCacheLockKey(name), new byte[0], expiration, SetOption.SET_IF_ABSENT);
+		while (!ObjectUtils.nullSafeEquals(connection.stringCommands().set(createCacheLockKey(name), new byte[0], expiration, SetOption.SET_IF_ABSENT),true)) {
+			checkAndPotentiallyWaitUntilUnlocked(name, connection);
+		}
+		return true;
 	}
 
 	/**
diff --git a/src/test/java/org/springframework/data/redis/cache/DefaultRedisCacheWriterTests.java b/src/test/java/org/springframework/data/redis/cache/DefaultRedisCacheWriterTests.java
index 3159376075..c44a9e85af 100644
--- a/src/test/java/org/springframework/data/redis/cache/DefaultRedisCacheWriterTests.java
+++ b/src/test/java/org/springframework/data/redis/cache/DefaultRedisCacheWriterTests.java
@@ -38,12 +38,14 @@
 import org.springframework.data.redis.test.condition.RedisDriver;
 import org.springframework.data.redis.test.extension.parametrized.MethodSource;
 import org.springframework.data.redis.test.extension.parametrized.ParameterizedRedisTest;
+import org.springframework.lang.Nullable;
 
 /**
  * Integration tests for {@link DefaultRedisCacheWriter}.
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author ChanYoung Joung
  */
 @MethodSource("testParams")
 public class DefaultRedisCacheWriterTests {
@@ -419,6 +421,45 @@ void noOpStatisticsCollectorReturnsEmptyStatsInstance() {
 		assertThat(stats.getPuts()).isZero();
 	}
 
+	@ParameterizedRedisTest
+	void doLockShouldGetLock() throws InterruptedException {
+
+		int threadCount = 3;
+		CountDownLatch beforeWrite = new CountDownLatch(threadCount);
+		CountDownLatch afterWrite = new CountDownLatch(threadCount);
+
+		DefaultRedisCacheWriter cw = new DefaultRedisCacheWriter(connectionFactory, Duration.ofMillis(50),
+				BatchStrategies.keys()){
+			@Nullable
+			protected Boolean doLock(String name, Object contextualKey, @Nullable Object contextualValue,
+									 RedisConnection connection) {
+				Boolean doLock = super.doLock(name, contextualKey, contextualValue, connection);
+				assertThat(doLock).isTrue();
+				return doLock;
+			}
+		};
+
+		cw.lock(CACHE_NAME);
+
+		for (int i = 0; i < threadCount; i++) {
+			Thread th = new Thread(() -> {
+				beforeWrite.countDown();
+				cw.putIfAbsent(CACHE_NAME, binaryCacheKey, binaryCacheValue, Duration.ZERO);
+				afterWrite.countDown();
+			});
+
+			th.start();
+		}
+
+		beforeWrite.await();
+
+		Thread.sleep(200);
+
+		cw.unlock(CACHE_NAME);
+		afterWrite.await();
+
+	}
+
 	private void doWithConnection(Consumer<RedisConnection> callback) {
 
 		try (RedisConnection connection = connectionFactory.getConnection()) {