org.apache.helix.HelixManager Java Examples

public AbstractYarnAppSecurityManager(Config config, HelixManager helixManager, FileSystem fs, Path tokenFilePath) {
  this.config = config;
  this.helixManager = helixManager;
  this.fs = fs;
  this.tokenFilePath = tokenFilePath;
  this.fs.makeQualified(tokenFilePath);
  this.loginIntervalInMinutes = ConfigUtils.getLong(config, GobblinYarnConfigurationKeys.LOGIN_INTERVAL_IN_MINUTES,
      GobblinYarnConfigurationKeys.DEFAULT_LOGIN_INTERVAL_IN_MINUTES);
  this.tokenRenewIntervalInMinutes = ConfigUtils.getLong(config, GobblinYarnConfigurationKeys.TOKEN_RENEW_INTERVAL_IN_MINUTES,
      GobblinYarnConfigurationKeys.DEFAULT_TOKEN_RENEW_INTERVAL_IN_MINUTES);

  this.loginExecutor = Executors.newSingleThreadScheduledExecutor(
      ExecutorsUtils.newThreadFactory(Optional.of(LOGGER), Optional.of("KeytabReLoginExecutor")));
  this.tokenRenewExecutor = Executors.newSingleThreadScheduledExecutor(
      ExecutorsUtils.newThreadFactory(Optional.of(LOGGER), Optional.of("TokenRenewExecutor")));
  this.isHelixClusterManaged = ConfigUtils.getBoolean(config, GobblinClusterConfigurationKeys.IS_HELIX_CLUSTER_MANAGED,
      GobblinClusterConfigurationKeys.DEFAULT_IS_HELIX_CLUSTER_MANAGED);
  this.helixInstanceName = ConfigUtils.getString(config, GobblinClusterConfigurationKeys.HELIX_INSTANCE_NAME_KEY,
      GobblinClusterManager.class.getSimpleName());
}
 

static boolean isJobFinished(String workflowName, String jobName, HelixManager helixManager) {
  WorkflowContext workflowContext = TaskDriver.getWorkflowContext(helixManager, workflowName);
  if (workflowContext == null) {
    // this workflow context doesn't exist, considered as finished.
    return true;
  }

  TaskState jobState = workflowContext.getJobState(TaskUtil.getNamespacedJobName(workflowName, jobName));
  switch (jobState) {
    case STOPPED:
    case FAILED:
    case COMPLETED:
    case ABORTED:
    case TIMED_OUT:
      return true;
    default:
      return false;
  }
}
 

@Test
public void testNoControllerLeader() {
  HelixManager helixManager = mock(HelixManager.class);
  HelixDataAccessor helixDataAccessor = mock(HelixDataAccessor.class);

  // Mock that there is no helix leader.
  when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  PropertyKey.Builder keyBuilder = mock(PropertyKey.Builder.class);
  when(helixDataAccessor.keyBuilder()).thenReturn(keyBuilder);
  PropertyKey controllerLeader = mock(PropertyKey.class);
  when(keyBuilder.controllerLeader()).thenReturn(controllerLeader);
  when(helixDataAccessor.getProperty(controllerLeader)).thenReturn(null);

  // Lead controller resource disabled.
  ConfigAccessor configAccessor = mock(ConfigAccessor.class);
  ResourceConfig resourceConfig = mock(ResourceConfig.class);
  when(helixManager.getConfigAccessor()).thenReturn(configAccessor);
  when(configAccessor.getResourceConfig(any(), any())).thenReturn(resourceConfig);
  when(resourceConfig.getSimpleConfig(anyString())).thenReturn("false");

  // Create Controller Leader Locator
  FakeControllerLeaderLocator.create(helixManager);
  ControllerLeaderLocator controllerLeaderLocator = FakeControllerLeaderLocator.getInstance();

  Assert.assertNull(controllerLeaderLocator.getControllerLeader(testTable));
}
 

