Java Code Examples for org.apache.helix.tools.ClusterSetup#rebalanceStorageCluster()

protected void setupStorageCluster(ClusterSetup setupTool, String clusterName, String dbName,
    int partitionNr, String prefix, int startPort, String stateModel, int replica,
    boolean rebalance) {
  setupTool.addResourceToCluster(clusterName, dbName, partitionNr, stateModel);
  for (int i = 0; i < NODE_NR; i++) {
    String instanceName = prefix + "_" + (startPort + i);
    setupTool.addInstanceToCluster(clusterName, instanceName);
  }
  if (rebalance) {
    setupTool.rebalanceStorageCluster(clusterName, dbName, replica);
  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  _setupTool = new ClusterSetup(ZK_ADDR);
  // setup storage cluster
  _setupTool.addCluster(CLUSTER_NAME, true);
  _setupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, NUM_PARTITIONS, STATE_MODEL);

  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _setupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  _setupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB, NUM_REPLICAS);

  // start dummy participants
  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].syncStart();
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  boolean result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.MasterNbInExtViewVerifier(ZK_ADDR, CLUSTER_NAME), 10000, 100);
  Assert.assertTrue(result);

  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME), 10000, 100);
  Assert.assertTrue(result);

  // create cluster manager
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
}
 

public static void setupCluster(String clusterName, String zkAddr, int startPort,
    String participantNamePrefix, String resourceNamePrefix, int resourceNb, int partitionNb,
    int nodesNb, int replica, String stateModelDef, RebalanceMode mode, boolean doRebalance)
    throws Exception {
  HelixZkClient zkClient = SharedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddr));
  if (zkClient.exists("/" + clusterName)) {
    LOG.warn("Cluster already exists:" + clusterName + ". Deleting it");
    zkClient.deleteRecursively("/" + clusterName);
  }

  ClusterSetup setupTool = new ClusterSetup(zkAddr);
  setupTool.addCluster(clusterName, true);

  for (int i = 0; i < nodesNb; i++) {
    int port = startPort + i;
    setupTool.addInstanceToCluster(clusterName, participantNamePrefix + "_" + port);
  }

  for (int i = 0; i < resourceNb; i++) {
    String resourceName = resourceNamePrefix + i;
    setupTool.addResourceToCluster(clusterName, resourceName, partitionNb, stateModelDef,
        mode.toString());
    if (doRebalance) {
      setupTool.rebalanceStorageCluster(clusterName, resourceName, replica);
    }
  }
  zkClient.close();
}
 

protected void createDBInSemiAuto(ClusterSetup clusterSetup, String clusterName, String dbName,
    List<String> preferenceList, String stateModelDef, int numPartition, int replica) {
  clusterSetup.addResourceToCluster(clusterName, dbName, numPartition, stateModelDef,
      IdealState.RebalanceMode.SEMI_AUTO.toString());
  clusterSetup.rebalanceStorageCluster(clusterName, dbName, replica);

  IdealState is =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(clusterName, dbName);
  for (String p : is.getPartitionSet()) {
    is.setPreferenceList(p, preferenceList);
  }
  clusterSetup.getClusterManagementTool().setResourceIdealState(clusterName, dbName, is);
}
 

protected void setupDBs(ClusterSetup clusterSetup) {
  // Set up target db
  if (_numDbs > 1) {
    for (int i = 0; i < _numDbs; i++) {
      int varyNum = _partitionVary ? 10 * i : 0;
      String db = WorkflowGenerator.DEFAULT_TGT_DB + i;
      clusterSetup.addResourceToCluster(CLUSTER_NAME, db, _numPartitions + varyNum,
          MASTER_SLAVE_STATE_MODEL, IdealState.RebalanceMode.FULL_AUTO.toString());
      clusterSetup.rebalanceStorageCluster(CLUSTER_NAME, db, _numReplicas);
      _testDbs.add(db);
    }
  } else {
    if (_instanceGroupTag) {
      clusterSetup.addResourceToCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
          _numPartitions, "OnlineOffline", IdealState.RebalanceMode.FULL_AUTO.name());
      IdealState idealState = clusterSetup.getClusterManagementTool()
          .getResourceIdealState(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB);
      idealState.setInstanceGroupTag("TESTTAG0");
      clusterSetup.getClusterManagementTool().setResourceIdealState(CLUSTER_NAME,
          WorkflowGenerator.DEFAULT_TGT_DB, idealState);
    } else {
      clusterSetup.addResourceToCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
          _numPartitions, MASTER_SLAVE_STATE_MODEL, IdealState.RebalanceMode.FULL_AUTO.name());
    }
    clusterSetup.rebalanceStorageCluster(CLUSTER_NAME, WorkflowGenerator.DEFAULT_TGT_DB,
        _numReplicas);
  }
}
 

