Java分布式锁之Zookeeper

上一篇有基于redis的分布式锁，大家有兴趣也可以去看看https://blog.csdn.net/lyz812672598/article/details/84994489

如果需要转载，请贴明出处，谢谢大家！

/**
 * <p>ZK分布式锁</p>
 *
 * @author liyz
 * @version 1.0.0
 * @date 2018/12/13 0013 16:27
 */
@Slf4j
@Configuration
@ConditionalOnExpression("'${zookeeper.lock.url}'.length() > 0")
public class ZookeeperLock {

    /**
     * zookeeper host
     */
    @Value("zookeeper.lock.url")
    private String zkUrl;

    /**
     * 锁节点默认父节点
     */
    private static final String DEFAULT_PATH = "/zookeeper.lock";

    /**
     * map默认容量
     */
    private static final int DEFAULT_MAP_CAPACITY = 2 << 14;

    /**
     * 默认session时间
     */
    private static final int DEFAULT_SESSION_TIME_OUT = 10000;

    /**
     * 默认connection时间
     */
    private static final int DEFAULT_CONNECTION_TIME_OUT = 10000;

    /**
     * 默认持有锁的有效时间
     */
    private static final int DEFAULT_EXPIRE = 60;

    /**
     * 节点信息容器:单号对应的节点的信息
     * <p>
     *     key:单号
     *     value:临时节点信息
     * </p>
     */
    private Map<String, LockModel> noteMap = new ConcurrentHashMap<>(DEFAULT_MAP_CAPACITY);

    /**
     * zk客户端
     */
    private ZkClient zkClient;

    @PostConstruct
    public void initLock() {
        zkClient = new ZkClient(zkUrl, DEFAULT_SESSION_TIME_OUT, DEFAULT_CONNECTION_TIME_OUT,
                new SerializableSerializer());
    }

    /**
     * <P>
     *     1.尝试获取锁，直到超时为止
     *     2.超时分为两种：没有获得锁超时；获得锁业务没有处理完超时
     *     3.方法 tryLock、lock适用于临时节点的父节点会经常重复，比如：账户-每个人下每个钱的类型就是一个父节点，
     *       不适用于类似订单锁，每个节点基本都不一致，这样会创建太多的永久节点
     *     4.如果是订单锁，可以用更加暴力的orderLock方法，舍弃了watch机制，公平特性
     * </P>
     *
     * @param lockKey
     * @param orderNo
     * @return
     */
    public void tryLock(String lockKey, String orderNo) {
        tryLock(lockKey, orderNo, DEFAULT_EXPIRE);
    }

    public void tryLock(String lockKey, String orderNo, String lockPath) {
        tryLock(lockKey, orderNo, lockPath, DEFAULT_EXPIRE);
    }

    public void tryLock(String lockKey, String orderNo, int expire) {
        tryLock(lockKey, orderNo, DEFAULT_PATH, expire);
    }

    public void tryLock(String lockKey, String orderNo, String lockPath, int expire) {
        tryLock(lockKey, orderNo, lockPath, expire, true);
    }

    public void tryLock(String lockKey, String orderNo, String lockPath, int expire, boolean isFirst) {
        log.info("*************当前有{}个节点正在等待", noteMap.size());
        initLockPath(lockPath);

        if (!noteMap.containsKey(orderNo)) {
            LockModel lockModel = LockModel.builder().expire(expire).build();
            Object obj = noteMap.putIfAbsent(orderNo, lockModel);
            if (obj != null) {
                throw new ZookeeperLockException(ZkLockEnums.ORDER_DEALING);
            }
        } else {
            if (isFirst) {
                throw new ZookeeperLockException(ZkLockEnums.ORDER_DEALING);
            }
        }

        /**
         * <p>
         *     如果需要加入尝试获取锁的超时时间，可以加一个参数，在这里判断
         * </p>
         */
        if (!firstChild(lockKey, orderNo, lockPath)) {
            configureZkWatch(lockKey, orderNo);
            tryLock(lockKey, orderNo, lockPath, expire, false);
        } else {
            log.info("**************线程:{},获得分布式锁:{}", Thread.currentThread().getName(), lockKey);
        }
    }

    /**
     * <p>
     *     1.获取zk锁，立马返回结果
     *     2.就算返回false，也要进行unlock操作
     * </p>
     *
     * @param lockKey
     * @param lockValue
     * @return true:success; false:fail
     */
    public boolean lock(String lockKey, String lockValue) {
        return lock(lockKey, lockValue, DEFAULT_PATH);
    }