private void syncSessionToController(HelixManager manager) {
  if (_lastSessionSyncTime == null ||
          System.currentTimeMillis() - _lastSessionSyncTime > SESSION_SYNC_INTERVAL) { // > delay since last sync
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    PropertyKey key = new Builder(manager.getClusterName()).controllerMessage(SESSION_SYNC);
    if (accessor.getProperty(key) == null) {
      LOG.info(String.format("Participant %s syncs session with controller", manager.getInstanceName()));
      Message msg = new Message(MessageType.PARTICIPANT_SESSION_CHANGE, SESSION_SYNC);
      msg.setSrcName(manager.getInstanceName());
      msg.setTgtSessionId("*");
      msg.setMsgState(MessageState.NEW);
      msg.setMsgId(SESSION_SYNC);

      Criteria cr = new Criteria();
      cr.setRecipientInstanceType(InstanceType.CONTROLLER);
      cr.setSessionSpecific(false);

      manager.getMessagingService().send(cr, msg);
      _lastSessionSyncTime = System.currentTimeMillis();
    }
  }
}
 

/**
 * @return A valid rebalancer object.
 *         If the rebalancer is no longer valid, it will be reset before returning.
 * TODO: Make rebalancer volatile or make it singleton, if this method is called in multiple
 * TODO: threads outside the controller object.
 */
synchronized T getRebalancer(HelixManager helixManager) {
  // Lazily initialize the stateful rebalancer instance since the GenericHelixController
  // instance is instantiated without the HelixManager information that is required.
  if (_rebalancer == null) {
    _rebalancer = createRebalancer(helixManager);
    _isRebalancerValid = true;
  }
  // If the rebalance exists but has been marked as invalid (due to leadership switch), it needs
  // to be reset before return.
  if (!_isRebalancerValid) {
    _rebalancer.reset();
    _isRebalancerValid = true;
  }
  return _rebalancer;
}
 

private static HelixManager getMockHelixManager(Config config) {
  HelixManager helixManager = Mockito.mock(HelixManager.class);
  HelixAdmin helixAdmin = Mockito.mock(HelixAdmin.class);
  HelixDataAccessor helixDataAccessor = Mockito.mock(HelixDataAccessor.class);
  PropertyKey propertyKey = Mockito.mock(PropertyKey.class);
  PropertyKey.Builder propertyKeyBuilder = Mockito.mock(PropertyKey.Builder.class);

  Mockito.when(helixManager.getInstanceName()).thenReturn("helixInstance1");
  Mockito.when(helixManager.getClusterName()).thenReturn(config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY));
  Mockito.doNothing().when(helixAdmin).enableInstance(Mockito.anyString(), Mockito.anyString(), Mockito.anyBoolean());
  Mockito.when(helixManager.getHelixDataAccessor()).thenReturn(helixDataAccessor);
  Mockito.when(helixDataAccessor.keyBuilder()).thenReturn(propertyKeyBuilder);
  Mockito.when(propertyKeyBuilder.liveInstance(Mockito.anyString())).thenReturn(propertyKey);
  Mockito.when(helixDataAccessor.getProperty(propertyKey)).thenReturn(null);

  return helixManager;
}
 

@Override
public void process(ClusterEvent event) throws Exception {
  _eventId = event.getEventId();
  CurrentStateOutput currentStateOutput =
      event.getAttribute(AttributeName.CURRENT_STATE.name());
  final Map<String, Resource> resourceMap =
      event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
  WorkflowControllerDataProvider cache = event.getAttribute(AttributeName.ControllerDataProvider.name());

  if (currentStateOutput == null || resourceMap == null || cache == null) {
    throw new StageException(
        "Missing attributes in event:" + event + ". Requires CURRENT_STATE|RESOURCES|DataCache");
  }

  // Reset current INIT/RUNNING tasks on participants for throttling
  cache.resetActiveTaskCount(currentStateOutput);

  buildQuotaBasedWorkflowPQsAndInitDispatchers(cache,
      (HelixManager) event.getAttribute(AttributeName.helixmanager.name()),
      (ClusterStatusMonitor) event.getAttribute(AttributeName.clusterStatusMonitor.name()));

  final BestPossibleStateOutput bestPossibleStateOutput =
      compute(event, resourceMap, currentStateOutput);
  event.addAttribute(AttributeName.BEST_POSSIBLE_STATE.name(), bestPossibleStateOutput);
}
 

