/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.job.master;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import alluxio.Constants;
import alluxio.conf.PropertyKey;
import alluxio.exception.status.ResourceExhaustedException;
import alluxio.job.SleepJobConfig;
import alluxio.job.plan.PlanDefinitionRegistryRule;
import alluxio.job.plan.SleepPlanDefinition;
import alluxio.job.util.JobTestUtils;
import alluxio.job.wire.JobInfo;
import alluxio.job.wire.JobWorkerHealth;
import alluxio.job.wire.Status;
import alluxio.job.workflow.composite.CompositeConfig;
import alluxio.master.LocalAlluxioJobCluster;
import alluxio.master.job.JobMaster;
import alluxio.testutils.BaseIntegrationTest;
import alluxio.testutils.LocalAlluxioClusterResource;
import alluxio.util.CommonUtils;
import alluxio.util.WaitForOptions;
import alluxio.wire.WorkerInfo;
import alluxio.worker.JobWorkerProcess;

import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Integration tests for the job master.
 */
public final class JobMasterIntegrationTest extends BaseIntegrationTest {
  private static final long WORKER_TIMEOUT_MS = 2000;
  private static final long LOST_WORKER_INTERVAL_MS = 2000;
  private JobMaster mJobMaster;
  private JobWorkerProcess mJobWorker;
  private LocalAlluxioJobCluster mLocalAlluxioJobCluster;

  @Rule
  public LocalAlluxioClusterResource mLocalAlluxioClusterResource =
      new LocalAlluxioClusterResource.Builder()
          .setProperty(PropertyKey.JOB_MASTER_WORKER_HEARTBEAT_INTERVAL, 20)
          .setProperty(PropertyKey.JOB_MASTER_WORKER_TIMEOUT, WORKER_TIMEOUT_MS)
          .setProperty(PropertyKey.JOB_MASTER_LOST_WORKER_INTERVAL, LOST_WORKER_INTERVAL_MS)
          .build();

  @Rule
  public PlanDefinitionRegistryRule mJobRule =
      new PlanDefinitionRegistryRule(SleepJobConfig.class, new SleepPlanDefinition());

  @Before
  public void before() throws Exception {
    mLocalAlluxioJobCluster = new LocalAlluxioJobCluster();
    mLocalAlluxioJobCluster.start();
    mJobMaster = mLocalAlluxioJobCluster.getMaster().getJobMaster();
    mJobWorker = mLocalAlluxioJobCluster.getWorker();
  }

  @After
  public void after() throws Exception {
    mLocalAlluxioJobCluster.stop();
  }

  @Test
  public void multipleTasksPerWorker() throws Exception {
    long jobId = mJobMaster.run(new SleepJobConfig(1, 2));

    JobInfo jobStatus = mJobMaster.getStatus(jobId);
    assertEquals(2, jobStatus.getChildren().size());

    JobTestUtils.waitForJobStatus(mJobMaster, jobId, Status.COMPLETED);

    jobStatus = mJobMaster.getStatus(jobId);
    assertEquals(2, jobStatus.getChildren().size());
  }

  @Test
  @LocalAlluxioClusterResource.Config(confParams = {PropertyKey.Name.JOB_MASTER_JOB_CAPACITY, "1",
      PropertyKey.Name.JOB_MASTER_FINISHED_JOB_RETENTION_TIME, "0"})
  public void flowControl() throws Exception {
    for (int i = 0; i < 10; i++) {
      while (true) {
        try {
          mJobMaster.run(new SleepJobConfig(100));
          break;
        } catch (ResourceExhaustedException e) {
          // sleep for a little before retrying the job
          CommonUtils.sleepMs(100);
        }
      }
    }
  }

  @Test
  public void restartMasterAndLoseWorker() throws Exception {
    long jobId = mJobMaster.run(new SleepJobConfig(1));
    JobTestUtils.waitForJobStatus(mJobMaster, jobId, Status.COMPLETED);
    mJobMaster.stop();
    mJobMaster.start(true);
    CommonUtils.waitFor("Worker to register with restarted job master",
        () -> !mJobMaster.getWorkerInfoList().isEmpty(),
        WaitForOptions.defaults().setTimeoutMs(10 * Constants.SECOND_MS));
    mJobWorker.stop();
    CommonUtils.sleepMs(WORKER_TIMEOUT_MS + LOST_WORKER_INTERVAL_MS);
    assertTrue(mJobMaster.getWorkerInfoList().isEmpty());
  }

  @Test
  @LocalAlluxioClusterResource.Config(
      confParams = {PropertyKey.Name.JOB_MASTER_LOST_WORKER_INTERVAL, "10000000"})
  public void restartMasterAndReregisterWorker() throws Exception {
    long jobId = mJobMaster.run(new SleepJobConfig(1));
    JobTestUtils.waitForJobStatus(mJobMaster, jobId, Status.COMPLETED);
    mJobMaster.stop();
    mJobMaster.start(true);
    CommonUtils.waitFor("Worker to register with restarted job master",
        () -> !mJobMaster.getWorkerInfoList().isEmpty(),
        WaitForOptions.defaults().setTimeoutMs(10 * Constants.SECOND_MS));
    final long firstWorkerId = mJobMaster.getWorkerInfoList().get(0).getId();
    mLocalAlluxioJobCluster.restartWorker();
    CommonUtils.waitFor("Restarted worker to register with job master", () -> {
      List<WorkerInfo> workerInfo = mJobMaster.getWorkerInfoList();
      return !workerInfo.isEmpty() && workerInfo.get(0).getId() != firstWorkerId;
    }, WaitForOptions.defaults().setTimeoutMs(10 * Constants.SECOND_MS));
    // The restarted worker should replace the original worker since they have the same address.
    assertEquals(1, mJobMaster.getWorkerInfoList().size());
  }

