Java Code Examples for org.apache.helix.HelixManagerFactory#getZKHelixManager()

private void initHelixManager() {
  String zkConnectionString =
      this.clusterConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  logger.info("Using ZooKeeper connection string: " + zkConnectionString);

  if (this.isTaskDriver && this.dedicatedTaskDriverCluster) {
    // This will create a Helix manager to receive the planning job
    this.taskDriverHelixManager = Optional.of(HelixManagerFactory.getZKHelixManager(
        ConfigUtils.getString(this.clusterConfig, GobblinClusterConfigurationKeys.TASK_DRIVER_CLUSTER_NAME_KEY, ""),
        this.helixInstanceName,
        InstanceType.PARTICIPANT,
        zkConnectionString));
    this.jobHelixManager = HelixManagerFactory.getZKHelixManager(
        this.clusterName,
        this.helixInstanceName,
        InstanceType.ADMINISTRATOR,
        zkConnectionString);
  } else {
    this.jobHelixManager = HelixManagerFactory.getZKHelixManager(
        this.clusterName,
        this.helixInstanceName,
        InstanceType.PARTICIPANT,
        zkConnectionString);
  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  String storageNodeName = PARTICIPANT_PREFIX + "_" + START_PORT;
  _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  _participant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, storageNodeName);
  _participant.syncStart();

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

@Test() public void testInstanceHistory() throws Exception {
  HelixManager manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  manager.connect();

  PropertyKey.Builder keyBuilder = new PropertyKey.Builder(CLUSTER_NAME);
  PropertyKey propertyKey = keyBuilder.participantHistory(_participants[0].getInstanceName());
  ParticipantHistory history = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(history);
  List<String> list = history.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 1);

  for (int i = 0; i <= 22; i++) {
    _participants[0].disconnect();
    _participants[0].connect();
  }

  history = manager.getHelixDataAccessor().getProperty(propertyKey);
  Assert.assertNotNull(history);
  list = history.getRecord().getListField("HISTORY");
  Assert.assertEquals(list.size(), 20);
  list = history.getRecord().getListField("OFFLINE");
  Assert.assertEquals(list.size(), 20);
  manager.disconnect();
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _participants = new MockParticipantManager[NUM_NODES];
  _gSetupTool.addCluster(CLUSTER_NAME, true);

  _participants = new MockParticipantManager[NUM_NODES];
  for (int i = 0; i < NUM_NODES; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

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

  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

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

@Override
public void onBecomeLeaderFromStandby(Message message, NotificationContext context)
    throws Exception {
  String clusterName = message.getPartitionName();
  String controllerName = message.getTgtName();

  logger.info(controllerName + " becoming leader from standby for " + clusterName);

  if (_controller == null) {
    _controller =
        HelixManagerFactory.getZKHelixManager(clusterName, controllerName,
            InstanceType.CONTROLLER, _zkAddr);
    _controller.setEnabledControlPipelineTypes(_enabledPipelineTypes);
    _controller.connect();
    _controller.startTimerTasks();
    logStateTransition("STANDBY", "LEADER", clusterName, controllerName);
  } else {
    logger.error("controller already exists:" + _controller.getInstanceName() + " for "
        + clusterName);
  }

}
 

public void start() throws Exception {
  manager =
      HelixManagerFactory.getZKHelixManager(clusterName, instanceName, InstanceType.PARTICIPANT,
          zkConnectString);

  if ("MasterSlave".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new MasterSlaveStateModelFactory(this.instanceName, delay);
  } else if ("OnlineOffline".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new OnlineOfflineStateModelFactory(this.instanceName, delay);
  } else if ("LeaderStandby".equalsIgnoreCase(stateModelType)) {
    stateModelFactory = new LeaderStandbyStateModelFactory(this.instanceName, delay);
  }
  // genericStateMachineHandler = new StateMachineEngine();
  // genericStateMachineHandler.registerStateModelFactory(stateModelType,
  // stateModelFactory);

  StateMachineEngine stateMach = manager.getStateMachineEngine();
  stateMach.registerStateModelFactory(stateModelType, stateModelFactory);
  manager.connect();
  manager.getMessagingService().registerMessageHandlerFactory(
      MessageType.STATE_TRANSITION.name(), stateMach);
}
 

@BeforeClass
public void setUp()
    throws Exception {
  Config jobConfigOverrides = ClusterIntegrationTestUtils.buildSleepingJob(JOB_ID, TASK_STATE_FILE);
  //Set up a Gobblin Helix cluster integration job
  suite = new IntegrationBasicSuite(jobConfigOverrides);

  helixConfig = suite.getManagerConfig();
  String clusterName = helixConfig.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);
  String zkConnectString = helixConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "TestManager",
      InstanceType.SPECTATOR, zkConnectString);
}
 

