Reactive distributed lock-Redis implementation

Background reason

• 1. The original authorization project integrates RedisLockRegistry in Spring to implement distributed locks. When migrating the authorization service to Reactive programming, it is necessary to implement the distributed lock implementation in the Reactive mode ( Reference[1] ).
• 2. The original RedisLockRegistry is processed based on Lua-Script and ThreadId
• 3. The main purpose is to keep the original business logic in the migrated project unchanged and to ensure concurrency problems.

Technical solutions and difficulties

• 1. Due to the changes of the Reactive programming model compared to the traditional programming model, in the Event-Loop environment of Reactor-Netty, the thread ID can no longer be used for logical distinction. However, Redis Lua-Script can still be used to achieve concurrency control
• 2. During concurrency, the traditional while(true) {... break} and Thread.sleep methods can no longer be used to wait for the lock to be acquired and to check the lock status. Need to change the way of thinking, use the Reactive way for processing.
• 3. The final realization of the lock processing scheme that is basically the same as RedisLockRegistry

Core dependency

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-webflux</artifactId>
    <optional>true</optional>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis-reactive</artifactId>
    <optional>true</optional>
</dependency>
<dependency>
    <groupId>io.projectreactor.addons</groupId>
    <artifactId>reactor-extra</artifactId>
</dependency>

Implementation logic

• 1. Lock processing Lua-Script

private static final String OBTAIN_LOCK_SCRIPT = "local lockSet = redis.call('SETNX', KEYS[1], ARGV[1])\n" +
            "if lockSet == 1 then\n" +
            "  redis.call('PEXPIRE', KEYS[1], ARGV[2])\n" +
            "  return true\n" +
            "else\n" +
            "  return false\n" +
            "end";

• 2. Core acquiring lock code snippet

/**
 * execute redis-script to obtain lock
 * @return if obtain success then return true otherwise return false
 */
private Mono<Boolean> obtainLock() {
    return Mono.from(ReactiveRedisDistributedLockRegistry.this.reactiveStringRedisTemplate
                    .execute(ReactiveRedisDistributedLockRegistry.this.obtainLockScript,
                            Collections.singletonList(this.lockKey),
                            List.of(this.lockId,String.valueOf(ReactiveRedisDistributedLockRegistry.this.expireAfter)))
            )
            .map(success -> {
                boolean result = Boolean.TRUE.equals(success);
                if (result) {
                    this.lockedAt = System.currentTimeMillis();
                }
                return result;
            });
}

• 3. Core release lock code snippet

/**
 * remove redis lock key
 * @return
 */
private Mono<Boolean> removeLockKey() {
    return Mono.just(ReactiveRedisDistributedLockRegistry.this.unlinkAvailable)
            .filter(unlink -> unlink)
            .flatMap(unlink -> ReactiveRedisDistributedLockRegistry.this.reactiveStringRedisTemplate
                    .unlink(this.lockKey)
                    .doOnError(throwable -> {
                        ReactiveRedisDistributedLockRegistry.this.unlinkAvailable = false;
                        if (log.isDebugEnabled()) {
                            log.debug("The UNLINK command has failed (not supported on the Redis server?); " +
                                    "falling back to the regular DELETE command", throwable);
                        } else {
                            log.warn("The UNLINK command has failed (not supported on the Redis server?); " +
                                    "falling back to the regular DELETE command: " + throwable.getMessage());
                        }
                    })
                    .onErrorResume(throwable -> ReactiveRedisDistributedLockRegistry.this.reactiveStringRedisTemplate.delete(this.lockKey))
            )
            .switchIfEmpty(ReactiveRedisDistributedLockRegistry.this.reactiveStringRedisTemplate.delete(this.lockKey))
            .then(Mono.just(true));
}

• 5. Check whether the lock is occupied by the code snippet

/**
 * check is the acquired is in this process
 * @return
 */
Mono<Boolean> isAcquiredInThisProcess() {
    return ReactiveRedisDistributedLockRegistry.this.reactiveStringRedisTemplate.opsForValue()
            .get(this.lockKey)
            .map(this.lockId::equals);
}

• 4. Basic lock interface definition

public interface ReactiveDistributedLock {

    /**
     * get lock Key
     * @return
     */
    String getLockKey();

    /**
     * Try to acquire the lock once. Lock is acquired for a pre configured duration.
     * @return if lock succeeded then return true otherwise return false
     * <strong>if flow is empty default return false</strong>
     */
    Mono<Boolean> acquireOnce();

