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

@Test()
public void testDisableNode() throws Exception {
  String command =
      "-zkSvr " + ZK_ADDR + " -enableInstance " + CLUSTER_NAME + " " + PARTICIPANT_PREFIX
          + "_12918" + " TestDB TestDB_0 false";
  ClusterSetup.processCommandLineArgs(command.split(" "));
  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_12918", true);

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

}
 

@BeforeClass()
public void beforeClass() throws Exception {
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addCluster "
      + clusterName));
  // ClusterSetup
  // .processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR +
  // " -addCluster relay-cluster-12345"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addResource "
      + clusterName + " db-12345 120 MasterSlave"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addNode " + clusterName
      + " localhost:8900"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addNode " + clusterName
      + " localhost:8901"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addNode " + clusterName
      + " localhost:8902"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addNode " + clusterName
      + " localhost:8903"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -addNode " + clusterName
      + " localhost:8904"));
  ClusterSetup.processCommandLineArgs(createArgs("-zkSvr " + ZK_ADDR + " -rebalance "
      + clusterName + " db-12345 3"));
}
 

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

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, _PARTITIONS, STATE_MODEL);

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

  // Set the timeout values
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, TEST_DB);
  String stateTransition = "SLAVE" + "-" + "MASTER" + "_" + Message.Attributes.TIMEOUT;
  idealState.getRecord().setSimpleField(stateTransition, "300");

  String command =
      "-zkSvr " + ZK_ADDR + " -addResourceProperty " + CLUSTER_NAME + " " + TEST_DB + " "
          + stateTransition + " 200";
  ClusterSetup.processCommandLineArgs(command.split(" "));
}
 

@Test()
public void testDropResource() throws Exception {
  // add a resource to be dropped
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "MyDB", 6, STATE_MODEL);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "MyDB", 3);

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

  String command = "-zkSvr " + ZK_ADDR + " -dropResource " + CLUSTER_NAME + " " + "MyDB";
  ClusterSetup.processCommandLineArgs(command.split(" "));

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB",
      TestHelper.<String> setOf("localhost_12918", "localhost_12919", "localhost_12920",
          "localhost_12921", "localhost_12922"), ZK_ADDR);
}
 

@Test()
public void testDropResourceAutoRebalance() throws Exception {
  // add a resource to be dropped
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "MyDB", _PARTITIONS, "OnlineOffline",
      RebalanceMode.FULL_AUTO + "");

  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "MyDB", 1);

  boolean result = ClusterStateVerifier
      .verifyByZkCallback(new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, "MyDB"));
  Assert.assertTrue(result);

  String command = "-zkSvr " + ZK_ADDR + " -dropResource " + CLUSTER_NAME + " " + "MyDB";
  ClusterSetup.processCommandLineArgs(command.split(" "));

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB",
      TestHelper.setOf("localhost_12918", "localhost_12919", "localhost_12920", "localhost_12921",
          "localhost_12922"),
      ZK_ADDR);

  // add a resource to be dropped
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "MyDB2", _PARTITIONS, "MasterSlave",
      RebalanceMode.FULL_AUTO + "");

  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "MyDB2", 1);

  result = ClusterStateVerifier
      .verifyByZkCallback(new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, "MyDB2"));
  Assert.assertTrue(result);

  command = "-zkSvr " + ZK_ADDR + " -dropResource " + CLUSTER_NAME + " " + "MyDB2";
  ClusterSetup.processCommandLineArgs(command.split(" "));

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB2",
      TestHelper.setOf("localhost_12918", "localhost_12919", "localhost_12920", "localhost_12921",
          "localhost_12922"),
      ZK_ADDR);
}
 

