package com.my.lionrpc.registry; import java.io.IOException; 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.ZooDefs; import org.apache.zookeeper.ZooKeeper; import org.apache.zookeeper.data.Stat; import org.slf4j.Logger; import org.slf4j.LoggerFactory; public class ServiceRegistry { private static final Logger LOGGER = LoggerFactory.getLogger(ServiceRegistry.class); private CountDownLatch latch = new CountDownLatch(1); private String registryAddress; private String data; public ServiceRegistry(String registryAddress) { this.registryAddress = registryAddress; } public void register(String data) { this.data = data; if (data != null) { ZooKeeper zk = connectServer(); if (zk != null) { addRootNode(zk); // Add root node if not exist createNode(zk, data); } } } public void reRegister() { if (data != null) { ZooKeeper zk = connectServer(); if (zk != null) { addRootNode(zk); // Add root node if not exist createNode(zk, data); } } } private ZooKeeper connectServer() { ZooKeeper zk = null; try { zk = new ZooKeeper(registryAddress, Constant.ZK_SESSION_TIMEOUT, new Watcher() { public void process(WatchedEvent event) { if (event.getState() == Event.KeeperState.SyncConnected) { latch.countDown(); }else{ if(event.getState() == Event.KeeperState.Expired || event.getState() == Event.KeeperState.Disconnected){ LOGGER.info("ServiceRegistry reRegister..."); reRegister(); } } } }); latch.await(); } catch (IOException e) { LOGGER.error("", e); }catch (InterruptedException ex){ LOGGER.error("", ex); } return zk; } private void addRootNode(ZooKeeper zk){ try { Stat s = zk.exists(Constant.ZK_REGISTRY_PATH, false); if (s == null) { zk.create(Constant.ZK_REGISTRY_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); } } catch (KeeperException e) { LOGGER.error(e.toString()); } catch (InterruptedException e) { LOGGER.error(e.toString()); } } private void createNode(ZooKeeper zk, String data) { try { byte[] bytes = data.getBytes(); String path = zk.create(Constant.ZK_DATA_PATH, bytes, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL); LOGGER.debug("create zookeeper node ({} => {})", path, data); } catch (KeeperException e) { LOGGER.error("", e); } catch (InterruptedException ex){ LOGGER.error("", ex); } } }