Java Code Examples for org.apache.helix.task.WorkflowConfig#Builder

public static JobQueue.Builder buildRecurrentJobQueue(String jobQueueName, int delayStart,
    int recurrenceInSeconds, TargetState targetState) {
  WorkflowConfig.Builder workflowCfgBuilder = new WorkflowConfig.Builder(jobQueueName);
  workflowCfgBuilder.setExpiry(120000);
  if (targetState != null) {
    workflowCfgBuilder.setTargetState(TargetState.STOP);
  }

  Calendar cal = Calendar.getInstance();
  cal.set(Calendar.MINUTE, cal.get(Calendar.MINUTE) + delayStart / 60);
  cal.set(Calendar.SECOND, cal.get(Calendar.SECOND) + delayStart % 60);
  cal.set(Calendar.MILLISECOND, 0);
  ScheduleConfig scheduleConfig =
      ScheduleConfig.recurringFromDate(cal.getTime(), TimeUnit.SECONDS, recurrenceInSeconds);
  workflowCfgBuilder.setScheduleConfig(scheduleConfig);
  return new JobQueue.Builder(jobQueueName).setWorkflowConfig(workflowCfgBuilder.build());
}
 

public static JobQueue.Builder buildJobQueue(String jobQueueName, int delayStart,
    int failureThreshold, int capacity) {
  WorkflowConfig.Builder workflowCfgBuilder = new WorkflowConfig.Builder(jobQueueName);
  workflowCfgBuilder.setExpiry(120000);
  workflowCfgBuilder.setCapacity(capacity);

  Calendar cal = Calendar.getInstance();
  cal.set(Calendar.MINUTE, cal.get(Calendar.MINUTE) + delayStart / 60);
  cal.set(Calendar.SECOND, cal.get(Calendar.SECOND) + delayStart % 60);
  cal.set(Calendar.MILLISECOND, 0);
  workflowCfgBuilder.setScheduleConfig(ScheduleConfig.oneTimeDelayedStart(cal.getTime()));

  if (failureThreshold > 0) {
    workflowCfgBuilder.setFailureThreshold(failureThreshold);
  }
  return new JobQueue.Builder(jobQueueName).setWorkflowConfig(workflowCfgBuilder.build());
}
 

@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);
}
 

@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);
}
 

/**
 * 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 explicitly allows overlap job assignment.
 * @throws InterruptedException
 */
@Test
// This test does NOT allow multiple jobs being assigned to an instance.
public void testMultipleJobAssignmentOverlapEnabled() throws InterruptedException {
  _runCounts.clear();
  failTask = false;
  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 < 40; 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);
}
 

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

/**
 * This test starts 4 jobs in job queue, the job all stuck, and verify that
 * (1) the number of running job does not exceed configured max allowed parallel jobs
 * (2) one instance can only be assigned to one job in the workflow
 */
@Test
public void testWhenDisallowOverlapJobAssignment() throws Exception {
  String queueName = TestHelper.getTestMethodName();

  WorkflowConfig.Builder cfgBuilder = new WorkflowConfig.Builder(queueName);
  cfgBuilder.setParallelJobs(PARALLEL_COUNT);
  cfgBuilder.setAllowOverlapJobAssignment(false);

  JobQueue.Builder queueBuild =
      new JobQueue.Builder(queueName).setWorkflowConfig(cfgBuilder.build());
  JobQueue queue = queueBuild.build();
  _driver.createQueue(queue);

  List<JobConfig.Builder> jobConfigBuilders = new ArrayList<JobConfig.Builder>();
  for (String testDbName : _testDbs) {
    jobConfigBuilders.add(
        new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND).setTargetResource(testDbName)
            .setTargetPartitionStates(Collections.singleton("SLAVE"))
            .setJobCommandConfigMap(Collections.singletonMap(MockTask.JOB_DELAY, "1000")));
  }

  _driver.stop(queueName);
  for (int i = 0; i < jobConfigBuilders.size(); ++i) {
    _driver.enqueueJob(queueName, "job_" + (i + 1), jobConfigBuilders.get(i));
  }
  _driver.resume(queueName);
  Thread.sleep(1000L);
  Assert.assertTrue(TaskTestUtil.pollForWorkflowParallelState(_driver, queueName));
}
 

/**
 * Tests whether jobs with undefined types (not found in ClusterConfig) run as DEFAULT.
 * @throws InterruptedException
 */
