zookeeper原生API实现分布式锁

原生API实现分布式锁

原理

每个客户端都往zookeeper上创建一个临时有序节点，注册完后，取/locks下所有节点中最小的节点，则创建该最小节点的客户端拿到锁。如果当前最小节点是/lock_seq1，那么/lock_seq2边监听/lock_seq1边阻塞等待，当/lock_seq1被删除后，/lock_seq2节点的客户端获得锁，以此类推。

代码

public class DistributedLock implements Lock,Watcher{

    private ZooKeeper zk=null;
    private String ROOT_LOCK="/locks";//定义根节点
    private String WAIT_LOCK;//等待前一个锁
    private String CURRENT_LOCK;//当前的锁

    private CountDownLatch countDownLatch;

    public DistributedLock() {
        try {
            zk=new ZooKeeper("192.168.3.140:2181",4000,this);
            //判断根节点是否存在
            Stat stat=zk.exists(ROOT_LOCK,false);
            if(stat==null){
                zk.create(ROOT_LOCK,"0".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

            }
        } catch (Exception e) {
            e.printStackTrace();
        }

    }

    public void lock() {
        if(this.tryLock()){//如果获得锁成功
            System.out.println(Thread.currentThread().getName()+"->"+CURRENT_LOCK+"->获得锁成功");
            return;
        }
        try {
            waitForLock(WAIT_LOCK);//没有获得锁，继续等待
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

    }

    private boolean waitForLock(String prev) throws KeeperException, InterruptedException {
        Stat stat=zk.exists(prev,true);//监听当前节点的上一个节点
        if(stat!=null){
            System.out.println(Thread.currentThread().getName()+"->等待锁"+ROOT_LOCK+"/"+prev+"释放");
            countDownLatch=new CountDownLatch(1);
            countDownLatch.await();
            System.out.println(Thread.currentThread().getName()+"->获得锁成功");
        }
        return true;
    }

    public void lockInterruptibly() throws InterruptedException {

    }

    public boolean tryLock() {
        try {
            //创建临时有序节点
            CURRENT_LOCK=zk.create(ROOT_LOCK+"/","0".getBytes(),
                    ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            System.out.println(Thread.currentThread().getName()+"->"+CURRENT_LOCK+" 尝试竞争锁");
            //例：children=[0000000260...000000000269]
            List<String> children=zk.getChildren(ROOT_LOCK,false);//获取根节点下所有子节点，这里的watcher设为false，因为不需要去监听
            //例：sortedSet=[/locks/0000000260.../locks/0000000260]
            SortedSet<String> sortedSet=new TreeSet();//定义一个有序集合进行排序
            for(String child:children){
                sortedSet.add(ROOT_LOCK+"/"+child);
            }
            //例：firstNode=/locks/0000000260
            String firstNode=sortedSet.first();//获得当前所有子节点中最小的节点
            //获取比当前节点更小的节点
            //例：CURRENT_LOCK=/locks/0000000262 lessThenMe=[/locks/0000000260,/locks/0000000261]
            SortedSet<String> lessThenMe=((TreeSet<String>)sortedSet).headSet(CURRENT_LOCK);
            if(CURRENT_LOCK.equals(firstNode)){//当前的节点和子节点中最小的节点比较，如果相等，则获得锁
                return true;
            }
            if(!lessThenMe.isEmpty()){
                //例：WAIT_LOCK=/locks/0000000261
                WAIT_LOCK=lessThenMe.last();//获得比当前节点更小的最后一个节点，设置给WAIT_LOCK
            }
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return false;
    }

    public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
        return false;
    }

    public void unlock() {
        System.out.println(Thread.currentThread().getName()+"->释放锁"+CURRENT_LOCK);
        try {
            zk.delete(CURRENT_LOCK,-1);
            CURRENT_LOCK=null;
            zk.close();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (KeeperException e) {
            e.printStackTrace();
        }
    }

    public Condition newCondition() {
        return null;
    }

    public void process(WatchedEvent watchedEvent) {
        if(this.countDownLatch!=null){
            this.countDownLatch.countDown();
        }
    }
}

测试运行

//开启10个线程，并且同时往zookeeper上注册节点
public class DistributedLockTest {

    public static void main(String[] args) throws IOException {
        final CountDownLatch countDownLatch=new CountDownLatch(10);
        for(int i=0;i<10;i++){
            new Thread(new Runnable() {
                public void run() {
                    try {
                        countDownLatch.await();
                        //10个线程一起并发执行以下代码
                        DistributedLock distributedLock=new DistributedLock();
                        distributedLock.lock();//当前线程获得锁
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            },"Thread-"+i).start();
            countDownLatch.countDown();
        }
        System.in.read();
    }
}

执行结果
Thread-5->/locks/0000000301 尝试竞争锁
Thread-0->/locks/0000000300 尝试竞争锁
Thread-3->/locks/0000000303 尝试竞争锁
Thread-7->/locks/0000000305 尝试竞争锁
Thread-9->/locks/0000000306 尝试竞争锁
Thread-6->/locks/0000000304 尝试竞争锁
Thread-1->/locks/0000000309 尝试竞争锁
Thread-2->/locks/0000000302 尝试竞争锁
Thread-4->/locks/0000000307 尝试竞争锁
Thread-8->/locks/0000000308 尝试竞争锁
Thread-0->/locks/0000000300->获得锁成功
Thread-4->等待锁/locks//locks/0000000306释放
Thread-3->等待锁/locks//locks/0000000302释放
Thread-7->等待锁/locks//locks/0000000304释放
Thread-5->等待锁/locks//locks/0000000300释放
Thread-1->等待锁/locks//locks/0000000308释放
Thread-8->等待锁/locks//locks/0000000307释放
Thread-9->等待锁/locks//locks/0000000305释放
Thread-2->等待锁/locks//locks/0000000301释放
Thread-6->等待锁/locks//locks/0000000303释放