com.taobao.pamirs.schedule.zk.ScheduleStrategyDataManager4ZK Maven / Gradle / Ivy

Go to download
package com.taobao.pamirs.schedule.zk;

import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.taobao.pamirs.schedule.strategy.ManagerFactoryInfo;
import com.taobao.pamirs.schedule.strategy.ScheduleStrategy;
import com.taobao.pamirs.schedule.strategy.ScheduleStrategyRunntime;
import com.taobao.pamirs.schedule.strategy.TBScheduleManagerFactory;
import java.io.Writer;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.UUID;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.ZooKeeper;

public class ScheduleStrategyDataManager4ZK {

  private ZKManager zkManager;
  private String PATH_Strategy;
  private String PATH_ManagerFactory;
  private Gson gson;

  //??Spring???󴴽???Ϻ󣬴????ڲ?????
  public ScheduleStrategyDataManager4ZK(ZKManager aZkManager) throws Exception {
    this.zkManager = aZkManager;
    gson = new GsonBuilder().registerTypeAdapter(Timestamp.class, new TimestampTypeAdapter())
        .setDateFormat("yyyy-MM-dd HH:mm:ss").create();
    this.PATH_Strategy = this.zkManager.getRootPath() + "/strategy";
    this.PATH_ManagerFactory = this.zkManager.getRootPath() + "/factory";

    if (this.getZooKeeper().exists(this.PATH_Strategy, false) == null) {
      ZKTools.createPath(getZooKeeper(), this.PATH_Strategy, CreateMode.PERSISTENT,
          this.zkManager.getAcl());
    }
    if (this.getZooKeeper().exists(this.PATH_ManagerFactory, false) == null) {
      ZKTools.createPath(getZooKeeper(), this.PATH_ManagerFactory, CreateMode.PERSISTENT,
          this.zkManager.getAcl());
    }
  }

  public ScheduleStrategy loadStrategy(String strategyName)
      throws Exception {
    String zkPath = this.PATH_Strategy + "/" + strategyName;
    if (this.getZooKeeper().exists(zkPath, false) == null) {
      return null;
    }
    String valueString = new String(this.getZooKeeper().getData(zkPath, false, null));
    ScheduleStrategy result = (ScheduleStrategy) this.gson
        .fromJson(valueString, ScheduleStrategy.class);
    return result;
  }

  public void createScheduleStrategy(ScheduleStrategy scheduleStrategy) throws Exception {
    String zkPath = this.PATH_Strategy + "/" + scheduleStrategy.getStrategyName();
    String valueString = this.gson.toJson(scheduleStrategy);
    if (this.getZooKeeper().exists(zkPath, false) == null) {
      this.getZooKeeper()
          .create(zkPath, valueString.getBytes(), this.zkManager.getAcl(), CreateMode.PERSISTENT);
    } else {
      throw new Exception("???Ȳ???" + scheduleStrategy.getStrategyName()
          + "?Ѿ?????,???ȷ????Ҫ?ؽ??????ȵ???deleteMachineStrategy(String taskType)ɾ??");
    }
  }

  public void updateScheduleStrategy(ScheduleStrategy scheduleStrategy)
      throws Exception {
    String zkPath = this.PATH_Strategy + "/" + scheduleStrategy.getStrategyName();
    String valueString = this.gson.toJson(scheduleStrategy);
    if (this.getZooKeeper().exists(zkPath, false) == null) {
      this.getZooKeeper()
          .create(zkPath, valueString.getBytes(), this.zkManager.getAcl(), CreateMode.PERSISTENT);
    } else {
      this.getZooKeeper().setData(zkPath, valueString.getBytes(), -1);
    }

  }

  public void deleteMachineStrategy(String taskType) throws Exception {
    deleteMachineStrategy(taskType, false);
  }

  public void pause(String strategyName) throws Exception {
    ScheduleStrategy strategy = this.loadStrategy(strategyName);
    strategy.setSts(ScheduleStrategy.STS_PAUSE);
    this.updateScheduleStrategy(strategy);
  }

  public void resume(String strategyName) throws Exception {
    ScheduleStrategy strategy = this.loadStrategy(strategyName);
    strategy.setSts(ScheduleStrategy.STS_RESUME);
    this.updateScheduleStrategy(strategy);
  }