@Test(dependsOnMethods = "testSchedulingWithoutQuota")
public void testSchedulingUndefinedTypes() throws InterruptedException {
  ClusterConfig clusterConfig = _manager.getConfigAccessor().getClusterConfig(CLUSTER_NAME);
  clusterConfig.resetTaskQuotaRatioMap();
  clusterConfig.setTaskQuotaRatio(DEFAULT_QUOTA_TYPE, 1);
  clusterConfig.setTaskQuotaRatio("A", 1);
  clusterConfig.setTaskQuotaRatio("B", 1);
  _manager.getConfigAccessor().setClusterConfig(CLUSTER_NAME, clusterConfig);
  _availableQuotaTypes = clusterConfig.getTaskQuotaRatioMap().keySet();

  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).setJobType("UNDEFINED");
    workflowBuilder.addJob("JOB" + i, jobConfigBulider);
  }

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

  // Check run counts for each quota type
  Assert.assertEquals((int) _quotaTypeExecutionCount.get("DEFAULT"), 10);
  Assert.assertFalse(_quotaTypeExecutionCount.containsKey("A"));
  Assert.assertFalse(_quotaTypeExecutionCount.containsKey("B"));
}
 

/**
 * Tests that stopping a workflow does result in its task ending up in STOPPED state.
 * @throws InterruptedException
 */
@Test(dependsOnMethods = "testStopWorkflow")
public void testStopTask() throws Exception {
  stopTestSetup(1);

  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 < 1; i++) {
    List<TaskConfig> taskConfigs = new ArrayList<>();
    taskConfigs.add(new TaskConfig("StopTask", new HashMap<>()));
    JobConfig.Builder jobConfigBulider = new JobConfig.Builder().setCommand("Dummy")
        .addTaskConfigs(taskConfigs).setJobCommandConfigMap(new HashMap<>());
    workflowBuilder.addJob("JOB" + i, jobConfigBulider);
  }

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

  // Stop the workflow
  _driver.stop(workflowName);
  _driver.pollForWorkflowState(workflowName, TaskState.STOPPED);

  Assert.assertEquals(TaskDriver.getWorkflowContext(_manager, workflowName).getWorkflowState(),
      TaskState.STOPPED);

  cleanupParticipants(1);
}
 

@Test
public void testWorkflowWithDependencies() throws InterruptedException {
  String workflowName = TestHelper.getTestMethodName();
  final int PARALLEL_NUM = 2;
  // Workflow setup
  WorkflowConfig.Builder workflowcfgBuilder =
      new WorkflowConfig.Builder().setWorkflowId(workflowName).setParallelJobs(PARALLEL_NUM);
  Workflow.Builder builder = new Workflow.Builder(workflowName);
  builder.setWorkflowConfig(workflowcfgBuilder.build());

  builder.addParentChildDependency("job" + _testDbs.get(0), "job" + _testDbs.get(1));
  for (int i = 0; i < 2; i++) {
    JobConfig.Builder jobConfig = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
        .setTargetResource(_testDbs.get(i))
        .setTargetPartitionStates(Sets.newHashSet("SLAVE", "MASTER"))
        .setJobCommandConfigMap(WorkflowGenerator.DEFAULT_COMMAND_CONFIG);
    String jobName = "job" + _testDbs.get(i);
    builder.addJob(jobName, jobConfig);
  }

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

  // Wait until the workflow completes
  _driver.pollForWorkflowState(workflowName, TaskState.COMPLETED);

  JobContext context1 =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, "job" + _testDbs.get(0)));
  JobContext context2 =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(workflowName, "job" + _testDbs.get(1)));
  Assert.assertTrue(context2.getStartTime() - context1.getFinishTime() >= 0L);
}
 

/**
 * This test starts 4 jobs in job queue, the job all stuck, and verify that
 * (1) the number of running job does not exceed configured max allowed parallel jobs
 * (2) one instance can be assigned to multiple jobs in the workflow when allow overlap assignment
 */
