org.apache.helix.task.JobConfig Java Examples

@Test
public void stopAndResume() throws Exception {
  Map<String, String> commandConfig = ImmutableMap.of(MockTask.JOB_DELAY, String.valueOf(100));

  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(commandConfig);
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(JOB_RESOURCE, jobBuilder).build();

  LOG.info("Starting flow " + flow.getName());
  _driver.start(flow);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.IN_PROGRESS);

  LOG.info("Pausing job");
  _driver.stop(JOB_RESOURCE);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.STOPPED);

  LOG.info("Resuming job");
  _driver.resume(JOB_RESOURCE);
  _driver.pollForWorkflowState(JOB_RESOURCE, TaskState.COMPLETED);
}
 

/**
 * Task type: generic
 * Rebalance running task: enabled
 * Story: new node added
 * NOTE: This test is disabled because this "load-balancing" would happen at the Task Assigner
 * level. In the legacy assignment strategy (Consistent Hashing) did not take instance's capacity
 * into account. However, the new quota-based scheduling takes capacity into account, and it will
 * generally assign to the most "free" instance, so load-balancing of tasks will happen at the
 * Assigner layer. Deprecating this test.
 */
@Deprecated
@Test(enabled = false)
public void testGenericTaskAndEnabledRebalanceAndNodeAdded() throws InterruptedException {
  WORKFLOW = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setWorkflow(WORKFLOW)
      .setNumberOfTasks(10).setNumConcurrentTasksPerInstance(100)
      .setCommand(MockTask.TASK_COMMAND).setRebalanceRunningTask(true)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999")); // task stuck

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW).addJob(JOB, jobBuilder);

  _driver.start(workflowBuilder.build());

  // All tasks stuck on the same instance
  Assert.assertTrue(checkTasksOnSameInstances());
  // Add a new instance, and some running tasks will be rebalanced to the new node
  startParticipant(_initialNumNodes);
  Assert.assertTrue(checkTasksOnDifferentInstances());
}
 

@Test public void testJobQueueNotCleanupRunningJobs() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2);
  for (int i = 0; i < 3; i++) {
    builder.enqueueJob("JOB" + i, jobBuilder);
  }
  builder.enqueueJob("JOB" + 3,
      jobBuilder.setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "1000000")));
  builder.enqueueJob("JOB" + 4, jobBuilder);
  _driver.start(builder.build());
  _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + 3),
      TaskState.IN_PROGRESS);
  _driver.cleanupQueue(queueName);
  Assert.assertEquals(_driver.getWorkflowConfig(queueName).getJobDag().size(), 2);
}
 

@Test
public void testJobQueueCleanUp() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.SUCCESS_COUNT_BEFORE_FAIL, "2"));
  for (int i = 0; i < 5; i++) {
    builder.enqueueJob("JOB" + i, jobBuilder);
  }
  _driver.start(builder.build());
  _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + 4),
      TaskState.FAILED);
  _driver.cleanupQueue(queueName);
  Assert.assertEquals(_driver.getWorkflowConfig(queueName).getJobDag().size(), 0);
}
 

@Test
public void testWorkflowAndJobTaskUserContentStore() throws InterruptedException {
  String jobName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(jobName);
  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
  Map<String, String> taskConfigMap = Maps.newHashMap();
  TaskConfig taskConfig1 = new TaskConfig("ContentStoreTask", taskConfigMap);
  taskConfigs.add(taskConfig1);
  Map<String, String> jobCommandMap = Maps.newHashMap();
  jobCommandMap.put("Timeout", "1000");

  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .addTaskConfigs(taskConfigs).setWorkflow(jobName)
      .setJobCommandConfigMap(jobCommandMap);
  workflowBuilder.addJob(jobName, jobBuilder);

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);
  Assert
      .assertEquals(_driver.getWorkflowContext(jobName).getWorkflowState(), TaskState.COMPLETED);
}
 

private List<String> createAndEnqueueJob(JobQueue.Builder queueBuild, int jobCount) {
  List<String> currentJobNames = new ArrayList<>();
  for (int i = 0; i < jobCount; i++) {
    String targetPartition = (i == 0) ? "MASTER" : "SLAVE";

    JobConfig.Builder jobConfig = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
        .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
        .setTargetPartitionStates(Sets.newHashSet(targetPartition))
        .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "100"));
    String jobName = targetPartition.toLowerCase() + "Job" + i;
    queueBuild.enqueueJob(jobName, jobConfig);
    currentJobNames.add(jobName);
  }
  Assert.assertEquals(currentJobNames.size(), jobCount);
  return currentJobNames;
}
 

