/*
 * 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.worker.job.task;

import alluxio.collections.Pair;
import alluxio.conf.Configuration;
import alluxio.conf.PropertyKey;
import alluxio.grpc.RunTaskCommand;
import alluxio.job.ErrorUtils;
import alluxio.job.RunTaskContext;
import alluxio.job.wire.Status;
import alluxio.job.wire.TaskInfo;
import alluxio.util.ThreadFactoryUtils;
import alluxio.util.logging.SamplingLogger;
import alluxio.wire.WorkerNetAddress;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import javax.annotation.concurrent.ThreadSafe;

/**
 * Manages the task executors.
 */
@ThreadSafe
public class TaskExecutorManager {
  private static final Logger LOG = LoggerFactory.getLogger(TaskExecutorManager.class);
  private static final SamplingLogger SAMPLING_LOGGER = new SamplingLogger(LOG, 30 * 1000);

  private static final int MAX_TASK_EXECUTOR_POOL_SIZE = 10000;

  private final PausableThreadPoolExecutor mTaskExecutionService;

  // These maps are all indexed by <Job ID, Task ID> pairs.
  /** Stores the futures for all running tasks. */
  private final Map<Pair<Long, Long>, Future<?>> mTaskFutures;
  /** Stores the task info for all running tasks. */
  private final Map<Pair<Long, Long>, TaskInfo> mUnfinishedTasks;
  /** Stores the updated tasks since the last call to {@link #getAndClearTaskUpdates()}. */
  private final Map<Pair<Long, Long>, TaskInfo> mTaskUpdates;

  private final WorkerNetAddress mAddress;

  private int mDefaultTaskExecutorPoolSize;
  private boolean mThrottled;

  /**
   * Constructs a new instance of {@link TaskExecutorManager}.
   * @param taskExecutorPoolSize number of task executors in the pool
   * @param address the worker address
   */
  public TaskExecutorManager(int taskExecutorPoolSize, WorkerNetAddress address) {
    mTaskFutures = Maps.newHashMap();
    mUnfinishedTasks = Maps.newHashMap();
    mTaskUpdates = Maps.newHashMap();
    mTaskExecutionService = new PausableThreadPoolExecutor(taskExecutorPoolSize,
        MAX_TASK_EXECUTOR_POOL_SIZE, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>(),
        ThreadFactoryUtils.build("task-execution-service-%d", true));
    mDefaultTaskExecutorPoolSize = taskExecutorPoolSize;
    mThrottled = false;
    mAddress = address;
  }

  /**
   * @return number of active tasks
   */
  public int getNumActiveTasks() {
    return mTaskExecutionService.getNumActiveTasks();
  }

  /**
   * @return task executor pool size
   */
  public int getTaskExecutorPoolSize() {
    return mTaskExecutionService.getCorePoolSize();
  }

  /**
   * @param taskExecutorPoolSize number of threads in the task executor pool
   */
  private synchronized void setTaskExecutorPoolSize(int taskExecutorPoolSize) {
    Preconditions.checkArgument(taskExecutorPoolSize >= 0);
    Preconditions.checkArgument(taskExecutorPoolSize <= MAX_TASK_EXECUTOR_POOL_SIZE);

    if (taskExecutorPoolSize == 0) {
      // treat 0 as a special case because ThreadedTaskExecutorService can't seem to have 0 threads
      mTaskExecutionService.pause();
    } else {
      mTaskExecutionService.resume();
    }

    mTaskExecutionService.setCorePoolSize(taskExecutorPoolSize);
  }

  /**
   * Sets the number of threads in the task executor pool when not throttled.
   *
   * @param defaultTaskExecutorPoolSize number of threads in the task executor pool
   */
  public synchronized void setDefaultTaskExecutorPoolSize(int defaultTaskExecutorPoolSize) {
    mDefaultTaskExecutorPoolSize = defaultTaskExecutorPoolSize;
    if (!mThrottled) {
      setTaskExecutorPoolSize(mDefaultTaskExecutorPoolSize);
    }
  }

  /**
   * Throttle the TaskExecutorManager. This sets the internal thread pool size to 0 until
   * unthrottled.
   */
  public synchronized void throttle() {
    mThrottled = true;
    setTaskExecutorPoolSize(0);
  }

  /**
   * Unthrottle the TaskExecutorManager. This sets the internal thread pool size back to the
   * default thread pool size.
   */
  public synchronized void unthrottle() {
    mThrottled = false;
    setTaskExecutorPoolSize(mDefaultTaskExecutorPoolSize);
  }

  /**
   * @return number of unfinished tasks
   */
  public int unfinishedTasks() {
    return mUnfinishedTasks.size();
  }