@Test (dependsOnMethods = {"testWhenDisallowOverlapJobAssignment"})
public void testWhenAllowOverlapJobAssignment() throws Exception {
  // Disable all participants except one to enforce all assignment to be on one host
  for (int i = 1; i < _numNodes; i++) {
    _participants[i].syncStop();
  }
  ClusterLiveNodesVerifier verifier = new ClusterLiveNodesVerifier(_gZkClient, CLUSTER_NAME,
      Collections.singletonList(_participants[0].getInstanceName()));
  Assert.assertTrue(verifier.verifyByPolling());

  String queueName = TestHelper.getTestMethodName();

  WorkflowConfig.Builder cfgBuilder = new WorkflowConfig.Builder(queueName);
  cfgBuilder.setParallelJobs(PARALLEL_COUNT);
  cfgBuilder.setAllowOverlapJobAssignment(true);

  JobQueue.Builder queueBuild = new JobQueue.Builder(queueName).setWorkflowConfig(cfgBuilder.build());
  JobQueue queue = queueBuild.build();
  _driver.createQueue(queue);

  // Create jobs that can be assigned to any instances
  List<JobConfig.Builder> jobConfigBuilders = new ArrayList<JobConfig.Builder>();
  for (int i = 0; i < PARALLEL_COUNT; i++) {
    List<TaskConfig> taskConfigs = new ArrayList<TaskConfig>();
    for (int j = 0; j < TASK_COUNT; j++) {
      taskConfigs.add(
          new TaskConfig.Builder().setTaskId("task_" + j).setCommand(MockTask.TASK_COMMAND)
              .build());
    }
    jobConfigBuilders.add(new JobConfig.Builder().addTaskConfigs(taskConfigs));
  }

  _driver.stop(queueName);
  for (int i = 0; i < jobConfigBuilders.size(); ++i) {
    _driver.enqueueJob(queueName, "job_" + (i + 1), jobConfigBuilders.get(i));
  }
  _driver.resume(queueName);
  Thread.sleep(2000);
  Assert.assertTrue(TaskTestUtil.pollForWorkflowParallelState(_driver, queueName));
}
 

@Test
public void testJobQueueAutoCleanUp() throws InterruptedException {
  int capacity = 10;
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName, capacity);
  WorkflowConfig.Builder cfgBuilder = new WorkflowConfig.Builder(builder.getWorkflowConfig());
  cfgBuilder.setJobPurgeInterval(1000);
  builder.setWorkflowConfig(cfgBuilder.build());

  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, String.valueOf(capacity / 2)))
          .setExpiry(200L);
  Set<String> deletedJobs = new HashSet<String>();
  Set<String> remainJobs = new HashSet<String>();
  for (int i = 0; i < capacity; i++) {
    builder.enqueueJob("JOB" + i, jobBuilder);
    if (i < capacity/2) {
      deletedJobs.add("JOB" + i);
    } else {
      remainJobs.add(TaskUtil.getNamespacedJobName(queueName, "JOB" + i));
    }
  }
  _driver.start(builder.build());
  _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + (capacity - 1)), TaskState.FAILED);
  Thread.sleep(2000);

  WorkflowConfig config = _driver.getWorkflowConfig(queueName);
  Assert.assertEquals(config.getJobDag().getAllNodes(), remainJobs);

  WorkflowContext context = _driver.getWorkflowContext(queueName);
  Assert.assertEquals(context.getJobStates().keySet(), remainJobs);
  Assert.assertTrue(remainJobs.containsAll(context.getJobStartTimes().keySet()));

  for (String job : deletedJobs) {
    JobConfig cfg = _driver.getJobConfig(job);
    JobContext ctx = _driver.getJobContext(job);
    Assert.assertNull(cfg);
    Assert.assertNull(ctx);
  }

}
 

