GH-8699: Enhancing unlock() Method of RedisLock with Atomic Redis Operation

Author: EddieChoCho

spring-integration-redis/src/main/java/org/springframework/integration/redis/util/RedisLockRegistry.java

@@ -343,12 +343,12 @@ public long getLockedAt() {
 		/**
 		 * Unlock the lock using the unlink method in redis.
 		 */
-		protected abstract void removeLockKeyInnerUnlink();
+		protected abstract Boolean removeLockKeyInnerUnlink();
 
 		/**
 		 * Unlock the lock using the delete method in redis.
 		 */
-		protected abstract void removeLockKeyInnerDelete();
+		protected abstract Boolean removeLockKeyInnerDelete();
 
 		@Override
 		public final void lock() {
@@ -454,11 +454,6 @@ public final void unlock() {
 				return;
 			}
 			try {
-				if (!isAcquiredInThisProcess()) {
-					throw new IllegalStateException("Lock was released in the store due to expiration. " +
-							"The integrity of data protected by this lock may have been compromised.");
-				}
-
 				if (Thread.currentThread().isInterrupted()) {
 					RedisLockRegistry.this.executor.execute(this::removeLockKey);
 				}
@@ -480,9 +475,9 @@ public final void unlock() {
 
 		private void removeLockKey() {
 			if (RedisLockRegistry.this.unlinkAvailable) {
+				Boolean unlinkResult = null;
 				try {
-					removeLockKeyInnerUnlink();
-					return;
+					unlinkResult = removeLockKeyInnerUnlink();
 				}
 				catch (Exception ex) {
 					RedisLockRegistry.this.unlinkAvailable = false;
@@ -495,8 +490,17 @@ private void removeLockKey() {
 								"falling back to the regular DELETE command: " + ex.getMessage());
 					}
 				}
+				if(Boolean.TRUE.equals(unlinkResult)){
+					return;
+				} else if(unlinkResult!= null){
+					throw new IllegalStateException("Lock was released in the store due to expiration. " +
+							"The integrity of data protected by this lock may have been compromised.");
+				}
+			}
+			if (!removeLockKeyInnerDelete()) {
+				throw new IllegalStateException("Lock was released in the store due to expiration. " +
+						"The integrity of data protected by this lock may have been compromised.");
 			}
-			removeLockKeyInnerDelete();
 		}
 
 		@Override
@@ -559,19 +563,23 @@ private RedisLockRegistry getOuterType() {
 
 	private final class RedisPubSubLock extends RedisLock {
 
-		private static final String UNLINK_UNLOCK_SCRIPT =
-				"if (redis.call('unlink', KEYS[1]) == 1) then " +
-						"redis.call('publish', ARGV[1], KEYS[1]) " +
-						"return true " +
-						"end " +
-						"return false";
+		private static final String UNLINK_UNLOCK_SCRIPT = """
+				local lockClientId = redis.call('GET', KEYS[1])
+				if (lockClientId == ARGV[1] and redis.call('UNLINK', KEYS[1]) == 1) then
+					redis.call('PUBLISH', ARGV[2], KEYS[1])
+					return true
+				end
+				return false
+				""";
 
-		private static final String DELETE_UNLOCK_SCRIPT =
-				"if (redis.call('del', KEYS[1]) == 1) then " +
-						"redis.call('publish', ARGV[1], KEYS[1]) " +
-						"return true " +
-						"end " +
-						"return false";
+		private static final String DELETE_UNLOCK_SCRIPT = """
+				local lockClientId = redis.call('GET', KEYS[1])
+				if (lockClientId == ARGV[1] and redis.call('DEL', KEYS[1]) == 1) then
+					redis.call('PUBLISH', ARGV[2], KEYS[1])
+					return true
+				end
+				return false
+				""";
 
 		private static final RedisScript<Boolean>
 				UNLINK_UNLOCK_REDIS_SCRIPT = new DefaultRedisScript<>(UNLINK_UNLOCK_SCRIPT, Boolean.class);
@@ -589,17 +597,17 @@ protected boolean tryRedisLockInner(long time) throws ExecutionException, Interr
 		}
 
 		@Override
-		protected void removeLockKeyInnerUnlink() {
-			RedisLockRegistry.this.redisTemplate.execute(
+		protected Boolean removeLockKeyInnerUnlink() {
+			return RedisLockRegistry.this.redisTemplate.execute(
 					UNLINK_UNLOCK_REDIS_SCRIPT, Collections.singletonList(this.lockKey),
-					RedisLockRegistry.this.unLockChannelKey);
+					RedisLockRegistry.this.clientId,RedisLockRegistry.this.unLockChannelKey);
 		}
 
 		@Override
-		protected void removeLockKeyInnerDelete() {
-			RedisLockRegistry.this.redisTemplate.execute(
+		protected Boolean removeLockKeyInnerDelete() {
+			return RedisLockRegistry.this.redisTemplate.execute(
 					DELETE_UNLOCK_REDIS_SCRIPT, Collections.singletonList(this.lockKey),
-					RedisLockRegistry.this.unLockChannelKey);
+					RedisLockRegistry.this.clientId,RedisLockRegistry.this.unLockChannelKey);
 
 		}
 
@@ -694,6 +702,30 @@ private void unlockNotify(String lockKey) {
 
 	private final class RedisSpinLock extends RedisLock {
 
+		private static final String UNLINK_UNLOCK_SCRIPT = """
+				local lockClientId = redis.call('GET', KEYS[1])
+				if lockClientId == ARGV[1] then
+					redis.call('UNLINK', KEYS[1])
+					return true
+				end
+				return false
+				""";
+
+		private static final String DELETE_UNLOCK_SCRIPT = """
+				local lockClientId = redis.call('GET', KEYS[1])
+				if lockClientId == ARGV[1] then
+					redis.call('DEL', KEYS[1])
+					return true
+				end
+				return false
+				""";
+
+		private static final RedisScript<Boolean>
+				UNLINK_UNLOCK_REDIS_SCRIPT = new DefaultRedisScript<>(UNLINK_UNLOCK_SCRIPT, Boolean.class);
+
+		private static final RedisScript<Boolean>
+				DELETE_UNLOCK_REDIS_SCRIPT = new DefaultRedisScript<>(DELETE_UNLOCK_SCRIPT, Boolean.class);
+
 		private RedisSpinLock(String path) {
 			super(path);
 		}
@@ -718,13 +750,17 @@ protected boolean tryRedisLockInner(long time) throws InterruptedException {
 		}
 
 		@Override
-		protected void removeLockKeyInnerUnlink() {
-			RedisLockRegistry.this.redisTemplate.unlink(this.lockKey);
+		protected Boolean removeLockKeyInnerUnlink() {
+			return RedisLockRegistry.this.redisTemplate.execute(
+					UNLINK_UNLOCK_REDIS_SCRIPT, Collections.singletonList(this.lockKey),
+					RedisLockRegistry.this.clientId);
 		}
 
 		@Override
-		protected void removeLockKeyInnerDelete() {
-			RedisLockRegistry.this.redisTemplate.delete(this.lockKey);
+		protected Boolean removeLockKeyInnerDelete() {
+			return RedisLockRegistry.this.redisTemplate.execute(
+					DELETE_UNLOCK_REDIS_SCRIPT, Collections.singletonList(this.lockKey),
+					RedisLockRegistry.this.clientId);
 		}
 
 	}