zookeeper实现分布式锁（学习zookeeper）

1、手写这种重复轮子的原因

学习zookeeper的过程中，用来熟悉zookeeper的操作及原理

2、上代码

该代码不可用在生产环境，因为是学习用途，逻辑并不谨慎（重要、重要、重要）

实现逻辑，代码注释已经写得比较清楚了，就不赘述了（我很懒）

package com.example.demo22.zookeerper;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.List;
import java.util.concurrent.locks.LockSupport;
import java.util.stream.Collectors;

/**
 * @author anxu
 * @date 2023/2/9 9:23
 */
public class Lock {
    private static final String NODE_PATH = "/anxu/lock";
    private static final String PARENT_NODE_PATH = "/anxu";

    /**
     * 当前获取到锁的Node（排在最前面的NODE）
     */
    private String currentNode;
    /**
     * 当前获取锁的线程
     */
    private Thread currentThread;

    /**
     * zooKeeper对象
     */
    private ZooKeeper zooKeeper;

    public Lock() {
        this.zooKeeper = ZookeeperFactory.getZookeeper();
    }

    /**
     * 加锁
     *
     * @return
     */
    public boolean lock() throws KeeperException, InterruptedException {
        /**
         * 加锁流程
         * 1、申请在 /anxu节点下创建有序节点 lock
         * 2、查看 /anxu下的所有子节点，判断自己的节点是否是最前面的
         *  2。1、如果是，提示获取锁成功
         *  2.2、如果不是，则watch排在自己前面的节点，并park自己，等待唤醒
         *
         */
        //检查是否为当前线程， 支持 可重入锁
        if (Thread.currentThread() == this.currentThread) {
            //获取重入次数
            int count = Integer.valueOf(new String(zooKeeper.getData(this.currentNode, true, new Stat())));
            //获取版本号
            int version = zooKeeper.exists(this.currentNode, true).getVersion();
            //更新重入次数
            zooKeeper.setData(this.currentNode, String.valueOf(count + 1).getBytes(), version);
            return true;
        }
        //创建有序节点
        String node = zooKeeper.create(NODE_PATH, "1".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        return this.tryLock(node);
    }

    private boolean tryLock(String node) throws KeeperException, InterruptedException {
        //查看父节点下的所有有序子节点 自己是否排在最前面
        List<String> children = zooKeeper.getChildren(PARENT_NODE_PATH, true);
        String minNode = getMinNode(children);
        //检查自己是否是排在最前面的节点
        if (node.equals(minNode)) {
            //设置当前占有锁的线程
            this.currentThread = Thread.currentThread();
            this.currentNode = node;
            return true;
        }

        //获取排在自己前面的节点名称（路径）
        String preNode = this.getPreNode(children, node);
        //记录当前线程，便于唤醒
        Thread currentThread = Thread.currentThread();
        //在前一个节点上注册watch监听
        zooKeeper.getData(preNode, watchedEvent -> {
            if (watchedEvent.getType() == Watcher.Event.EventType.NodeDeleted) { //监听节点删除
                LockSupport.unpark(currentThread); //如果前一个节点被删除，说明锁已经放开，我们唤醒线程继续尝试获取锁
            }
        }, new Stat());
        LockSupport.park();
        return tryLock(node);
    }

    /**
     * 解锁
     *
     * @return
     */
    public boolean unlock() throws KeeperException, InterruptedException {
        //获取重入次数
        int count = Integer.valueOf(new String(zooKeeper.getData(this.currentNode, true, new Stat())));
        count = count - 1;
        //获取版本号
        int version = zooKeeper.exists(this.currentNode, true).getVersion();
        if (count > 0) {
            //更新 重入次数
            zooKeeper.setData(this.currentNode, String.valueOf(count).getBytes(), version);
        } else {
            //删除节点
            String currentNode = this.currentNode;
            this.currentNode = null;
            this.currentThread = null;
            zooKeeper.delete(currentNode, version);
        }
        return true;
    }

    /**
     * 获取集合中的排在最前面的节点
     *
     * @param list
     * @return
     */
    private String getMinNode(List<String> list) {
        return PARENT_NODE_PATH + "/" + list.stream().sorted().collect(Collectors.toList()).get(0);
    }

    /**
     * 获取节点的前一个节点
     *
     * @param list
     * @param currentNode
     * @return
     */
    private String getPreNode(List<String> list, String currentNode) {
        List<String> sortList = list.stream().map(it -> PARENT_NODE_PATH + "/" + it).sorted().collect(Collectors.toList());
        int index = sortList.indexOf(currentNode);
        return sortList.get(index - 1);
    }
}

3、运行效果

测试代码

package com.example.demo22.zookeerper;

import org.apache.zookeeper.KeeperException;

/**
 * @author anxu
 * @date 2023/1/10 17:21
 */
public class Test {
    private static int sum = 0;
    public static void main(String[] args) throws InterruptedException, KeeperException {
        Lock lock = new Lock();
        Thread.sleep(1000);
        for (int j = 0; j < 3; j++) {
            new Thread(() -> {
                for (int k = 0; k < 1000; k++) {
                    try {
                        lock.lock();
                        sum++;
                    } catch (KeeperException e) {
                        e.printStackTrace();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    } finally {
                        try {
                            lock.unlock();
                        } catch (KeeperException e) {
                            e.printStackTrace();
                        } catch (InterruptedException e) {
                            e.printStackTrace();
                        }
                    }
                }
            }).start();
        }
        Thread.sleep(20000);
        System.out.println(sum);
    }
}

执行结果