RedisLockRegistry implements RenewableLockRegistry

Author: NaccOll

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

@@ -27,15 +27,12 @@
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.locks.Condition;
@@ -58,6 +55,7 @@
 import org.springframework.data.redis.listener.RedisMessageListenerContainer;
 import org.springframework.data.redis.listener.Topic;
 import org.springframework.integration.support.locks.ExpirableLockRegistry;
+import org.springframework.integration.support.locks.RenewableLockRegistry;
 import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
 import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
@@ -97,7 +95,7 @@
  * @since 4.0
  *
  */
-public final class RedisLockRegistry implements ExpirableLockRegistry, DisposableBean {
+public final class RedisLockRegistry implements ExpirableLockRegistry, DisposableBean, RenewableLockRegistry {
 
 	private static final Log LOGGER = LogFactory.getLog(RedisLockRegistry.class);
 
@@ -142,6 +140,11 @@ protected boolean removeEldestEntry(Entry<String, RedisLock> eldest) {
 	private Executor executor =
 			Executors.newCachedThreadPool(new CustomizableThreadFactory("redis-lock-registry-"));
 
+	/**
+	 * renewal scheduleExecutor
+	 */
+	private ScheduledExecutorService renewalScheduledExecutor;
+
 	/**
 	 * Flag to denote whether the {@link ExecutorService} was provided via the setter and
 	 * thus should not be shutdown when {@link #destroy()} is called
@@ -211,6 +214,14 @@ public void setExecutor(Executor executor) {
 		this.executorExplicitlySet = true;
 	}
 
+	/**
+	 * Set the {@link ScheduledExecutorService}, where is not provided then renewal not work
+	 * @param renewalScheduledExecutor the executor service
+	 */
+	public void setRenewalScheduledExecutor(ScheduledExecutorService renewalScheduledExecutor) {
+		this.renewalScheduledExecutor = renewalScheduledExecutor;
+	}
+
 	/**
 	 * Set the capacity of cached locks.
 	 * @param cacheCapacity The capacity of cached lock, (default 100_000).
@@ -293,6 +304,29 @@ public void destroy() {
 				throw new IllegalStateException(ex);
 			}
 		}
+		if (renewalScheduledExecutor != null) {
+			renewalScheduledExecutor.shutdown();;
+		}
+	}
+
+	@Override
+	public void renewLock(Object lockKey) {
+		String path = (String) lockKey;
+		RedisLock redisLock;
+		this.lock.lock();
+		try {
+			redisLock = this.locks.computeIfAbsent(path, getRedisLockConstructor(this.redisLockType));
+		}
+		finally {
+			this.lock.unlock();
+		}
+		if (redisLock == null) {
+			throw new IllegalStateException("Could not renew mutex at " + path);
+		}
+
+		if (!redisLock.renew()) {
+			throw new IllegalStateException("Could not renew mutex at " + path);
+		}
 	}
 
 	/**
@@ -305,11 +339,6 @@ public enum RedisLockType {
 		 */
 		SPIN_LOCK,
 
-		/**
-		 * The lock is acquired by periodically(idleBetweenTries property) checking whether the lock can be acquired and schedule renewal expire time
-		 */
-		SPIN_LOCK_WITH_RENEWAL,
-
 		/**
 		 * The lock is acquired by redis pub-sub subscription.
 		 */
@@ -320,7 +349,6 @@ private Function<String, RedisLock> getRedisLockConstructor(RedisLockType redisL
 		return switch (redisLockType) {
 			case SPIN_LOCK -> RedisSpinLock::new;
 			case PUB_SUB_LOCK -> RedisPubSubLock::new;
-			case SPIN_LOCK_WITH_RENEWAL -> RedisSpinLockWithRenewal::new;
 		};
 	}
 
@@ -338,9 +366,20 @@ private abstract class RedisLock implements Lock {
 				return false
 				""";
 
+		private static final String RENEW_SCRIPT = """
+				if (redis.call('GET', KEYS[1]) == ARGV[2]) then
+					redis.call('PEXPIRE', KEYS[1], ARGV[1])
+					return true
+				end
+				return false
+				""";
+
 		protected static final RedisScript<Boolean>
 				OBTAIN_LOCK_REDIS_SCRIPT = new DefaultRedisScript<>(OBTAIN_LOCK_SCRIPT, Boolean.class);
 
+		public static final RedisScript<Boolean>
+				RENEW_REDIS_SCRIPT = new DefaultRedisScript<>(RENEW_SCRIPT, Boolean.class);
+
 		protected final String lockKey;
 
 		private final ReentrantLock localLock = new ReentrantLock();
@@ -464,6 +503,13 @@ private boolean tryRedisLock(long time) throws ExecutionException, InterruptedEx
 					LOGGER.debug("Acquired lock; " + this);
 				}
 				this.lockedAt = System.currentTimeMillis();
+				if (RedisLockRegistry.this.renewalScheduledExecutor != null) {
+					long delay = RedisLockRegistry.this.expireAfter / 3;
+					Runnable task = () -> {
+						renewLock(this.lockKey.substring((RedisLockRegistry.this.registryKey + ':').length()));
+					};
+					RedisLockRegistry.this.renewalScheduledExecutor.scheduleWithFixedDelay(task, delay, delay, TimeUnit.MILLISECONDS);
+				}
 			}
 			return acquired;
 		}
@@ -538,6 +584,12 @@ else if (Boolean.FALSE.equals(unlinkResult)) {
 			}
 		}
 
+		public final boolean renew() {
+			return Boolean.TRUE.equals(RedisLockRegistry.this.redisTemplate.execute(
+					RENEW_REDIS_SCRIPT, Collections.singletonList(this.lockKey), String.valueOf(RedisLockRegistry.this.expireAfter),
+					RedisLockRegistry.this.clientId));
+		}
+
 		@Override
 		public final Condition newCondition() {
 			throw new UnsupportedOperationException("Conditions are not supported");
@@ -734,7 +786,7 @@ private void unlockNotify(String lockKey) {
 
 	}
 
-	private class RedisSpinLock extends RedisLock {
+	private final class RedisSpinLock extends RedisLock {
 
 		private static final String UNLINK_UNLOCK_SCRIPT = """
 				local lockClientId = redis.call('GET', KEYS[1])
@@ -801,176 +853,4 @@ private boolean removeLockKeyWithScript(RedisScript<Boolean> redisScript) {
 
 	}
 
-	public final class RedisSpinLockWithRenewal extends RedisSpinLock {
-
-		private static final ConcurrentMap<String, ExpirationEntry> EXPIRATION_RENEWAL_MAP = new ConcurrentHashMap<>();
-
-		private static final ScheduledExecutorService scheduledExecutorService = new ScheduledThreadPoolExecutor(10,
-				new CustomizableThreadFactory("redis-lock-renewal-schedule"));
-
-		private static final String RENEW_SCRIPT = """
-				if (redis.call('GET', KEYS[1]) == ARGV[2]) then
-					redis.call('PEXPIRE', KEYS[1], ARGV[1])
-					return true
-				end
-				return false
-				""";
-
-		public static final RedisScript<Boolean>
-				RENEW_REDIS_SCRIPT = new DefaultRedisScript<>(RENEW_SCRIPT, Boolean.class);
-
-		private RedisSpinLockWithRenewal(String path) {
-			super(path);
-		}
-
-		@Override
-		protected boolean tryRedisLockInner(long time) throws InterruptedException {
-			boolean res = super.tryRedisLockInner(time);
-			if (res) {
-				scheduleExpirationRenewal(Thread.currentThread().getId());
-			}
-			return res;
-		}
-
-		@Override
-		protected boolean removeLockKeyInnerUnlink() {
-			boolean result = super.removeLockKeyInnerUnlink();
-			if (result) {
-				cancelExpirationRenewal(Thread.currentThread().getId());
-			}
-			return result;
-		}
-
-		@Override
-		protected boolean removeLockKeyInnerDelete() {
-			boolean result = super.removeLockKeyInnerDelete();
-			if (result) {
-				cancelExpirationRenewal(Thread.currentThread().getId());
-			}
-			return result;
-		}
-
-		protected void scheduleExpirationRenewal(long threadId) {
-			ExpirationEntry entry = new ExpirationEntry();
-			ExpirationEntry oldEntry = EXPIRATION_RENEWAL_MAP.putIfAbsent(lockKey, entry);
-			if (oldEntry != null) {
-				oldEntry.addThreadId(threadId);
-			}
-			else {
-				entry.addThreadId(threadId);
-				renewExpiration();
-			}
-		}
-
-		private void renewExpiration() {
-			ExpirationEntry ee = EXPIRATION_RENEWAL_MAP.get(lockKey);
-			if (ee == null) {
-				return;
-			}
-			Runnable task = () -> {
-				ExpirationEntry ent = EXPIRATION_RENEWAL_MAP.get(lockKey);
-				if (ent == null) {
-					return;
-				}
-				Long threadId = ent.getFirstThreadId();
-				if (threadId == null) {
-					return;
-				}
-
-				boolean res = renewKeyWithScript(RENEW_REDIS_SCRIPT);
-				if (res) {
-					renewExpiration();
-				}
-				else {
-					EXPIRATION_RENEWAL_MAP.remove(lockKey);
-				}
-
-			};
-			ScheduledFuture<?> timeout = scheduledExecutorService.schedule(task, RedisLockRegistry.this.expireAfter / 3,
-					TimeUnit.MILLISECONDS);
-			ee.setTimeout(timeout);
-		}
-
-		private void cancelExpirationRenewal(Long threadId) {
-			ExpirationEntry task = EXPIRATION_RENEWAL_MAP.get(lockKey);
-			if (task == null) {
-				return;
-			}
-
-			if (threadId != null) {
-				task.removeThreadId(threadId);
-			}
-
-			if (threadId == null || task.hasNoThreads()) {
-				ScheduledFuture<?> timeout = task.getTimeout();
-				if (timeout != null) {
-					timeout.cancel(true);
-				}
-				EXPIRATION_RENEWAL_MAP.remove(lockKey);
-			}
-		}
-
-		public boolean renewKeyWithScript(RedisScript<Boolean> redisScript) {
-			return Boolean.TRUE.equals(RedisLockRegistry.this.redisTemplate.execute(
-					redisScript, Collections.singletonList(this.lockKey), String.valueOf(RedisLockRegistry.this.expireAfter),
-					RedisLockRegistry.this.clientId));
-		}
-
-		static class ExpirationEntry {
-
-			private final Map<Long, Integer> threadIds = new LinkedHashMap<>();
-
-			private ScheduledFuture<?> timeout;
-
-			ExpirationEntry() {
-				super();
-			}
-
-			public synchronized void addThreadId(long threadId) {
-				Integer counter = this.threadIds.get(threadId);
-				if (counter == null) {
-					counter = 1;
-				}
-				else {
-					counter++;
-				}
-				this.threadIds.put(threadId, counter);
-			}
-
-			public synchronized boolean hasNoThreads() {
-				return this.threadIds.isEmpty();
-			}
-
-			public synchronized Long getFirstThreadId() {
-				if (this.threadIds.isEmpty()) {
-					return null;
-				}
-				return this.threadIds.keySet().iterator().next();
-			}
-
-			public synchronized void removeThreadId(long threadId) {
-				Integer counter = this.threadIds.get(threadId);
-				if (counter == null) {
-					return;
-				}
-				counter--;
-				if (counter == 0) {
-					this.threadIds.remove(threadId);
-				}
-				else {
-					this.threadIds.put(threadId, counter);
-				}
-			}
-
-			public ScheduledFuture<?> getTimeout() {
-				return this.timeout;
-			}
-
-			public void setTimeout(ScheduledFuture<?> timeout) {
-				this.timeout = timeout;
-			}
-		}
-
-	}
-
 }