public class SnowflakeZookeeperHolder {
private static final Logger LOGGER = LoggerFactory.getLogger(SnowflakeZookeeperHolder.class);
private String zk_AddressNode = null;
private String listenAddress = null;
private int workerID;
private static final String PREFIX_ZK_PATH = "/snowflake";
private static String PATH_FOREVER;
private String connectionString;
private long lastUpdateTime;
private String port;
private int maxWordId;
private String lockPath;
public SnowflakeZookeeperHolder(String hostName,String port, String connectionString,String snowTye,int maxWordId,String app) {
this.listenAddress = hostName+":"+port;
this.port = port;
this.connectionString = connectionString;
PATH_FOREVER = PREFIX_ZK_PATH + "/"+snowTye+"/"+app;
this.maxWordId = maxWordId;
this.lockPath = PREFIX_ZK_PATH + "/"+snowTye;
}
public boolean init() {
try {
CuratorFramework curator = createWithOptions(connectionString, new RetryUntilElapsed(1000, 4), 10000, 6000);
curator.start();
InterProcessMutex lock = new InterProcessMutex(curator,lockPath);
try {
if(lock.acquire(2,TimeUnit.SECONDS)) {
Stat stat = curator.checkExists().forPath(PATH_FOREVER);
if (stat == null) {
//不存在根节点,机器第一次启动,并上传数据
zk_AddressNode = createNode(curator);
//保存workID
updateNewData(curator, zk_AddressNode);
}else {
ArrayList<Integer> usedWorkIddSet= new ArrayList<>();
//存在根节点,先检查是否有属于自己的根节点
List<String> keys = curator.getChildren().forPath(PATH_FOREVER);
//临时节点未找到当前节点
if(!keys.contains(listenAddress)) {
for (String key : keys) {
Endpoint endpoint = getNodeData(curator, PATH_FOREVER + "/" + key);
if (null != endpoint) {
usedWorkIddSet.add(endpoint.getWorkId());
}
}
if(keys.size() == 0){
zk_AddressNode = createNode(curator);
updateNewData(curator, zk_AddressNode);
}else {
ArrayList<Integer> allWordId = new ArrayList<>(maxWordId);
for (int i = 0; i <= maxWordId; i++) {
allWordId.add(i);
}
allWordId.removeAll(usedWorkIddSet);
if(allWordId.size()<=0){
LOGGER.error("workId全部被占用");
return false;
}else{
workerID = allWordId.get(0);
zk_AddressNode = createNode(curator);
//保存workID
updateNewData(curator, zk_AddressNode);
}
}
}else{
zk_AddressNode = PATH_FOREVER + "/" + listenAddress;
Endpoint endpoint = getNodeData(curator, zk_AddressNode);
workerID = endpoint.getWorkId();
updateNewData(curator, zk_AddressNode);
}
}
}else{
LOGGER.error("get zk lock fail!!!");
return false;
}
}catch (Exception e){
LOGGER.error("get zk lock ERROR {}", e);
return false;
}finally {
lock.release();
}
} catch (Exception e) {
LOGGER.error("Start node ERROR {}", e);
return false;
}
LOGGER.info(PATH_FOREVER+" workId="+workerID);
return true;
}
private Endpoint getNodeData(CuratorFramework curator, String zk_AddressNode) throws Exception {
byte[] bytes = curator.getData().forPath(zk_AddressNode);
Endpoint endPoint = deBuildData(new String(bytes));
return endPoint;
}
/**
* 创建持久顺序节点 ,并把节点数据放入 value
*
* @param curator
* @return
* @throws Exception
*/
private String createNode(CuratorFramework curator) throws Exception {
try {
return curator.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL).forPath(PATH_FOREVER + "/" + listenAddress, buildData().getBytes());
} catch (Exception e) {
LOGGER.error("create node error msg {} ", e.getMessage());
throw e;
}
}
private void updateNewData(CuratorFramework curator, String path){
try {
if (System.currentTimeMillis() < lastUpdateTime) {
return;
}
curator.setData().forPath(path, buildData().getBytes());
lastUpdateTime = System.currentTimeMillis();
} catch (Exception e) {
try {
init();
} catch (Exception e1) {
e1.printStackTrace();
}
LOGGER.info("update init data error path is {} error is {}", path, e);
}
}
/**
* 构建需要上传的数据
*
* @return
*/
private String buildData() throws JsonProcessingException {
Endpoint endpoint = new Endpoint(workerID, System.currentTimeMillis());
ObjectMapper mapper = new ObjectMapper();
String json = mapper.writeValueAsString(endpoint);
return json;
}
private Endpoint deBuildData(String json) throws IOException {
ObjectMapper mapper = new ObjectMapper();
Endpoint endpoint = mapper.readValue(json, Endpoint.class);
return endpoint;
}
private CuratorFramework createWithOptions(String connectionString, RetryPolicy retryPolicy, int connectionTimeoutMs, int sessionTimeoutMs) {
return CuratorFrameworkFactory.builder().connectString(connectionString)
.retryPolicy(retryPolicy)
.connectionTimeoutMs(connectionTimeoutMs)
.sessionTimeoutMs(sessionTimeoutMs)
.build();
}
/**
* 上报数据结构
*/
static class Endpoint {
private int workId;
private long timestamp;
public Endpoint() {
}
public Endpoint(int workId, long timestamp) {
this.workId = workId;
this.timestamp = timestamp;
}
public int getWorkId() {
return workId;
}
public void setWorkId(int workId) {
this.workId = workId;
}
public long getTimestamp() {
return timestamp;
}
public void setTimestamp(long timestamp) {
this.timestamp = timestamp;
}
}
public String getZk_AddressNode() {
return zk_AddressNode;
}
public void setZk_AddressNode(String zk_AddressNode) {
this.zk_AddressNode = zk_AddressNode;
}
public String getListenAddress() {
return listenAddress;
}
public void setListenAddress(String listenAddress) {
this.listenAddress = listenAddress;
}
public int getWorkerID() {
return workerID;
}
public void setWorkerID(int workerID) {
this.workerID = workerID;
}
}