    /**
     * Try to acquire the lock. Lock is acquired for a pre configured duration.
     * @return
     * <strong>if flow is empty then throw an excpetion {@link CannotAcquireLockException}</strong>
     */
    Mono<Boolean> acquire();

    /**
     * Try to acquire the lock for a given duration.
     * @param duration duration in used
     * @return
     * <strong>the given duration must less than the default duration.Otherwise the lockKey well be expire by redis with default expire duration</strong>
     * <strong>if flow is empty then throw an excpetion {@link CannotAcquireLockException}</strong>
     */
    Mono<Boolean> acquire(Duration duration);

    /**
     * Release the lock.
     * @return
     * <strong>if lock key doesn't exist in the redis,then throw an exception {@link IllegalStateException}</strong>
     */
    Mono<Boolean> release();
    
}    

• 5. Basic lock interface implementation

private final class ReactiveRedisDistributedLock implements ReactiveDistributedLock {

        @Override
        public String getLockKey() {
            return this.lockKey;
        }

        @Override
        public Mono<Boolean> acquireOnce() {
            log.debug("Acquire Lock Once,LockKey:{}",this.lockKey);
            return this.obtainLock()
                    .doOnNext(lockResult -> log.info("Obtain Lock Once,LockKey:{},Result:{}",this.lockKey,lockResult))
                    .doOnError(this::rethrowAsLockException);
        }

        @Override
        public Mono<Boolean> acquire() {
            log.debug("Acquire Lock By Default Duration :{}" ,expireDuration);
            // 这里使用默认配置的最大等待时间获取锁
            return this.acquire(ReactiveRedisDistributedLockRegistry.this.expireDuration);
        }

        @Override
        public Mono<Boolean> acquire(Duration duration) {
          //尝试获取锁
            return this.obtainLock()
                 //过滤获取锁成功
                    .filter(result -> result)
                    //如果是Empty，则重试
                    .repeatWhenEmpty(Repeat.onlyIf(repeatContext -> true)
                           //重试超时时间
                            .timeout(duration)
                            //重试间隔
                            .fixedBackoff(Duration.ofMillis(100))
                            .//重试日志记录
                            .doOnRepeat(objectRepeatContext -> {
                                if (log.isTraceEnabled()) {
                                    log.trace("Repeat Acquire Lock Repeat Content:{}",objectRepeatContext);
                                }
                            })
                    )
                    //这里必须使用 `defaultIfEmpty`,在repeat超时后，整个流的信号会变为empty，如果不处理empty则整个留就中断了或者由最外层的empty处理方法处理
                    .defaultIfEmpty(false)
                    //记录上锁结果日志
                    .doOnNext(lockResult -> log.info("Obtain Lock,Lock Result :{},Lock Info:{}",lockResult,this))
                    //如果出错，则抛出异常信息
                    .doOnError(this::rethrowAsLockException);
        }

        @Override
        public Mono<Boolean> release() {
          //检查当前锁是否是自己占用
            return this.isAcquiredInThisProcess()
                 //占用的锁
                    .filter(isThisProcess -> isThisProcess)
                    //释放锁
                    .flatMap(isThisProcess -> this.removeLockKey()
                           //记录日志
                            .doOnNext(releaseResult -> log.info("Released Lock:{},Lock Info:{}",releaseResult,this))
                            //出现未知异常，则重新抛出
                            .onErrorResume(throwable -> Mono.fromRunnable(() -> ReflectionUtils.rethrowRuntimeException(throwable)))
                            //如果流是empty，则表示，锁已经不存在了，被Redis配置的最大过期时间释放
                            .switchIfEmpty(Mono.error(new IllegalStateException("Lock was released in the store due to expiration. " + "The integrity of data protected by this lock may have been compromised.")))
                    );
        }
}

• 6. Built-in timing tasks, used to detect expired RedisLock that is not in use, and release the memory cache. The timed task is mounted to the SpringBean's declaration cycle, and the startup and shutdown of the timed task have been completed. ( InitializingBean , DisposableBean )

private Scheduler scheduler = Schedulers.newSingle("redis-lock-evict",true);