@Override
@PreFetch(enabled = false)
public void onIdealStateChange(List<IdealState> idealStates, NotificationContext changeContext) {
  logger.info(
      "START: Generic GenericClusterController.onIdealStateChange() for cluster " + _clusterName);
  notifyCaches(changeContext, ChangeType.IDEAL_STATE);
  pushToEventQueues(ClusterEventType.IdealStateChange, changeContext,
      Collections.<String, Object>emptyMap());

  if (changeContext.getType() != NotificationContext.Type.FINALIZE) {
    HelixManager manager = changeContext.getManager();
    if (manager != null) {
      HelixDataAccessor dataAccessor = changeContext.getManager().getHelixDataAccessor();
      checkRebalancingTimer(changeContext.getManager(), idealStates,
          (ClusterConfig) dataAccessor.getProperty(dataAccessor.keyBuilder().clusterConfig()));
    }
  }

  logger.info("END: GenericClusterController.onIdealStateChange() for cluster " + _clusterName);
}
 

@Test()
public void testParticipant() throws Exception {
  String className = getShortClassName();
  LOG.info("RUN " + className + " at " + new Date(System.currentTimeMillis()));

  final String clusterName = CLUSTER_PREFIX + "_" + className + "_" + "testParticipant";
  TestHelper.setupEmptyCluster(_gZkClient, clusterName);

  final String controllerName = "participant_0";
  HelixManager manager =
      new MockZKHelixManager(clusterName, controllerName, InstanceType.PARTICIPANT, _gZkClient);
  GenericHelixController participant0 = new GenericHelixController();
  List<HelixTimerTask> timerTasks = Collections.emptyList();

  try {
    DistributedLeaderElection election =
        new DistributedLeaderElection(manager, participant0, timerTasks);
    Assert.fail(
        "Should not be able construct DistributedLeaderElection object using participant manager.");
  } catch (HelixException ex) {
    // expected
  }

  TestHelper.dropCluster(clusterName, _gZkClient);
}
 

public ParticipantManager(HelixManager manager, RealmAwareZkClient zkclient, int sessionTimeout,
    LiveInstanceInfoProvider liveInstanceInfoProvider,
    List<PreConnectCallback> preConnectCallbacks, final String sessionId,
    HelixManagerProperty helixManagerProperty) {
  _zkclient = zkclient;
  _manager = manager;
  _clusterName = manager.getClusterName();
  _instanceName = manager.getInstanceName();
  _keyBuilder = new PropertyKey.Builder(_clusterName);
  _sessionId = sessionId;
  _sessionTimeout = sessionTimeout;
  _configAccessor = manager.getConfigAccessor();
  _instanceType = manager.getInstanceType();
  _helixAdmin = manager.getClusterManagmentTool();
  _dataAccessor = (ZKHelixDataAccessor) manager.getHelixDataAccessor();
  _messagingService = (DefaultMessagingService) manager.getMessagingService();
  _stateMachineEngine = manager.getStateMachineEngine();
  _liveInstanceInfoProvider = liveInstanceInfoProvider;
  _preConnectCallbacks = preConnectCallbacks;
  _helixManagerProperty = helixManagerProperty;
}
 

private Message generateCancellationMessageForPendingMessage(final String desiredState, final String currentState,
    final String nextState, final Message pendingMessage, final HelixManager manager,
    final Resource resource, final Partition partition, final Map<String, String> sessionIdMap,
    final String instanceName, final StateModelDefinition stateModelDef,
    final Message cancellationMessage, final boolean isCancellationEnabled) {

  Message message = null;

  if (pendingMessage != null) {
    String pendingState = pendingMessage.getToState();
    if (nextState.equalsIgnoreCase(pendingState)) {
      LogUtil.logInfo(logger, _eventId,
          "Message already exists for " + instanceName + " to transit " + resource
              .getResourceName() + "." + partition.getPartitionName() + " from "
              + currentState + " to " + nextState + ", isRelay: " + pendingMessage.isRelayMessage());
    } else if (currentState.equalsIgnoreCase(pendingState)) {
      LogUtil.logInfo(logger, _eventId,
          "Message hasn't been removed for " + instanceName + " to transit " + resource
              .getResourceName() + "." + partition.getPartitionName() + " to "
              + pendingState + ", desiredState: " + desiredState + ", isRelay: " + pendingMessage.isRelayMessage());
    } else {
      LogUtil.logInfo(logger, _eventId,
          "IdealState changed before state transition completes for " + resource
              .getResourceName() + "." + partition.getPartitionName() + " on "
              + instanceName + ", pendingState: " + pendingState + ", currentState: "
              + currentState + ", nextState: " + nextState + ", isRelay: " + pendingMessage.isRelayMessage());

      message = createStateTransitionCancellationMessage(manager, resource,
          partition.getPartitionName(), instanceName, sessionIdMap.get(instanceName),
          stateModelDef.getId(), pendingMessage.getFromState(), pendingState, nextState,
          cancellationMessage, isCancellationEnabled, currentState);
    }
  }
  return message;
}
 