  public void deleteMachineStrategy(String taskType, boolean isForce) throws Exception {
    String zkPath = this.PATH_Strategy + "/" + taskType;
    if (isForce == false && this.getZooKeeper().getChildren(zkPath, null).size() > 0) {
      throw new Exception("????ɾ??" + taskType + "?????в??ԣ??ᵼ?±???????????Ӧ?ò???ֹͣʧȥ???Ƶĵ??Ƚ??̡?" +
          "?????????IP??ַ???????еĵ???????ֹͣ????ɾ?????Ȳ???");
    }
    ZKTools.deleteTree(this.getZooKeeper(), zkPath);
  }

  public List loadAllScheduleStrategy() throws Exception {
    String zkPath = this.PATH_Strategy;
    List result = new ArrayList();
    List names = this.getZooKeeper().getChildren(zkPath, false);
    Collections.sort(names);
    for (String name : names) {
      result.add(this.loadStrategy(name));
    }
    return result;
  }

  /**
   * ע??ManagerFactory
   *
   * @return ??Ҫȫ??ע???ĵ??ȣ????統IP?????б???
   */
  public List registerManagerFactory(TBScheduleManagerFactory managerFactory)
      throws Exception {

    if (managerFactory.getUuid() == null) {
      String uuid =
          managerFactory.getIp() + "$" + managerFactory.getHostName() + "$" + UUID.randomUUID()
              .toString().replaceAll("-", "").toUpperCase();
      String zkPath = this.PATH_ManagerFactory + "/" + uuid + "$";
      zkPath = this.getZooKeeper()
          .create(zkPath, null, this.zkManager.getAcl(), CreateMode.EPHEMERAL_SEQUENTIAL);
      managerFactory.setUuid(zkPath.substring(zkPath.lastIndexOf("/") + 1));
    } else {
      String zkPath = this.PATH_ManagerFactory + "/" + managerFactory.getUuid();
      if (this.getZooKeeper().exists(zkPath, false) == null) {
        zkPath = this.getZooKeeper()
            .create(zkPath, null, this.zkManager.getAcl(), CreateMode.EPHEMERAL);
      }
    }

    List result = new ArrayList();
    for (ScheduleStrategy scheduleStrategy : loadAllScheduleStrategy()) {
      boolean isFind = false;
      //??ͣ???߲???IP??Χ
      if (ScheduleStrategy.STS_PAUSE.equalsIgnoreCase(scheduleStrategy.getSts()) == false
          && scheduleStrategy.getIPList() != null) {
        for (String ip : scheduleStrategy.getIPList()) {
          if (ip.equals("127.0.0.1") || ip.equalsIgnoreCase("localhost") || ip
              .equals(managerFactory.getIp()) || ip
              .equalsIgnoreCase(managerFactory.getHostName())) {
            //???ӿɹ???TaskType
            String zkPath =
                this.PATH_Strategy + "/" + scheduleStrategy.getStrategyName() + "/" + managerFactory
                    .getUuid();
            if (this.getZooKeeper().exists(zkPath, false) == null) {
              zkPath = this.getZooKeeper()
                  .create(zkPath, null, this.zkManager.getAcl(), CreateMode.EPHEMERAL);
            }
            isFind = true;
            break;
          }
        }
      }
      if (isFind == false) {//???ԭ??ע???Factory
        String zkPath =
            this.PATH_Strategy + "/" + scheduleStrategy.getStrategyName() + "/" + managerFactory
                .getUuid();
        if (this.getZooKeeper().exists(zkPath, false) != null) {
          ZKTools.deleteTree(this.getZooKeeper(), zkPath);
          result.add(scheduleStrategy.getStrategyName());
        }
      }
    }
    return result;
  }

  /**
   * ע??????ֹͣ????
   */
  public void unRregisterManagerFactory(TBScheduleManagerFactory managerFactory) throws Exception {
    for (String taskName : this.getZooKeeper().getChildren(this.PATH_Strategy, false)) {
      String zkPath = this.PATH_Strategy + "/" + taskName + "/" + managerFactory.getUuid();
      if (this.getZooKeeper().exists(zkPath, false) != null) {
        ZKTools.deleteTree(this.getZooKeeper(), zkPath);
      }
    }
  }

  public ScheduleStrategyRunntime loadScheduleStrategyRunntime(String strategyName, String uuid)
      throws Exception {
    String zkPath = this.PATH_Strategy + "/" + strategyName + "/" + uuid;
    ScheduleStrategyRunntime result = null;
    if (this.getZooKeeper().exists(zkPath, false) != null) {
      byte[] value = this.getZooKeeper().getData(zkPath, false, null);
      if (value != null) {
        String valueString = new String(value);
        result = (ScheduleStrategyRunntime) this.gson
            .fromJson(valueString, ScheduleStrategyRunntime.class);
        if (null == result) {
          throw new Exception("gson ?????л??쳣,????Ϊnull");
        }
        if (null == result.getStrategyName()) {
          throw new Exception("gson ?????л??쳣,????????Ϊnull");
        }
        if (null == result.getUuid()) {
          throw new Exception("gson ?????л??쳣,uuidΪnull");
        }
      } else {
        result = new ScheduleStrategyRunntime();
        result.setStrategyName(strategyName);
        result.setUuid(uuid);
        result.setRequestNum(0);
        result.setMessage("");
      }
    }
    return result;
  }