@BeforeClass
public void beforeClass() throws Exception {
  super.beforeClass();
  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, _replica);

  // start dummy participants
  for (int i = 0; i < NODE_NR; 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();

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

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

@BeforeClass
public void beforeClass() throws Exception {
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _participants = new MockParticipantManager[NUM_NODES];
  for (int i = 0; i < NUM_NODES; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

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

  _manager = HelixManagerFactory
      .getZKHelixManager(CLUSTER_NAME, "Admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

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

  ConfigAccessor _configAccessor = _manager.getConfigAccessor();
  ClusterConfig clusterConfig = _configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.enableTargetExternalView(true);
  _configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);
}
 

public void start() throws Exception {
  System.out.println("STARTING " + instanceName);
  configureInstance(instanceName);
  participantManager =
      HelixManagerFactory.getZKHelixManager(clusterName, instanceName, InstanceType.PARTICIPANT,
          zkAddress);
  participantManager.getStateMachineEngine().registerStateModelFactory("OnlineOffline",
      new LockFactory());
  participantManager.connect();
  if (startController) {
    startController();
  }
  System.out.println("STARTED " + instanceName);

}
 

public void start(String stateModelName, StateModelFactory<T> stateModelFactory) {
    try {
        manager = HelixManagerFactory.getZKHelixManager(clusterName, instanceName,
                InstanceType.PARTICIPANT, zkAddr);

        StateMachineEngine stateMachine = manager.getStateMachineEngine();
        stateMachine.registerStateModelFactory(stateModelName, stateModelFactory);

        manager.connect();
    } catch (Exception e) {
        logger.error("failed to connect manager", e);
        throw new RuntimeException("failed to start HelixPartitionManager");
    }
}
 

public void start() throws Exception {
  try {
    workerInstance.start(srcCluster, dstCluster, routeId, federatedDeploymentName);
    KafkaUReplicatorMetricsReporter.get()
        .registerMetric(METRIC_ADD_TOPIC_PARTITION_FAILURE, addTopicPartitionFailureMeter);
    KafkaUReplicatorMetricsReporter.get()
        .registerMetric(METRIC_DELETE_TOPIC_PARTITION_FAILURE, deleteTopicPartitionFailureMeter);

    if (helixManager != null && helixManager.isConnected()) {
      LOGGER.warn("ControllerWorkerHelixManager already connected");
      return;
    }
    LOGGER.info("Starting ControllerWorkerHelixHandler");
    helixManager = HelixManagerFactory.getZKHelixManager(helixClusterName,
        workerInstanceId,
        InstanceType.PARTICIPANT,
        helixZkURL);
    helixManager.connect();
    helixManager.getStateMachineEngine()
        .registerStateModelFactory(OnlineOfflineStateFactory.STATE_MODE_DEF,
            new OnlineOfflineStateFactory(this));
  } catch (Exception e) {
    LOGGER.error("Add instance to helix cluster failed. helixCluster: {}",
        helixClusterName, e);
    throw e;
  }
  LOGGER.info("Register ControllerWorkerHelixHandler finished");
}
 

public void start() throws Exception {
  _helixZkManager = HelixManagerFactory.getZKHelixManager(_helixClusterName, _instanceId,
      InstanceType.PARTICIPANT, _zkServer);
  StateMachineEngine stateMachineEngine = _helixZkManager.getStateMachineEngine();
  // create a stateModelFactory that returns a statemodel object for each partition.
  TestOnlineOfflineStateModelFactory stateModelFactory =
      new TestOnlineOfflineStateModelFactory(_instanceId, 0);
  stateMachineEngine.registerStateModelFactory("OnlineOffline", stateModelFactory);
  _helixZkManager.connect();
}
 

@BeforeClass
public void setUp() throws Exception {
  // Use a random ZK port
  this.testingZKServer = new TestingServer(-1);
  LOG.info("Testing ZK Server listening on: " + testingZKServer.getConnectString());

  URL url = GobblinClusterManagerTest.class.getClassLoader().getResource(
      GobblinClusterManager.class.getSimpleName() + ".conf");
  Assert.assertNotNull(url, "Could not find resource " + url);

  Config config = ConfigFactory.parseURL(url)
      .withValue("gobblin.cluster.zk.connection.string",
                 ConfigValueFactory.fromAnyRef(testingZKServer.getConnectString()))
      .withValue(GobblinClusterConfigurationKeys.HELIX_TASK_QUOTA_CONFIG_KEY,
          ConfigValueFactory.fromAnyRef("DEFAULT:1,OTHER:10"))
      .withValue(GobblinClusterConfigurationKeys.HADOOP_CONFIG_OVERRIDES_PREFIX + "." + HADOOP_OVERRIDE_PROPERTY_NAME,
          ConfigValueFactory.fromAnyRef("value"))
      .withValue(GobblinClusterConfigurationKeys.HADOOP_CONFIG_OVERRIDES_PREFIX + "." + "fs.file.impl.disable.cache",
          ConfigValueFactory.fromAnyRef("true"))
      .resolve();

  String zkConnectionString = config.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  HelixUtils.createGobblinHelixCluster(zkConnectionString,
      config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY));

  this.helixManager = HelixManagerFactory
      .getZKHelixManager(config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY),
          TestHelper.TEST_HELIX_INSTANCE_NAME, InstanceType.PARTICIPANT, zkConnectionString);
  this.helixManager.connect();
  this.helixManager.getMessagingService().registerMessageHandlerFactory(GobblinHelixConstants.SHUTDOWN_MESSAGE_TYPE,
      new TestShutdownMessageHandlerFactory(this));

  this.gobblinClusterManager =
      new GobblinClusterManager(GobblinClusterManagerTest.class.getSimpleName(), TestHelper.TEST_APPLICATION_ID, config,
          Optional.<Path>absent());
  this.gobblinClusterManager.getEventBus().register(this.gobblinClusterManager);
  this.gobblinClusterManager.connectHelixManager();
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _participants = new MockParticipantManager[_numNodes];
  String namespace = "/" + CLUSTER_NAME;
  if (_gZkClient.exists(namespace)) {
    _gZkClient.deleteRecursively(namespace);
  }

  // Setup cluster and instances
  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  setupTool.addCluster(CLUSTER_NAME, true);
  for (int i = 0; i < _numNodes; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (_startPort + i);
    setupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

  // start dummy participants
  for (int i = 0; i < _numNodes; i++) {
    final String instanceName = PARTICIPANT_PREFIX + "_" + (_startPort + i);

    // Set task callbacks
    Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
    TaskFactory shortTaskFactory = WriteTask::new;
    taskFactoryReg.put("WriteTask", shortTaskFactory);

    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);

    // Register a Task state model factory.
    StateMachineEngine stateMachine = _participants[i].getStateMachineEngine();
    stateMachine.registerStateModelFactory("Task",
        new TaskStateModelFactory(_participants[i], taskFactoryReg));
    _participants[i].syncStart();
  }

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

  // Start an admin connection
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _driver = new TaskDriver(_manager);

  _jobCommandMap = new HashMap<>();
}
 