public GobblinServiceJobScheduler(String serviceName, Config config, Optional<HelixManager> helixManager,
    Optional<FlowCatalog> flowCatalog, Optional<TopologyCatalog> topologyCatalog, Orchestrator orchestrator,
    SchedulerService schedulerService, Optional<Logger> log) throws Exception {
  super(ConfigUtils.configToProperties(config), schedulerService);

  _log = log.isPresent() ? log.get() : LoggerFactory.getLogger(getClass());
  this.serviceName = serviceName;
  this.flowCatalog = flowCatalog;
  this.helixManager = helixManager;
  this.orchestrator = orchestrator;
  this.scheduledFlowSpecs = Maps.newHashMap();
  this.isNominatedDRHandler = config.hasPath(GOBBLIN_SERVICE_SCHEDULER_DR_NOMINATED)
      && config.hasPath(GOBBLIN_SERVICE_SCHEDULER_DR_NOMINATED);
}
 

public AutoRebalanceLiveInstanceChangeListener(HelixMirrorMakerManager helixMirrorMakerManager,
    HelixManager helixManager, ControllerConf controllerConf) {
  _helixMirrorMakerManager = helixMirrorMakerManager;
  _helixManager = helixManager;
  _maxWorkingInstances = controllerConf.getMaxWorkingInstances();
  _delayedAutoReblanceTimeInSeconds = controllerConf.getAutoRebalanceDelayInSeconds();
  _overloadedRatioThreshold = controllerConf.getAutoRebalanceWorkloadRatioThreshold();
  _maxDedicatedInstancesRatio = controllerConf.getMaxDedicatedLaggingInstancesRatio();
  _maxStuckPartitionMovements = controllerConf.getMaxStuckPartitionMovements();
  _movePartitionAfterStuckMillis = TimeUnit.MINUTES
      .toMillis(controllerConf.getMoveStuckPartitionAfterMinutes());
  LOGGER.info("Delayed Auto Reblance Time In Seconds: {}", _delayedAutoReblanceTimeInSeconds);
  registerMetrics();

  int autoRebalancePeriodInSeconds = controllerConf.getAutoRebalancePeriodInSeconds();
  final int minIntervalInSeconds = controllerConf.getAutoRebalanceMinIntervalInSeconds();
  if (autoRebalancePeriodInSeconds > 0) {
    LOGGER.info("Trying to schedule auto rebalancing at rate " + autoRebalancePeriodInSeconds
        + " seconds");
    _delayedScheuler.scheduleWithFixedDelay(
        new Runnable() {
          @Override
          public void run() {
            try {
              if (_helixMirrorMakerManager.getWorkloadInfoRetriever().isInitialized()
                  && System.currentTimeMillis() - _lastRebalanceTimeMillis
                  > 1000L * minIntervalInSeconds) {
                rebalanceCurrentCluster(_helixMirrorMakerManager.getCurrentLiveInstances(),
                    _helixMirrorMakerManager.getBlacklistedInstances(), false, false);
              }
            } catch (Exception e) {
              LOGGER
                  .error("Got exception during periodically rebalancing the whole cluster! ", e);
            }
          }
        }, Math.max(_delayedAutoReblanceTimeInSeconds, autoRebalancePeriodInSeconds),
        autoRebalancePeriodInSeconds,
        TimeUnit.SECONDS);
  }
}
 

public SubmitJobTask(String zkAddress, int index) throws Exception {
  HelixManager manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Administrator",
      InstanceType.ADMINISTRATOR, zkAddress);
  manager.connect();
  _driver = new TaskDriver(manager);
  _jobPrefixName = "JOB_" + index + "#";
}
 

@Override
public void init(HelixManager helixManager, TableConfig tableConfig) {
  _realtimeTableName = tableConfig.getTableName();
  _replication = tableConfig.getValidationConfig().getReplicasPerPartitionNumber();

  LOGGER.info("Initialized RealtimeSegmentAssignment with replication: {} for table: {}", _replication,
      _realtimeTableName);
}
 

