使用zookeeper实现分布式锁

package com.concurrent;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;
 
public class ConcurrentTest {
    private CountDownLatch startSignal = new CountDownLatch(1);//开始阀门
    private CountDownLatch doneSignal = null;//结束阀门
    private CopyOnWriteArrayList<Long> list = new CopyOnWriteArrayList<Long>();
    private AtomicInteger err = new AtomicInteger();//原子递增
    private ConcurrentTask[] task = null;
     
    public ConcurrentTest(ConcurrentTask... task){
        this.task = task;
        if(task == null){
            System.out.println("task can not null");
            System.exit(1);
        }
        doneSignal = new CountDownLatch(task.length);
        start();
    }
    /**
     * @param args
     * @throws ClassNotFoundException 
     */
    private void start(){
        //创建线程，并将所有线程等待在阀门处
        createThread();
        //打开阀门
        startSignal.countDown();//递减锁存器的计数，如果计数到达零，则释放所有等待的线程
        try {
            doneSignal.await();//等待所有线程都执行完毕
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        //计算执行时间
        getExeTime();
    }
    /**
     * 初始化所有线程，并在阀门处等待
     */
    private void createThread() {
        long len = doneSignal.getCount();
        for (int i = 0; i < len; i++) {
            final int j = i;
            new Thread(new Runnable(){
                public void run() {
                    try {
                        startSignal.await();//使当前线程在锁存器倒计数至零之前一直等待
                        long start = System.currentTimeMillis();
                        task[j].run();
                        long end = (System.currentTimeMillis() - start);
                        list.add(end);
                    } catch (Exception e) {
                        err.getAndIncrement();//相当于err++
                    } 
                    doneSignal.countDown();
                }
            }).start();
        }
    }
    /**
     * 计算平均响应时间
     */
    private void getExeTime() {
        int size = list.size();
        List<Long> _list = new ArrayList<Long>(size);
        _list.addAll(list);
        Collections.sort(_list);
        long min = _list.get(0);
        long max = _list.get(size-1);
        long sum = 0L;
        for (Long t : _list) {
            sum += t;
        }
        long avg = sum/size;
        System.out.println("min: " + min);
        System.out.println("max: " + max);
        System.out.println("avg: " + avg);
        System.out.println("err: " + err.get());
    }
     
    public interface ConcurrentTask {
        void run();
    }
 
}

package com.concurrent;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
 
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;
 
public class DistributedLock implements Lock, Watcher{
    private ZooKeeper zk;
    private String root = "/locks";//根
    private String lockName;//竞争资源的标志
    private String waitNode;//等待前一个锁
    private String myZnode;//当前锁
    private CountDownLatch latch;//计数器
    private int sessionTimeout = 30000;
    private List<Exception> exception = new ArrayList<Exception>();
     
    /**
     * 创建分布式锁,使用前请确认config配置的zookeeper服务可用
     * @param config 127.0.0.1:2181
     * @param lockName 竞争资源标志,lockName中不能包含单词lock
     */
    public DistributedLock(String config, String lockName){
        this.lockName = lockName;
        // 创建一个与服务器的连接
         try {
            zk = new ZooKeeper(config, sessionTimeout, this);
            Stat stat = zk.exists(root, false);
            if(stat == null){
                // 创建根节点
                zk.create(root, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE,CreateMode.PERSISTENT); 
            }
        } catch (IOException e) {
            exception.add(e);
        } catch (KeeperException e) {
            exception.add(e);
        } catch (InterruptedException e) {
            exception.add(e);
        }
    }
 
    /**
     * zookeeper节点的监视器
     */
    public void process(WatchedEvent event) {
        if(this.latch != null) {  
            this.latch.countDown();  
        }
    }
     
    public void lock() {
        if(exception.size() > 0){
            throw new LockException(exception.get(0));
        }
        try {
            if(this.tryLock()){
                System.out.println("Thread " + Thread.currentThread().getId() + " " +myZnode + " get lock true");
                return;
            }
            else{
                waitForLock(waitNode, sessionTimeout);//等待锁
            }
        } catch (KeeperException e) {
            throw new LockException(e);
        } catch (InterruptedException e) {
            throw new LockException(e);
        } 
    }
 
