Spring Boot + Lua 手写分布式锁（支持自动续期 / 可重入）

1、简介

在分布式系统环境中，多个服务或节点可能并发地访问和修改同一资源，这种情况极易导致数据不一致或死锁问题。为解决这一问题，分布式锁机制应运而生。相较于直接使用现成的分布式锁解决方案，通过自己动手实践，我们能够更深刻地理解其内部的运作机制与核心原理。
通过Spring Boot集成Redis，并使用Lua脚本，我们可以实现一个支持自动续期和可重入的分布式锁。Lua脚本的原子性执行确保了获取和释放锁的操作是不可分割的，从而避免了竞态条件。自动续期功能则通过守护线程或看门狗机制实现，确保锁在业务处理过程中不会因过期而被其他客户端获取。

2、实战案例

2.1 环境准备

这里我们只需要引入一个redis相关的依赖即可

<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-data-redis</artifactId>
</dependency>

2.2 定义加锁/解锁LUA脚本

加锁Lua脚本

local key = KEYS[1]
local lockId = ARGV[1]
local expireTime = ARGV[2]
-- 判断锁是否存在
if (redis.call('exists', key) == 0) then
  redis.call('hset', key, lockId, 1)
  redis.call('pexpire', key, expireTime)
  return 1
end
-- 判断是否是当前线程持有锁
if (redis.call('hexists', key, lockId) == 1) then
  -- 如果当前线程已经获取锁了，则进行累加计数器
  redis.call('hincrby', key, lockId, 1)
  -- 重置过期时间
  redis.call('pexpire', key, expireTime)
  return 1
end
return 0

释放锁Lua脚本

local key = KEYS[1]
local lockId = ARGV[1]
-- 当前线程未获取锁
if (redis.call('hexists', key, lockId) == 0) then
  return nil
end
local count = redis.call('hincrby', key, lockId, -1)
if (count > 0) then
  redis.call('pexpire', key, ARGV[2])
  return 0
else
  redis.call('del', key)
  return 1
end

锁续期Lua脚本

local key = KEYS[1]
local lockId = ARGV[1]
local expireTime = ARGV[2]
-- 判断是否是当前线程持有锁
if (redis.call('hexists', key, lockId) == 1) then
  -- 重置过期时间
  redis.call('pexpire', key, expireTime)
  return 1
end
return 0

接下来配置配置上面3个脚本

2.3 Lua脚本配置

@Configuration
public class RedisLockConfig {
  @Bean
  DefaultRedisScript<Long> lockScript() {
    DefaultRedisScript<Long> script = new DefaultRedisScript<>();
    script.setLocation(new ClassPathResource("lua/lock.lua"));
    script.setResultType(Long.class);
    return script;
  }
  @Bean
  DefaultRedisScript<Long> unlockScript() {
    DefaultRedisScript<Long> script = new DefaultRedisScript<>();
    script.setLocation(new ClassPathResource("lua/unlock.lua"));
    script.setResultType(Long.class);
    return script;
  }
  @Bean
  DefaultRedisScript<Long> renewalScript() {
    DefaultRedisScript<Long> script = new DefaultRedisScript<>();
    script.setLocation(new ClassPathResource("lua/renewal.lua"));
    script.setResultType(Long.class);
    return script;
  }
}

通过如上的定义方便我们在分布式锁的实现中使用。

2.4 分布式锁实现

@Component
public class RedisDistributedLock {
  /**默认锁前缀*/
  private static final String LOCK_PREFIX = "pack:lock:";
  /**默认锁超时时间*/
  private static final long LOCK_EXPIRE_TIME = 30 * 1000 ;
  /**记录锁key与ID*/
  private static final ThreadLocal<Map<String, String>> lockHolder = ThreadLocal.withInitial(HashMap::new);

  private final StringRedisTemplate redisTemplate;
  private final DefaultRedisScript<Long> lockScript;
  private final DefaultRedisScript<Long> unlockScript;
  private final DefaultRedisScript<Long> renewalScript ;
  private final String id ;
  private final String lockKey ;
  private final ScheduledExecutorService scheduler ;
  public RedisDistributedLock(StringRedisTemplate redisTemplate,
      DefaultRedisScript<Long> lockScript,
      DefaultRedisScript<Long> unlockScript,
      DefaultRedisScript<Long> renewalScript,
      String lockKey) {
    this.redisTemplate = redisTemplate;
    this.lockScript = lockScript;
    this.unlockScript = unlockScript;
    this.renewalScript = renewalScript ;
    this.id = UUID.randomUUID().toString().replaceAll("-", "") ;
    this.lockKey = lockKey ;
    this.scheduler = Executors.newSingleThreadScheduledExecutor() ;
  }
  public void lock() {
    /**TODO*/
  }
  public void unlock() {
    /**TODO*/
  }
  private String generateLockId() {
    return this.id + ":" + Thread.currentThread().threadId();
  }
}

