Zookeeper实现简易的分布式锁

Zookeeper实现简易的分布式锁

利用Zookeeper实现分布式锁的主要思路就是利用了Zookeeper提供的znode创建临时序列功能

首页使用Zookeeper客户端创建节点

这里我们利用Zookeeper官方提供的客户端代码，下面就是maven依赖,这里使用的是3.5.5这个版本

<!-- https://mvnrepository.com/artifact/org.apache.zookeeper/zookeeper -->
<dependency>
    <groupId>org.apache.zookeeper</groupId>
    <artifactId>zookeeper</artifactId>
    <version>3.5.5</version>
</dependency>

我们来看看正常情况下如何创建一个节点的

@Test
public void creatTest()throws Exception{
    /**
     * 下面创建Zookeeper 客户端代码貌似是一个定式保证一定获取到了Zookeeper的连接
     */
    CountDownLatch latch = new CountDownLatch(1);
    ZooKeeper zk = new ZooKeeper("172.16.12.255:2181", 3000, new Watcher() {
        @Override
        public void process(WatchedEvent event) {
            if(event.getState() == Watcher.Event.KeeperState.SyncConnected){
                latch.countDown();
            }
        }
    });
    latch.await();
    //这里代码为创建一个路径为 /helloworld 值为 helloworld 的持久化节点
    String nodeID = zk.create("/helloworld", "helloworld".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,CreateMode.PERSISTENT);
    System.out.println(nodeID);
    
}

接着我们利用Zookeeper自带的脚本 在Zookeeper bin目录下

sh bin/zkCli.sh

#下面是使用脚本执行结果

[zk: localhost:2181(CONNECTED) 103] ls /
[helloworld, cluster, MyLocks, MyFirstZnode, brokers, zookeeper, admin, isr_change_notification, log_dir_event_notification, controller_epoch, MyFirstLock, consumers, latest_producer_id_block, config, zk_test]
[zk: localhost:2181(CONNECTED) 104] get /helloworld
helloworld
cZxid = 0x11a
ctime = Mon Jul 01 03:34:24 EDT 2019
mZxid = 0x11a
mtime = Mon Jul 01 03:34:24 EDT 2019
pZxid = 0x11a
cversion = 0
dataVersion = 0
aclVersion = 0
ephemeralOwner = 0x0
dataLength = 10
numChildren = 0
[zk: localhost:2181(CONNECTED) 105]

Zookeeper 节点类型

org.apache.zookeeper.CreateMode 这是一个枚举类型包含以下几个类型

PERSISTENT  //持久化节点，当客户端端失联时不会自动删除
PERSISTENT_SEQUENTIAL //持久化节点，当客户端端失联时不会自动删除,并且创建时带有自增长序列
EPHEMERAL  //临时性节点，当客户端失联时会自动删除
EPHEMERAL_SEQUENTIAL //临时性节点，当客户端失联时会自动删除，并且创建时带有自增长序
CONTAINER  //容器节点，当容器的最后一个子级被删除后容器节点就将作为被删除的候选选项
PERSISTENT_WITH_TTL  //改节点当客户端失联时不会自动删除，但是在给定的TTL内未修改，当米有子级节点时会被删除
PERSISTENT_SEQUENTIAL_WITH_TTL //同上多一个序列

进入正文编写一个简易的锁

编写的思路大致为：

1. 创建一个Zookeeper路径作为存放我们的锁资源，该资源为PERSISTENT持久化类型的，不能因为客户端失联就删除了
2. 当申请锁资源时会去创建一个EPHEMERAL_SEQUENTIAL类型的节点，临时性节点当客户端失联自动释放资源确保其他锁等待者可以获取到锁
3. 创建临时锁资源后，去获取锁资源下的所有子级，判断自己创建的资源是不是编号最小的那个,如果自己是最小值则获取到锁
4. 判断结果自己不是最小值则获取到排在自己前方的一个锁，并且向Zookeeper往该节点注册一个监听器当其释放时通知
5. 释放锁，将当前线程的锁ID拿到并且使用Zookeeper删除该节点

以上就是实现的思路，下面就是实现代码了

public class MyDistributedLock {

    private static final String rootPath = "/MyLocks";

    private static final int sessionTimeout = 3000;

    private ThreadLocal<LockNode> lockNodes = new ThreadLocal<>();

    private final String lockName;