@Override
public void doTransition(Message message, NotificationContext context) {
  HelixManager manager = context.getManager();
  String clusterName = manager.getClusterName();

  String instance = message.getTgtName();
  String partitionName = message.getPartitionName();
  String fromState = message.getFromState();
  String toState = message.getToState();
  if (instance.equals("localhost_12919") && partitionName.equals("TestDB0_0")) {
    if (fromState.equals("SLAVE") && toState.equals("OFFLINE")) {
      slaveToOfflineCnt++;

      try {
        String command = "--zkSvr " + ZK_ADDR + " --enablePartition true " + clusterName
            + " localhost_12919 TestDB0 TestDB0_0";

        ClusterSetup.processCommandLineArgs(command.split("\\s+"));
      } catch (Exception e) {
        LOG.error("Exception in cluster setup", e);
      }

    } else if (slaveToOfflineCnt > 0 && fromState.equals("OFFLINE")
        && toState.equals("SLAVE")) {
      offlineToSlave++;
    }
  }
}
 

@Test()
public void testDropResourceWhileNodeDead() throws Exception {
  // add a resource to be dropped
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "MyDB2", 16, STATE_MODEL);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "MyDB2", 3);

  boolean verifyResult =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(verifyResult);

  String hostToKill = "localhost_12920";

  _participants[2].syncStop();
  Thread.sleep(1000);

  String command = "-zkSvr " + ZK_ADDR + " -dropResource " + CLUSTER_NAME + " " + "MyDB2";
  ClusterSetup.processCommandLineArgs(command.split(" "));

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB2",
      TestHelper.<String> setOf("localhost_12918", "localhost_12919",
      /* "localhost_12920", */"localhost_12921", "localhost_12922"), ZK_ADDR);

  _participants[2] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, hostToKill);
  _participants[2].syncStart();

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB2",
      TestHelper.<String> setOf("localhost_12918", "localhost_12919", "localhost_12920",
          "localhost_12921", "localhost_12922"), ZK_ADDR);
}
 

void assertClusterSetupException(String command) {
  boolean exceptionThrown = false;
  try {
    ClusterSetup.processCommandLineArgs(command.split(" "));
  } catch (Exception e) {
    exceptionThrown = true;
  }
  Assert.assertTrue(exceptionThrown);
}
 