//挂载Spring 声明周期
@Override
public void afterPropertiesSet() {
    log.debug("Initialize Auto Remove Unused Lock Execution");
    //使用Flux的特性来实现定时任务
    Flux.interval(expireEvictIdle, scheduler)
            .flatMap(value -> {
                long now = System.currentTimeMillis();
                log.trace("Auto Remove Unused Lock ,Evict Triggered");
                return Flux.fromIterable(this.locks.entrySet())
                        //过滤已经过期的锁对象
                        .filter(entry -> now - entry.getValue().getLockedAt() > expireAfter)
                        //将没有被占用的锁删除
                        .flatMap(entry -> entry.getValue()
                                .isAcquiredInThisProcess()
                                .filter(inProcess -> !inProcess)
                                .doOnNext(inProcess -> {
                                    this.locks.remove(entry.getKey());
                                    log.debug("Auto Remove Unused Lock,Lock Info:{}", entry);
                                })
                                //错误记录日志
                                .onErrorResume(throwable -> {
                                    log.error("Auto Remove Unused Locks Occur Exception,Lock Info: " + entry, throwable);
                                    return Mono.empty();
                                })
                        );
            })
            //Scheduler 需要订阅才能执行
            .subscribe();
}

@Override
public void destroy() {
    log.debug("Shutdown Auto Remove Unused Lock Execution");
    //挂载SpringBean声明周期，销毁Scheduler
    this.scheduler.dispose();
}

• 7. Optimize the processing logic of the lock interface and increase the default method of the interface to facilitate lock control and processing. Wrap the downstream execution logic of the lock into Supplier for easy calling and processing

/**
 * Acquire a lock and release it after action is executed or fails.
 *
 * @param <T>  type od value emitted by the action
 * @param monoExecution to be executed subscribed to when lock is acquired
 * @return true if lock is acquired.
 * @see ReactiveDistributedLock#acquire()
 */
default <T> Mono<T> acquireAndExecute(Mono<T> monoExecution) {
    return acquire()
            .flatMap(acquireResult -> Mono.just(acquireResult)
                        .filter(result -> result)
                        //这里配合上锁逻辑，如果是空，则表示无法获取锁
                        .switchIfEmpty(Mono.error(new CannotAcquireLockException("Failed to Obtain Lock ,LockKey: " + getLockKey())))
                        .flatMap(lockResult -> monoExecution
                                .flatMap(result -> this.release()
                                        .flatMap(releaseResult -> Mono.just(result))
                                )
                                .switchIfEmpty(this.release().then(Mono.empty()))
                                .onErrorResume(throwable -> this.release().flatMap(r -> Mono.error(throwable)))
                        )
            );
}

/**
 * Acquire a lock for a given duration and release it after action is executed.
 *
 * @param <T>      type od value emitted by the action
 * @param duration how much time must pass for the acquired lock to expire
 * @param monoExecution to be executed subscribed to when lock is acquired
 * @return true, if lock is acquired
 * @see ReactiveDistributedLock#acquire(Duration)
 */
default <T> Mono<T> acquireAndExecute(Duration duration, Mono<T> monoExecution) {
    return acquire(duration)
            .flatMap(acquireResult -> Mono.just(acquireResult)
                    .filter(result -> result)
                    .switchIfEmpty(Mono.error(new CannotAcquireLockException("Failed to Obtain Lock ,LockKey: " + getLockKey())))
                    .flatMap(lockResult -> monoExecution
                            .flatMap(result -> this.release()
                                    .flatMap(releaseResult -> Mono.just(result))
                            )
                            .switchIfEmpty(this.release().then(Mono.empty()))
                            .onErrorResume(throwable -> this.release().flatMap(r -> Mono.error(throwable)))
                    )
            );
}

/**
 * Acquire a lock and release it after action is executed or fails.
 *
 * @param <T>  type od value emitted by the action
 * @param fluxExecution  to be executed subscribed to when lock is acquired
 * @return true if lock is acquired.
 * @see ReactiveDistributedLock#acquire()
 */
default <T> Flux<T> acquireAndExecuteMany(Flux<T> fluxExecution) {
    return acquire()
            .flatMapMany(acquireResult -> Mono.just(acquireResult)
                    .filter(result -> result)
                    .switchIfEmpty(Mono.error(new CannotAcquireLockException("Failed to Obtain Lock ,LockKey: " + getLockKey())))
                    .flatMapMany(lockResult -> fluxExecution
                            .flatMap(result -> this.release()
                                    .flatMap(releaseResult -> Mono.just(result))
                            )
                            .switchIfEmpty(this.release().thenMany(Flux.empty()))
                            .onErrorResume(throwable -> this.release().flatMap(r -> Mono.error(throwable)))
                    )
            );
}

/**
 * Acquire a lock for a given duration and release it after action is executed.
 *
 * @param <T>      type od value emitted by the action
 * @param duration how much time must pass for the acquired lock to expire
 * @param fluxExecution  to be executed subscribed to when lock is acquired
 * @return true, if lock is acquired
 * @see ReactiveDistributedLock#acquire(Duration)
 */