@Test
public void testNoSlaveToRunTask() throws InterruptedException {
  // first job can't be assigned to any instance and stuck, timeout, second job runs and workflow
  // succeed.
  final String FIRST_JOB = "first_job";
  final String SECOND_JOB = "second_job";
  final String WORKFLOW_NAME = TestHelper.getTestMethodName();
  final String DB_NAME = WorkflowGenerator.DEFAULT_TGT_DB;

  JobConfig.Builder firstJobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          // since replica=1, there's no Slave, but target only Slave, the task can't be assigned,
          // job stuck and timeout
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.SLAVE.name()))
          .setCommand(MockTask.TASK_COMMAND).setTimeout(1000);

  JobConfig.Builder secondJobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND).setIgnoreDependentJobFailure(true); // ignore first
                                                                                 // job's timeout

  WorkflowConfig.Builder workflowConfigBuilder =
      new WorkflowConfig.Builder(WORKFLOW_NAME).setFailureThreshold(1); // workflow ignores first
                                                                        // job's timeout and
                                                                        // schedule second job and
                                                                        // succeed.

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW_NAME)
      .setWorkflowConfig(workflowConfigBuilder.build()).addJob(FIRST_JOB, firstJobBuilder)
      .addJob(SECOND_JOB, secondJobBuilder).addParentChildDependency(FIRST_JOB, SECOND_JOB);

  // Check that there are no SLAVE replicas
  BestPossibleExternalViewVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkClient(_gZkClient)
          .setZkAddr(ZK_ADDR).build();
  Assert.assertTrue(verifier.verifyByPolling());
  ExternalView view =
      _gSetupTool.getClusterManagementTool().getResourceExternalView(CLUSTER_NAME, DB_NAME);
  view.getPartitionSet().forEach(partition -> {
    for (Map.Entry<String, String> entry : view.getStateMap(partition).entrySet()) {
      if (!entry.getValue().equals("MASTER")) {
        try {
          Thread.sleep(2000L);
        } catch (InterruptedException e) {
          // OK
        }
      }
    }
  });

  _driver.start(workflowBuilder.build());

  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, FIRST_JOB),
      TaskState.TIMED_OUT);
  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, SECOND_JOB),
      TaskState.COMPLETED);
  _driver.pollForWorkflowState(WORKFLOW_NAME, TaskState.COMPLETED);

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(WORKFLOW_NAME, FIRST_JOB));
  for (int pId : jobContext.getPartitionSet()) {
    // No task assigned for first job
    Assert.assertNull(jobContext.getPartitionState(pId));
  }
}
 

@Test
public void testTaskRunningIndefinitely() throws InterruptedException {
  // first job runs indefinitely and timeout, the second job runs successfully, the workflow
  // succeed.
  final String FIRST_JOB = "first_job";
  final String SECOND_JOB = "second_job";
  final String WORKFLOW_NAME = TestHelper.getTestMethodName();
  final String DB_NAME = WorkflowGenerator.DEFAULT_TGT_DB;

  JobConfig.Builder firstJobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND)
          .setJobCommandConfigMap(ImmutableMap.of(MockTask.JOB_DELAY, "99999999")) // task stuck
          .setTimeout(10);

  JobConfig.Builder secondJobBuilder =
      new JobConfig.Builder().setWorkflow(WORKFLOW_NAME).setTargetResource(DB_NAME)
          .setTargetPartitionStates(Sets.newHashSet(MasterSlaveSMD.States.MASTER.name()))
          .setCommand(MockTask.TASK_COMMAND).setIgnoreDependentJobFailure(true); // ignore first
                                                                                 // job's timeout

  WorkflowConfig.Builder workflowConfigBuilder =
      new WorkflowConfig.Builder(WORKFLOW_NAME).setFailureThreshold(1); // workflow ignores first
                                                                        // job's timeout and
                                                                        // schedule second job and
                                                                        // succeed.

  Workflow.Builder workflowBuilder = new Workflow.Builder(WORKFLOW_NAME)
      .setWorkflowConfig(workflowConfigBuilder.build()).addJob(FIRST_JOB, firstJobBuilder)
      .addJob(SECOND_JOB, secondJobBuilder).addParentChildDependency(FIRST_JOB, SECOND_JOB);

  _driver.start(workflowBuilder.build());

  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, FIRST_JOB),
      TaskState.TIMED_OUT);
  _driver.pollForJobState(WORKFLOW_NAME, TaskUtil.getNamespacedJobName(WORKFLOW_NAME, SECOND_JOB),
      TaskState.COMPLETED);
  _driver.pollForWorkflowState(WORKFLOW_NAME, TaskState.COMPLETED);

  JobContext jobContext =
      _driver.getJobContext(TaskUtil.getNamespacedJobName(WORKFLOW_NAME, FIRST_JOB));
  for (int pId : jobContext.getPartitionSet()) {
    // All tasks aborted because of job timeout
    Assert.assertEquals(jobContext.getPartitionState(pId), TaskPartitionState.TASK_ABORTED);
  }
}
 