@Test
public void testDropSchemataResource() 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
      8, // partitions per resource
      n, // 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 < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

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

  ZkHelixClusterVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());

  // add schemata resource group
  String command = "--zkSvr " + ZK_ADDR + " --addResource " + clusterName
      + " schemata 1 STORAGE_DEFAULT_SM_SCHEMATA";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  command = "--zkSvr " + ZK_ADDR + " --rebalance " + clusterName + " schemata " + n;
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  Assert.assertTrue(verifier.verifyByPolling());

  // drop schemata resource group
  // System.out.println("Dropping schemata resource group...");
  command = "--zkSvr " + ZK_ADDR + " --dropResource " + clusterName + " schemata";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(verifier.verifyByPolling());

  Thread.sleep(400);

  assertEmptyCSandEV(clusterName, "schemata", participants);

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

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

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

  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
      8, // partitions per resource
      n, // number of nodes
      3, // replicas
      "MasterSlave", true); // do rebalance

  // set ideal state to customized mode
  ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<>(_gZkClient);
  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setRebalanceMode(RebalanceMode.CUSTOMIZED);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

  // start controller
  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();
  }

  ZkHelixClusterVerifier _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // disable localhost_12918
  String command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " false";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is in OFFLINE state
  Map<String, Map<String, String>> expectStateMap = new HashMap<>();
  Map<String, String> expectInstanceStateMap = new HashMap<>();
  expectInstanceStateMap.put(disableNode, "OFFLINE");
  expectStateMap.put(".*", expectInstanceStateMap);
  boolean result =
      ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "==");
  Assert.assertTrue(result, disableNode + " should be in OFFLINE");

  // re-enable localhost_12918
  command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " true";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is NOT in OFFLINE state
  result = ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "!=");
  Assert.assertTrue(result, disableNode + " should NOT be in OFFLINE");

  // 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 testDisableNodeAutoIS() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final int n = 5;
  String disableNode = "localhost_12919";

  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
      8, // partitions per resource
      n, // 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 < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

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

  ZkHelixClusterVerifier _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // disable localhost_12919
  String command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " false";
  ClusterSetup.processCommandLineArgs(command.split(" "));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12919 is in OFFLINE state
  Map<String, Map<String, String>> expectStateMap = new HashMap<>();
  Map<String, String> expectInstanceStateMap = new HashMap<>();
  expectInstanceStateMap.put(disableNode, "OFFLINE");
  expectStateMap.put(".*", expectInstanceStateMap);
  boolean result =
      ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "==");
  Assert.assertTrue(result, disableNode + " should be in OFFLINE");

  // re-enable localhost_12919
  command =
      "--zkSvr " + ZK_ADDR + " --enableInstance " + clusterName + " " + disableNode + " true";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  result = ClusterStateVerifier
      .verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // make sure localhost_12919 is NOT in OFFLINE state
  result = ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "!=");
  Assert.assertTrue(result, disableNode + " should NOT be in OFFLINE");

  // 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 testDisablePartitionCustomIS() 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
      8, // partitions per resource
      n, // number of nodes
      3, // replicas
      "MasterSlave", true); // do rebalance

  // set ideal state to customized mode
  ZkBaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
  ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setRebalanceMode(RebalanceMode.CUSTOMIZED);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

  // start controller
  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();
  }

  BestPossibleExternalViewVerifier _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR)
          .setZkClient(_gZkClient).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // disable [TestDB0_0, TestDB0_5] on localhost_12919
  String command = "--zkSvr " + ZK_ADDR + " --enablePartition false " + clusterName
      + " localhost_12919 TestDB0 TestDB0_0 TestDB0_5";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is in OFFLINE state for [TestDB0_0, TestDB0_5]
  Map<String, Map<String, String>> expectStateMap = new HashMap<>();
  Map<String, String> expectInstanceStateMap = new HashMap<>();
  expectInstanceStateMap.put("localhost_12919", "OFFLINE");
  expectStateMap.put("TestDB0_0", expectInstanceStateMap);
  expectStateMap.put("TestDB0_5", expectInstanceStateMap);
  boolean result =
      ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "==");
  Assert.assertTrue(result,
      "localhost_12919" + " should be in OFFLINE for [TestDB0_0, TestDB0_5]");

  // re-enable localhost_12919 for [TestDB0_0, TestDB0_5]
  command = "--zkSvr " + ZK_ADDR + " --enablePartition true " + clusterName
      + " localhost_12919 TestDB0 TestDB0_0 TestDB0_5";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  result = ClusterStateVerifier
      .verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR, clusterName));
  Assert.assertTrue(result);

  // make sure localhost_12919 is NOT in OFFLINE state for [TestDB0_0, TestDB0_5]
  result = ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "!=");
  Assert.assertTrue(result, "localhost_12919" + " should NOT be in OFFLINE");

  // 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 testDisablePartitionAutoIS() 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
      8, // partitions per resource
      n, // 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 < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

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

  ZkHelixClusterVerifier _clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // disable [TestDB0_0, TestDB0_5] on localhost_12919
  String command = "--zkSvr " + ZK_ADDR + " --enablePartition false " + clusterName
      + " localhost_12919 TestDB0 TestDB0_0 TestDB0_5";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12918 is in OFFLINE state for [TestDB0_0, TestDB0_5]
  Map<String, Map<String, String>> expectStateMap = new HashMap<>();
  Map<String, String> expectInstanceStateMap = new HashMap<>();
  expectInstanceStateMap.put("localhost_12919", "OFFLINE");
  expectStateMap.put("TestDB0_0", expectInstanceStateMap);
  expectStateMap.put("TestDB0_5", expectInstanceStateMap);
  boolean result =
      ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "==");
  Assert.assertTrue(result,
      "localhost_12919" + " should be in OFFLINE for [TestDB0_0, TestDB0_5]");

  // re-enable localhost_12919 for [TestDB0_0, TestDB0_5]
  command = "--zkSvr " + ZK_ADDR + " --enablePartition true " + clusterName
      + " localhost_12919 TestDB0 TestDB0_0 TestDB0_5";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // make sure localhost_12919 is NOT in OFFLINE state for [TestDB0_0, TestDB0_5]
  result = ZkTestHelper.verifyState(_gZkClient, clusterName, "TestDB0", expectStateMap, "!=");
  Assert.assertTrue(result, "localhost_12919" + " should NOT be in OFFLINE");

  // 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 testDropErrorPartitionCustomIS() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final int n = 2;

  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
      2, // partitions per resource
      n, // number of nodes
      2, // replicas
      "MasterSlave", false); // do rebalance

  // set custom ideal-state
  CustomModeISBuilder isBuilder = new CustomModeISBuilder("TestDB0");
  isBuilder.setNumPartitions(2);
  isBuilder.setNumReplica(2);
  isBuilder.setStateModel("MasterSlave");
  isBuilder.assignInstanceAndState("TestDB0_0", "localhost_12918", "MASTER");
  isBuilder.assignInstanceAndState("TestDB0_0", "localhost_12919", "SLAVE");
  isBuilder.assignInstanceAndState("TestDB0_1", "localhost_12919", "MASTER");
  isBuilder.assignInstanceAndState("TestDB0_1", "localhost_12918", "SLAVE");

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), isBuilder.build());

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

  // start participants
  Map<String, Set<String>> errTransitions = new HashMap<>();
  errTransitions.put("SLAVE-MASTER", TestHelper.setOf("TestDB0_0"));

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

    if (i == 0) {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
      participants[i].setTransition(new ErrTransition(errTransitions));
    } else {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    }
    participants[i].syncStart();
  }

  Map<String, Map<String, String>> errStateMap = new HashMap<>();
  errStateMap.put("TestDB0", new HashMap<>());
  errStateMap.get("TestDB0").put("TestDB0_0", "localhost_12918");

  ZkHelixClusterVerifier verifier = new BestPossibleExternalViewVerifier.Builder(clusterName)
      .setZkAddr(ZK_ADDR).setErrStates(errStateMap).build();
  Assert.assertTrue(verifier.verifyByPolling());

  // drop resource containing error partitions should drop the partition successfully
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--dropResource", clusterName, "TestDB0"
  });

  // make sure TestDB0_0 partition is dropped
  verifier = new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling(), "Should be empty exeternal-view");
  Thread.sleep(400);

  assertEmptyCSandEV(clusterName, "TestDB0", participants);

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

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