  /**
   * װ?????еĲ???????״̬
   */
  public List loadAllScheduleStrategyRunntime() throws Exception {
    List result = new ArrayList();
    String zkPath = this.PATH_Strategy;
    for (String taskType : this.getZooKeeper().getChildren(zkPath, false)) {
      for (String uuid : this.getZooKeeper().getChildren(zkPath + "/" + taskType, false)) {
        result.add(loadScheduleStrategyRunntime(taskType, uuid));
      }
    }
    return result;
  }

  public List loadAllScheduleStrategyRunntimeByUUID(
      String managerFactoryUUID) throws Exception {
    List result = new ArrayList();
    String zkPath = this.PATH_Strategy;

    List taskTypeList = this.getZooKeeper().getChildren(zkPath, false);
    Collections.sort(taskTypeList);
    for (String taskType : taskTypeList) {
      if (this.getZooKeeper().exists(zkPath + "/" + taskType + "/" + managerFactoryUUID, false)
          != null) {
        result.add(loadScheduleStrategyRunntime(taskType, managerFactoryUUID));
      }
    }
    return result;
  }

  public List loadAllScheduleStrategyRunntimeByTaskType(
      String strategyName) throws Exception {
    List result = new ArrayList();
    String zkPath = this.PATH_Strategy;
    if (this.getZooKeeper().exists(zkPath + "/" + strategyName, false) == null) {
      return result;
    }
    List uuidList = this.getZooKeeper().getChildren(zkPath + "/" + strategyName, false);
    //????
    Collections.sort(uuidList, new Comparator() {
      public int compare(String u1, String u2) {
        return u1.substring(u1.lastIndexOf("$") + 1).compareTo(
            u2.substring(u2.lastIndexOf("$") + 1));
      }
    });

    for (String uuid : uuidList) {
      result.add(loadScheduleStrategyRunntime(strategyName, uuid));
    }
    return result;
  }

  /**
   * ????????????
   */
  public void updateStrategyRunntimeReqestNum(String strategyName, String manangerFactoryUUID,
      int requestNum) throws Exception {
    String zkPath = this.PATH_Strategy + "/" + strategyName + "/" + manangerFactoryUUID;
    ScheduleStrategyRunntime result = null;
    if (this.getZooKeeper().exists(zkPath, false) != null) {
      result = this.loadScheduleStrategyRunntime(strategyName, manangerFactoryUUID);
    } else {
      result = new ScheduleStrategyRunntime();
      result.setStrategyName(strategyName);
      result.setUuid(manangerFactoryUUID);
      result.setRequestNum(requestNum);
      result.setMessage("");
    }
    result.setRequestNum(requestNum);
    String valueString = this.gson.toJson(result);
    this.getZooKeeper().setData(zkPath, valueString.getBytes(), -1);
  }

  /**
   * ???µ??ȹ????е???Ϣ
   */
  public void updateStrategyRunntimeErrorMessage(String strategyName, String manangerFactoryUUID,
      String message) throws Exception {
    String zkPath = this.PATH_Strategy + "/" + strategyName + "/" + manangerFactoryUUID;
    ScheduleStrategyRunntime result = null;
    if (this.getZooKeeper().exists(zkPath, false) != null) {
      result = this.loadScheduleStrategyRunntime(strategyName, manangerFactoryUUID);
    } else {
      result = new ScheduleStrategyRunntime();
      result.setStrategyName(strategyName);
      result.setUuid(manangerFactoryUUID);
      result.setRequestNum(0);
    }
    result.setMessage(message);
    String valueString = this.gson.toJson(result);
    this.getZooKeeper().setData(zkPath, valueString.getBytes(), -1);
  }

  public void updateManagerFactoryInfo(String uuid, boolean isStart) throws Exception {
    String zkPath = this.PATH_ManagerFactory + "/" + uuid;
    if (this.getZooKeeper().exists(zkPath, false) == null) {
      throw new Exception("???????????????:" + uuid);
    }
    this.getZooKeeper().setData(zkPath, Boolean.toString(isStart).getBytes(), -1);
  }