@BeforeClass
public void setUp() throws Exception {

  // Mock AWS SDK calls
  MockitoAnnotations.initMocks(this);

  PowerMockito.whenNew(AWSSdkClient.class).withAnyArguments().thenReturn(awsSdkClient);

  Mockito.doNothing()
      .when(awsSdkClient)
      .createSecurityGroup(Mockito.anyString(), Mockito.anyString());
  Mockito.doReturn(Lists.<AvailabilityZone>newArrayList(availabilityZone))
      .when(awsSdkClient)
      .getAvailabilityZones();
  Mockito.doReturn("dummy")
      .when(awsSdkClient)
      .createKeyValuePair(Mockito.anyString());
  Mockito.doReturn(Lists.<AutoScalingGroup>newArrayList(masterASG, workerASG))
      .when(awsSdkClient)
      .getAutoScalingGroupsWithTag(Mockito.any(Tag.class));
  Mockito.doReturn(Lists.<Instance>newArrayList(instance))
      .when(awsSdkClient)
      .getInstancesForGroup(Mockito.anyString(), Mockito.anyString());
  Mockito.doReturn(Lists.<S3ObjectSummary>newArrayList())
      .when(awsSdkClient)
      .listS3Bucket(Mockito.anyString(), Mockito.anyString());
  Mockito.doNothing()
      .when(awsSdkClient)
      .addPermissionsToSecurityGroup(Mockito.any(String.class), Mockito.any(String.class), Mockito.any(String.class),
          Mockito.any(Integer.class), Mockito.any(Integer.class));
  Mockito.doNothing()
      .when(awsSdkClient)
      .createAutoScalingGroup(Mockito.any(String.class), Mockito.any(String.class), Mockito.any(Integer.class),
          Mockito.any(Integer.class), Mockito.any(Integer.class), Mockito.any(Optional.class),
          Mockito.any(Optional.class), Mockito.any(Optional.class), Mockito.any(Optional.class),
          Mockito.any(Optional.class), Mockito.any(Optional.class), Mockito.any(List.class));
  Mockito.doNothing()
      .when(awsSdkClient)
      .createLaunchConfig(Mockito.any(String.class), Mockito.any(String.class), Mockito.any(String.class),
          Mockito.any(String.class), Mockito.any(String.class), Mockito.any(Optional.class),
          Mockito.any(Optional.class), Mockito.any(Optional.class), Mockito.any(Optional.class),
          Mockito.any(Optional.class), Mockito.any(String.class));
  Mockito
      .doNothing()
      .when(awsSdkClient)
      .deleteAutoScalingGroup(Mockito.any(String.class), Mockito.any(boolean.class));
  Mockito
      .doNothing()
      .when(awsSdkClient)
      .deleteLaunchConfiguration(Mockito.any(String.class));
  Mockito.doNothing()
      .when(awsSdkClient)
      .addPermissionsToSecurityGroup(Mockito.any(String.class), Mockito.any(String.class), Mockito.any(String.class),
          Mockito.any(Integer.class), Mockito.any(Integer.class));

  // Local test Zookeeper
  final TestingServer testingZKServer = this.closer.register(new TestingServer(-1));
  LOG.info("Testing ZK Server listening on: " + testingZKServer.getConnectString());
  this.curatorFramework = TestHelper.createZkClient(testingZKServer, this.closer);

  // Load configuration
  final URL url = GobblinAWSClusterLauncherTest.class.getClassLoader().getResource(
      GobblinAWSClusterLauncherTest.class.getSimpleName() + ".conf");
  Assert.assertNotNull(url, "Could not find resource " + url);
  this.config = ConfigFactory.parseURL(url)
      .withValue("gobblin.cluster.zk.connection.string",
                 ConfigValueFactory.fromAnyRef(testingZKServer.getConnectString()))
      .resolve();
  this.helixClusterName = this.config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);

  final String zkConnectionString = this.config.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  this.helixManager = HelixManagerFactory
      .getZKHelixManager(this.config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY),
          TestHelper.TEST_HELIX_INSTANCE_NAME, InstanceType.CONTROLLER, zkConnectionString);

  // Gobblin AWS Cluster Launcher to test
  this.gobblinAwsClusterLauncher = new GobblinAWSClusterLauncher(this.config);
}
 