    public boolean lock(String lockKey, String lockValue, String lockPath) {
        return lock(lockKey, lockValue, lockPath, DEFAULT_EXPIRE);
    }

    public boolean lock(String lockKey, String lockValue, int expire) {
        return lock(lockKey, lockValue, DEFAULT_PATH, expire);
    }

    public boolean lock(String lockKey, String orderNo, String lockPath, int expire) {
        log.info("*************当前有{}个节点正在等待", noteMap.size());
        initLockPath(lockPath);

        if (!noteMap.containsKey(orderNo)) {
            LockModel lockModel = LockModel.builder().expire(expire).build();
            Object obj = noteMap.putIfAbsent(orderNo, lockModel);
            if (obj != null) {
                throw new ZookeeperLockException(ZkLockEnums.ORDER_DEALING);
            }
        }

        return firstChild(lockKey, orderNo, lockPath);
    }

    /**
     * 订单锁，简单暴力版，大家其实可以优化优化的
     * 注：最后也需要调用unLock()方法
     * 
     * @param orderNo
     */
    public void orderLock(String orderNo) {
        orderLock(orderNo, DEFAULT_PATH);
    }
    
    public void orderLock(String orderNo, String lockPath) {
        orderLock(orderNo, lockPath, DEFAULT_EXPIRE);
    }
    
    public void orderLock(String orderNo, int expire) {
        orderLock(orderNo, DEFAULT_PATH, expire);
    }

    public void orderLock(String orderNo, String lockPath, int expire) {
        log.info("***************线程:{}，尝试获取锁:{}", Thread.currentThread().getName(), orderNo);
        initLockPath(lockPath);
        
        long timeOut = expire * 1000000;
        long now = System.nanoTime();
        String zkPath = lockPath + "/" + orderNo;
        while ((System.nanoTime() - now) < timeOut) {
            if (!zkClient.exists(zkPath)) {
                try {
                    zkClient.createEphemeral(zkPath);
                    log.info("***************线程:{}，获取锁成功:{}", Thread.currentThread().getName(), orderNo);
                    return;
                } catch (ZkNodeExistsException e) {}
            }

            //休眠一小下
            seleep(10, 50000);
        }

        throw new ZookeeperLockException(ZkLockEnums.TIME_OUT);
    }

    /**
     * 休眠
     *
     * @param millis
     * @param nanos
     */
    private void seleep(int millis, int nanos) {
        try {
            Thread.sleep(millis, new Random().nextInt(nanos));
        } catch (Exception e) {
            log.info("***************获取锁休眠失败:{}", e.getMessage());
        }
    }

    /**
     * 判断该临时节点是否是第一个子节点
     *
     * @param lockKey
     * @param orderNo
     * @param lockPath
     * @return
     */
    private boolean firstChild(String lockKey, String orderNo, String lockPath) {
        String currentPath = noteMap.get(orderNo).getCurrentPath();
        String zkPath = lockPath + "/" + lockKey;
        if (StringUtils.isBlank(currentPath)) {
            if (!zkClient.exists(zkPath)) {
                /**
                 * <p>
                 *     1.这里写入时间是用来回收这里创建的永久节点
                 *     2.定时任务扫描可以删除的节点
                 *     3.如果有更好的方法可以提供出来
                 *     注：也可以根据节点zk提供的节点信息获取时间戳，这里根据自己的实际情况可以适当修改
                 * </p>
                 */
                try {
                    zkClient.createPersistent(zkPath, System.currentTimeMillis());
                } catch (ZkNodeExistsException e) {
                    zkClient.writeData(zkPath, System.currentTimeMillis());
                }
            } else {
                zkClient.writeData(zkPath, System.currentTimeMillis());
            }
            currentPath = zkClient.createEphemeralSequential(zkPath + "/", orderNo);
            noteMap.get(orderNo).setCurrentPath(currentPath);
            log.info("*************线程:{},创建了临时节点:{}", Thread.currentThread().getName(), currentPath);
        }

        List<String> childrenList = zkClient.getChildren(zkPath);
        //排序为了更快获得最小的节点
        Collections.sort(childrenList);
        if (currentPath.equals(zkPath + "/" + childrenList.get(0))) {
            return true;
        }
        //得到该临时节点的前一个节点
        int position = Collections.binarySearch(childrenList, currentPath.substring(zkPath.length() + 1));
        noteMap.get(orderNo).setBeforePath(zkPath + "/" + childrenList.get(position - 1));
        return false;
    }