@Test
public void testJobQueueAddingJobsOneByOne() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  WorkflowConfig.Builder workflowCfgBuilder = new WorkflowConfig.Builder().setWorkflowId(queueName).setParallelJobs(1);
  _driver.start(builder.setWorkflowConfig(workflowCfgBuilder.build()).build());
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2);
  _driver.enqueueJob(queueName, "JOB0", jobBuilder);
  for (int i = 1; i < 5; i++) {
    _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + (i - 1)),
        10000L, TaskState.COMPLETED);
    _driver.waitToStop(queueName, 5000L);
    _driver.enqueueJob(queueName, "JOB" + i, jobBuilder);
    _driver.resume(queueName);
  }

  _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + 4),
      TaskState.COMPLETED);
}
 

/**
 * Task type: fixed target
 * Rebalance running task: disabled
 * Story: 1 node is down
 */
@Test
public void testFixedTargetTaskAndDisabledRebalanceAndNodeDown() throws InterruptedException {
  WORKFLOW = TestHelper.getTestMethodName();
  startParticipant(_initialNumNodes);

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW).setTargetResource(DATABASE)
          .setNumConcurrentTasksPerInstance(100).setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999"));

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW).addJob(JOB, jobBuilder);
  _driver.start(workflowBuilder.build());
  Assert.assertTrue(checkTasksOnDifferentInstances());
  // Stop a participant and partitions will be moved to the same instance,
  // and tasks rebalanced accordingly
  stopParticipant(_initialNumNodes);
  Assert.assertTrue(checkTasksOnSameInstances());
}
 

@Test
public void testJobAndWorkflowType() throws InterruptedException {
  LOG.info("Start testing job and workflow type");
  String jobName = TestHelper.getTestMethodName();
  JobConfig.Builder jobConfig = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG)
      .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG).setJobType(DEFAULT_TYPE);

  Map<String, String> tmp = new HashMap<>();
  tmp.put("WorkflowType", DEFAULT_TYPE);
  Workflow.Builder builder =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobName, jobConfig).fromMap(tmp);

  // Start workflow
  _driver.start(builder.build());

  _driver.pollForWorkflowState(jobName, TaskState.COMPLETED);
  String fetchedJobType =
      _driver.getJobConfig(String.format("%s_%s", jobName, jobName)).getJobType();
  String fetchedWorkflowType = _driver.getWorkflowConfig(jobName).getWorkflowType();

  Assert.assertEquals(fetchedJobType, DEFAULT_TYPE);
  Assert.assertEquals(fetchedWorkflowType, DEFAULT_TYPE);
}
 

private List<String> createAndEnqueueJob(JobQueue.Builder queueBuild, int jobCount) {
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < jobCount; i++) {
    String targetPartition = (i == 0) ? "MASTER" : "SLAVE";

    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
            .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
            .setTargetPartitionStates(Sets.newHashSet(targetPartition));
    String jobName = targetPartition.toLowerCase() + "Job" + i;
    queueBuild.enqueueJob(jobName, jobConfig);
    currentJobNames.add(jobName);
  }
  Assert.assertEquals(currentJobNames.size(), jobCount);
  return currentJobNames;
}
 

@Test(dependsOnMethods = "testGetJobs")
public void testGetJob() throws IOException {
  System.out.println("Start test :" + TestHelper.getTestMethodName());

  String body =
      get("clusters/" + CLUSTER_NAME + "/workflows/" + WORKFLOW_NAME + "/jobs/" + JOB_NAME, null,
          Response.Status.OK.getStatusCode(), true);
  JsonNode node = OBJECT_MAPPER.readTree(body);
  Assert.assertNotNull(node.get(JobAccessor.JobProperties.JobConfig.name()));
  Assert.assertNotNull(node.get(JobAccessor.JobProperties.JobContext.name()));
  String workflowId =
      node.get(JobAccessor.JobProperties.JobConfig.name()).get("simpleFields").get("WorkflowID")
          .getTextValue();
  Assert.assertEquals(workflowId, WORKFLOW_NAME);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

@Test
public void testTaskAssignment() throws InterruptedException {
  _gSetupTool.getClusterManagementTool()
      .enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_" + (_startPort + 0), false);
  String jobResource = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB);

  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait 1 sec. The task should not be complete since it is not assigned.
  Thread.sleep(1000L);

  // The task is not assigned so the task state should be null in this case.
  Assert.assertNull(
      _driver.getJobContext(TaskUtil.getNamespacedJobName(jobResource)).getPartitionState(0));
}
 