public GobblinYarnAppLauncher(Config config, YarnConfiguration yarnConfiguration) throws IOException {
  this.config = config;

  this.applicationName = config.getString(GobblinYarnConfigurationKeys.APPLICATION_NAME_KEY);
  this.appQueueName = config.getString(GobblinYarnConfigurationKeys.APP_QUEUE_KEY);

  String zkConnectionString = config.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  LOGGER.info("Using ZooKeeper connection string: " + zkConnectionString);

  this.helixManager = HelixManagerFactory.getZKHelixManager(
      config.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY), GobblinClusterUtils.getHostname(),
      InstanceType.SPECTATOR, zkConnectionString);

  this.yarnConfiguration = yarnConfiguration;
  YarnHelixUtils.setAdditionalYarnClassPath(config, this.yarnConfiguration);
  this.yarnConfiguration.set("fs.automatic.close", "false");
  this.yarnClient = YarnClient.createYarnClient();
  this.yarnClient.init(this.yarnConfiguration);

  this.fs = GobblinClusterUtils.buildFileSystem(config, this.yarnConfiguration);
  this.closer.register(this.fs);

  this.applicationStatusMonitor = Executors.newSingleThreadScheduledExecutor(
      ExecutorsUtils.newThreadFactory(Optional.of(LOGGER), Optional.of("GobblinYarnAppStatusMonitor")));
  this.appReportIntervalMinutes = config.getLong(GobblinYarnConfigurationKeys.APP_REPORT_INTERVAL_MINUTES_KEY);

  this.appMasterJvmArgs = config.hasPath(GobblinYarnConfigurationKeys.APP_MASTER_JVM_ARGS_KEY) ?
      Optional.of(config.getString(GobblinYarnConfigurationKeys.APP_MASTER_JVM_ARGS_KEY)) :
      Optional.<String>absent();

  this.sinkLogRootDir = new Path(config.getString(GobblinYarnConfigurationKeys.LOGS_SINK_ROOT_DIR_KEY));

  this.maxGetApplicationReportFailures = config.getInt(GobblinYarnConfigurationKeys.MAX_GET_APP_REPORT_FAILURES_KEY);

  this.emailNotificationOnShutdown =
      config.getBoolean(GobblinYarnConfigurationKeys.EMAIL_NOTIFICATION_ON_SHUTDOWN_KEY);

  this.appMasterMemoryMbs = this.config.getInt(GobblinYarnConfigurationKeys.APP_MASTER_MEMORY_MBS_KEY);

  this.jvmMemoryXmxRatio = ConfigUtils.getDouble(this.config,
      GobblinYarnConfigurationKeys.APP_MASTER_JVM_MEMORY_XMX_RATIO_KEY,
      GobblinYarnConfigurationKeys.DEFAULT_APP_MASTER_JVM_MEMORY_XMX_RATIO);

  Preconditions.checkArgument(this.jvmMemoryXmxRatio >= 0 && this.jvmMemoryXmxRatio <= 1,
      GobblinYarnConfigurationKeys.APP_MASTER_JVM_MEMORY_XMX_RATIO_KEY + " must be between 0 and 1 inclusive");

  this.jvmMemoryOverheadMbs = ConfigUtils.getInt(this.config,
      GobblinYarnConfigurationKeys.APP_MASTER_JVM_MEMORY_OVERHEAD_MBS_KEY,
      GobblinYarnConfigurationKeys.DEFAULT_APP_MASTER_JVM_MEMORY_OVERHEAD_MBS);

  Preconditions.checkArgument(this.jvmMemoryOverheadMbs < this.appMasterMemoryMbs * this.jvmMemoryXmxRatio,
      GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_OVERHEAD_MBS_KEY + " cannot be more than "
          + GobblinYarnConfigurationKeys.CONTAINER_MEMORY_MBS_KEY + " * "
          + GobblinYarnConfigurationKeys.CONTAINER_JVM_MEMORY_XMX_RATIO_KEY);

  this.appViewAcl = ConfigUtils.getString(this.config, GobblinYarnConfigurationKeys.APP_VIEW_ACL,
      GobblinYarnConfigurationKeys.DEFAULT_APP_VIEW_ACL);
  this.containerTimezone = ConfigUtils.getString(this.config, GobblinYarnConfigurationKeys.GOBBLIN_YARN_CONTAINER_TIMEZONE,
      GobblinYarnConfigurationKeys.DEFAULT_GOBBLIN_YARN_CONTAINER_TIMEZONE);

  this.isHelixClusterManaged = ConfigUtils.getBoolean(this.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());

  this.detachOnExitEnabled = ConfigUtils
      .getBoolean(config, GobblinYarnConfigurationKeys.GOBBLIN_YARN_DETACH_ON_EXIT_ENABLED,
          GobblinYarnConfigurationKeys.DEFAULT_GOBBLIN_YARN_DETACH_ON_EXIT);
  this.appLauncherMode = ConfigUtils.getString(config, GOBBLIN_YARN_APP_LAUNCHER_MODE, DEFAULT_GOBBLIN_YARN_APP_LAUNCHER_MODE);

  this.messagingService = new GobblinHelixMessagingService(this.helixManager);

  try {
    config = addDynamicConfig(config);
    outputConfigToFile(config);
  } catch (SchemaRegistryException e) {
    throw new IOException(e);
  }
}
 