    public boolean tryLock() {
        try {
            String splitStr = "_lock_";
            if(lockName.contains(splitStr))
                throw new LockException("lockName can not contains \\u000B");
            //创建临时子节点
            myZnode = zk.create(root + "/" + lockName + splitStr, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE,CreateMode.EPHEMERAL_SEQUENTIAL);
            System.out.println(myZnode + " is created ");
            //取出所有子节点
            List<String> subNodes = zk.getChildren(root, false);
            //取出所有lockName的锁
            List<String> lockObjNodes = new ArrayList<String>();
            for (String node : subNodes) {
                String _node = node.split(splitStr)[0];
                if(_node.equals(lockName)){
                    lockObjNodes.add(node);
                }
            }
            Collections.sort(lockObjNodes);
            System.out.println(myZnode + "==" + lockObjNodes.get(0));
            if(myZnode.equals(root+"/"+lockObjNodes.get(0))){
                //如果是最小的节点,则表示取得锁
                return true;
            }
            //如果不是最小的节点，找到比自己小1的节点
            String subMyZnode = myZnode.substring(myZnode.lastIndexOf("/") + 1);
            waitNode = lockObjNodes.get(Collections.binarySearch(lockObjNodes, subMyZnode) - 1);
        } catch (KeeperException e) {
            throw new LockException(e);
        } catch (InterruptedException e) {
            throw new LockException(e);
        }
        return false;
    }
 
    public boolean tryLock(long time, TimeUnit unit) {
        try {
            if(this.tryLock()){
                return true;
            }
            return waitForLock(waitNode,time);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }
 
    private boolean waitForLock(String lower, long waitTime) throws InterruptedException, KeeperException {
        Stat stat = zk.exists(root + "/" + lower,true);
        //判断比自己小一个数的节点是否存在,如果不存在则无需等待锁,同时注册监听
        if(stat != null){
            System.out.println("Thread " + Thread.currentThread().getId() + " waiting for " + root + "/" + lower);
            this.latch = new CountDownLatch(1);
            this.latch.await(waitTime, TimeUnit.MILLISECONDS);
            this.latch = null;
        }
        return true;
    }
 
    public void unlock() {
        try {
            System.out.println("unlock " + myZnode);
            zk.delete(myZnode,-1);
            myZnode = null;
            zk.close();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (KeeperException e) {
            e.printStackTrace();
        }
    }
 
    public void lockInterruptibly() throws InterruptedException {
        this.lock();
    }
 
    public Condition newCondition() {
        return null;
    }
     
    public class LockException extends RuntimeException {
        private static final long serialVersionUID = 1L;
        public LockException(String e){
            super(e);
        }
        public LockException(Exception e){
            super(e);
        }
    }
 
}

package com.concurrent;

import com.concurrent.ConcurrentTest.ConcurrentTask;

public class ZkTest {
    public static void main(String[] args) {
        Runnable task1 = new Runnable(){
            public void run() {
                DistributedLock lock = null;
                try {
                    lock = new DistributedLock("127.0.0.1:2182","test1");
                    //lock = new DistributedLock("127.0.0.1:2182","test2");
                    lock.lock();
                    Thread.sleep(3000);
                    System.out.println("===Thread " + Thread.currentThread().getId() + " running");
                } catch (Exception e) {
                    e.printStackTrace();
                } 
                finally {
                    if(lock != null)
                        lock.unlock();
                }
                 
            }
             
        };
        new Thread(task1).start();
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e1) {
            e1.printStackTrace();
        }
        ConcurrentTask[] tasks = new ConcurrentTask[10];
        for(int i=0;i<tasks.length;i++){
            ConcurrentTask task3 = new ConcurrentTask(){
                public void run() {
                    DistributedLock lock = null;
                    try {
                        lock = new DistributedLock("127.0.0.1:2183","test2");
                        lock.lock();
                        System.out.println("Thread " + Thread.currentThread().getId() + " running");
                    } catch (Exception e) {
                        e.printStackTrace();
                    } 
                    finally {
                        lock.unlock();
                    }
                     
                }
            };
            tasks[i] = task3;
        }
        new ConcurrentTest(tasks);
    }
}

package com.concurrent2;

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

import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.Watcher.Event.KeeperState;

public class ConnectionWatcher implements Watcher 
{
    private static final int SESSION_TIMEOUT = 5000;
    protected ZooKeeper zk;
    private CountDownLatch connectedSignal = new CountDownLatch(1);

    public void connect(String hosts) throws IOException, InterruptedException 
    {
        zk = new ZooKeeper(hosts, SESSION_TIMEOUT, this);
        connectedSignal.await();
    }

    @Override
    public void process(WatchedEvent event) 
    {
        if (event.getState() == KeeperState.SyncConnected) 
        {
            connectedSignal.countDown();
        }
    }

