Java Code Examples for org.apache.helix.model.IdealState#setBatchMessageMode()

private static void createBrokerResourceIfNeeded(String helixClusterName, HelixAdmin helixAdmin,
    boolean enableBatchMessageMode) {
  // Add state model definition if needed
  String stateModel =
      PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.PINOT_BROKER_RESOURCE_ONLINE_OFFLINE_STATE_MODEL;
  StateModelDefinition stateModelDef = helixAdmin.getStateModelDef(helixClusterName, stateModel);
  if (stateModelDef == null) {
    LOGGER.info("Adding state model: {}", stateModel);
    helixAdmin.addStateModelDef(helixClusterName, stateModel,
        PinotHelixBrokerResourceOnlineOfflineStateModelGenerator.generatePinotStateModelDefinition());
  }

  // Add broker resource if needed
  if (helixAdmin.getResourceIdealState(helixClusterName, BROKER_RESOURCE_INSTANCE) == null) {
    LOGGER.info("Adding resource: {}", BROKER_RESOURCE_INSTANCE);
    IdealState idealState = new CustomModeISBuilder(BROKER_RESOURCE_INSTANCE).setStateModel(stateModel).build();
    idealState.setBatchMessageMode(enableBatchMessageMode);
    helixAdmin.addResource(helixClusterName, BROKER_RESOURCE_INSTANCE, idealState);
  }
}
 

@Override
public void enableBatchMessageMode(String clusterName, String resourceName, boolean enabled) {
  logger.info("{} batch message mode for resource {} in cluster {}.",
      enabled ? "Enable" : "Disable", resourceName, clusterName);
  // TODO: Change IdealState to ResourceConfig when configs are migrated to ResourceConfig
  IdealState idealState = getResourceIdealState(clusterName, resourceName);
  if (idealState == null) {
    throw new HelixException("Cluster " + clusterName + ", resource: " + resourceName
        + ", ideal-state does not exist");
  }

  idealState.setBatchMessageMode(enabled);
  setResourceIdealState(clusterName, resourceName, idealState);
}
 

@Test (dependsOnMethods = "TestPerStateTransitionTypeThreadPool")
public void testBatchMessageThreadPoolSize() throws InterruptedException {
  int customizedPoolSize = 5;
  _participants[0].getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new TestOnlineOfflineStateModelFactory(customizedPoolSize, 2000), "TestFactory");
  for (int i = 1; i < _participants.length; i++) {
    _participants[i].syncStop();
  }
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  // Add 10 dbs with batch message enabled. Each db has 10 partitions.
  // So it will have 10 batch messages and each batch message has 10 sub messages.
  int numberOfDbs = 10;
  for (int i = 0; i < numberOfDbs; i++) {
    String dbName = "TestDBABatch" + i;
    IdealState idealState = new FullAutoModeISBuilder(dbName).setStateModel("OnlineOffline")
        .setStateModelFactoryName("TestFactory").setNumPartitions(10).setNumReplica(1).build();
    idealState.setBatchMessageMode(true);
    _gSetupTool.getClusterManagementTool().addResource(CLUSTER_NAME, dbName, idealState);
    _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, dbName, 1);
  }
  Assert.assertTrue(_clusterVerifier.verifyByPolling());

  DefaultMessagingService svc =
      (DefaultMessagingService) (_participants[0].getMessagingService());
  HelixTaskExecutor helixExecutor = svc.getExecutor();
  ThreadPoolExecutor executor = (ThreadPoolExecutor) (helixExecutor._batchMessageExecutorService);
  Assert.assertNotNull(executor);
  Assert.assertTrue(executor.getPoolSize() >= numberOfDbs);

  BestPossibleExternalViewVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());
}
 

/**
 *
 * Building an empty idealState for a given table.
 * Used when creating a new table.
 *
 * @param tableName resource name
 * @param numCopies is the number of replicas
 * @return
 */