@BeforeClass
public void setUp() throws Exception {
  TestingServer testingZKServer = this.closer.register(new TestingServer(-1));
  LOG.info("Testing ZK Server listening on: " + testingZKServer.getConnectString());

  URL url = GobblinHelixJobLauncherTest.class.getClassLoader().getResource(
      GobblinHelixJobLauncherTest.class.getSimpleName() + ".conf");
  Assert.assertNotNull(url, "Could not find resource " + url);

  this.appWorkDir = new Path(GobblinHelixJobLauncherTest.class.getSimpleName());

  // Prepare the source Json file
  File sourceJsonFile = new File(this.appWorkDir.toString(), TestHelper.TEST_JOB_NAME + ".json");
  TestHelper.createSourceJsonFile(sourceJsonFile);

  baseConfig = ConfigFactory.parseURL(url)
      .withValue("gobblin.cluster.zk.connection.string",
                 ConfigValueFactory.fromAnyRef(testingZKServer.getConnectString()))
      .withValue(ConfigurationKeys.SOURCE_FILEBASED_FILES_TO_PULL,
          ConfigValueFactory.fromAnyRef(sourceJsonFile.getAbsolutePath()))
      .withValue(ConfigurationKeys.JOB_STATE_IN_STATE_STORE, ConfigValueFactory.fromAnyRef("true"))
      .resolve();

  String zkConnectingString = baseConfig.getString(GobblinClusterConfigurationKeys.ZK_CONNECTION_STRING_KEY);
  String helixClusterName = baseConfig.getString(GobblinClusterConfigurationKeys.HELIX_CLUSTER_NAME_KEY);

  HelixUtils.createGobblinHelixCluster(zkConnectingString, helixClusterName);

  this.helixManager = HelixManagerFactory
      .getZKHelixManager(helixClusterName, TestHelper.TEST_HELIX_INSTANCE_NAME, InstanceType.CONTROLLER,
          zkConnectingString);
  this.closer.register(new Closeable() {
    @Override
    public void close() throws IOException {
      helixManager.disconnect();
    }
  });
  this.helixManager.connect();

  this.localFs = FileSystem.getLocal(new Configuration());

  this.closer.register(new Closeable() {
    @Override
    public void close() throws IOException {
      if (localFs.exists(appWorkDir)) {
        localFs.delete(appWorkDir, true);
      }
    }
  });

  this.gobblinTaskRunner =
      new GobblinTaskRunner(TestHelper.TEST_APPLICATION_NAME, TestHelper.TEST_HELIX_INSTANCE_NAME,
          TestHelper.TEST_APPLICATION_ID, TestHelper.TEST_TASK_RUNNER_ID, baseConfig, Optional.of(appWorkDir));

  String stateStoreType = ConfigUtils.getString(baseConfig, ConfigurationKeys.STATE_STORE_TYPE_KEY,
      ConfigurationKeys.DEFAULT_STATE_STORE_TYPE);

  ClassAliasResolver<DatasetStateStore.Factory> resolver =
      new ClassAliasResolver<>(DatasetStateStore.Factory.class);

  DatasetStateStore.Factory stateStoreFactory =
      resolver.resolveClass(stateStoreType).newInstance();

  this.datasetStateStore = stateStoreFactory.createStateStore(baseConfig);

  this.thread = new Thread(new Runnable() {
    @Override
    public void run() {
      gobblinTaskRunner.start();
    }
  });
  this.thread.start();
}
 