@Test
public void testUpdateStoppedQueue() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();

  // Create a queue
  LOG.info("Starting job-queue: " + queueName);
  JobQueue queue = createDefaultRecurrentJobQueue(queueName, 2);
  _driver.start(queue);

  WorkflowContext wCtx = TaskTestUtil.pollForWorkflowContext(_driver, queueName);

  // ensure current schedule is started
  String scheduledQueue = wCtx.getLastScheduledSingleWorkflow();
  _driver.pollForWorkflowState(scheduledQueue, TaskState.IN_PROGRESS);

  _driver.stop(queueName);

  WorkflowConfig workflowConfig = _driver.getWorkflowConfig(queueName);
  Assert.assertEquals(workflowConfig.getTargetState(), TargetState.STOP);

  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowConfig);
  Calendar startTime = Calendar.getInstance();
  startTime.set(Calendar.SECOND, startTime.get(Calendar.SECOND) + 1);

  ScheduleConfig scheduleConfig =
      ScheduleConfig.recurringFromDate(startTime.getTime(), TimeUnit.MINUTES, 2);

  configBuilder.setScheduleConfig(scheduleConfig);

  _driver.updateWorkflow(queueName, configBuilder.build());

  workflowConfig = _driver.getWorkflowConfig(queueName);
  Assert.assertEquals(workflowConfig.getTargetState(), TargetState.STOP);

  _driver.resume(queueName);

  // ensure current schedule is completed
  _driver.pollForWorkflowState(scheduledQueue, TaskState.COMPLETED);

  Thread.sleep(1000);

  wCtx = TaskTestUtil.pollForWorkflowContext(_driver, queueName);
  scheduledQueue = wCtx.getLastScheduledSingleWorkflow();
  WorkflowConfig wCfg = _driver.getWorkflowConfig(scheduledQueue);

  Calendar configStartTime = Calendar.getInstance();
  configStartTime.setTime(wCfg.getStartTime());

  Assert.assertTrue(
      (startTime.get(Calendar.HOUR_OF_DAY) == configStartTime.get(Calendar.HOUR_OF_DAY) &&
          startTime.get(Calendar.MINUTE) == configStartTime.get(Calendar.MINUTE) &&
          startTime.get(Calendar.SECOND) == configStartTime.get(Calendar.SECOND)));
}
 

@Test
public void testUpdateRunningQueue() throws InterruptedException {
  String queueName = TestHelper.getTestMethodName();

  // Create a queue
  LOG.info("Starting job-queue: " + queueName);
  JobQueue queue = createDefaultRecurrentJobQueue(queueName, 2);
  _driver.start(queue);

  WorkflowContext wCtx = TaskTestUtil.pollForWorkflowContext(_driver, queueName);

  WorkflowConfig workflowConfig = _driver.getWorkflowConfig(queueName);
  WorkflowConfig.Builder configBuilder = new WorkflowConfig.Builder(workflowConfig);

  Calendar startTime = Calendar.getInstance();
  startTime.set(Calendar.SECOND, startTime.get(Calendar.SECOND) + 1);

  ScheduleConfig scheduleConfig =
      ScheduleConfig.recurringFromDate(startTime.getTime(), TimeUnit.MINUTES, 2);

  configBuilder.setScheduleConfig(scheduleConfig);

  // ensure current schedule is started
  String scheduledQueue = wCtx.getLastScheduledSingleWorkflow();
  _driver.pollForWorkflowState(scheduledQueue, TaskState.IN_PROGRESS);

  _driver.updateWorkflow(queueName, configBuilder.build());

  // ensure current schedule is completed
  _driver.pollForWorkflowState(scheduledQueue, TaskState.COMPLETED);

  Thread.sleep(1000);

  wCtx = TaskTestUtil.pollForWorkflowContext(_driver, queueName);
  scheduledQueue = wCtx.getLastScheduledSingleWorkflow();
  WorkflowConfig wCfg = _driver.getWorkflowConfig(scheduledQueue);

  Calendar configStartTime = Calendar.getInstance();
  configStartTime.setTime(wCfg.getStartTime());

  Assert.assertTrue(
      (startTime.get(Calendar.HOUR_OF_DAY) == configStartTime.get(Calendar.HOUR_OF_DAY) &&
          startTime.get(Calendar.MINUTE) == configStartTime.get(Calendar.MINUTE) &&
          startTime.get(Calendar.SECOND) == configStartTime.get(Calendar.SECOND)));
}
 