  @Test
  public void getAllWorkerHealth() throws Exception {
    final AtomicReference<List<JobWorkerHealth>> singleton = new AtomicReference<>();
    CommonUtils.waitFor("allWorkerHealth", () -> {
      List<JobWorkerHealth> allWorkerHealth = mJobMaster.getAllWorkerHealth();
      singleton.set(allWorkerHealth);
      return allWorkerHealth.size() == 1;
    });
    List<JobWorkerHealth> allWorkerHealth = singleton.get();

    JobWorkerHealth workerHealth = allWorkerHealth.get(0);
    assertNotNull(workerHealth.getHostname());
    assertEquals(3, workerHealth.getLoadAverage().size());
  }

  @Test
  @LocalAlluxioClusterResource.Config(confParams = {PropertyKey.Name.JOB_MASTER_JOB_CAPACITY, "20"})
  public void stopJobWorkerTasks() throws Exception {
    long jobId0 = mJobMaster.run(new SleepJobConfig(5000));
    long jobId1 = mJobMaster.run(new SleepJobConfig(5000));
    long jobId2 = mJobMaster.run(new SleepJobConfig(1));
    long jobId3 = mJobMaster.run(new SleepJobConfig(1));

    JobTestUtils.waitForJobStatus(mJobMaster, jobId2, Status.COMPLETED);
    JobTestUtils.waitForJobStatus(mJobMaster, jobId3, Status.COMPLETED);

    assertFalse(mJobMaster.getStatus(jobId1).getStatus().isFinished());
    assertFalse(mJobMaster.getStatus(jobId0).getStatus().isFinished());
    assertEquals(2, mJobMaster.getAllWorkerHealth().get(0).getNumActiveTasks());

    mJobMaster.setTaskPoolSize(0);
    long jobId5 = mJobMaster.run(new SleepJobConfig(1));

    // wait to make sure that this job is not completing any time soon
    CommonUtils.sleepMs(300);
    assertFalse(mJobMaster.getStatus(jobId5).getStatus().isFinished());
    assertEquals(0, mJobMaster.getAllWorkerHealth().get(0).getTaskPoolSize());

    // existing running tasks will continue to run
    assertEquals(2, mJobMaster.getAllWorkerHealth().get(0).getNumActiveTasks());
  }

  @Test
  @LocalAlluxioClusterResource.Config(confParams = {PropertyKey.Name.JOB_MASTER_JOB_CAPACITY, "20"})
  public void throttleJobWorkerTasks() throws Exception {
    mJobMaster.setTaskPoolSize(1);

    long jobId0 = mJobMaster.run(new SleepJobConfig(1));
    JobTestUtils.waitForJobStatus(mJobMaster, jobId0,
        Sets.newHashSet(Status.RUNNING, Status.COMPLETED));

    long jobId1 = mJobMaster.run(new SleepJobConfig(50000));
    JobTestUtils.waitForJobStatus(mJobMaster, jobId1, Status.RUNNING);
    JobTestUtils.waitForJobStatus(mJobMaster, jobId0, Status.COMPLETED);

    long jobId2 = mJobMaster.run(new SleepJobConfig(1));
    long jobId3 = mJobMaster.run(new SleepJobConfig(1));

    // wait a bit more to make sure other jobs aren't completing
    CommonUtils.sleepMs(300);

    assertEquals(Status.RUNNING, mJobMaster.getStatus(jobId1).getStatus());
    assertEquals(Status.CREATED, mJobMaster.getStatus(jobId2).getStatus());
    assertEquals(Status.CREATED, mJobMaster.getStatus(jobId3).getStatus());

    assertEquals(1, mJobMaster.getAllWorkerHealth().get(0).getTaskPoolSize());
    assertEquals(1, mJobMaster.getAllWorkerHealth().get(0).getNumActiveTasks());

    mJobMaster.cancel(jobId1);

    JobTestUtils.waitForJobStatus(mJobMaster, jobId2, Status.COMPLETED);
    JobTestUtils.waitForJobStatus(mJobMaster, jobId3, Status.COMPLETED);
  }

  @Test
  public void cancel() throws Exception {
    SleepJobConfig childJob1 = new SleepJobConfig(50000);
    SleepJobConfig childJob2 = new SleepJobConfig(45000);
    SleepJobConfig childJob3 = new SleepJobConfig(40000);

    CompositeConfig jobConfig = new CompositeConfig(
        Lists.newArrayList(childJob1, childJob2, childJob3), false);

    long jobId = mJobMaster.run(jobConfig);

    JobInfo status = mJobMaster.getStatus(jobId);
    List<JobInfo> children = status.getChildren();

    assertEquals(3, children.size());

    long child0 = children.get(0).getId();
    long child1 = children.get(1).getId();
    long child2 = children.get(2).getId();

    mJobMaster.cancel(jobId);
    JobTestUtils.waitForJobStatus(mJobMaster, jobId, Status.CANCELED);
    JobTestUtils.waitForJobStatus(mJobMaster, child0, Status.CANCELED);
    JobTestUtils.waitForJobStatus(mJobMaster, child1, Status.CANCELED);
    JobTestUtils.waitForJobStatus(mJobMaster, child2, Status.CANCELED);
  }
}
