package zoo.test.use;
import java.io.IOException;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.apache.zookeeper.ZooDefs.Ids;
import org.apache.zookeeper.data.Stat;
public class ZooKeeperLock {
private String dirName;
private static final String prefix="lock-";
private ZooKeeper zooKeeper;
private String holdVersion;
private AtomicInteger count = new AtomicInteger();
public ZooKeeperLock(String connectStr,String dirName) throws IOException{
this.dirName = dirName;
initConnect(connectStr,Integer.MAX_VALUE);
}
private void initConnect(String connectStr,int timeout) throws IOException{
this.zooKeeper = new ZooKeeper(connectStr, timeout, null);
}
public void lock() throws KeeperException, InterruptedException{
if(count.incrementAndGet()>1){
return;
}
if(holdVersion==null){
String nodeName = null;
while(true){
try{
nodeName = zooKeeper.create(dirName+"/"+prefix, null, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
break;
}catch(KeeperException.NoNodeException e){
zooKeeper.create(dirName, null, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
}
}
holdVersion = nodeName.substring((dirName+"/"+prefix).length());
}
Set<String> orderChildren = this.getOrderChildren();
if(orderChildren.isEmpty()){
throw new KeeperException.NoNodeException("创建的节点被其他的客户端删除");
}
String smallestVersion = orderChildren.iterator().next();
if(compare(holdVersion, smallestVersion)!=0){
CountDownLatchWatcher latchWatcher = new CountDownLatchWatcher();
String watcherNodeVersion = getWatcherNodeVersion(orderChildren);
Stat stat = zooKeeper.exists(dirName+"/"+prefix+watcherNodeVersion, latchWatcher);
if(stat!=null){
latchWatcher.await();
}
}
}
public void unLock() throws InterruptedException, KeeperException{
if(count.decrementAndGet()==0){
zooKeeper.delete(dirName+"/"+prefix+holdVersion, -1);
holdVersion=null;
}
}
private class CountDownLatchWatcher implements Watcher{
private CountDownLatch latch = new CountDownLatch(1);
@Override
public void process(WatchedEvent event) {
// TODO Auto-generated method stub
if(event.getType()==EventType.NodeDeleted){
latch.countDown();
}
}
public void await() throws InterruptedException{
latch.await();
}
}
private long compare(String o1,String o2){
Long l1 = Long.valueOf(o1);
Long l2 = Long.valueOf(o2);
return l1-l2;
}
private String getWatcherNodeVersion(Set<String> orderChildren){
String watcherNodeVersion = null;
for(String child:orderChildren){
if(child.equals(holdVersion)){
break;
}
watcherNodeVersion = child;
}
return watcherNodeVersion;
}
private Set<String> getOrderChildren(){
Set<String> orderChildren = new TreeSet<String>(new Comparator<String>() {
@Override
public int compare(String o1, String o2) {
// TODO Auto-generated method stub
int l1 = Integer.valueOf(o1);
int l2 = Integer.valueOf(o2);
return l1-l2;
}
});
try {
List<String> children = zooKeeper.getChildren(dirName, null);
if(children!=null&&!children.isEmpty()){
for(String child:children){
if(child.regionMatches(0, prefix, 0, prefix.length())){
orderChildren.add(child.substring(prefix.length()));
}
}
}
} catch (KeeperException e) {
// TODO Auto-generated catch block
e.printStackTrace();
} catch (InterruptedException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
return orderChildren;
}
public void close() throws InterruptedException{
zooKeeper.close();
}
}