@Test(dependsOnMethods = "testInstanceAutoJoin")
public void testAutoRegistrationShouldFailWhenWaitingResponse() throws Exception {
  // Create CloudConfig object and add to config
  CloudConfig.Builder cloudConfigBuilder = new CloudConfig.Builder();
  cloudConfigBuilder.setCloudEnabled(true);
  cloudConfigBuilder.setCloudProvider(CloudProvider.AZURE);
  cloudConfigBuilder.setCloudID("TestID");
  CloudConfig cloudConfig = cloudConfigBuilder.build();

  HelixManager manager = _participants[0];
  HelixDataAccessor accessor = manager.getHelixDataAccessor();

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, db3, 60, "OnlineOffline",
      RebalanceMode.FULL_AUTO + "");
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db3, 1);
  String instance3 = "localhost_279700";

  ConfigScope scope = new ConfigScopeBuilder().forCluster(CLUSTER_NAME).build();

  manager.getConfigAccessor().set(scope, ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, "true");
  // Write the CloudConfig to Zookeeper
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.cloudConfig(), cloudConfig);

  MockParticipantManager autoParticipant =
      new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instance3);
  autoParticipant.syncStart();

  Assert.assertTrue(null == manager.getHelixDataAccessor()
      .getProperty(accessor.keyBuilder().liveInstance(instance3)));
  try {
    manager.getConfigAccessor().getInstanceConfig(CLUSTER_NAME, instance3);
    fail(
        "Exception should be thrown because the instance cannot be added to the cluster due to the disconnection with Azure endpoint");
  } catch (HelixException e) {

  }
}
 

/**
 * Return a {@link MockHelixManager}
 * @param clusterName the name of the cluster for which the manager is to be gotten.
 * @param instanceName the name of the instance on whose behalf the manager is to be gotten.
 * @param instanceType the {@link InstanceType} of the requester.
 * @param zkAddr the address identifying the zk service to which this request is to be made.
 * @return the {@link MockHelixManager}
 */
public HelixManager getZKHelixManager(String clusterName, String instanceName, InstanceType instanceType,
    String zkAddr) {
  if (helixCluster.getZkAddrs().contains(zkAddr)) {
    return new MockHelixManager(instanceName, instanceType, zkAddr, helixCluster, znRecordMap, beBadException);
  } else {
    throw new IllegalArgumentException("Invalid ZkAddr");
  }
}
 

/**
 * Helix participant cannot pre-configure tags before it connects to ZK. So this method can only be invoked after
 * {@link HelixManager#connect()}. However this will still work because tagged jobs won't be sent to a non-tagged instance. Hence
 * the job with EXAMPLE_INSTANCE_TAG will remain in the ZK until an instance with EXAMPLE_INSTANCE_TAG was found.
 */
private void addInstanceTags() {
  List<String> tags = ConfigUtils.getStringList(this.clusterConfig, GobblinClusterConfigurationKeys.HELIX_INSTANCE_TAGS_KEY);
  HelixManager receiverManager = getReceiverManager();
  if (receiverManager.isConnected()) {
    if (!tags.isEmpty()) {
      logger.info("Adding tags binding " + tags);
      tags.forEach(tag -> receiverManager.getClusterManagmentTool()
          .addInstanceTag(this.clusterName, this.helixInstanceName, tag));
      logger.info("Actual tags binding " + receiverManager.getClusterManagmentTool()
          .getInstanceConfig(this.clusterName, this.helixInstanceName).getTags());
    }
  }
}
 

public ServiceDiscovery(String zkAddress, String cluster, Mode mode) {
  this.zkAddress = zkAddress;
  this.cluster = cluster;
  this.mode = mode;
  serviceMap = new HashMap<String, HelixManager>();
  cache = Collections.emptyList();
}
 

public static HelixManager startHelixController(final String zkConnectString,
    final String clusterName, final String controllerName, final String controllerMode) {
  HelixManager manager = null;
  try {
    if (controllerMode.equalsIgnoreCase(STANDALONE)) {
      manager =
          HelixManagerFactory.getZKHelixManager(clusterName, controllerName,
              InstanceType.CONTROLLER, zkConnectString);
      manager.connect();
    } else if (controllerMode.equalsIgnoreCase(DISTRIBUTED)) {
      manager =
          HelixManagerFactory.getZKHelixManager(clusterName, controllerName,
              InstanceType.CONTROLLER_PARTICIPANT, zkConnectString);

      DistClusterControllerStateModelFactory stateModelFactory =
          new DistClusterControllerStateModelFactory(zkConnectString);

      StateMachineEngine stateMach = manager.getStateMachineEngine();
      stateMach.registerStateModelFactory("LeaderStandby", stateModelFactory);
      manager.connect();
    } else {
      logger.error("cluster controller mode:" + controllerMode + " NOT supported");
    }
  } catch (Exception e) {
    logger.error("Exception while starting controller", e);
  }

  return manager;
}
 

