zookeeper 分布式队列初实现
package zoo.test.use;
import java.io.IOException;
import java.util.Comparator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
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.KeeperException.NoNodeException;
import org.apache.zookeeper.ZooDefs.Ids;
public class ZooKeeperQueue {
private String dirName;
private ZooKeeper zooKeeper;
private static final String prefix="queue-";
public ZooKeeperQueue(String connectStr,String dirName) throws IOException{
this.dirName = dirName;
initConnect(connectStr,30*60*1000);
}
private void initConnect(String connectStr,int timeout) throws IOException{
this.zooKeeper = new ZooKeeper(connectStr, timeout, null);
}
public byte[] take() throws InterruptedException, KeeperException{
Set<String> versions = null;
while(true){
CountDownLatchWatcher latchWatcher = new CountDownLatchWatcher();
try{
versions = this.getOrderChildrenVersion(latchWatcher);
}catch(KeeperException.NoNodeException e){
try{
zooKeeper.create(dirName, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
}catch(KeeperException.NodeExistsException ex){
// node create by another client try again
continue;
}
}
if(versions.isEmpty()){
latchWatcher.await();
System.out.println("wait for offer");
continue;
}
for(String version:versions){
String dirName = this.getDireNameByVersion(version);
try{
byte[] data = zooKeeper.getData(dirName, false, null);
zooKeeper.delete(dirName, -1);
return data;
}catch(KeeperException.NoNodeException e){
// node remove by another client try next
}
}
}
}
private class CountDownLatchWatcher implements Watcher{
private CountDownLatch latch = new CountDownLatch(1);
@Override
public void process(WatchedEvent event) {
// TODO Auto-generated method stub
latch.countDown();
}
public void await() throws InterruptedException{
latch.await();
}
}
public boolean offer(byte[] data){
for(;;){
try{
zooKeeper.create(dirName+"/"+prefix, data, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT_SEQUENTIAL);
return true;
}catch(KeeperException.NoNodeException e){
try{
zooKeeper.create(dirName, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
}catch(KeeperException.NodeExistsException ex){
// this node create by another client try again
}catch(Exception ex){
ex.printStackTrace();
return false;
}
}catch(Exception e){
e.printStackTrace();
return false;
}
}
}
public byte[] pool(){
try{
return this.remove();
}catch(Exception e){
return null;
}
}
public byte[] peer(){
try{
return this.element();
}catch(Exception e){
return null;
}
}
private byte[] remove() throws KeeperException, InterruptedException{
Set<String> versions = null;
while(true){
versions = this.getOrderChildrenVersion(null);
if(versions.isEmpty()){
throw new NoSuchElementException();
}
for(String version:versions){
String dirName = this.getDireNameByVersion(version);
try{
byte[] data = zooKeeper.getData(dirName, false, null);
zooKeeper.delete(dirName, -1);
return data;
}catch(KeeperException.NoNodeException e){
// node remove by another client try next
}
}
}
}
private byte[] element() throws KeeperException, InterruptedException{
Set<String> versions = null;
while(true){
versions = this.getOrderChildrenVersion(null);
if(versions.isEmpty()){
throw new NoSuchElementException();
}
for(String version:versions){
String dirName = this.getDireNameByVersion(version);
try {
return zooKeeper.getData(dirName, false, null);
} catch (KeeperException.NoNodeException e) {
// TODO Auto-generated catch block
//node remove by another client try next
}
}
}
}
private Set<String> getOrderChildrenVersion(Watcher watcher) throws NoNodeException{
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, watcher);
if(children!=null&&!children.isEmpty()){
for(String str:children){
if(!str.regionMatches(0, str, 0, prefix.length())){
System.out.println("node not belong the queue:"+str);
continue;
}
orderChildren.add(str.substring(prefix.length()));
}
}
} catch(KeeperException.NoNodeException e){
throw new KeeperException.NoNodeException();
} catch (KeeperException e) {
// TODO Auto-generated catch block
e.printStackTrace();
} catch (InterruptedException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
return orderChildren;
}
private String getDireNameByVersion(String version){
return this.dirName+"/"+prefix+version;
}
}