Redis的用武之处--由短信限流引出的分布式锁和漏斗限流（下）

分布式锁

有些人应该用过Redission这个redis中间件框架，它以 使用者忘记redis本身命令，而更多关注业务为目标，所以它的api不同于jedis，redission就原生提供了分布式锁，限流器等现成的工具类。我以重复"造轮子"为宗旨，试着写写这个分布锁。

上一篇我们知道光一个漏斗限流在生产环境是不行的，容易因为并发导致出现问题，我们需要给这个限流器上一把锁，先贴流程图/代码：

这把锁的入参我都是仿照Redission来弄，实现也比Redission简单一些

/**
     * 尝试获取锁
     *
     * @param lockKey
     * @param waitTime  最多等待时间(秒)
     * @param leaseTime 上锁后自动释放锁时间(秒)
     * @return
     */
    public boolean tryLock(String lockKey, long waitTime, int leaseTime) {
        long currentThreadid = Thread.currentThread().getId();
        final String finalLockKey = packageLockKey(lockKey);
        try (Jedis jedis = jedisPool.getResource()) {
            long time;
            long begin = System.currentTimeMillis();
            log.info("获取锁开始 {} {}", currentThreadid, begin);
            String result = jedis.set(finalLockKey, "1", "NX", "EX", leaseTime);
            if (OK.equalsIgnoreCase(result)) {
                log.info("锁为空直接拿到锁，Thread {} 拿到锁", currentThreadid);
                return true;
            }

//            到目前为止已经超时，则返回false
            time = System.currentTimeMillis() - begin;
            if (time > TimeUnit.SECONDS.toMillis(waitTime)) {
                return false;
            }
            CountDownLatch l = new CountDownLatch(1);
            ScheduledFuture scheduledFuture = executorService.scheduleAtFixedRate(() -> {
                long id = Thread.currentThread().getId();
                String waitResult = jedis.set(finalLockKey, "1", "NX", "EX", leaseTime);
                if (OK.equalsIgnoreCase(waitResult)) {
                    log.info("轮询阶段拿到锁,Thread {} 拿到锁", id);
                    l.countDown();
                    throw new RuntimeException();
                }
            }, 0, 500, TimeUnit.MILLISECONDS);
            boolean await = l.await(TimeUnit.SECONDS.toMillis(waitTime) - time, TimeUnit.MILLISECONDS);
            if (await) {
                log.info("拿锁阶段,Thread {} 拿到锁", currentThreadid);
            } else {
                scheduledFuture.cancel(true);
            }
            return await;
        } catch (InterruptedException e) {
            log.error("FunnelRateLimiter InterruptedException {}", e);
            return false;
        }
    }

这句代码其实是保证原子性的核心，因为redis这条命令 将以前 设置对象&&设置过期时间 两条命令 合并成一条去执行。至于NX与EX，NX|XX决定设置前提，EX|PX 表示过期时间的单位
nxxx参数:

• NX – 只有key不存在才设置成功
• XX – 只有key存在才设置成功

expx参数:

• EX = 秒
• PX = 毫秒

jedis.set(finalLockKey, "1", "NX", "EX", 5);

CountDownLatch

当第一次设置失败以后('NX’表示只有key为空才能设置成功)，接下来就使用了CountDownLatch，如果有多线程编程经验，这个类就不会陌生，CountDownLatch这个类能够使一个线程等待其他线程完成各自的工作后再执行，在这个代码里我们给了对象里的变量count 初始值为1，当在其它线程中执行 l.countDown(), 1降为0，await方法返回 true ，被await方法阻塞的线程会继续向下执行。

await方法也有两个参数，表示当多长时间以后如果没有线程执行l.countDown()，也就是对象内部的count值仍大于0，则返回false，线程也不再阻塞。

CountDownLatch l = new CountDownLatch(1);
l.await(TimeUnit.SECONDS.toMillis(waitTime) - time, TimeUnit.MILLISECONDS);

ScheduledExecutorService

那第一次设置失败后怎样再次尝试获取锁？我们借用ScheduledExecutorService定时服务，它的scheduleAtFixedRate方法可以定时轮询，我们就可以写业务逻辑再次执行jedis.set方法尝试获取锁。下图代码以 500毫秒/次 的频率去反复执行。当获取锁成功，及时将CountDownLatch对象里的变量count 减1，并且抛出异常停止轮询。

ScheduledFuture scheduledFuture = executorService.scheduleAtFixedRate(() -> {
                long id = Thread.currentThread().getId();
                String waitResult = jedis.set(finalLockKey, "1", "NX", "EX", leaseTime);
                if (OK.equalsIgnoreCase(waitResult)) {
                    log.info("轮询阶段拿到锁,Thread {} 拿到锁", id);
                    l.countDown();
                    throw new RuntimeException();
                }
}, 0, 500, TimeUnit.MILLISECONDS);

这时，主线程唤醒继续执行。当然还有另外一种情况，就是之前说的CountDownLatch超时，这时取消定时服务，获取锁失败，返回false。

当tryLock方法返回true或false，我们也就知道获取锁是否成功，在业务代码里也就可以决定是否执行接下来的方法。

unlock

unlock方法就简单一点了，直接del这个key即可。

 /**
     * 释放锁
     *
     * @param lockKey
     */
    public void unlock(String lockKey) {
        final String finalLockKey = packageLockKey(lockKey);
        try (Jedis jedis = jedisPool.getResource()) {
            Long del = jedis.del(finalLockKey);
            if (del > 0) {
                long currentThreadid = Thread.currentThread().getId();
                log.info("FunnelRateLimiter 锁已经释放 Thread {}", currentThreadid);
            }
        } catch (Exception e) {
            log.error("FunnelRateLimiter unlock error {}", e);
        }
    }

总结

当 漏斗限流 和 分布式锁 结合在一起，这个工具类才算可以上生产环境（完整结合代码在上一篇文章）。通过两个篇幅我们由短信限流这个实际需求 引出了分布式锁和漏斗限流的简单实现，要说简单，因为这个确实只满足了基本的需求，但是麻雀虽小五脏俱全，我已经将其运用在生产环境，只要redis不挂，暂时没什么问题。代码后续肯定也有优化的地方，继续努力学习。

最后分享一下掘金的小册，我也是从小册中学习获得灵感，对于公司而言，鼓励大家不要重复造轮子，但是对于个人学习，还是要鼓励大家多造轮子，自己才能进步。