  public ManagerFactoryInfo loadManagerFactoryInfo(String uuid) throws Exception {
    String zkPath = this.PATH_ManagerFactory + "/" + uuid;
    if (this.getZooKeeper().exists(zkPath, false) == null) {
      throw new Exception("???????????????:" + uuid);
    }
    byte[] value = this.getZooKeeper().getData(zkPath, false, null);
    ManagerFactoryInfo result = new ManagerFactoryInfo();
    result.setUuid(uuid);
    if (value == null) {
      result.setStart(true);
    } else {
      result.setStart(Boolean.parseBoolean(new String(value)));
    }
    return result;
  }

  /**
   * ??????????Ϣ??Ŀǰ֧??baseTaskType??strategy???ݡ?
   */
  public void importConfig(String config, Writer writer, boolean isUpdate)
      throws Exception {
    ConfigNode configNode = gson.fromJson(config, ConfigNode.class);
    if (configNode != null) {
      String path = configNode.getRootPath() + "/"
          + configNode.getConfigType();
      ZKTools.createPath(getZooKeeper(), path, CreateMode.PERSISTENT, zkManager.getAcl());
      String y_node = path + "/" + configNode.getName();
      if (getZooKeeper().exists(y_node, false) == null) {
        writer.append("?ɹ?????????????Ϣ\n");
        getZooKeeper().create(y_node, configNode.getValue().getBytes(),
            zkManager.getAcl(), CreateMode.PERSISTENT);
      } else if (isUpdate) {
        writer.append("????????Ϣ?Ѿ????ڣ?????ǿ?Ƹ?????\n");
        getZooKeeper().setData(y_node,
            configNode.getValue().getBytes(), -1);
      } else {
        writer.append("????????Ϣ?Ѿ????ڣ??????Ҫ???£???????ǿ?Ƹ???\n");
      }
    }
    writer.append(configNode.toString());
  }

  /**
   * ?????????Ϣ??Ŀǰ????baseTaskType??strategy???ݡ?
   */
  public StringBuffer exportConfig(String rootPath, Writer writer)
      throws Exception {
    StringBuffer buffer = new StringBuffer();
    for (String type : new String[]{"baseTaskType", "strategy"}) {
      if (type.equals("baseTaskType")) {
        writer.write("?????????????б???\n");
      } else {
        writer.write("?????????????б???\n");
      }
      String bTTypePath = rootPath + "/" + type;
      List fNodeList = getZooKeeper().getChildren(bTTypePath,
          false);
      for (int i = 0; i < fNodeList.size(); i++) {
        String fNode = fNodeList.get(i);
        ConfigNode configNode = new ConfigNode(rootPath, type, fNode);
        configNode.setValue(
            new String(this.getZooKeeper().getData(bTTypePath + "/" + fNode, false, null)));
        buffer.append(gson.toJson(configNode));
        buffer.append("\n");
        writer.write(configNode.toString());
      }
      writer.write("\n\n");
    }
    if (buffer.length() > 0) {
      String str = buffer.toString();
      return new StringBuffer(str.substring(0, str.length() - 1));
    }
    return buffer;
  }

  public List loadAllManagerFactoryInfo() throws Exception {
    String zkPath = this.PATH_ManagerFactory;
    List result = new ArrayList();
    List names = this.getZooKeeper().getChildren(zkPath, false);
    Collections.sort(names, new Comparator() {
      public int compare(String u1, String u2) {
        return u1.substring(u1.lastIndexOf("$") + 1).compareTo(
            u2.substring(u2.lastIndexOf("$") + 1));
      }
    });
    for (String name : names) {
      ManagerFactoryInfo info = new ManagerFactoryInfo();
      info.setUuid(name);
      byte[] value = this.getZooKeeper().getData(zkPath + "/" + name, false, null);
      if (value == null) {
        info.setStart(true);
      } else {
        info.setStart(Boolean.parseBoolean(new String(value)));
      }
      result.add(info);
    }
    return result;
  }

  public void printTree(String path, Writer writer, String lineSplitChar)
      throws Exception {
    ZKTools.printTree(this.getZooKeeper(), path, writer, lineSplitChar);
  }

  public void deleteTree(String path) throws Exception {
    ZKTools.deleteTree(this.getZooKeeper(), path);
  }

  public ZooKeeper getZooKeeper() throws Exception {
    return this.zkManager.getZooKeeper();
  }

  public String getRootPath() {
    return this.zkManager.getRootPath();
  }
}