@BeforeClass
public void beforeClass() throws Exception {
  _numNodes = 2; // For easier debugging by inspecting ZNodes

  _participants = new MockParticipantManager[_numNodes];
  String namespace = "/" + CLUSTER_NAME;
  if (_gZkClient.exists(namespace)) {
    _gZkClient.deleteRecursively(namespace);
  }

  // Setup cluster and instances
  ClusterSetup setupTool = new ClusterSetup(ZK_ADDR);
  setupTool.addCluster(CLUSTER_NAME, true);
  for (int i = 0; i < _numNodes; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (_startPort + i);
    setupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }

  // start dummy participants
  for (int i = 0; i < _numNodes; i++) {
    final String instanceName = PARTICIPANT_PREFIX + "_" + (_startPort + i);

    // Set task callbacks
    Map<String, TaskFactory> taskFactoryReg = new HashMap<>();
    TaskFactory shortTaskFactory = context -> new ShortTask(context, instanceName);
    TaskFactory longTaskFactory = context -> new LongTask(context, instanceName);
    TaskFactory failTaskFactory = context -> new FailTask(context, instanceName);
    taskFactoryReg.put("ShortTask", shortTaskFactory);
    taskFactoryReg.put("LongTask", longTaskFactory);
    taskFactoryReg.put("FailTask", failTaskFactory);

    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);

    // Register a Task state model factory.
    StateMachineEngine stateMachine = _participants[i].getStateMachineEngine();
    stateMachine.registerStateModelFactory("Task",
        new TaskStateModelFactory(_participants[i], taskFactoryReg));
    _participants[i].syncStart();
  }

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

  // Start an admin connection
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
  _driver = new TaskDriver(_manager);

  _jobCommandMap = Maps.newHashMap();
}
 