@Test
public void testStoppingQueueFailToStop() throws Exception {
  String jobQueueName = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder0 =
      new JobConfig.Builder().setWorkflow(jobQueueName).setTargetResource(DATABASE)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "100000"));

  JobQueue.Builder jobQueue = TaskTestUtil.buildJobQueue(jobQueueName);
  jobQueue.enqueueJob("JOB0", jobBuilder0);
  _driver.start(jobQueue.build());
  _driver.pollForJobState(jobQueueName, TaskUtil.getNamespacedJobName(jobQueueName, "JOB0"),
      TaskState.IN_PROGRESS);
  boolean exceptionHappened = false;
  try {
    _driver.waitToStop(jobQueueName, 5000L);
  } catch (HelixException e) {
    exceptionHappened = true;
  }
  _driver.pollForWorkflowState(jobQueueName, TaskState.STOPPING);
  Assert.assertTrue(exceptionHappened);
  latch.countDown();
}
 

@Test
public void testJobQueueAddingJobsAtSametime() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  WorkflowConfig.Builder workflowCfgBuilder =
      new WorkflowConfig.Builder().setWorkflowId(queueName).setParallelJobs(1);
  _driver.start(builder.setWorkflowConfig(workflowCfgBuilder.build()).build());

  // Adding jobs
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2);
  _driver.waitToStop(queueName, 5000L);
  for (int i = 0; i < 5; i++) {
    _driver.enqueueJob(queueName, "JOB" + i, jobBuilder);
  }
  _driver.resume(queueName);

  _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + 4),
      TaskState.COMPLETED);
}
 

/**
 * Helper method for creating custom workflows.
 * @param workflowName
 * @param shouldOverlapJobAssign
 * @param quotaType
 * @param numJobs
 * @param numTasks
 * @param taskType
 * @return a workflow per parameters given
 */
private Workflow createWorkflow(String workflowName, boolean shouldOverlapJobAssign,
    String quotaType, int numJobs, int numTasks, String taskType) {

  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(shouldOverlapJobAssign);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int jobIndex = 0; jobIndex < numJobs; jobIndex++) {
    String jobName = workflowName + "_" + jobIndex;
    List<TaskConfig> taskConfigs = new ArrayList<>();
    for (int taskIndex = 0; taskIndex < numTasks; taskIndex++) {
      Map<String, String> taskConfigMap = new HashMap<>();
      taskConfigs.add(new TaskConfig(taskType, taskConfigMap));
    }
    JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand(JOB_COMMAND)
        .setJobCommandConfigMap(_jobCommandMap).addTaskConfigs(taskConfigs).setJobType(quotaType);
    workflowBuilder.addJob(jobName, jobBuilder);
  }
  return workflowBuilder.build();
}
 

/**
 * This test does NOT allow multiple jobs being assigned to an instance.
 * @throws InterruptedException
 */
@Test
public void testMultipleJobAssignment() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);

  for (int i = 0; i < 20; i++) {
    List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(1);
    taskConfigs.add(new TaskConfig("TaskOne", new HashMap<>()));
    JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(_jobCommandMap);
    workflowBuilder.addJob("JOB" + i, jobBuilder);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  Assert.assertEquals(_runCounts.size(), 5);
}
 