public static IdealState buildEmptyIdealStateFor(String tableName, int numCopies, boolean enableBatchMessageMode) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(tableName);
  final int replicas = numCopies;
  customModeIdealStateBuilder
      .setStateModel(PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL)
      .setNumPartitions(0).setNumReplica(replicas).setMaxPartitionsPerNode(1);
  final IdealState idealState = customModeIdealStateBuilder.build();
  idealState.setInstanceGroupTag(tableName);
  idealState.setBatchMessageMode(enableBatchMessageMode);
  return idealState;
}
 

public static IdealState buildEmptyRealtimeIdealStateFor(String realtimeTableName, int replicaCount,
    boolean enableBatchMessageMode) {
  final CustomModeISBuilder customModeIdealStateBuilder = new CustomModeISBuilder(realtimeTableName);
  customModeIdealStateBuilder
      .setStateModel(PinotHelixSegmentOnlineOfflineStateModelGenerator.PINOT_SEGMENT_ONLINE_OFFLINE_STATE_MODEL)
      .setNumPartitions(0).setNumReplica(replicaCount).setMaxPartitionsPerNode(1);
  final IdealState idealState = customModeIdealStateBuilder.build();
  idealState.setInstanceGroupTag(realtimeTableName);
  idealState.setBatchMessageMode(enableBatchMessageMode);

  return idealState;
}
 

@Test
public void testSubMessageFailed() throws Exception {
  TestOnlineOfflineStateModel._numOfSuccessBeforeFailure.set(6);

  // Let one instance handle all the batch messages.
  _participants[0].getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new TestOnlineOfflineStateModelFactory(), "TestFactory");
  for (int i = 1; i < _participants.length; i++) {
    _participants[i].syncStop();
  }

  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  // Check that the Participants really stopped
  boolean result = TestHelper.verify(() -> {
    List<String> liveInstances =
        dataAccessor.getChildNames(dataAccessor.keyBuilder().liveInstances());
    for (int i = 1; i < _participants.length; i++) {
      if (_participants[i].isConnected()
          || liveInstances.contains(_participants[i].getInstanceName())) {
        return false;
      }
    }
    return true;
  }, TestHelper.WAIT_DURATION);
  Assert.assertTrue(result);

  // Add 1 db with batch message enabled. Each db has 10 partitions.
  // So it will have 1 batch message and 10 sub messages.
  String dbName = "TestDBSubMessageFail";
  IdealState idealState = new FullAutoModeISBuilder(dbName).setStateModel("OnlineOffline")
      .setStateModelFactoryName("TestFactory").setNumPartitions(10).setNumReplica(1).build();
  idealState.setBatchMessageMode(true);
  _gSetupTool.getClusterManagementTool().addResource(CLUSTER_NAME, dbName, idealState);

  // Check that IdealState has really been added
  result = TestHelper.verify(
      () -> dataAccessor.getPropertyStat(dataAccessor.keyBuilder().idealStates(dbName)) != null,
      TestHelper.WAIT_DURATION);
  Assert.assertTrue(result);

  for (int i = 0; i < 5; i++) {
    IdealState is =
        _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, dbName);
    if (!idealState.equals(is)) {
      Thread.sleep(1000L);
    }
  }
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, dbName, 1);
  Assert.assertTrue(_clusterVerifier.verifyByPolling());
  Thread.sleep(2000L);

  int numOfOnlines = 0;
  int numOfErrors = 0;
  ExternalView externalView =
      _gSetupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, dbName);
  for (String partition : externalView.getPartitionSet()) {
    if (externalView.getStateMap(partition).values().contains("ONLINE")) {
      numOfOnlines++;
    }
    if (externalView.getStateMap(partition).values().contains("ERROR")) {
      numOfErrors++;
    }
  }
  Assert.assertEquals(numOfErrors, 4);
  Assert.assertEquals(numOfOnlines, 6);
}
 

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

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // startPort
      "localhost", // participantNamePrefix
      "TestDB", // resourceNamePrefix
      1, // resourceNb
      2, // partitionNb
      n, // nodesNb
      2, // replica
      "MasterSlave", true);

  // enable batch-message
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setBatchMessageMode(true);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

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

  // start participants
  MockParticipantManager[] participants = new MockParticipantManager[n];
  TestMockMSModelFactory[] ftys = new TestMockMSModelFactory[n];

  for (int i = 0; i < n; i++) {
    String instanceName = "localhost_" + (12918 + i);
    ftys[i] = new TestMockMSModelFactory();
    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].getStateMachineEngine().registerStateModelFactory("MasterSlave", ftys[i]);
    participants[i].syncStart();
    int finalI = i;
    TestHelper.verify(() -> participants[finalI].isConnected()
        && accessor.getChildNames(keyBuilder.liveInstances())
            .contains(participants[finalI].getInstanceName()),
        TestHelper.WAIT_DURATION);

    // wait for each participant to complete state transitions, so we have deterministic results
    ZkHelixClusterVerifier _clusterVerifier =
        new BestPossibleExternalViewVerifier.Builder(clusterName).setZkAddr(ZK_ADDR).build();
    Assert.assertTrue(_clusterVerifier.verifyByPolling(),
        "participant: " + instanceName + " fails to complete all transitions");
  }

  // check batch-msg-wrapper counts
  MockBatchMsgWrapper wrapper = (MockBatchMsgWrapper) ftys[0].getBatchMessageWrapper("TestDB0");
  // System.out.println("startCount: " + wrapper._startCount);
  Assert.assertEquals(wrapper._startCount, 3,
      "Expect 3 batch.start: O->S, S->M, and M->S for 1st participant");
  // System.out.println("endCount: " + wrapper._endCount);
  Assert.assertEquals(wrapper._endCount, 3,
      "Expect 3 batch.end: O->S, S->M, and M->S for 1st participant");

  wrapper = (MockBatchMsgWrapper) ftys[1].getBatchMessageWrapper("TestDB0");
  // System.out.println("startCount: " + wrapper._startCount);
  Assert.assertEquals(wrapper._startCount, 2,
      "Expect 2 batch.start: O->S and S->M for 2nd participant");
  // System.out.println("endCount: " + wrapper._endCount);
  Assert.assertEquals(wrapper._startCount, 2,
      "Expect 2 batch.end: O->S and S->M for 2nd participant");

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

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