@Test
public void testBasic() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final int n = 5;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  MockParticipantManager[] participants = new MockParticipantManager[n];

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      10, // partitions per resource
      n, // number of nodes
      3, // replicas
      "MasterSlave", true); // do rebalance

  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
  controller.syncStart();

  // start participants
  for (int i = 0; i < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);

  // add a new idealState without registering message handling factory
  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  setupTool.addResourceToCluster(clusterName, "TestDB1", 16, "MasterSlave");

  ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB1"));
  idealState.setStateModelFactoryName("TestDB1_Factory");
  accessor.setProperty(keyBuilder.idealStates("TestDB1"), idealState);
  setupTool.rebalanceStorageCluster(clusterName, "TestDB1", 3);

  // assert that we have received OFFLINE->SLAVE messages for all partitions
  int totalMsgs = 0;
  for (int retry = 0; retry < 5; retry++) {
    Thread.sleep(100);
    totalMsgs = 0;
    for (int i = 0; i < n; i++) {
      List<Message> msgs =
          accessor.getChildValues(keyBuilder.messages(participants[i].getInstanceName()), true);
      totalMsgs += msgs.size();
    }

    if (totalMsgs == 48) // partition# x replicas
      break;
  }

  Assert
      .assertEquals(
          totalMsgs,
          48,
          "Should accumulated 48 unprocessed messages (1 O->S per partition per replica) because TestDB1 is added without state-model-factory but was "
              + totalMsgs);

  // register "TestDB1_Factory" state model factory
  // Logger.getRootLogger().setLevel(Level.INFO);
  for (int i = 0; i < n; i++) {
    participants[i].getStateMachineEngine()
        .registerStateModelFactory("MasterSlave", new MockMSModelFactory(), "TestDB1_Factory");
  }

  result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);

  // clean up
  // wait for all zk callbacks done
  controller.syncStop();
  for (int i = 0; i < 5; i++) {
    participants[i].syncStop();
  }
  deleteCluster(clusterName);

  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));

}
 

@Test()
public void testPauseSignal() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  final String clusterName = className + "_" + methodName;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  MockParticipantManager[] participants = new MockParticipantManager[5];

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      10, // partitions per resource
      5, // number of nodes
      3, // replicas
      "MasterSlave", true); // do rebalance

  // start controller
  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
  controller.syncStart();

  // start participants
  for (int i = 0; i < 5; i++) {
    String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  boolean result = ClusterStateVerifier
      .verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // pause the cluster and make sure pause is persistent
  final HelixDataAccessor tmpAccessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));

  String cmd = "-zkSvr " + ZK_ADDR + " -enableCluster " + clusterName + " false";
  ClusterSetup.processCommandLineArgs(cmd.split(" "));

  tmpAccessor.setProperty(tmpAccessor.keyBuilder().pause(), new PauseSignal("pause"));

  // wait for controller to be signaled by pause
  Thread.sleep(1000);

  // add a new resource group
  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  setupTool.addResourceToCluster(clusterName, "TestDB1", 10, "MasterSlave");
  setupTool.rebalanceStorageCluster(clusterName, "TestDB1", 3);

  // make sure TestDB1 external view is empty
  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 1000, clusterName, "TestDB1",
      TestHelper.setOf("localhost_12918", "localhost_12919", "localhost_12920",
          "localhost_12921", "localhost_12922"),
      ZK_ADDR);

  // resume controller
  final HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));

  cmd = "-zkSvr " + ZK_ADDR + " -enableCluster " + clusterName + " true";
  ClusterSetup.processCommandLineArgs(cmd.split(" "));
  result = ClusterStateVerifier
      .verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // clean up
  controller.syncStop();
  for (int i = 0; i < 5; i++) {
    participants[i].syncStop();
  }

  deleteCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