接下来，我们具体来实现里面的逻辑。

加锁操作

public void lock() {
  final long expireMillis = LOCK_EXPIRE_TIME ; 
  String actualKey = LOCK_PREFIX + lockKey;
  try {
    while (true) {
      /**生成唯一的锁ID*/
      String lockId = generateLockId();
      Long result = redisTemplate.execute(lockScript, List.of(actualKey), lockId,
          String.valueOf(expireMillis));
      if (result != null && result == 1) {
        /**成功获取分布式锁*/
        lockHolder.get().put(actualKey, lockId);
        startRenewalThread(actualKey, lockId, expireMillis);
        return; 
      }
      // 为了解决惊群效应；但是这里是有明显缺陷的，这导致始终会再这里等待一定的时间；我们应该采用发布/订阅消息的方式来实现
      TimeUnit.MILLISECONDS.sleep((long)(Math.random() * 50) + 50) ;
    }
  } catch (Exception e) {
    System.err.println("加锁失败: " + e.getMessage()) ;
    throw new RuntimeException(e) ;
  }
}
/**根据UUID+当前线程id生成唯一id*/
private String generateLockId() {
  return this.id + ":" + Thread.currentThread().threadId();
}
/**锁续期*/
private void startRenewalThread(String key, String lockId, long expireTime) {
  long delay = expireTime / 3;
  scheduler.scheduleAtFixedRate(() -> {
    Long ret = this.redisTemplate.execute(renewalScript, List.of(key), lockId, String.valueOf(expireTime)) ;
    if (ret == null || ret == 0) {
      scheduler.shutdownNow();
    }
  }, delay, delay, TimeUnit.MILLISECONDS);
}

释放锁操作

public void unlock() {
  final long expireMillis = LOCK_EXPIRE_TIME ; 
  String actualKey = LOCK_PREFIX + lockKey;
  String lockId = lockHolder.get().get(actualKey);
  if (lockId == null) {
    return;
  }
  Long ret = redisTemplate.execute(unlockScript, List.of(actualKey), lockId,
      String.valueOf(expireMillis)) ;
  if (ret != null && ret == 1) {
    lockHolder.get().remove(actualKey) ;
    stopRenewalThread() ;
  }
}
private void stopRenewalThread() {
  if (scheduler != null && !scheduler.isShutdown()) {
    scheduler.shutdownNow();
  }
}

以上就是分布式锁的所有代码了；下面定义获取锁的工具类。

@Component
public class PackLock {
  private final StringRedisTemplate stringRedisTemplate ;
  private final DefaultRedisScript<Long> lockScript;
  private final DefaultRedisScript<Long> unlockScript;
  private final DefaultRedisScript<Long> renewalScript;
  public PackLock(StringRedisTemplate stringRedisTemplate,
      DefaultRedisScript<Long> lockScript,
      DefaultRedisScript<Long> unlockScript,
      DefaultRedisScript<Long> renewalScript) {
    this.stringRedisTemplate = stringRedisTemplate;
    this.lockScript = lockScript ;
    this.unlockScript = unlockScript ;
    this.renewalScript = renewalScript ;
  }
  /**每次都需要创建新的锁对象*/
  public RedisDistributedLock getLock(String lockKey) {
    RedisDistributedLock lock = new RedisDistributedLock(
      stringRedisTemplate, this.lockScript, 
      this.unlockScript, renewalScript, lockKey) ;
    return lock ;
  }
}

2.5 锁的使用

@Service
public class ProductService {
  private int count = 20 ;
  
  private final PackLock packLock ;
  public ProductService(PackLock packLock) {
    this.packLock = packLock;
  }
  public void calc() {
    if (count <= 0) {
      return ;
    }
    RedisDistributedLock lock = this.packLock.getLock("xxx") ;
    lock.lock() ;
    try {
      if (count > 0) {
        count-- ;
      }
    } finally {
      lock.unlock() ;
    }
  }
}