    private final String fullLockName;

    private final String lockPath;

    private ZooKeeper zk;

    private CountDownLatch connectedSignal = new CountDownLatch(1);

    private final static byte[] data = new byte[0];

    public MyDistributedLock(String lockName, String zkUrl) {
        this.lockName = lockName;
        this.lockPath = rootPath + "/" + lockName;
        this.fullLockName = lockPath + "/id-";
        try {
            zk = new ZooKeeper(zkUrl, sessionTimeout, event -> {
                if (event.getState() == Watcher.Event.KeeperState.SyncConnected) {
                    connectedSignal.countDown();
                }
            });
            connectedSignal.await();
            Stat stat = zk.exists(rootPath, false);
            if (stat == null) {
                zk.create(rootPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            stat = zk.exists(lockPath, false);
            if (stat == null) {
                zk.create(lockPath, data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (Exception e) {
            throw new RuntimeException();
        }
    }

    public void lock() {
        LockNode lockNode = lockNodes.get();
        if (lockNode == null) {
            lockNew();
        } else {
            lockNode.count++;
        }
    }

    private void lockNew(){
        try {
            String currentNodeId = zk.create(fullLockName, data,
                    ZooDefs.Ids.OPEN_ACL_UNSAFE,CreateMode.EPHEMERAL_SEQUENTIAL);

            LockNode lockNode = new LockNode();
            lockNode.id = currentNodeId;


            List<String> children = zk.getChildren(this.lockPath, false);
            TreeSet<String> sortedNodes = new TreeSet<>();
            for(String node :children) {
                sortedNodes.add(lockPath +"/" +node);
            }
            String first = sortedNodes.first();
            if(first.equals(currentNodeId)){
                lockNode.count ++;
                this.lockNodes.set(lockNode);
                return;
            }else {
                String pre = sortedNodes.lower(currentNodeId);
                CountDownLatch currentLatch = new CountDownLatch(1);
                LockWatcher lockWatcher = new LockWatcher(currentLatch);
                Stat exists = zk.exists(pre, lockWatcher);
                if(exists != null){
                    currentLatch.await();
                }
                lockNode.count ++;
                this.lockNodes.set(lockNode);

            }
        }catch (Exception e){
            throw new RuntimeException(e);
        }
    }

    public void unlock(){
        LockNode currentNode = this.lockNodes.get();
        if(currentNode == null){
            throw new RuntimeException("No lock to release");
        }
        currentNode.count --;
        if(currentNode.count == 0){
            this.lockNodes.remove();
            try {
                System.out.println(currentNode.id);
                zk.delete(currentNode.id,-1);
            }catch (Exception e){
                e.printStackTrace();
            }
        }

    }

    @Data
    static class LockNode {
        String id;

        int count = 0;
    }

    static class LockWatcher implements Watcher {

        private final CountDownLatch latch;

        LockWatcher(CountDownLatch latch) {
            this.latch = latch;
        }

        @Override
        public void process(WatchedEvent event) {
            if (event.getType() == Event.EventType.NodeDeleted) {
                latch.countDown();
            }
        }
    }

}

下面是针对以上代码的测试代码

public class MyDistributedLockTest {

    private static MyDistributedLock LOCK = new MyDistributedLock("test","172.16.12.255:2181");

    @Test
    public void test1(){
        LOCK.lock();
        System.out.println("getLock");
        LOCK.unlock();
        System.out.println("release");
    }

    @Test
    public void test2()throws Exception{
        ZooKeeper zooKeeper = new ZooKeeper("172.16.12.255:2181", 30000, event -> System.out.println(event));
        zooKeeper.create("/MyLocks/test/ids-",new byte[0],ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        Thread.sleep(3000);
    }

    private static class R1 implements Runnable{

        private final CountDownLatch latch;

        public R1(CountDownLatch latch) {
            this.latch = latch;
        }

        @Override
        public void run() {
            LOCK.lock();
            System.out.println(Thread.currentThread().getName() + ":getLock");
            LOCK.unlock();
            latch.countDown();
        }
    }

    @Test
    public void test3()throws Exception{
        CountDownLatch latch = new CountDownLatch(2);
        Thread t1 = new Thread(new R1(latch));
        Thread t2 = new Thread(new R1(latch));
        t1.start();
        t2.start();
        latch.await();
        System.out.println("get it");
    }
}