@Test()
public void testStandAloneCMSessionExpiry() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, PARTICIPANT_PREFIX, "TestDB", 1, 20, 5, 3,
      "MasterSlave", true);

  MockParticipantManager[] participants = new MockParticipantManager[5];
  for (int i = 0; i < 5; i++) {
    String instanceName = "localhost_" + (12918 + i);
    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
  controller.syncStart();

  boolean result;
  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // participant session expiry
  MockParticipantManager participantToExpire = participants[1];

  // System.out.println("Expire participant session");
  String oldSessionId = participantToExpire.getSessionId();

  ZkTestHelper.expireSession(participantToExpire.getZkClient());
  String newSessionId = participantToExpire.getSessionId();
  // System.out.println("oldSessionId: " + oldSessionId + ", newSessionId: " + newSessionId);
  Assert.assertTrue(newSessionId.compareTo(oldSessionId) > 0,
      "Session id should be increased after expiry");

  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  setupTool.addResourceToCluster(clusterName, "TestDB1", 10, "MasterSlave");
  setupTool.rebalanceStorageCluster(clusterName, "TestDB1", 3);

  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // controller session expiry
  // System.out.println("Expire controller session");
  oldSessionId = controller.getSessionId();
  ZkTestHelper.expireSession(controller.getZkClient());
  newSessionId = controller.getSessionId();
  // System.out.println("oldSessionId: " + oldSessionId + ", newSessionId: " + newSessionId);
  Assert.assertTrue(newSessionId.compareTo(oldSessionId) > 0,
      "Session id should be increased after expiry");

  setupTool.addResourceToCluster(clusterName, "TestDB2", 8, "MasterSlave");
  setupTool.rebalanceStorageCluster(clusterName, "TestDB2", 3);

  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // clean up
  System.out.println("Clean up ...");
  controller.syncStop();
  for (int i = 0; i < 5; i++) {
    participants[i].syncStop();
  }

  deleteCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

public static void setupCluster(String uniqClusterName, String zkAddr, int numResources,
    int numPartitionsPerResource, int numInstances, int replica, boolean doRebalance)
    throws Exception {
  HelixZkClient zkClient = SharedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(ZK_ADDR));

  try {
    zkClient.setZkSerializer(new ZNRecordSerializer());

    // String clusterName = CLUSTER_PREFIX + "_" + uniqClusterName;
    String clusterName = uniqClusterName;
    if (zkClient.exists("/" + clusterName)) {
      LOG.warn("test cluster already exists:" + clusterName + ", test name:" + uniqClusterName + " is not unique or test has been run without cleaning up zk; deleting it");
      zkClient.deleteRecursively("/" + clusterName);
    }

    if (_testInfoMap.containsKey(uniqClusterName)) {
      LOG.warn("test info already exists:" + uniqClusterName + " is not unique or test has been run without cleaning up test info map; removing it");
      _testInfoMap.remove(uniqClusterName);
    }
    TestInfo testInfo =
        new TestInfo(clusterName, zkClient, numResources, numPartitionsPerResource, numInstances,
            replica);
    _testInfoMap.put(uniqClusterName, testInfo);

    ClusterSetup setupTool = new ClusterSetup(zkAddr);
    setupTool.addCluster(clusterName, true);

    for (int i = 0; i < numInstances; i++) {
      int port = START_PORT + i;
      setupTool.addInstanceToCluster(clusterName, PARTICIPANT_PREFIX + "_" + port);
    }

    for (int i = 0; i < numResources; i++) {
      String dbName = TEST_DB_PREFIX + i;
      setupTool.addResourceToCluster(clusterName, dbName, numPartitionsPerResource, STATE_MODEL);
      if (doRebalance) {
        setupTool.rebalanceStorageCluster(clusterName, dbName, replica);

        // String idealStatePath = "/" + clusterName + "/" +
        // PropertyType.IDEALSTATES.toString() + "/"
        // + dbName;
        // ZNRecord idealState = zkClient.<ZNRecord> readData(idealStatePath);
        // testInfo._idealStateMap.put(dbName, idealState);
      }
    }
  } finally {
    zkClient.close();
  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  String className = TestHelper.getTestClassName();
  _clusterNamePrefix = className;

  System.out
      .println("START " + _clusterNamePrefix + " at " + new Date(System.currentTimeMillis()));

  // setup 10 clusters
  for (int i = 0; i < clusterNb; i++) {
    String clusterName = _clusterNamePrefix + "0_" + i;
    String participantName = "localhost" + i;
    String resourceName = "TestDB" + i;
    TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
        participantName, // participant name prefix
        resourceName, // resource name prefix
        1, // resources
        8, // partitions per resource
        n, // number of nodes
        3, // replicas
        "MasterSlave", true); // do rebalance

    _clusters.add(clusterName);
  }

  // setup controller cluster
  _controllerClusterName = "CONTROLLER_" + _clusterNamePrefix;
  TestHelper.setupCluster(_controllerClusterName, ZK_ADDR, // controller
      0, // port
      "controller", // participant name prefix
      _clusterNamePrefix, // resource name prefix
      1, // resources
      clusterNb, // partitions per resource
      n, // number of nodes
      3, // replicas
      "LeaderStandby", true); // do rebalance

  // start distributed cluster controllers
  _controllers = new ClusterDistributedController[n + n];
  for (int i = 0; i < n; i++) {
    _controllers[i] =
        new ClusterDistributedController(ZK_ADDR, _controllerClusterName, "controller_" + i);
    _controllers[i].syncStart();
  }

  ZkHelixClusterVerifier controllerClusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(_controllerClusterName).setZkClient(_gZkClient)
          .build();

  Assert.assertTrue(controllerClusterVerifier.verifyByPolling(),
      "Controller cluster NOT in ideal state");

  // start first cluster
  _participants = new MockParticipantManager[n];
  _firstClusterName = _clusterNamePrefix + "0_0";
  for (int i = 0; i < n; i++) {
    String instanceName = "localhost0_" + (12918 + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, _firstClusterName, instanceName);
    _participants[i].syncStart();
  }

  ZkHelixClusterVerifier firstClusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(_firstClusterName).setZkClient(_gZkClient)
          .build();
  Assert.assertTrue(firstClusterVerifier.verifyByPolling(), "first cluster NOT in ideal state");

  // add more controllers to controller cluster
  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  for (int i = 0; i < n; i++) {
    String controller = "controller_" + (n + i);
    setupTool.addInstanceToCluster(_controllerClusterName, controller);
  }
  setupTool.rebalanceStorageCluster(_controllerClusterName, _clusterNamePrefix + "0", 6);
  for (int i = n; i < 2 * n; i++) {
    _controllers[i] =
        new ClusterDistributedController(ZK_ADDR, _controllerClusterName, "controller_" + i);
    _controllers[i].syncStart();
  }

  // verify controller cluster
  Assert.assertTrue(controllerClusterVerifier.verifyByPolling(),
      "Controller cluster NOT in ideal state");

  // verify first cluster
  Assert.assertTrue(firstClusterVerifier.verifyByPolling(), "first cluster NOT in ideal state");
  // verify all the rest clusters
  for (int i = 1; i < clusterNb; i++) {
    ZkHelixClusterVerifier clusterVerifier =
        new BestPossibleExternalViewVerifier.Builder(_clusterNamePrefix + "0_" + i)
            .setZkClient(_gZkClient).build();
    Assert.assertTrue(clusterVerifier.verifyByPolling(), "Cluster NOT in ideal state.");
  }
}