public TaskRunner(Task task, String taskName, String taskPartition, String instance,
    HelixManager manager, String sessionId) {
  _task = task;
  _taskName = taskName;
  _taskPartition = taskPartition;
  _instance = instance;
  _manager = manager;
  _sessionId = sessionId;
}
 

private HelixManager initializeHelixManager(String clusterName, String instanceName) {
  HelixManager manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
      InstanceType.PARTICIPANT, org.apache.helix.common.ZkTestBase.ZK_ADDR);

  MasterSlaveStateModelFactory stateModelFactory = new MasterSlaveStateModelFactory(instanceName);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory("id1", stateModelFactory);
  return manager;
}
 

public static IdealState buildInitialHighLevelRealtimeIdealStateFor(String realtimeTableName,
    TableConfig realtimeTableConfig, HelixManager helixManager, ZkHelixPropertyStore<ZNRecord> zkHelixPropertyStore,
    boolean enableBatchMessageMode) {
  List<String> realtimeInstances = HelixHelper.getInstancesWithTag(helixManager,
      TagNameUtils.extractConsumingServerTag(realtimeTableConfig.getTenantConfig()));
  IdealState idealState = buildEmptyRealtimeIdealStateFor(realtimeTableName, 1, enableBatchMessageMode);
  if (realtimeInstances.size() % Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()) != 0) {
    throw new RuntimeException(
        "Number of instance in current tenant should be an integer multiples of the number of replications");
  }
  setupInstanceConfigForHighLevelConsumer(realtimeTableName, realtimeInstances.size(),
      Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()),
      realtimeTableConfig.getIndexingConfig().getStreamConfigs(), zkHelixPropertyStore, realtimeInstances);
  return idealState;
}
 

@Test()
public void testNoRetry() throws InterruptedException {
  // String p = "test_";
  // System.out.println(p.substring(p.lastIndexOf('_')+1));
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

  CancellableHandlerFactory factory = new CancellableHandlerFactory();
  for (String type : factory.getMessageTypes()) {
    executor.registerMessageHandlerFactory(type, factory);
  }
  NotificationContext changeContext = new NotificationContext(manager);

  List<Message> msgList = new ArrayList<Message>();
  int nMsgs2 = 4;
  // Test the case in which retry = 0
  for (int i = 0; i < nMsgs2; i++) {
    Message msg = new Message(factory.getMessageTypes().get(0), UUID.randomUUID().toString());
    msg.setTgtSessionId("*");
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setExecutionTimeout((i + 1) * 600);
    msgList.add(msg);
  }
  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage("someInstance", msgList, changeContext);

  Thread.sleep(4000);

  AssertJUnit.assertTrue(factory._handlersCreated == nMsgs2);
  AssertJUnit.assertEquals(factory._timedOutMsgIds.size(), 2);
  // AssertJUnit.assertFalse(msgList.get(0).getRecord().getSimpleFields().containsKey("TimeOut"));
  for (int i = 0; i < nMsgs2 - 2; i++) {
    if (factory.getMessageTypes().contains(msgList.get(i).getMsgType())) {
      AssertJUnit.assertTrue(msgList.get(i).getRecord().getSimpleFields()
          .containsKey("Cancelcount"));
      AssertJUnit.assertTrue(factory._timedOutMsgIds.containsKey(msgList.get(i).getId()));
    }
  }
}
 