default <T> Flux<T> acquireAndExecuteMany(Duration duration, Flux<T> fluxExecution) {
    return acquire(duration)
            .flatMapMany(acquireResult -> Mono.just(acquireResult)
                    .filter(result -> result)
                    .switchIfEmpty(Mono.error(new CannotAcquireLockException("Failed to Obtain Lock ,LockKey: " + getLockKey())))
                    .flatMapMany(lockResult -> fluxExecution
                            .flatMap(result -> this.release()
                                    .flatMap(releaseResult -> Mono.just(result))
                            )
                            .switchIfEmpty(this.release().thenMany(Flux.empty()))
                            .onErrorResume(throwable -> this.release().flatMap(r -> Mono.error(throwable)))
                    )
            );
}

Instructions

• Parameter configuration in application.yml

lock:
  redis:
    reactive:
      expire-after: 10s
      expire-evict-idle: 1s

• Inject Bean

@Autowired
private ReactiveRedisDistributedLockRegistry reactiveRedisDistributedLockRegistry;

• 1. Lock once, fail quickly

@Test
public void testAcquireOnce() throws Exception {
    ProcessFunctions processFunctions = new ProcessFunctions();
    String key = "LOCK_ONCE";
    Flux<String> flux = Flux.range(0, 5)
            .flatMap(value -> this.reactiveRedisDistributedLockRegistry.obtain(key)
                    .acquireOnce()
                    .filter(acquireResult -> acquireResult)
                    .flatMap(acquireResult -> processFunctions.processFunction())
                    .switchIfEmpty(Mono.just(FAILED))
            )
            .doOnNext(System.out::println);
    StepVerifier.create(flux)
            .expectNext(OK)
            .expectNext(FAILED)
            .expectNext(FAILED)
            .expectNext(FAILED)
            .expectNext(FAILED)
            .verifyComplete();

}

• 2. The default timeout period of waiting for lock

@Test
public void testAcquireDefaultDurationAndProcessDuringTheExpireDuration() throws Exception {
    //default lock expire is 10S
    ProcessFunctions processFunctions = new ProcessFunctions();
    String key = "LOCK_DEFAULT";
    Flux<String> flux = Flux.range(0, 3)
            .flatMap(value -> this.reactiveRedisDistributedLockRegistry.obtain(key)
                    .acquireAndExecute(
                           processFunctions.processDelayFunction(Duration.ofSeconds(2))
                    )
                    .doOnNext(System.out::println)
                    .onErrorResume(throwable -> CannotAcquireLockException.class.isAssignableFrom(throwable.getClass()),throwable -> {
                        System.out.println("Lock Error");
                        return Mono.just(FAILED);
                    })
            );
    StepVerifier.create(flux)
            .expectNext(OK)
            .expectNext(OK)
            .expectNext(OK)
            .verifyComplete();

}

• 3. Lock for specified time

@Test
public void testAcquireDuration() throws Exception {
    ProcessFunctions processFunctions = new ProcessFunctions();
    String key = "LOCK_GIVEN_DURATION";
    Flux<String> flux = Flux.range(0, 3)
            .subscribeOn(Schedulers.parallel())
            .flatMap(value -> this.reactiveRedisDistributedLockRegistry.obtain(key)
                    .acquireAndExecute(Duration.ofSeconds(3), processFunctions.processDelayFunction(Duration.ofSeconds(2))
                    )
                    .doOnNext(System.out::println)
                    .onErrorResume(throwable -> CannotAcquireLockException.class.isAssignableFrom(throwable.getClass()), throwable -> {
                        System.out.println("Lock Error");
                        return Mono.just(FAILED);
                    })
            );
    StepVerifier.create(flux)
            .expectNext(OK)
            .expectNext(FAILED)
            .expectNext(OK)
            .verifyComplete();

}

Reference documents

• 1 . RedisLockRegistry: https://docs.spring.io/spring-integration/docs/5.3.6.RELEASE/reference/html/redis.html#redis-lock-registry
• 2 . Trigger Mono Execution After Another Mono Terminates: https://stackoverflow.com/questions/50686524/how-to-trigger-mono-execution-after-another-mono-terminates

Source code related

• Maintained on GitHub, welcome to Issue and Star reactive-redis-distributed-lock
• It is currently written in the form of SpringBoot scaffolding and has not been published to the Maven central warehouse. You can package it yourself if necessary.