public static void submitJobToWorkFlow(JobConfig.Builder jobConfigBuilder,
    String workFlowName,
    String jobName,
    TaskDriver helixTaskDriver,
    HelixManager helixManager,
    long workFlowExpiryTime) throws Exception {

  WorkflowConfig workFlowConfig = new WorkflowConfig.Builder().setExpiry(workFlowExpiryTime, TimeUnit.SECONDS).build();
  // Create a work flow for each job with the name being the queue name
  Workflow workFlow = new Workflow.Builder(workFlowName).setWorkflowConfig(workFlowConfig).addJob(jobName, jobConfigBuilder).build();
  // start the workflow
  helixTaskDriver.start(workFlow);
  log.info("Created a work flow {}", workFlowName);

  waitJobInitialization(helixManager, workFlowName, jobName, Long.MAX_VALUE);
}
 

@Test
public void testTaskRetryWithDelay() throws Exception {
  String jobResource = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(2).setCommand(MockTask.TASK_COMMAND).setWorkflow(jobResource)
      .setFailureThreshold(Integer.MAX_VALUE).setTaskRetryDelay(2000L)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.FAILURE_COUNT_BEFORE_SUCCESS, "2"));
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job completes.
  _driver.pollForWorkflowState(jobResource, TaskState.COMPLETED);

  long startTime = _driver.getWorkflowContext(jobResource).getStartTime();
  long finishedTime = _driver.getWorkflowContext(jobResource).getFinishTime();

  // It should finish and take more than 2 secs with the retry delay built in
  Assert.assertTrue(finishedTime - startTime >= 2000L);
}
 

@Override
public DistributeJobResult call()
    throws Exception {
  String planningId = getPlanningJobId(this.jobPlanningProps);
  JobConfig.Builder builder = createJobBuilder(this.jobPlanningProps);
  try {
    long submitStartTime = System.currentTimeMillis();
    GobblinHelixDistributeJobExecutionLauncher.this.helixMetrics.submitMeter.mark();
    submitJobToHelix(planningId, planningId, builder);
    GobblinHelixDistributeJobExecutionLauncher.this.helixMetrics.updateTimeForHelixSubmit(submitStartTime);
    long waitStartTime = System.currentTimeMillis();

    // we should not wait if in non-blocking mode.
    DistributeJobResult rst = new DistributeJobResult();
    if (!GobblinHelixDistributeJobExecutionLauncher.this.nonBlockingMode) {
      rst = waitForJobCompletion(planningId, planningId);
      GobblinHelixDistributeJobExecutionLauncher.this.helixMetrics.updateTimeForHelixWait(waitStartTime);
    }

    return rst;
  } catch (Exception e) {
    log.error(planningId + " is not able to submit.");
    return new DistributeJobResult();
  }
}
 

@Test
public void testTaskRetryWithoutDelay() throws Exception {
  String jobResource = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(2).setCommand(MockTask.TASK_COMMAND)
      .setFailureThreshold(Integer.MAX_VALUE)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.FAILURE_COUNT_BEFORE_SUCCESS, "1"));
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job completes.
  _driver.pollForWorkflowState(jobResource, TaskState.COMPLETED);

  long startTime = _driver.getWorkflowContext(jobResource).getStartTime();
  long finishedTime = _driver.getWorkflowContext(jobResource).getFinishTime();

  // It should have finished within less than 2 sec
  Assert.assertTrue(finishedTime - startTime <= 2000L);
}
 

/**
 * Tests whether jobs can run successfully without quotaTypes or quota configuration defined in
 * ClusterConfig. This test is to ensure backward-compatibility. This test must go first because
 * we want to make sure there is no quota config set anywhere.
 * @throws InterruptedException
 */
@Test
public void testSchedulingWithoutQuota() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowName);
  configBuilder.setAllowOverlapJobAssignment(true);
  workflowBuilder.setWorkflowConfig(configBuilder.build());

  for (int i = 0; i < 10; i++) {
    List<TaskConfig> taskConfigs = new ArrayList<>();
    taskConfigs.add(new TaskConfig("ShortTask", new HashMap<>()));
    JobConfig.Builder jobConfigBulider = new JobConfig.Builder().setCommand(JOB_COMMAND)
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(_jobCommandMap);
    workflowBuilder.addJob("JOB" + i, jobConfigBulider);
  }

  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  for (int i = 0; i < 10; i++) {
    String jobName = workflowName + "_" + "JOB" + i;
    TaskState jobState = _driver.getWorkflowContext(workflowName).getJobState(jobName);
    Assert.assertEquals(jobState, TaskState.COMPLETED);
  }
}
 