    /**
     * 对前一个节点增加一个watch机制
     * 利用发令枪和zk的watch机制，实现redis没有的轮询方式
     *
     * @param lockKey
     * @param orderNo
     */
    private void configureZkWatch(String lockKey, String orderNo) {
        LockModel lockModel = noteMap.get(orderNo);
        IZkDataListener listener = new IZkDataListener() {

            @Override
            public void handleDataChange(String dataPath, Object data) throws Exception {

            }

            @Override
            public void handleDataDeleted(String dataPath) throws Exception {
                log.info("*****************线程:{},捕获到节点删除:{}", Thread.currentThread().getName(), dataPath);
                if (lockModel.getCdl() != null) {
                    lockModel.getCdl().countDown();
                }
            }
        };
        //对前一个节点加入watch机制
        zkClient.subscribeDataChanges(lockModel.getBeforePath(), listener);

        if(zkClient.exists(lockModel.getBeforePath())) {
            lockModel.setCdl(new CountDownLatch(1));
            try {
                lockModel.getCdl().await(lockModel.getExpire(), TimeUnit.SECONDS);
            } catch (Exception e) {
                log.error("****************线程:{},发令枪发生异常:{}", Thread.currentThread().getName(), e.getMessage());
            }
            if (zkClient.exists(lockModel.getBeforePath())) {
                unlock(orderNo);
                throw new ZookeeperLockException(ZkLockEnums.TIME_OUT);
            }
        }

        zkClient.unsubscribeDataChanges(lockModel.getBeforePath(), listener);
    }

    /**
     * <p>
     *      由于我们只用了该单例，所以我们要主动删除该临时节点，但是我们并没有删除其父永久节点，如果有人需要删除，
     *      可以进行修改，优化，如果不删，也可以用别的方式删除，上述提到定时任务，有好的想法可以提出来
     * </p>
     *
     * @param orderNo
     */
    public void unlock(String orderNo) {
        if (noteMap.containsKey(orderNo)) {
            LockModel lockModel = noteMap.get(orderNo);
            String currentPath = lockModel.getCurrentPath();
            if (zkClient.exists(currentPath)) {
                zkClient.delete(currentPath);
            }
            noteMap.remove(orderNo);
            log.info("******************线程:{}--锁释放:{}", Thread.currentThread().getName(), currentPath);
        }
    }

    /**
     * 该方法只能释放orderLock方法产生的锁
     * 
     * @param orderNo
     */
    public void orderUnlock(String orderNo) {
        orderUnlock(orderNo, DEFAULT_PATH);
    }
    
    public void orderUnlock(String orderNo, String lockPath) {
        String zkPath = lockPath + "/" + orderNo;
        if (zkClient.exists(zkPath)) {
            zkClient.delete(zkPath);
            log.info("******************线程:{}--锁释放:{}", Thread.currentThread().getName(), orderNo);
        }
    }

    /**
     * 初始化zk锁根节点
     *
     * @param lockPath
     */
    private void initLockPath(String lockPath) {
        if (zkClient.exists(lockPath)) {
            return;
        }
        try {
            zkClient.createPersistent(lockPath);
        } catch (ZkNodeExistsException e) {}
    }

    /**
     * lock内容
     */
    @Data
    @Builder
    class LockModel implements Serializable {
        private static final long serialVersionUID = -5182247768216621018L;

        private CountDownLatch cdl;

        private String beforePath;

        private String currentPath;

        private int expire;
    }

    @Data
    class ZookeeperLockException extends RuntimeException implements Serializable {
        private static final long serialVersionUID = -2749187337703750042L;

        private String errorMsg;

        private ZkLockEnums zkLockEnums;

        public ZookeeperLockException(String errorMsg) {
            super(errorMsg);
            this.errorMsg = errorMsg;
        }

        public ZookeeperLockException(String errorMsg, Throwable cause) {
            super(errorMsg, cause);
            this.errorMsg = errorMsg;
        }

        public ZookeeperLockException(ZkLockEnums zkLockEnums) {
            super(zkLockEnums.getMsg());
            this.zkLockEnums = zkLockEnums;
            this.errorMsg = zkLockEnums.getMsg();
        }
    }

    enum ZkLockEnums {

        ORDER_DEALING("订单正在处理"),
        TIME_OUT("超时请重新尝试");

        ZkLockEnums(String msg) {
            this.msg = msg;
        }

        private String msg;

        public String getMsg() {
            return msg;
        }
    }
}

两者各有优点，各有缺点，大家可以根据自己的实际需求，并发量来选取对应的锁。