@Test
public void testQueueParallelJobs() throws InterruptedException {
  final int parallelJobs = 3;
  final int numberOfJobsAddedBeforeControllerSwitch = 4;
  final int totalNumberOfJobs = 7;
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  WorkflowConfig.Builder workflowCfgBuilder = new WorkflowConfig.Builder()
      .setWorkflowId(queueName).setParallelJobs(parallelJobs).setAllowOverlapJobAssignment(true);
  _driver.start(builder.setWorkflowConfig(workflowCfgBuilder.build()).build());
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(Collections.singletonMap(MockTask.JOB_DELAY, "10000"));

  // Add 4 jobs to the queue
  for (int i = 0; i < numberOfJobsAddedBeforeControllerSwitch; i++) {
    _driver.enqueueJob(queueName, "JOB" + i, jobBuilder);
  }

  // Wait until all of the enqueued jobs (Job0 to Job3) are finished
  for (int i = 0; i < numberOfJobsAddedBeforeControllerSwitch; i++) {
    _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + i),
        TaskState.COMPLETED);
  }

  // Stop the Controller
  _controller.syncStop();

  // Add 3 more jobs to the queue which should run in parallel after the Controller is started
  for (int i = numberOfJobsAddedBeforeControllerSwitch; i < totalNumberOfJobs; i++) {
    _driver.enqueueJob(queueName, "JOB" + i, jobBuilder);
  }

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

  // Wait until all of the newly added jobs (Job4 to Job6) are finished
  for (int i = numberOfJobsAddedBeforeControllerSwitch; i < totalNumberOfJobs; i++) {
    _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + i),
        TaskState.COMPLETED);
  }

  // Make sure the jobs have been running in parallel by checking the jobs start time and finish
  // time
  long maxStartTime = Long.MIN_VALUE;
  long minFinishTime = Long.MAX_VALUE;

  for (int i = numberOfJobsAddedBeforeControllerSwitch; i < totalNumberOfJobs; i++) {
    JobContext jobContext =
        _driver.getJobContext(TaskUtil.getNamespacedJobName(queueName, "JOB" + i));
    maxStartTime = Long.max(maxStartTime, jobContext.getStartTime());
    minFinishTime = Long.min(minFinishTime, jobContext.getFinishTime());
  }
  Assert.assertTrue(minFinishTime > maxStartTime);
}
 

@Test
public void testQueueJobsMaxCapacity() throws InterruptedException {
  final int numberOfJobsAddedInitially = 4;
  final int queueCapacity = 5;
  final String newJobName = "NewJob";
  String queueName = TestHelper.getTestMethodName();
  JobQueue.Builder builder = TaskTestUtil.buildJobQueue(queueName);
  WorkflowConfig.Builder workflowCfgBuilder =
      new WorkflowConfig.Builder().setWorkflowId(queueName).setParallelJobs(1)
          .setAllowOverlapJobAssignment(true).setCapacity(queueCapacity);
  _driver.start(builder.setWorkflowConfig(workflowCfgBuilder.build()).build());
  JobConfig.Builder jobBuilder =
      new JobConfig.Builder().setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB)
          .setCommand(MockTask.TASK_COMMAND).setMaxAttemptsPerTask(2)
          .setJobCommandConfigMap(Collections.singletonMap(MockTask.JOB_DELAY, "1000"));

  // Add 4 jobs to the queue
  for (int i = 0; i < numberOfJobsAddedInitially; i++) {
    _driver.enqueueJob(queueName, "JOB" + i, jobBuilder);
  }

  // Wait until all of the enqueued jobs (Job0 to Job3) are finished
  for (int i = 0; i < numberOfJobsAddedInitially; i++) {
    _driver.pollForJobState(queueName, TaskUtil.getNamespacedJobName(queueName, "JOB" + i),
        TaskState.COMPLETED);
  }

  boolean exceptionHappenedWhileAddingNewJob = false;
  try {
    // This call will produce the exception because 4 jobs have been already added
    // By adding the new job the queue will hit its capacity limit
    _driver.enqueueJob(queueName, newJobName, jobBuilder);
  } catch (Exception e) {
    exceptionHappenedWhileAddingNewJob = true;
  }
  Assert.assertTrue(exceptionHappenedWhileAddingNewJob);

  // Make sure that jobConfig has not been created
  JobConfig jobConfig =
      _driver.getJobConfig(TaskUtil.getNamespacedJobName(queueName, newJobName));
  Assert.assertNull(jobConfig);
}