Zookeeper实现分布式锁(基于回调和Watch)

实现思路

1. 创建Zookeeper的连接工具类

package com.zhf.model.zookeeper;
import org.apache.commons.lang.StringUtils;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;

import java.io.IOException;
import java.util.concurrent.CountDownLatch;

public class ZKConnection {

    private static final String ADDRESS = "120.79.165.112:2181," +
            "                              120.79.218.19:2181," +
            "                              47.106.85.178:2181," +
            "                              39.108.13.179:2181";

    private static volatile ZooKeeper zooKeeper;

    public static ZooKeeper getConnection(String path){
        String address = ADDRESS;
        if(!StringUtils.isBlank(path)){
            address += "/" + path;
        }
        CountDownLatch zkConnection = new CountDownLatch(1);
        try {
            if(null == zooKeeper){
                synchronized (ZKConnection.class){
                    if(null == zooKeeper){
                        zooKeeper = new ZooKeeper(address, 5000, watchedEvent ->  {
                            switch (watchedEvent.getState()) {
                                case NoSyncConnected:
                                    System.out.println("Zookeeper NO Sync Connection：" + watchedEvent.getPath());
                                    zkConnection.countDown();
                                    break;
                                case SyncConnected:
                                    //Sync连接上
                                    System.out.println("Zookeeper Sync Connection：" + watchedEvent.getPath());
                                    zkConnection.countDown();
                                    break;
                            }
                        });
                        zkConnection.await();

                    }
                }
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        return zooKeeper;
    }
}

2. 创建Watch和回调类

此处偷懒，Callback和Watch放在同一个类中

package com.zhf.model.zookeeper.lock.v2;

import com.zhf.model.zookeeper.ZKConnection;
import com.zhf.model.zookeeper.lock.v1.ZkLock;
import lombok.Data;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

@Data
public class ZooLock implements AsyncCallback.StringCallback, AsyncCallback.Children2Callback,Watcher, AsyncCallback.StatCallback {

    //Zookeeper连接的跟路径
    public static final String LOCK_PATH = "lock";

    //创建的临时有序节点
    public static final String LOCK_VALUE = "/testLock";

    //执行回调方法的线程
    private Thread thread;

    //创建的路径名称
    private String pathName;

    //获取锁的时候需要使用CountDownLatch将线程阻塞
    private CountDownLatch count = new CountDownLatch(1);

    //zk连接对象
    private ZooKeeper zk = ZKConnection.getConnection(LOCK_PATH);

    public ZooLock(Thread thread){
        this.thread = thread;
    }

    /**
     * 获取锁的方法，现使用的方案为异步获取
     */
    public void tryLock(){
        //创建临时有序节点，然后注册回调方法回调方法1
        zk.create(LOCK_VALUE,thread.getName().getBytes(),
                ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL, this,thread.getName());
        try {
            //阻塞住，等待回调方法取获取锁，获取成功后取消阻塞
            count.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    /**
     * ----------回调方法1------------
     * tryLock方法的回调方法(StringCallBack)
     * 获取所有线程创建的临时有序节点，然后注册获取锁的回调方法
     * @param i index
     * @param s 在哪个节点下创建
     * @param o 回调的提示
     * @param s1 创建的节点名称
     */
    @Override
    public void processResult(int i, String s, Object o, String s1) {
        System.out.println("StringCallBack回调方法参数：i:" + i + "\ts:" + s + "\to:" + o + "\ts1:" + s1);
        //设置创建的pathName,为/testLock0000000032,所以需要去掉斜杠
        pathName = s1.substring(s1.lastIndexOf("/")+1);
        //获取/lock节点下面创建的所有/testLock的有序节点，加入回调方法2
        zk.getChildren("/",false, this,"Get Children");

    }


    /**
     * ----------回调方法2------------
     * 获取临时有序节点的回调方法
     * 获取锁的主要方法
     * @param i index
     * @param s 在哪个节点下创建
     * @param o 回调的提示
     * @param list 获取的临时有序节点的列表
     * @param stat 状态
     */
    @Override
    public void processResult(int i, String s, Object o, List<String> list, Stat stat) {
        System.out.println("Children2CallBack回调方法参数：i:" + i + "\t s:" + s + "\to:" + o + "\tstat:" + stat + "\t list size:" + list.size());
        if(null == list || (null != list && list.size() == 0)){
            return;
        }
        //1.对所有的/testLock节点排序
        Collections.sort(list);
        System.out.println("list:" + list +";\t pathName:" + pathName);
        //2.看当前创建的节点的index
        int index = list.indexOf(pathName);
        //2.1 如果当前节点为第一个，获取锁成功
        if(index == 0){
            try {
                zk.setData("/"+pathName, thread.getName().getBytes(),-1);
                //2.2解除当前线程的阻塞状态，回调成功
                count.countDown();
            } catch (KeeperException e) {
                e.printStackTrace();
                throw new RuntimeException(e);
            } catch (InterruptedException e) {
                e.printStackTrace();
                throw new RuntimeException(e);
            }
            //2.2 如果不是最第一个节点，监听它的前一个节点
        }else{
            System.out.println("index:" + index);
            //看它前面一个节点是否存在，加入回调方法3，然后监控前一个节点的NodeDeleted状态
            //!!!此处有一个问题，就是前一个线程执行完毕，直接结束了，当前线程监控不到前一个线程，目前解决方法为，在回调方法里面获取锁
            zk.exists("/"+list.get(index-1),  this,  this,"Listen pre node");
        }
    }

    /**
     * --------Watch方法-------
     * @param watchedEvent
     */
    @Override
    public void process(WatchedEvent watchedEvent) {
        Event.EventType type = watchedEvent.getType();
        switch (type) {
            case NodeDeleted:
                //如果前一个节点被删除，会通知到后一个节点
                System.out.println("Watch到节点被删除：" + watchedEvent.getPath());
                zk.getChildren("/",false, this,"Notify Pre Node");
                break;
        }
    }

    public void unLock() {
        try {
            System.out.println("删除节点：" + pathName);
            zk.delete("/"+pathName,-1);
        } catch (InterruptedException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        } catch (KeeperException e) {
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    /**
     * ----------回调方法3------------
     * @param i
     * @param s
     * @param o
     * @param stat
     */
    @Override
    public void processResult(int i, String s, Object o, Stat stat) {
        //不管监控成功或者失败，都继续获取锁
        zk.getChildren("/",false, this,"Notify Pre Node");
    }
}

3. 测试类

package com.zhf.model.zookeeper.lock;

import com.zhf.model.zookeeper.lock.v1.ZkLock;
import com.zhf.model.zookeeper.lock.v2.ZooLock;

import java.util.concurrent.CountDownLatch;

public class Test {

    public int num = 0;

    public void testLock(){
        CountDownLatch count = new CountDownLatch(20);
        for (int i = 0; i < 20; i++) {
            new Thread(()->{
                ZooLock lock = new ZooLock(Thread.currentThread());
                System.out.println("Thread start: " + Thread.currentThread().getName());
                lock.tryLock();
                num ++;
                System.out.println("Thread: " +Thread.currentThread().getName()+ " run------------");
                lock.unLock();
                count.countDown();
            }).start();
        }
        try {
            count.await();
            System.out.println("Count Num: = "+num);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }

    }

    public static void main(String[] args) throws InterruptedException {
        Test test = new Test();
        test.testLock();
        Thread.sleep(10000);

    }
}