  /**
   * Noitfy the start of the task.
   *
   * @param jobId the job id
   * @param taskId the task id
   */
  public synchronized void notifyTaskRunning(long jobId, long taskId) {
    Pair<Long, Long> id = new Pair<>(jobId, taskId);
    TaskInfo taskInfo = mUnfinishedTasks.get(id);
    taskInfo.setStatus(Status.RUNNING);
    LOG.info("Task {} for job {} started", taskId, jobId);
  }

  /**
   * Notifies the completion of the task.
   *
   * @param jobId the job id
   * @param taskId the task id
   * @param result the task execution result
   */
  public synchronized void notifyTaskCompletion(long jobId, long taskId, Serializable result) {
    Pair<Long, Long> id = new Pair<>(jobId, taskId);
    TaskInfo taskInfo = mUnfinishedTasks.get(id);
    taskInfo.setStatus(Status.COMPLETED);
    taskInfo.setResult(result);
    finishTask(id);
    LOG.info("Task {} for job {} completed.", taskId, jobId);
  }

  /**
   * Notifies the failure of the task.
   *
   * @param jobId the job id
   * @param taskId the task id
   * @param t the thrown exception
   */
  public synchronized void notifyTaskFailure(long jobId, long taskId, Throwable t) {
    Pair<Long, Long> id = new Pair<>(jobId, taskId);
    TaskInfo taskInfo = mUnfinishedTasks.get(id);
    taskInfo.setStatus(Status.FAILED);

    String errorMessage;
    if (Configuration.getBoolean(PropertyKey.DEBUG)) {
      errorMessage = Throwables.getStackTraceAsString(t);
    } else {
      errorMessage = t.getMessage();
    }

    taskInfo.setErrorType(ErrorUtils.getErrorType(t));
    if (errorMessage != null) {
      taskInfo.setErrorMessage(errorMessage);
    }
    finishTask(id);
    LOG.info("Task {} for job {} failed: {}", taskId, jobId, errorMessage);
    SAMPLING_LOGGER.info("Stack trace for taskId: {} jobId: {} : {}", taskId, jobId,
        Throwables.getStackTraceAsString(t));
  }

  /**
   * Executes the given task.
   *
   * @param jobId the job id
   * @param taskId the task id
   * @param runTaskCommand the run task command
   * @param context the context of the worker
   */
  public synchronized void executeTask(long jobId, long taskId, RunTaskCommand runTaskCommand,
      RunTaskContext context) {
    Future<?> future = mTaskExecutionService
        .submit(new TaskExecutor(jobId, taskId, runTaskCommand, context, this));
    Pair<Long, Long> id = new Pair<>(jobId, taskId);
    mTaskFutures.put(id, future);
    TaskInfo taskInfo = new TaskInfo(jobId, taskId, Status.CREATED, mAddress, null);

    mUnfinishedTasks.put(id, taskInfo);
    mTaskUpdates.put(id, taskInfo);
    LOG.info("Task {} for job {} received", taskId, jobId);
  }

  /**
   * Cancels the given task.
   *
   * @param jobId the job id
   * @param taskId the task id
   */
  public synchronized void cancelTask(long jobId, long taskId) {
    Pair<Long, Long> id = new Pair<>(jobId, taskId);
    TaskInfo taskInfo = mUnfinishedTasks.get(id);
    if (!mTaskFutures.containsKey(id) || taskInfo.getStatus().equals(Status.CANCELED)) {
      // job has finished, or failed, or canceled
      return;
    }
    Future<?> future = mTaskFutures.get(id);
    if (!future.cancel(true)) {
      taskInfo.setStatus(Status.FAILED);
      taskInfo.setErrorType("FailedCancel");
      taskInfo.setErrorMessage("Failed to cancel the task");
      LOG.info("Failed to cancel task {} for job {}", taskId, jobId);
    } else {
      taskInfo.setStatus(Status.CANCELED);
      LOG.info("Task {} for job {} canceled", taskId, jobId);
    }
    finishTask(id);
  }

  /**
   * @return the list of task information
   */
  public synchronized List<TaskInfo> getAndClearTaskUpdates() {
    try {
      return ImmutableList.copyOf(mTaskUpdates.values());
    } finally {
      mTaskUpdates.clear();
    }
  }

  /**
   * Adds the given tasks to the task updates data structure. If there is already an update for the
   * specified task, it is not changed.
   *
   * @param tasks the tasks to restore
   */
  public synchronized void restoreTaskUpdates(List<TaskInfo> tasks) {
    for (TaskInfo task : tasks) {
      Pair<Long, Long> id = new Pair<>(task.getParentId(), task.getId());
      if (!mTaskUpdates.containsKey(id)) {
        mTaskUpdates.put(id, task);
      }
    }
  }

  private void finishTask(Pair<Long, Long> id) {
    TaskInfo taskInfo = mUnfinishedTasks.get(id);
    mTaskFutures.remove(id);
    mUnfinishedTasks.remove(id);
    mTaskUpdates.put(id, taskInfo);
  }
}