@Override
public Task createTask(String id, Set<String> parentIds, HelixManager helixManager,
    TaskResultStore taskResultStore) {
  if (id.equalsIgnoreCase("filterImps")) {
    return new FilterTask(id, parentIds, helixManager, taskResultStore, FilterTask.IMPRESSIONS);
  } else if (id.equalsIgnoreCase("filterClicks")) {
    return new FilterTask(id, parentIds, helixManager, taskResultStore, FilterTask.CLICKS);
  } else if (id.equalsIgnoreCase("impClickJoin")) {
    return new JoinTask(id, parentIds, helixManager, taskResultStore,
        FilterTask.FILTERED_IMPRESSIONS, FilterTask.FILTERED_CLICKS);
  } else if (id.equalsIgnoreCase("impCountsByGender")) {
    return new CountTask(id, parentIds, helixManager, taskResultStore,
        FilterTask.FILTERED_IMPRESSIONS, "gender");
  } else if (id.equalsIgnoreCase("impCountsByCountry")) {
    return new CountTask(id, parentIds, helixManager, taskResultStore,
        FilterTask.FILTERED_IMPRESSIONS, "country");
  } else if (id.equalsIgnoreCase("clickCountsByGender")) {
    return new CountTask(id, parentIds, helixManager, taskResultStore, JoinTask.JOINED_CLICKS,
        "gender");
  } else if (id.equalsIgnoreCase("clickCountsByCountry")) {
    return new CountTask(id, parentIds, helixManager, taskResultStore, JoinTask.JOINED_CLICKS,
        "country");
  } else if (id.equalsIgnoreCase("report")) {
    return new ReportTask(id, parentIds, helixManager, taskResultStore);
  }

  throw new IllegalArgumentException("Cannot create task for " + id);
}
 

@Override
protected void createHelixCluster() throws Exception {
  super.createHelixCluster();
  String clusterName = super.getManagerConfig().getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);
  String zkConnectString = super.getManagerConfig().getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(clusterName, IntegrationBasicSuite.WORKER_INSTANCE_0, InstanceType.PARTICIPANT, zkConnectString);

  //Create a partial instance setup
  ClusterIntegrationTestUtils.createPartialInstanceStructure(helixManager, zkConnectString);
}
 

@Override
public void execute(final ClusterEvent event) throws Exception {
  ResourceControllerDataProvider cache = event.getAttribute(AttributeName.ControllerDataProvider.name());
  ClusterConfig clusterConfig = cache.getClusterConfig();

  if (!clusterConfig.isPersistBestPossibleAssignment() && !clusterConfig
      .isPersistIntermediateAssignment()) {
    return;
  }

  BestPossibleStateOutput bestPossibleAssignment =
      event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());

  HelixManager helixManager = event.getAttribute(AttributeName.helixmanager.name());
  HelixDataAccessor accessor = helixManager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Map<String, Resource> resourceMap = event.getAttribute(AttributeName.RESOURCES.name());

  for (String resourceId : bestPossibleAssignment.resourceSet()) {
    try {
      persistAssignment(resourceMap.get(resourceId), cache, event, bestPossibleAssignment,
          clusterConfig, accessor, keyBuilder);
    } catch (HelixException ex) {
      LogUtil
          .logError(LOG, _eventId, "Failed to persist assignment for resource " + resourceId, ex);
    }
  }
}
 

private static void printStatus(final HelixManager manager) {
  System.out.println("CLUSTER STATUS");
  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  Builder keyBuilder = helixDataAccessor.keyBuilder();
  System.out.println("External View \n"
      + helixDataAccessor.getProperty(keyBuilder.externalView("repository")));
}
 

public static SegmentAssignment getSegmentAssignment(HelixManager helixManager, TableConfig tableConfig) {
  SegmentAssignment segmentAssignment;
  if (tableConfig.getTableType() == TableType.OFFLINE) {
    segmentAssignment = new OfflineSegmentAssignment();
  } else {
    segmentAssignment = new RealtimeSegmentAssignment();
  }
  segmentAssignment.init(helixManager, tableConfig);
  return segmentAssignment;
}
 

private void relinquishLeadership(HelixManager manager,
    ControllerManagerHelper controllerHelper) {
  long start = System.currentTimeMillis();
  LOG.info(manager.getInstanceName() + " tries to relinquish leadership for cluster: " + manager
      .getClusterName());
  controllerHelper.stopControllerTimerTasks();
  controllerHelper.removeListenersFromController(_controller);
  // clear write-through cache
  manager.getHelixDataAccessor().getBaseDataAccessor().reset();
  LOG.info("{} relinquishes leadership for cluster: {}, took: {}ms", manager.getInstanceName(),
      manager.getClusterName(), System.currentTimeMillis() - start);
}