@Test
public void testMultipleTaskAssignment() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  String workflowName = TestHelper.getTestMethodName();
  Workflow.Builder workflowBuilder = new Workflow.Builder(workflowName);

  List<TaskConfig> taskConfigs = Lists.newArrayListWithCapacity(20);
  for (int i = 0; i < 20; i++) {
    Map<String, String> taskConfigMap = Maps.newHashMap();
    taskConfigs.add(new TaskConfig("TaskOne", taskConfigMap));
  }
  JobConfig.Builder jobBuilder = new JobConfig.Builder().setCommand("DummyCommand")
      .setJobCommandConfigMap(_jobCommandMap).addTaskConfigs(taskConfigs);

  workflowBuilder.addJob("JOB", jobBuilder);
  _driver.start(workflowBuilder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  Assert.assertEquals(_runCounts.size(), 5);
}
 

@Test
public void testTaskNumAttemptsWithDelay() throws Exception {
  int maxAttempts = Integer.MAX_VALUE; // Allow the count to increase infinitely
  // Use a delay that's long enough for multiple rounds of pipeline
  long retryDelay = 4000L;

  String workflowName = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(maxAttempts).setCommand(MockTask.TASK_COMMAND)
      .setWorkflow(workflowName).setFailureThreshold(Integer.MAX_VALUE)
      .setTaskRetryDelay(retryDelay).setJobCommandConfigMap(
          ImmutableMap.of(MockTask.FAILURE_COUNT_BEFORE_SUCCESS, String.valueOf(maxAttempts)));
  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(workflowName, jobBuilder).build();
  _driver.start(flow);

  // Wait until the job is running
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, workflowName));
  int expectedAttempts = jobContext.getPartitionNumAttempts(0);

  // Check 3 times to see if maxAttempts match up
  for (int i = 0; i < 3; i++) {
    // Add a small delay to make sure the check falls in the middle of the scheduling timeline
    Thread.sleep(retryDelay + 1000L);
    JobContext ctx =
        _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, workflowName));
    expectedAttempts++;
    Assert.assertEquals(ctx.getPartitionNumAttempts(0), expectedAttempts);
  }
}
 

@Test
public void testStopWorkflowInStoppingState() throws InterruptedException {
  final String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  Workflow.Builder builder = new Workflow.Builder(workflowName);

  // Add 2 jobs
  Map<String, String> jobCommandConfigMap = new HashMap<String, String>();
  jobCommandConfigMap.put(MockTask.JOB_DELAY, "1000000");
  jobCommandConfigMap.put(MockTask.NOT_ALLOW_TO_CANCEL, String.valueOf(true));
  List<TaskConfig> taskConfigs = ImmutableList.of(
      new TaskConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTaskId("testTask").build());
  JobConfig.Builder job1 = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .addTaskConfigs(taskConfigs).setJobCommandConfigMap(jobCommandConfigMap);
  String job1Name = "Job1";

  JobConfig.Builder job2 =
      new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).addTaskConfigs(taskConfigs);
  String job2Name = "Job2";

  builder.addJob(job1Name, job1);
  builder.addJob(job2Name, job2);

  _driver.start(builder.build());
  _driver.pollForWorkflowState(workflowName, TaskState.IN_PROGRESS);
  _driver.pollForJobState(workflowName, TaskUtil.getNamespacedJobName(workflowName, job1Name), TaskState.IN_PROGRESS);
  _driver.stop(workflowName);
  _driver.pollForWorkflowState(workflowName, TaskState.STOPPING);

  // Expect job and workflow stuck in STOPPING state.
  WorkflowContext workflowContext = _driver.getWorkflowContext(workflowName);
  Assert.assertEquals(
      workflowContext.getJobState(TaskUtil.getNamespacedJobName(workflowName, job1Name)),
      TaskState.STOPPING);
}
 