@Test
public void testEnablePartitionDuringDisable() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  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 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

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

  // start participants
  EnablePartitionTransition transition = new EnablePartitionTransition();
  MockParticipantManager[] participants = new MockParticipantManager[5];
  for (int i = 0; i < 5; i++) {
    String instanceName = "localhost_" + (12918 + i);

    if (instanceName.equals("localhost_12919")) {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
      participants[i].setTransition(transition);
    } else {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    }
    participants[i].syncStart();
  }

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

  // disable partitions
  String command = "--zkSvr " + ZK_ADDR + " --enablePartition false " + clusterName
      + " localhost_12919 TestDB0 TestDB0_0";

  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  // ensure we get 1 slaveToOffline and 1 offlineToSlave after disable partition
  long startT = System.currentTimeMillis();
  while (System.currentTimeMillis() - startT < 10000) // retry in 5s
  {
    if (transition.slaveToOfflineCnt > 0 && transition.offlineToSlave > 0) {
      break;
    }

    Thread.sleep(100);
  }
  long endT = System.currentTimeMillis();
  System.out.println("1 disable and re-enable took: " + (endT - startT) + "ms");
  Assert.assertEquals(transition.slaveToOfflineCnt, 1, "should get 1 slaveToOffline transition");
  Assert.assertEquals(transition.offlineToSlave, 1, "should get 1 offlineToSlave transition");

  // 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 testResetInstance() throws Exception {
  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()));

  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

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

  Map<String, Set<String>> errPartitions = new HashMap<String, Set<String>>() {
    {
      put("SLAVE-MASTER", TestHelper.setOf("TestDB0_4"));
      put("OFFLINE-SLAVE", TestHelper.setOf("TestDB0_8"));
    }
  };

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

    if (i == 0) {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
      participants[i].setTransition(new ErrTransition(errPartitions));
    } else {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    }
    participants[i].syncStart();
  }

  // verify cluster
  Map<String, Map<String, String>> errStateMap = new HashMap<String, Map<String, String>>();
  errStateMap.put("TestDB0", new HashMap<String, String>());
  errStateMap.get("TestDB0").put("TestDB0_4", "localhost_12918");
  errStateMap.get("TestDB0").put("TestDB0_8", "localhost_12918");
  boolean result =
      ClusterStateVerifier
          .verifyByZkCallback((new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR,
              clusterName, errStateMap)));
  Assert.assertTrue(result, "Cluster verification fails");

  // reset node "localhost_12918"
  participants[0].setTransition(null);
  String command = "--zkSvr " + ZK_ADDR + " --resetInstance " + clusterName + " localhost_12918";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  result =
      ClusterStateVerifier
          .verifyByZkCallback((new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR,
              clusterName)));
  Assert.assertTrue(result, "Cluster verification fails");

  // 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 testParticipantIncompatibleWithBatchMsg() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  int n = 2;

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

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

  // enable batch message
  // --addResourceProperty <clusterName resourceName propertyName propertyValue>
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--addResourceProperty", clusterName, "TestDB0",
      HelixPropertyAttribute.BATCH_MESSAGE_MODE.toString(), "true"
  });

  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();

  // register a message listener so we know how many message generated
  TestZkChildListener listener = new TestZkChildListener();
  _gZkClient.subscribeChildChanges(keyBuilder.messages("localhost_12918").getPath(), listener);

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

  // pause controller
  // --enableCluster <clusterName true/false>
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--enableCluster", clusterName, "false"
  });

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

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

  // change localhost_12918 version to 0.5, so batch-message-mode will be ignored
  LiveInstance liveInstance = accessor.getProperty(keyBuilder.liveInstance("localhost_12918"));
  liveInstance.setHelixVersion("0.5");
  accessor.setProperty(keyBuilder.liveInstance("localhost_12918"), liveInstance);

  // resume controller
  // --enableCluster <clusterName true/false>
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--enableCluster", clusterName, "true"
  });

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);
  Assert.assertTrue(listener._maxNumberOfChildren > 16,
      "Should see more than 16 messages at the same time (32 O->S and 32 S->M)");

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

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