/**
 * Creates a setup where the health API can be tested.
 * @param clusterName
 * @param resourceName
 * @param idealStateParams
 * @param partitionReplicaStates maps partitionName to its replicas' states
 * @throws Exception
 */
private void createDummyMapping(String clusterName, String resourceName,
    Map<String, String> idealStateParams, Map<String, List<String>> partitionReplicaStates)
    throws Exception {
  IdealState idealState = new IdealState(resourceName);
  idealState.setMinActiveReplicas(Integer.parseInt(idealStateParams.get("MinActiveReplicas"))); // 2
  idealState.setStateModelDefRef(idealStateParams.get("StateModelDefRef")); // MasterSlave
  idealState.setMaxPartitionsPerInstance(
      Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance"))); // 3
  idealState.setReplicas(idealStateParams.get("Replicas")); // 3
  idealState.setNumPartitions(Integer.parseInt(idealStateParams.get("NumPartitions"))); // 3
  idealState.enable(false);

  Map<String, List<String>> partitionNames = new LinkedHashMap<>();
  List<String> dummyPrefList = new ArrayList<>();

  for (int i = 0; i < Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance")); i++) {
    dummyPrefList.add(ANY_INSTANCE);
    partitionNames.put("p" + i, dummyPrefList);
  }
  idealState.getRecord().getListFields().putAll(partitionNames);

  if (!_gSetupTool.getClusterManagementTool().getClusters().contains(clusterName)) {
    _gSetupTool.getClusterManagementTool().addCluster(clusterName);
  }
  _gSetupTool.getClusterManagementTool().setResourceIdealState(clusterName, resourceName,
      idealState);

  // Set ExternalView's replica states for a given parameter map
  ExternalView externalView = new ExternalView(resourceName);

  Map<String, Map<String, String>> mappingCurrent = new LinkedHashMap<>();

  List<String> partitionReplicaStatesList = new ArrayList<>(partitionReplicaStates.keySet());
  for (int k = 0; k < partitionReplicaStatesList.size(); k++) {
    Map<String, String> replicaStatesForPartition = new LinkedHashMap<>();
    List<String> replicaStateList = partitionReplicaStates.get(partitionReplicaStatesList.get(k));
    for (int i = 0; i < replicaStateList.size(); i++) {
      replicaStatesForPartition.put("r" + i, replicaStateList.get(i));
    }
    mappingCurrent.put("p" + k, replicaStatesForPartition);
  }

  externalView.getRecord().getMapFields().putAll(mappingCurrent);

  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "p1",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  helixManager.connect();
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().externalView(resourceName),
      externalView);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}