@Test
public void testBasic() 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
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setBatchMessageMode(true);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

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

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

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          clusterName));
  Assert.assertTrue(result);
  // Change to three is because there is an extra factory registered
  // So one extra NO_OP message send
  Assert.assertTrue(listener._maxNumberOfChildren <= 3,
      "Should get no more than 2 messages (O->S and 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 testSubMsgExecutionFail() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;

  final int n = 5;
  MockParticipantManager[] participants = new MockParticipantManager[n];

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

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, "localhost", "TestDB", 1, // resource#
      6, // partition#
      n, // nodes#
      3, // replicas#
      "MasterSlave", true);

  // enable batch message
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, _baseAccessor);
  Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  idealState.setBatchMessageMode(true);
  accessor.setProperty(keyBuilder.idealStates("TestDB0"), idealState);

  // get MASTER for errPartition
  String errPartition = "TestDB0_0";
  String masterOfPartition0 = null;
  for (Map.Entry<String, String> entry : idealState.getInstanceStateMap(errPartition).entrySet()) {
    if (entry.getValue().equals("MASTER")) {
      masterOfPartition0 = entry.getKey();
      break;
    }
  }
  Assert.assertNotNull(masterOfPartition0);

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

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

    if (instanceName.equals(masterOfPartition0)) {
      Map<String, Set<String>> errPartitions = new HashMap<String, Set<String>>();
      errPartitions.put("SLAVE-MASTER", TestHelper.setOf("TestDB0_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();
  }

  Map<String, Map<String, String>> errStates = new HashMap<String, Map<String, String>>();
  errStates.put("TestDB0", new HashMap<String, String>());
  errStates.get("TestDB0").put(errPartition, masterOfPartition0);
  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, clusterName, errStates));
  Assert.assertTrue(result);

  Map<String, Set<String>> errorStateMap = new HashMap<String, Set<String>>();
  errorStateMap.put(errPartition, TestHelper.setOf(masterOfPartition0));

  // verify "TestDB0_0", masterOfPartition0 is in ERROR state
  TestHelper.verifyState(clusterName, ZK_ADDR, errorStateMap, "ERROR");

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