@Test
public void testDropErrorPartitionAutoIS() 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

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

  // start participants
  Map<String, Set<String>> errTransitions = new HashMap<>();
  errTransitions.put("SLAVE-MASTER", TestHelper.setOf("TestDB0_4"));
  errTransitions.put("OFFLINE-SLAVE", TestHelper.setOf("TestDB0_8"));

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

    if (i == 0) {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
      participants[i].setTransition(new ErrTransition(errTransitions));
    } else {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    }
    participants[i].syncStart();
  }

  Map<String, Map<String, String>> errStateMap = new HashMap<>();
  errStateMap.put("TestDB0", new HashMap<>());
  errStateMap.get("TestDB0").put("TestDB0_4", "localhost_12918");
  errStateMap.get("TestDB0").put("TestDB0_8", "localhost_12918");

  ZkHelixClusterVerifier verifier = new BestPossibleExternalViewVerifier.Builder(clusterName)
      .setZkAddr(ZK_ADDR).setErrStates(errStateMap).build();
  Assert.assertTrue(verifier.verifyByPolling());

  // drop resource containing error partitions should drop the partition successfully
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--dropResource", clusterName, "TestDB0"
  });

  // make sure TestDB0_4 and TestDB0_8 partitions are dropped
  verifier = new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());

  Thread.sleep(400);

  assertEmptyCSandEV(className, "TestDB0", participants);

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

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