@Test
public void testDeleteWorkflow() throws InterruptedException {
  String jobQueueName = TestHelper.getTestMethodName();
  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG)
      .setMaxAttemptsPerTask(1).setWorkflow(jobQueueName)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "100000"));

  JobQueue.Builder jobQueue = TaskTestUtil.buildJobQueue(jobQueueName);
  jobQueue.enqueueJob("job1", jobBuilder);
  _driver.start(jobQueue.build());
  _driver.pollForJobState(jobQueueName, TaskUtil.getNamespacedJobName(jobQueueName, "job1"),
      TaskState.IN_PROGRESS);

  // Check that WorkflowConfig, WorkflowContext, and IdealState are indeed created for this job
  // queue
  Assert.assertNotNull(_driver.getWorkflowConfig(jobQueueName));
  Assert.assertNotNull(_driver.getWorkflowContext(jobQueueName));
  Assert.assertNotNull(admin.getResourceIdealState(CLUSTER_NAME, jobQueueName));

  // Pause the Controller so that the job queue won't get deleted
  admin.enableCluster(CLUSTER_NAME, false);
  Thread.sleep(1000);
  // Attempt the deletion and time out
  try {
    _driver.deleteAndWaitForCompletion(jobQueueName, DELETE_DELAY);
    Assert.fail(
        "Delete must time out and throw a HelixException with the Controller paused, but did not!");
  } catch (HelixException e) {
    // Pass
  }

  // Resume the Controller and call delete again
  admin.enableCluster(CLUSTER_NAME, true);
  _driver.deleteAndWaitForCompletion(jobQueueName, DELETE_DELAY);

  // Check that the deletion operation completed
  Assert.assertNull(_driver.getWorkflowConfig(jobQueueName));
  Assert.assertNull(_driver.getWorkflowContext(jobQueueName));
  Assert.assertNull(admin.getResourceIdealState(CLUSTER_NAME, jobQueueName));
}
 

@Test
public void testJobFailsWithMissingTarget() throws Exception {
  String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  LOG.info("Starting job-queue: " + workflowName);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  // Create and Enqueue jobs
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < _numDbs; i++) {
    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(
            _testDbs.get(i))
            .setTargetPartitionStates(Sets.newHashSet("SLAVE"));
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
    if (i > 0) {
      builder.addParentChildDependency("job" + _testDbs.get(i - 1), "job" + _testDbs.get(i));
    }
    currentJobNames.add(jobName);
  }

  _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, _testDbs.get(1));
  _driver.start(builder.build());

  String namedSpaceJob = String.format("%s_%s", workflowName, currentJobNames.get(1));
  _driver.pollForJobState(workflowName, namedSpaceJob, TaskState.FAILED);
  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);

  _driver.delete(workflowName);
}
 

@Test(dependsOnMethods = "testJobContinueUponParentJobFailure")
public void testJobFailsWithMissingTargetInRunning() throws Exception {
  String workflowName = TestHelper.getTestMethodName();

  // Create a workflow
  LOG.info("Starting job-queue: " + workflowName);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  // Create and add jobs
  List<String> currentJobNames = new ArrayList<String>();
  for (int i = 0; i < _numDbs; i++) {
    JobConfig.Builder jobConfig =
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(_testDbs.get(i))
            .setTargetPartitionStates(Sets.newHashSet("SLAVE"));
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
    if (i > 0) {
      builder.addParentChildDependency("job" + _testDbs.get(i - 1), "job" + _testDbs.get(i));
    }      currentJobNames.add(jobName);
  }

  _driver.start(builder.build());
  _gSetupTool.dropResourceFromCluster(CLUSTER_NAME, _testDbs.get(0));

  String namedSpaceJob1 = String.format("%s_%s", workflowName, currentJobNames.get(0));
  _driver.pollForJobState(workflowName, namedSpaceJob1, TaskState.FAILED);
  _driver.pollForWorkflowState(workflowName, TaskState.FAILED);

  _driver.delete(workflowName);
}
 

@Test
public void test() throws Exception {
  String jobResource = TestHelper.getTestMethodName();

  JobConfig.Builder jobBuilder = JobConfig.Builder.fromMap(WorkflowGenerator.DEFAULT_JOB_CONFIG);
  jobBuilder.setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG)
      .setMaxAttemptsPerTask(2).setCommand(MockTask.TASK_COMMAND)
      .setFailureThreshold(Integer.MAX_VALUE)
      .setJobCommandConfigMap(ImmutableMap.of(MockTask.THROW_EXCEPTION, "true"));

  Workflow flow =
      WorkflowGenerator.generateSingleJobWorkflowBuilder(jobResource, jobBuilder).build();

  _driver.start(flow);

  // Wait until the job completes.
  _driver.pollForWorkflowState(jobResource, TaskState.COMPLETED);

  JobContext ctx = _driver.getJobContext(TaskUtil.getNamespacedJobName(jobResource));
  for (int i = 0; i < _numPartitions; i++) {
    TaskPartitionState state = ctx.getPartitionState(i);
    if (state != null) {
      Assert.assertEquals(state, TaskPartitionState.TASK_ERROR);
      Assert.assertEquals(ctx.getPartitionNumAttempts(i), 2);
    }
  }
}
 

