code: package resa.evaluation.scheduler.plan;
  
  import backtype.storm.Config;
  import backtype.storm.generated.Nimbus;
  import backtype.storm.generated.RebalanceOptions;
  import backtype.storm.scheduler.ExecutorDetails;
  import backtype.storm.utils.NimbusClient;
  import backtype.storm.utils.Utils;
  import org.apache.curator.framework.CuratorFramework;
  import org.apache.zookeeper.ZKUtil;
  import org.junit.Before;
  import org.junit.Test;
  import resa.evaluation.migrate.ConsistentHashing;
  import resa.evaluation.migrate.MigPlan;
  import resa.util.TopologyHelper;
  
  import java.nio.file.Files;
  import java.nio.file.Paths;
  import java.util.*;
  import java.util.stream.Collectors;
  
  /**
   * Created by ding on 14-8-12.
   */
  public class MigrateSimulateForCH {
  
      private CuratorFramework zk;
      private Map<String, Object> conf;
      private int[] allocations;
      private List<String> slots;
      private double[] dataSizes;
      private ExecutorDetails[] executors;
      private String topoName = "fpt";
  
      @Before
      public void init() throws Exception {
          conf = Utils.readStormConfig();
          conf.put(Config.NIMBUS_HOST, "192.168.0.30");
          conf.put(Config.STORM_ZOOKEEPER_SERVERS, Arrays.asList("192.168.0.30"));
          zk = Utils.newCuratorStarted(conf, (List<String>) conf.get(Config.STORM_ZOOKEEPER_SERVERS),
                  conf.get(Config.STORM_ZOOKEEPER_PORT));
          allocations = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/ch-migrate.txt")).stream()
                  .map(String::trim).filter(s -> !s.isEmpty())
                  .mapToInt(Integer::parseInt).toArray();
          dataSizes = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/fp-data-sizes-064-70.txt")).stream()
                  .map(String::trim).filter(s -> !s.isEmpty()).mapToDouble(Double::valueOf).toArray();
          slots = getSlots();
          String topoId = TopologyHelper.getTopologyId(topoName, conf);
          Map<String, Map<ExecutorDetails, String>> assignment = TopologyHelper.getAssignment(conf, topoId);
          Map<ExecutorDetails, String> compAssignment = assignment.get("detector");
          executors = compAssignment.keySet().stream().sorted(Comparator.comparing(ExecutorDetails::getStartTask))
                  .toArray(ExecutorDetails[]::new);
      }
  
      @Test
      public void deleteZkNode() throws Exception {
          ZKUtil.deleteRecursive(zk.getZookeeperClient().getZooKeeper(), "/resa");
      }
  
      @Test
      public void run() throws Exception {
          int state = 0;
          List<int[]> from = new ArrayList<>();
          List<int[]> to = new ArrayList<>();
          String[] workers = slots.stream().limit(allocations[state]).toArray(String[]::new);
          assignment2State(from, workers);
          doRebalance(topoName, allocations[state]);
          System.out.println("Init state, curr alloc=" + allocations[state]);
          Scanner scanner = new Scanner(System.in);
          String line;
          int count = 0;
          while ((line = scanner.nextLine()) != null) {
              if (line.equals("c")) {
                  count++;
                  System.out.println("run " + count + ", from " + allocations[state] + " to " + allocations[state + 1]);
                  System.out.println("old workers:" + Arrays.toString(workers));
                  workers = migrateTo(from, workers, to, allocations[state + 1]);
                  System.out.println("new workers:" + Arrays.toString(workers));
                  assignment2State(to, workers);
                  doRebalance(topoName, allocations[state + 1]);
                  state++;
                  from.clear();
                  from.addAll(to);
                  to.clear();
              }
          }
      }
  
      private String[] migrateTo(List<int[]> from, String[] workers, List<int[]> to, int toNumWorkers) {
          MigPlan res = ConsistentHashing.calc(dataSizes, from, toNumWorkers);
          to.addAll(res.tasks);
          String[] newWorkers = new String[toNumWorkers];
          Set<String> selected = new HashSet<>();
          int[][] map = res.matching;
          for (int i = 0; i < map.length; i++) {
              newWorkers[map[i][0]] = workers[map[i][1]];
              selected.add(newWorkers[map[i][0]]);
          }
          for (int i = 0; i < newWorkers.length; i++) {
              if (newWorkers[i] == null) {
                  for (String w : slots) {
                      if (!selected.contains(w)) {
                          selected.add(w);
                          newWorkers[i] = w;
                          break;
                      }
                  }
              }
          }
          return newWorkers;
      }
  
      public void assignment2State(List<int[]> alloc, String[] workers) throws Exception {
          Map<String, String> exe2WorkerSlot = new HashMap<>();
          for (int i = 0; i < alloc.size(); i++) {
              for (int task : alloc.get(i)) {
                  exe2WorkerSlot.put(executors[task].getStartTask() + "-" + executors[task].getEndTask(), workers[i]);
              }
          }
  //        System.out.println(exe2WorkerSlot);
          String path = "/resa/" + TopologyHelper.getTopologyId(topoName, conf);
          if (zk.checkExists().forPath(path) != null) {
              zk.setData().forPath(path, Utils.serialize(new HashMap<>(Collections.singletonMap("assignment",
                      exe2WorkerSlot))));
          } else {
              zk.create().creatingParentsIfNeeded().forPath(path, Utils.serialize(
                      new HashMap<>(Collections.singletonMap("assignment", exe2WorkerSlot))));
          }
      }
  
  
      private List<String> getSlots() throws Exception {
          NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
          Nimbus.Client nimbus = nimbusClient.getClient();
          List<String> slots = nimbus.getClusterInfo().get_supervisors().stream()
                  .flatMap(s -> Arrays.asList(s.get_host() + ":6700", s.get_host() + ":6701").stream())
                  .collect(Collectors.toList());
          Collections.shuffle(slots);
          System.out.println(slots.size());
          return slots;
      }
  
      private void doRebalance(String topoName, int numWorkers) throws Exception {
          NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
          Nimbus.Client nimbus = nimbusClient.getClient();
          RebalanceOptions options = new RebalanceOptions();
          options.set_num_workers(numWorkers);
          options.set_wait_secs(1);
          System.out.println("Reassigning to " + numWorkers + " workers");
          nimbus.rebalance(topoName, options);
      }
  
      public static void main(String[] args) throws Exception {
          MigrateSimulateForCH simulate = new MigrateSimulateForCH();
          simulate.init();
          simulate.run();
      }
  
  }
  