@Test
public void testExpandCluster() throws Exception {
  String DB2 = "TestDB2";
  int partitions = 30;
  int replica = 3;
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, DB2, partitions, STATE_MODEL);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, DB2, replica, "keyX");

  String DB3 = "TestDB3";

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, DB3, partitions, STATE_MODEL);

  IdealState testDB0 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, TEST_DB);
  IdealState testDB2 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, DB2);
  IdealState testDB3 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, DB3);

  for (int i = 0; i < 5; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (27960 + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  String command = "-zkSvr localhost:2183 -expandCluster " + CLUSTER_NAME;
  ClusterSetup.processCommandLineArgs(command.split(" "));

  IdealState testDB0_1 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, TEST_DB);
  IdealState testDB2_1 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, DB2);
  IdealState testDB3_1 =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, DB3);

  Map<String, Object> resultOld2 = RebalanceUtil.buildInternalIdealState(testDB2);
  Map<String, Object> result2 = RebalanceUtil.buildInternalIdealState(testDB2_1);

  TestEspressoStorageClusterIdealState.Verify(result2, partitions, replica - 1);

  Double masterKeepRatio = 0.0, slaveKeepRatio = 0.0;
  double[] result = TestEspressoStorageClusterIdealState.compareResult(resultOld2, result2);
  masterKeepRatio = result[0];
  slaveKeepRatio = result[1];
  Assert.assertTrue(masterKeepRatio > 0.49 && masterKeepRatio < 0.51);

  Assert.assertTrue(testDB3_1.getRecord().getListFields().size() == 0);

  // partitions should stay as same
  Assert.assertTrue(testDB0_1.getRecord().getListFields().keySet()
      .containsAll(testDB0.getRecord().getListFields().keySet()));
  Assert.assertTrue(testDB0_1.getRecord().getListFields().size() == testDB0.getRecord()
      .getListFields().size());
  Assert.assertTrue(testDB2_1.getRecord().getMapFields().keySet()
      .containsAll(testDB2.getRecord().getMapFields().keySet()));
  Assert.assertTrue(testDB2_1.getRecord().getMapFields().size() == testDB2.getRecord()
      .getMapFields().size());
  Assert.assertTrue(testDB3_1.getRecord().getMapFields().keySet()
      .containsAll(testDB3.getRecord().getMapFields().keySet()));
  Assert.assertTrue(testDB3_1.getRecord().getMapFields().size() == testDB3.getRecord()
      .getMapFields().size());

  Map<String, Object> resultOld = RebalanceUtil.buildInternalIdealState(testDB0);
  Map<String, Object> resultNew = RebalanceUtil.buildInternalIdealState(testDB0_1);

  result = TestEspressoStorageClusterIdealState.compareResult(resultOld, resultNew);
  masterKeepRatio = result[0];
  slaveKeepRatio = result[1];
  Assert.assertTrue(masterKeepRatio > 0.49 && masterKeepRatio < 0.51);
}
 

@Test
public void testResetNode() throws Exception {
  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()));

  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

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

  Map<String, Set<String>> errPartitions = new HashMap<String, Set<String>>() {
    {
      put("SLAVE-MASTER", TestHelper.setOf("TestDB0_4"));
      put("OFFLINE-SLAVE", TestHelper.setOf("TestDB0_8"));
    }
  };

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

    if (i == 0) {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
      participants[i].setTransition(new ErrTransition(errPartitions));
    } else {
      participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    }
    participants[i].syncStart();
  }

  // verify cluster
  Map<String, Map<String, String>> errStateMap = new HashMap<String, Map<String, String>>();
  errStateMap.put("TestDB0", new HashMap<String, String>());
  errStateMap.get("TestDB0").put("TestDB0_4", "localhost_12918");
  errStateMap.get("TestDB0").put("TestDB0_8", "localhost_12918");
  boolean result =
      ClusterStateVerifier
          .verifyByZkCallback((new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR,
              clusterName, errStateMap)));
  Assert.assertTrue(result, "Cluster verification fails");

  // reset resource "TestDB0"
  participants[0].setTransition(null);
  String command = "--zkSvr " + ZK_ADDR + " --resetResource " + clusterName + " TestDB0";
  ClusterSetup.processCommandLineArgs(command.split("\\s+"));

  result =
      ClusterStateVerifier
          .verifyByZkCallback((new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR,
              clusterName)));
  Assert.assertTrue(result, "Cluster verification fails");

  // 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()));

}