/**
 * Create a job from a given batch of {@link WorkUnit}s.
 */
JobConfig.Builder createHelixJob(List<WorkUnit> workUnits) throws IOException {
  Map<String, TaskConfig> taskConfigMap = Maps.newHashMap();

  try (ParallelRunner stateSerDeRunner = new ParallelRunner(this.stateSerDeRunnerThreads, this.fs)) {
    int multiTaskIdSequence = 0;
    for (WorkUnit workUnit : workUnits) {
      if (workUnit instanceof MultiWorkUnit) {
        workUnit.setId(JobLauncherUtils.newMultiTaskId(this.jobContext.getJobId(), multiTaskIdSequence++));
      }
      addWorkUnit(workUnit, stateSerDeRunner, taskConfigMap);
    }

    Path jobStateFilePath;

    // write the job.state using the state store if present, otherwise serialize directly to the file
    if (this.stateStores.haveJobStateStore()) {
      jobStateFilePath = GobblinClusterUtils.getJobStateFilePath(true, this.appWorkDir, this.jobContext.getJobId());
      this.stateStores.getJobStateStore().put(jobStateFilePath.getParent().getName(), jobStateFilePath.getName(),
          this.jobContext.getJobState());
    } else {
      jobStateFilePath = GobblinClusterUtils.getJobStateFilePath(false, this.appWorkDir, this.jobContext.getJobId());
      SerializationUtils.serializeState(this.fs, jobStateFilePath, this.jobContext.getJobState());
    }

    // Block on persistence of all workunits to be finished.
    // It is necessary when underlying storage being slow and Helix activate task-execution before the workunit being persisted.
    stateSerDeRunner.waitForTasks(Long.MAX_VALUE);

    LOGGER.debug("GobblinHelixJobLauncher.createHelixJob: jobStateFilePath {}, jobState {} jobProperties {}",
        jobStateFilePath, this.jobContext.getJobState().toString(), this.jobContext.getJobState().getProperties());

    return translateGobblinJobConfigToHelixJobConfig(this.jobContext.getJobState(), workUnits, taskConfigMap);
  }
}
 

/**
 * Number of instance is equal to number of failure threshold, thus job failure mechanism needs to
 * consider given up
 * tasks that no longer exist on the instance, not only given up tasks currently reported on
 * CurrentState.
 */
@Test
public void testHighThreshold() throws InterruptedException {
  final String WORKFLOW_NAME = "testWorkflow";
  final String JOB_NAME = "testJob";

  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(
              ImmutableMap.of(MockTask.TASK_RESULT_STATUS, TaskResult.Status.FATAL_FAILED.name()))
          .setFailureThreshold(1);
  Workflow.Builder workflowBuilder =
      new Workflow.Builder(WORKFLOW_NAME).addJob(JOB_NAME, jobBuilder);
  _driver.start(workflowBuilder.build());

  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, JOB_NAME),
      TaskState.FAILED);
  _driver.pollForWorkflowState(WORKFLOW_NAME, TaskState.FAILED);

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(WORKFLOW_NAME, JOB_NAME));
  int countAborted = 0;
  int countNoState = 0;
  for (int pId : jobContext.getPartitionSet()) {
    TaskPartitionState state = jobContext.getPartitionState(pId);
    if (state == TaskPartitionState.TASK_ABORTED) {
      countAborted++;
    } else if (state == null) {
      countNoState++;
    } else {
      throw new HelixException(String.format("State %s is not expected.", state));
    }
  }
  Assert.assertEquals(countAborted, 2); // Failure threshold is 1, so 2 tasks aborted.
  Assert.assertEquals(countNoState, 3); // Other 3 tasks are not scheduled at all.
}