    public void close() throws InterruptedException 
    {
        zk.close();
    }
}

package com.concurrent2;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.data.Stat;

public class DistributedLock extends ConnectionWatcher 
{
    public String join(String groupPath) throws KeeperException,InterruptedException 
    {
        String path = groupPath + "/lock-" + zk.getSessionId() + "-";
        // 建立一个顺序临时节点
        String createdPath = zk.create(path, null/* data */, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        System.out.println("Created " + createdPath);

        return createdPath;
    }

    /**
     * 检查本客户端是否得到了分布式锁
     * 
     * @param groupPath
     * @param myName
     * @return 
     * @throws KeeperException
     * @throws InterruptedException
     */
    public boolean checkState(String groupPath, String myName) throws KeeperException, InterruptedException 
    {
        System.out.println("groupPath = " + groupPath);
        System.out.println("myName = " + myName);
        List<String> childList = zk.getChildren(groupPath, false);
        // myName = /zkRoot/_locknode_/lock-94103680129368072-0000000003
        String[] myStr = myName.split("-");
        long myId = Long.parseLong(myStr[2]);

        boolean minId = true;
        int index = 0;
        for (String childName : childList)
        {
            System.out.println(index + " \t " + childName);
            String[] str = childName.split("-");
            long id = Long.parseLong(str[2]);
            if (id < myId)
            {
                minId = false;
                break;
            }
            index ++;
        }

        if (minId) 
        {
            System.out.println(new Date() + "我得到了分布锁，哈哈！ myId:" + myId);
            return true;
        }
        else 
        {
            System.out.println(new Date() + "继续努力吧，  myId:" + myId);
            return false;
        }
    }

    /**
     * 若本客户端没有得到分布式锁，则进行监听本节点前面的节点（避免羊群效应）
     * 
     * @param groupPath
     * @param myName
     * @throws KeeperException
     * @throws InterruptedException
     */
    public void listenNode(final String groupPath, final String myName) throws KeeperException, InterruptedException 
    {
        List<String> childList = zk.getChildren(groupPath, false);

        String[] myStr = myName.split("-");
        long myId = Long.parseLong(myStr[2]);

        List<Long> idList = new ArrayList<Long>();
        Map<Long, String> sessionMap = new HashMap<Long, String>();

        for (String childName : childList) 
        {
            String[] str = childName.split("-");
            long id = Long.parseLong(str[2]);
            idList.add(id);
            sessionMap.put(id, str[1] + "-" + str[2]);
        }

        Collections.sort(idList);
        int i = idList.indexOf(myId);
        if (i <= 0) 
        {
            throw new IllegalArgumentException("数据错误！");
        }

        // 得到前面的一个节点
        long headId = idList.get(i - 1);

        String headPath = groupPath + "/lock-" + sessionMap.get(headId);
        // 添加监听：/zkRoot/_locknode_/lock-94103680129368071-0000000002
        System.out.println("添加监听：" + headPath);

        Stat stat = zk.exists(headPath, new Watcher() {
            @Override
            public void process(WatchedEvent event)
            {
                System.out.println("已经触发了" + event.getType() + "事件！");
                try 
                {
                    while (true) 
                    {
                        if (checkState(groupPath, myName)) 
                        {
                            Thread.sleep(3000);
                            System.out.println(new Date() + " 系统关闭！");
                            System.exit(0);
                        }
                        Thread.sleep(3000);
                    }
                } catch (KeeperException e) 
                {
                    e.printStackTrace();
                } catch (InterruptedException e) 
                {
                    e.printStackTrace();
                }
            }
        });
        System.out.println(stat);
    }
    /**
     * 1.
     * Exception in thread "main" org.apache.zookeeper.KeeperException$NoNodeException: KeeperErrorCode = NoNode for /zkRoot/_locknode_/lock-94103680129368068-
     * 2.
     * Exception in thread "main" org.apache.zookeeper.KeeperException$NoChildrenForEphemeralsException: KeeperErrorCode = NoChildrenForEphemerals for /zkRoot/_locknode_
     * 
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception 
    {
        DistributedLock joinGroup = new DistributedLock();
        joinGroup.connect("localhost:" + "2181");
        
        // zookeeper的根节点；运行本程序前，需要提前生成
        String groupName = "zkRoot";
        String memberName = "_locknode_";
        Stat stat = joinGroup.zk.exists("/" + groupName, true);
        if(stat == null)
        {
            joinGroup.zk.create("/" + groupName, groupName.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        stat = joinGroup.zk.exists("/" + groupName + "/" + memberName, true);
        if(stat == null)
        {
            joinGroup.zk.create("/" + groupName + "/" + memberName, memberName.getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }
        String path = "/" + groupName + "/" + memberName;

        String myName = joinGroup.join(path);
        if (!joinGroup.checkState(path, myName)) {
            joinGroup.listenNode(path, myName);
        }

        Thread.sleep(Integer.MAX_VALUE);
        joinGroup.close();
    }
}