package com.imadcn.framework.idworker.register.zookeeper;
import java.io.File;
import java.io.IOException;
import java.net.UnknownHostException;
import java.nio.charset.StandardCharsets;
import java.util.Date;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.apache.commons.io.FileUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import com.imadcn.framework.idworker.config.ApplicationConfiguration;
import com.imadcn.framework.idworker.exception.RegException;
import com.imadcn.framework.idworker.register.WorkerRegister;
import com.imadcn.framework.idworker.registry.CoordinatorRegistryCenter;
import com.imadcn.framework.idworker.util.HostUtils;
/**
* 机器信息注册
*
* @author imadcn
* @since 1.0.0
*/
public class ZookeeperWorkerRegister implements WorkerRegister {
private static final Logger logger = LoggerFactory.getLogger(ZookeeperWorkerRegister.class);
/**
* 最大机器数
*/
private static final long MAX_WORKER_NUM = 1024;
/**
* 加锁最大等待时间
*/
private static final int MAX_LOCK_WAIT_TIME_MS = 30 * 1000;
/**
* 注册中心工具
*/
private final CoordinatorRegistryCenter regCenter;
/**
* 注册文件
*/
private String registryFile;
/**
* zk节点信息
*/
private final NodePath nodePath;
/**
* zk节点是否持久化存储
*/
private boolean durable;
public ZookeeperWorkerRegister(CoordinatorRegistryCenter regCenter,
ApplicationConfiguration applicationConfiguration) {
this.regCenter = regCenter;
this.nodePath = new NodePath(applicationConfiguration.getGroup());
this.durable = applicationConfiguration.isDurable();
if (StringUtils.isEmpty(applicationConfiguration.getRegistryFile())) {
this.registryFile = getDefaultFilePath(nodePath.getGroupName());
} else {
this.registryFile = applicationConfiguration.getRegistryFile();
}
}
/**
* 向zookeeper注册workerId
*
* @return workerId workerId
*/
@Override
public long register() {
InterProcessMutex lock = null;
try {
CuratorFramework client = (CuratorFramework) regCenter.getRawClient();
lock = new InterProcessMutex(client, nodePath.getGroupPath());
int numOfChildren = regCenter.getNumChildren(nodePath.getWorkerPath());
if (numOfChildren < MAX_WORKER_NUM) {
if (!lock.acquire(MAX_LOCK_WAIT_TIME_MS, TimeUnit.MILLISECONDS)) {
String message = String.format("acquire lock failed after %s ms.", MAX_LOCK_WAIT_TIME_MS);
throw new TimeoutException(message);
}
NodeInfo localNodeInfo = getLocalNodeInfo();
List<String> children = regCenter.getChildrenKeys(nodePath.getWorkerPath());
// 有本地缓存的节点信息，同时ZK也有这条数据
if (localNodeInfo != null && children.contains(String.valueOf(localNodeInfo.getWorkerId()))) {
String key = getNodePathKey(nodePath, localNodeInfo.getWorkerId());
String zkNodeInfoJson = regCenter.get(key);
NodeInfo zkNodeInfo = createNodeInfoFromJsonStr(zkNodeInfoJson);
if (checkNodeInfo(localNodeInfo, zkNodeInfo)) {
// 更新ZK节点信息，保存本地缓存，开启定时上报任务
nodePath.setWorkerId(zkNodeInfo.getWorkerId());
zkNodeInfo.setUpdateTime(new Date());
updateZookeeperNodeInfo(key, zkNodeInfo);
saveLocalNodeInfo(zkNodeInfo);
executeUploadNodeInfoTask(key, zkNodeInfo);
return zkNodeInfo.getWorkerId();
}
}
// 无本地信息或者缓存数据不匹配，开始向ZK申请节点机器ID
for (int workerId = 0; workerId < MAX_WORKER_NUM; workerId++) {
String workerIdStr = String.valueOf(workerId);
if (!children.contains(workerIdStr)) { // 申请成功
NodeInfo applyNodeInfo = createNodeInfo(nodePath.getGroupName(), workerId);
nodePath.setWorkerId(applyNodeInfo.getWorkerId());
// 保存ZK节点信息，保存本地缓存，开启定时上报任务
saveZookeeperNodeInfo(nodePath.getWorkerIdPath(), applyNodeInfo);
saveLocalNodeInfo(applyNodeInfo);
executeUploadNodeInfoTask(nodePath.getWorkerIdPath(), applyNodeInfo);
return applyNodeInfo.getWorkerId();
}
}
}
throw new RegException("max worker num reached. register failed");
} catch (RegException e) {
throw e;
} catch (Exception e) {
logger.error("", e);
throw new IllegalStateException(e.getMessage(), e);
} finally {
try {
if (lock != null) {
lock.release();
}
} catch (Exception ignored) {
logger.error("", ignored);
}
}
}
/**
* 添加连接监听
*
* @param listener zk状态监听listener
*/
@Deprecated
public void addConnectionListener(ConnectionStateListener listener) {
// CuratorFramework client = (CuratorFramework)
// regCenter.getRawClient();
// client.getConnectionStateListenable().addListener(listener);
}
/**
* 关闭注册
*/
@Override
public synchronized void logout() {
CuratorFramework client = (CuratorFramework) regCenter.getRawClient();
if (client != null && client.getState() == CuratorFrameworkState.STARTED) {
// 移除注册节点（最大程度的自动释放资源）
regCenter.remove(nodePath.getWorkerIdPath());
// 关闭连接
regCenter.close();
}
}
/**
* 检查节点信息
*
* @param localNodeInfo 本地缓存节点信息
* @param zkNodeInfo zookeeper节点信息
* @return
*/
private boolean checkNodeInfo(NodeInfo localNodeInfo, NodeInfo zkNodeInfo) {
try {
// NodeId、IP、HostName、GroupName 相等（本地缓存==ZK数据）
if (!zkNodeInfo.getNodeId().equals(localNodeInfo.getNodeId())) {
return false;
}
if (!zkNodeInfo.getIp().equals(localNodeInfo.getIp())) {
return false;
}
if (!zkNodeInfo.getHostName().equals(localNodeInfo.getHostName())) {
return false;
}
if (!zkNodeInfo.getGroupName().equals(localNodeInfo.getGroupName())) {
return false;
}
return true;
} catch (Exception e) {
logger.error("check node info error, {}", e);
return false;
}
}
/**
* 更新节点信息Task
*
* @param key zk path
* @param nodeInfo 节点信息
*/
private void executeUploadNodeInfoTask(final String key, final NodeInfo nodeInfo) {
Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
@Override
public Thread newThread(Runnable r) {
Thread thread = new Thread(r, "upload node info task thread");
thread.setDaemon(true);
return thread;
}
}).scheduleWithFixedDelay(new Runnable() {
@Override
public void run() {
updateZookeeperNodeInfo(key, nodeInfo);
}
}, 3L, 3L, TimeUnit.SECONDS);
}
/**
* 获取节点ZK Path KEy
*
* @param nodePath 节点路径信息
* @param workerId 节点机器ID
* @return
*/
private String getNodePathKey(NodePath nodePath, Integer workerId) {
StringBuilder builder = new StringBuilder();
builder.append(nodePath.getWorkerPath()).append("/");
builder.append(workerId);
return builder.toString();
}
/**
* 保存ZK节点信息
*
* @param key
* @param nodeInfo
*/
private void saveZookeeperNodeInfo(String key, NodeInfo nodeInfo) {
if (durable) {
regCenter.persist(key, jsonizeNodeInfo(nodeInfo));
} else {
regCenter.persistEphemeral(key, jsonizeNodeInfo(nodeInfo));
}
}
/**
* 刷新ZK节点信息（修改updateTime）
*
* @param key
* @param nodeInfo
*/
private void updateZookeeperNodeInfo(String key, NodeInfo nodeInfo) {
try {
nodeInfo.setUpdateTime(new Date());
if (durable) {
regCenter.persist(key, jsonizeNodeInfo(nodeInfo));
} else {
regCenter.persistEphemeral(key, jsonizeNodeInfo(nodeInfo));
}
} catch (Exception e) {
logger.debug("update zookeeper node info error, {}", e);
}
}
/**
* 缓存机器节点信息至本地
*
* @param nodeInfo 机器节点信息
*/
private void saveLocalNodeInfo(NodeInfo nodeInfo) {
try {
File nodeInfoFile = new File(registryFile);
String nodeInfoJson = jsonizeNodeInfo(nodeInfo);
FileUtils.writeStringToFile(nodeInfoFile, nodeInfoJson, StandardCharsets.UTF_8);
} catch (IOException e) {
logger.error("save node info cache error, {}", e);
}
}
/**
* 读取本地缓存机器节点
*
* @return 机器节点信息
*/
private NodeInfo getLocalNodeInfo() {
try {
File nodeInfoFile = new File(registryFile);
if (nodeInfoFile.exists()) {
String nodeInfoJson = FileUtils.readFileToString(nodeInfoFile, StandardCharsets.UTF_8);
NodeInfo nodeInfo = createNodeInfoFromJsonStr(nodeInfoJson);
return nodeInfo;
}
} catch (Exception e) {
logger.error("read node info cache error, {}", e);
}
return null;
}
/**
* 初始化节点信息
*
* @param groupName 分组名
* @param workerId 机器号
* @return 节点信息
* @throws UnknownHostException
*/
private NodeInfo createNodeInfo(String groupName, Integer workerId) throws UnknownHostException {
NodeInfo nodeInfo = new NodeInfo();
nodeInfo.setNodeId(genNodeId());
nodeInfo.setGroupName(groupName);
nodeInfo.setWorkerId(workerId);
nodeInfo.setIp(HostUtils.getLocalIP());
nodeInfo.setHostName(HostUtils.getLocalHostName());
nodeInfo.setCreateTime(new Date());
nodeInfo.setUpdateTime(new Date());
return nodeInfo;
}
/**
* 通过节点信息JSON字符串反序列化节点信息
*
* @param jsonStr 节点信息JSON字符串
* @return 节点信息
*/
private NodeInfo createNodeInfoFromJsonStr(String jsonStr) {
NodeInfo nodeInfo = JSON.parseObject(jsonStr, NodeInfo.class);
return nodeInfo;
}
/**
* 节点信息转json字符串
*
* @param nodeInfo 节点信息
* @return json字符串
*/
private String jsonizeNodeInfo(NodeInfo nodeInfo) {
String dateFormat = "yyyy-MM-dd HH:mm:ss";
return JSON.toJSONStringWithDateFormat(nodeInfo, dateFormat, SerializerFeature.WriteDateUseDateFormat);
}
/**
* 获取本地节点缓存文件路径
*
* @param groupName 分组名
* @return 文件路径
*/
private String getDefaultFilePath(String groupName) {
StringBuilder builder = new StringBuilder();
builder.append(".").append(File.separator).append("tmp");
builder.append(File.separator).append("idworker");
builder.append(File.separator).append(groupName).append(".cache");
return builder.toString();
}
/**
* 获取节点唯一ID （基于UUID）
*
* @return 节点唯一ID
*/
private String genNodeId() {
return UUID.randomUUID().toString().replace("-", "").toLowerCase();
}
}
