zookeeper的分布式锁

使用curator库的InterProcessLock类

package cn.sky.zookeepertest;

import java.util.concurrent.TimeUnit;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessLock;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;

public class Main {
    private static CuratorFramework client;
    private static String connectString = "127.0.0.1:2181";
    private static InterProcessLock lock1;
    private static InterProcessLock lock2;
    private static String lockpath = "/mylock";
    static {
        client = CuratorFrameworkFactory.newClient(connectString, new ExponentialBackoffRetry(1000, 3));
        lock1 = new InterProcessMutex(client, lockpath);
        lock2 = new InterProcessMutex(client, lockpath);
    }

    public static void main(String[] args) throws Exception {
        client.start();
        testLock();
    }

    private static void doSomething() throws Exception {
        System.out.println(Thread.currentThread().getName()+" begin.");
        Thread.sleep(5000);
        System.out.println(Thread.currentThread().getName()+" end.");
    }

    /**
     * 测试分布式锁
     * 用多线程模拟分布式环境
     */
    static void testLock() {
        int getLockTimeout = 6;
        Thread t1 = new Thread() {
            public void run() {
                try {
                    if(lock1.acquire(getLockTimeout,TimeUnit.SECONDS)) {
                        doSomething();
                    };
                } catch (Exception e) {
                    e.printStackTrace();
                }finally {
                    if(lock1.isAcquiredInThisProcess()) {
                        try {
                            lock1.release();
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                }
            };
        };
        Thread t2 = new Thread() {
            public void run() {
                try {
                    if(lock2.acquire(getLockTimeout,TimeUnit.SECONDS)) {
                        doSomething();
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    if(lock2.isAcquiredInThisProcess()) {
                        try {
                            lock2.release();
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                }
            };
        };
        t1.start();
        t2.start();
    }
}

自定义分布式锁的简单实现

package cn.sky.zookeepertest;

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

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.KeeperException.NoNodeException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.ZooKeeper;
/**
 * 对原生zookeeper的api进行简单封装
 * @author lenovo
 *
 */
public class MyZKClient {
    private String connectString = "127.0.0.1:2181";
    private int sessionTimeout_ms;
    private CountDownLatch cdl = new CountDownLatch(1);
    private ZooKeeper zkCli;

    private MyZKClient(){}

    public void start() {
        try {
            zkCli = new ZooKeeper(connectString, sessionTimeout_ms, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    KeeperState state = event.getState();
                    if(state.equals(KeeperState.SyncConnected)) {
                        System.out.println("zookeeper client SyncConnected");
                        cdl.countDown();
                    }
                }
            });
            cdl.await();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public void close(){
        if(zkCli!=null) {
            try {
                zkCli.close();
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
    }

    public static MyZKClient newClient(String connectString, int sessionTimeout_ms) {
        MyZKClient client = new MyZKClient();
        client.connectString = connectString;
        client.sessionTimeout_ms = sessionTimeout_ms;
        return client;
    }

    public String createPathWithMode(String path,byte[] data,CreateMode createMode) throws Exception {
        if(data==null) {
            data = new byte[0];
        }
        return zkCli.create(path, data, Ids.OPEN_ACL_UNSAFE, createMode);
    }

    public void deleteNode(String path) {
        try {
            zkCli.delete(path, -1);
        } catch (InterruptedException | KeeperException e) {
            throw new RuntimeException(e);
        }
    }

    public boolean existsNode(String path) throws Exception {
        return zkCli.exists(path, false)!=null;
    }

    public List<String> getChildren(String parentPath) throws Exception{
        return zkCli.getChildren(parentPath, null);
    }

    /**
     * 如果注册一次，则只会监听一次
     * @param path
     * @param watcher
     * @throws Exception
     */
    public void watch(String path,Watcher watcher) throws Exception{
        if(zkCli.exists(path, watcher)==null){
            throw new RuntimeException(new NoNodeException(path));
        }
    }
}

package cn.sky.zookeepertest;

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

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.EventType;
/**
 * 分布式锁的实现类
 * @author lenovo
 *
 */
public class MyZKLock {
    private String basePath;
    private static final String lockName = "mylock-";
    private MyZKClient client;
    private String myLockPath;
    private boolean required = false;

    public MyZKLock(String basePath,MyZKClient client) {
        this.basePath = basePath;
        this.client = client;
    }

    public boolean require(long timeout, TimeUnit unit) throws Exception {
        if(!client.existsNode(basePath)) {
            client.createPathWithMode(basePath, null, CreateMode.PERSISTENT);
        }
        myLockPath = client.createPathWithMode(basePath+"/"+lockName, null, CreateMode.EPHEMERAL_SEQUENTIAL);
        String prePath = null;
        while((prePath=getPreviousPath())!=null) {
            CountDownLatch cdl = new CountDownLatch(1);
            client.watch(prePath, new Watcher() {
                @Override
                public void process(WatchedEvent event) {
                    if(event.getType().equals(EventType.NodeDeleted)){
                        cdl.countDown();
                    }
                }
            });
            if(unit==null||timeout<0) {
                cdl.await();
            }else {
                if(!cdl.await(timeout, unit)) {
                    client.deleteNode(myLockPath);
                    required = false;
                    return required;
                }
            }
        }
        required = true;
        return required;
    }

    public boolean require() throws Exception {
        return require(-1, null);
    }

    private String getPreviousPath() throws Exception {
        List<String> children = client.getChildren(basePath);
        if(children==null||children.size()==0) {
            throw new RuntimeException("get lock fail!");
        }
        Collections.sort(children);
        for(int i=0;i<children.size();i++) {
            String path = basePath+"/"+children.get(i);
            if(path.equals(myLockPath)) {
                if(i==0) {
                    return null;
                }else {
                    return basePath+"/"+children.get(i-1);
                }
            }
        }
        return null;
    }

    public void release() {
        if(!required) {
            throw new RuntimeException("the lock is not required!");
        }
        client.deleteNode(myLockPath);
    }

    public boolean isRequired() {
        return required;
    }

}

package cn.sky.zookeepertest;

import java.util.concurrent.TimeUnit;
/**
 * 测试类
 * @author lenovo
 *
 */
public class Main {
    private static String basePath = "/mylock";
    private static String connectString = "127.0.0.1:2181";
    private static MyZKClient client = MyZKClient.newClient(connectString, 1000);

    public static void main(String[] args) throws Exception {
        client.start();
        for(int i=1;i<=3;i++) {
            new Thread(new MyRunnable(i)).start();
        }
        Thread.sleep(9999999);
    }


    private static class MyRunnable implements Runnable{
        private int i;
        private MyRunnable(int i){
            this.i = i;
        }
        @Override
        public void run() {
            MyZKLock lock = new MyZKLock(basePath, client);
            try {
                if(!lock.require(3,TimeUnit.SECONDS)) {
                    System.out.println("获取锁"+i+"失败");
                    return;
                }
                {
                    System.out.println("begin"+i);
                    Thread.sleep(2900);
                    System.out.println("end"+i);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }finally {
                if(lock.isRequired()) {
                    lock.release();
                }
            }
        }
    }
}