com.hubspot.singularity.data.TaskManager Maven / Gradle / Ivy

Go to download
package com.hubspot.singularity.data;

import static java.nio.charset.StandardCharsets.UTF_8;

import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.utils.ZKPaths;
import org.apache.mesos.Protos.TaskStatus;

import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.google.inject.name.Named;
import com.hubspot.singularity.ExtendedTaskState;
import com.hubspot.singularity.LoadBalancerRequestType;
import com.hubspot.singularity.SingularityCreateResult;
import com.hubspot.singularity.SingularityDeleteResult;
import com.hubspot.singularity.SingularityKilledTaskIdRecord;
import com.hubspot.singularity.SingularityLoadBalancerUpdate;
import com.hubspot.singularity.SingularityMainModule;
import com.hubspot.singularity.SingularityPendingRequest.PendingType;
import com.hubspot.singularity.SingularityPendingTask;
import com.hubspot.singularity.SingularityPendingTaskId;
import com.hubspot.singularity.SingularitySlave;
import com.hubspot.singularity.SingularityTask;
import com.hubspot.singularity.SingularityTaskCleanup;
import com.hubspot.singularity.SingularityTaskHealthcheckResult;
import com.hubspot.singularity.SingularityTaskHistory;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;
import com.hubspot.singularity.SingularityTaskIdHolder;
import com.hubspot.singularity.SingularityTaskStatusHolder;
import com.hubspot.singularity.config.SingularityConfiguration;
import com.hubspot.singularity.data.transcoders.IdTranscoder;
import com.hubspot.singularity.data.transcoders.StringTranscoder;
import com.hubspot.singularity.data.transcoders.Transcoder;

@Singleton
public class TaskManager extends CuratorAsyncManager {

  private static final String TASKS_ROOT = "/tasks";

  private static final String ACTIVE_PATH_ROOT = TASKS_ROOT + "/active";
  private static final String LAST_ACTIVE_TASK_STATUSES_PATH_ROOT = TASKS_ROOT + "/statuses";
  public static final String PENDING_PATH_ROOT = TASKS_ROOT + "/scheduled";
  private static final String CLEANUP_PATH_ROOT = TASKS_ROOT + "/cleanup";
  private static final String LB_CLEANUP_PATH_ROOT = TASKS_ROOT + "/lbcleanup";
  private static final String DRIVER_KILLED_PATH_ROOT = TASKS_ROOT + "/killed";

  private static final String HISTORY_PATH_ROOT = TASKS_ROOT + "/history";

  private static final String LAST_HEALTHCHECK_KEY = "LAST_HEALTHCHECK";
  private static final String DIRECTORY_KEY = "DIRECTORY";
  private static final String TASK_KEY = "TASK";
  private static final String NOTIFIED_OVERDUE_TO_FINISH_KEY = "NOTIFIED_OVERDUE_TO_FINISH";

  private static final String LOAD_BALANCER_PRE_KEY = "LOAD_BALANCER_";

  private static final String HEALTHCHECKS_PATH = "/healthchecks";
  private static final String UPDATES_PATH = "/updates";

  private final Transcoder healthcheckResultTranscoder;
  private final Transcoder taskCleanupTranscoder;
  private final Transcoder taskTranscoder;
  private final Transcoder taskStatusTranscoder;
  private final Transcoder killedTaskIdRecordTranscoder;
  private final Transcoder taskHistoryUpdateTranscoder;
  private final Transcoder taskLoadBalancerUpdateTranscoder;

  private final IdTranscoder pendingTaskIdTranscoder;
  private final IdTranscoder taskIdTranscoder;

  private final Function pendingTaskIdToPendingTaskFunction;

  private final WebhookManager webhookManager;
  private final String serverId;

  @Inject
  public TaskManager(SingularityConfiguration configuration, CuratorFramework curator, WebhookManager webhookManager, IdTranscoder pendingTaskIdTranscoder,
      IdTranscoder taskIdTranscoder, Transcoder taskLoadBalancerHistoryUpdateTranscoder,
      Transcoder taskStatusTranscoder, Transcoder healthcheckResultTranscoder, Transcoder taskTranscoder,
      Transcoder taskCleanupTranscoder, Transcoder taskHistoryUpdateTranscoder,
      Transcoder killedTaskIdRecordTranscoder, @Named(SingularityMainModule.SERVER_ID_PROPERTY) String serverId) {
    super(curator, configuration.getZookeeperAsyncTimeout());

    this.healthcheckResultTranscoder = healthcheckResultTranscoder;
    this.taskTranscoder = taskTranscoder;
    this.taskStatusTranscoder = taskStatusTranscoder;
    this.killedTaskIdRecordTranscoder = killedTaskIdRecordTranscoder;
    this.taskCleanupTranscoder = taskCleanupTranscoder;
    this.taskHistoryUpdateTranscoder = taskHistoryUpdateTranscoder;
    this.taskIdTranscoder = taskIdTranscoder;
    this.pendingTaskIdTranscoder = pendingTaskIdTranscoder;
    this.taskLoadBalancerUpdateTranscoder = taskLoadBalancerHistoryUpdateTranscoder;
    this.webhookManager = webhookManager;

    this.serverId = serverId;

    this.pendingTaskIdToPendingTaskFunction = new Function() {

      @Override
      public SingularityPendingTask apply(SingularityPendingTaskId input) {
        Optional maybeCmdLineArgs = Optional.absent();

        if (input.getPendingType() == PendingType.ONEOFF || input.getPendingType() == PendingType.IMMEDIATE) {
          maybeCmdLineArgs = getStringData(ZKPaths.makePath(PENDING_PATH_ROOT, input.getId()));
        }

        return new SingularityPendingTask(input, maybeCmdLineArgs);
      }
    };
  }

  private String getLastHealthcheckPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), LAST_HEALTHCHECK_KEY);
  }

  private String getHealthcheckParentPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), HEALTHCHECKS_PATH);
  }

  private String getLastActiveTaskStatusPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(LAST_ACTIVE_TASK_STATUSES_PATH_ROOT, taskId.getId());
  }

  private String getHealthcheckPath(SingularityTaskHealthcheckResult healthcheck) {
    return ZKPaths.makePath(getHealthcheckParentPath(healthcheck.getTaskId()), Long.toString(healthcheck.getTimestamp()));
  }

  private String getTaskPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), TASK_KEY);
  }

  private String getLoadBalancerStatePath(SingularityTaskId taskId, LoadBalancerRequestType requestType) {
    return ZKPaths.makePath(getHistoryPath(taskId), LOAD_BALANCER_PRE_KEY + requestType.name());
  }

  private String getDirectoryPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), DIRECTORY_KEY);
  }

  private String getNotifiedOverduePath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), NOTIFIED_OVERDUE_TO_FINISH_KEY);
  }

  private String getUpdatesPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), UPDATES_PATH);
  }

  private String getUpdatePath(SingularityTaskId taskId, ExtendedTaskState state) {
    return ZKPaths.makePath(getUpdatesPath(taskId), state.name());
  }

  private String getRequestPath(String requestId) {
    return ZKPaths.makePath(HISTORY_PATH_ROOT, requestId);
  }

  private String getHistoryPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getRequestPath(taskId.getRequestId()), taskId.getId());
  }

  private String getActivePath(String taskId) {
    return ZKPaths.makePath(ACTIVE_PATH_ROOT, taskId);
  }

  private String getPendingPath(SingularityPendingTaskId pendingTaskId) {
    return ZKPaths.makePath(PENDING_PATH_ROOT, pendingTaskId.getId());
  }

  private String getCleanupPath(String taskId) {
    return ZKPaths.makePath(CLEANUP_PATH_ROOT, taskId);
  }

  public int getNumCleanupTasks() {
    return getNumChildren(CLEANUP_PATH_ROOT);
  }

  public int getNumLbCleanupTasks() {
    return getNumChildren(LB_CLEANUP_PATH_ROOT);
  }

  public int getNumActiveTasks() {
    return getNumChildren(ACTIVE_PATH_ROOT);
  }

  public int getNumScheduledTasks() {
    return getNumChildren(PENDING_PATH_ROOT);
  }

  public void saveLoadBalancerState(SingularityTaskId taskId, LoadBalancerRequestType requestType, SingularityLoadBalancerUpdate lbUpdate) {
    Preconditions.checkState(requestType != LoadBalancerRequestType.DEPLOY);

    save(getLoadBalancerStatePath(taskId, requestType), lbUpdate, taskLoadBalancerUpdateTranscoder);
  }

  public void saveTaskDirectory(SingularityTaskId taskId, String directory) {
    save(getDirectoryPath(taskId), Optional.of(directory.getBytes(UTF_8)));
  }

  public void saveLastActiveTaskStatus(SingularityTaskStatusHolder taskStatus) {
    save(getLastActiveTaskStatusPath(taskStatus.getTaskId()), taskStatus, taskStatusTranscoder);
  }

  public Optional getDirectory(SingularityTaskId taskId) {
    return getData(getDirectoryPath(taskId), StringTranscoder.INSTANCE);
  }

  public void saveHealthcheckResult(SingularityTaskHealthcheckResult healthcheckResult) {
    final Optional bytes = Optional.of(healthcheckResultTranscoder.toBytes(healthcheckResult));

    save(getHealthcheckPath(healthcheckResult), bytes);
    save(getLastHealthcheckPath(healthcheckResult.getTaskId()), bytes);
  }

  public void createPendingTasks(List tasks) {
    try {
      for (SingularityPendingTask task : tasks) {
        createPendingTask(task);
      }
    } catch (Throwable t) {
      throw Throwables.propagate(t);
    }
  }

  private void createPendingTask(SingularityPendingTask task) throws Exception {
    final String pendingPath = getPendingPath(task.getPendingTaskId());

    Optional data = Optional.absent();

    if (task.getMaybeCmdLineArgs().isPresent()) {
      data = Optional.of(task.getMaybeCmdLineArgs().get().getBytes(UTF_8));
    }

    create(pendingPath, data);
  }

  public List getAllTaskIds() {
    final List requestIds = getChildren(HISTORY_PATH_ROOT);
    final List paths = Lists.newArrayListWithCapacity(requestIds.size());

    for (String requestId : requestIds) {
      paths.add(getRequestPath(requestId));
    }

    return getChildrenAsIdsForParents(HISTORY_PATH_ROOT, paths, taskIdTranscoder);
  }

  private List getTaskIds(String root) {
    return getChildrenAsIds(root, taskIdTranscoder);
  }

  public List getActiveTaskIds() {
    return getTaskIds(ACTIVE_PATH_ROOT);
  }

  public List getCleanupTaskIds() {
    return getTaskIds(CLEANUP_PATH_ROOT);
  }

  public List getCleanupTasks() {
    return getAsyncChildren(CLEANUP_PATH_ROOT, taskCleanupTranscoder);
  }

  public List getActiveTasks() {
    List children = Lists.transform(getChildrenAsIds(ACTIVE_PATH_ROOT, taskIdTranscoder), new Function() {

      @Override
      public String apply(SingularityTaskId taskId) {
        return getTaskPath(taskId);
      }

    });

    return getAsync("active_tasks", children, taskTranscoder);
  }

  public List getLastActiveTaskStatuses() {
    return getAsyncChildren(LAST_ACTIVE_TASK_STATUSES_PATH_ROOT, taskStatusTranscoder);
  }

  public Optional getLastActiveTaskStatus(SingularityTaskId taskId) {
    return getData(getLastActiveTaskStatusPath(taskId), taskStatusTranscoder);
  }

  public List getLastActiveTaskStatusesFor(Collection activeTaskIds) {
    List paths = Lists.newArrayListWithExpectedSize(activeTaskIds.size());
    for (SingularityTaskId taskId : activeTaskIds) {
      paths.add(getLastActiveTaskStatusPath(taskId));
    }
    return getAsync(LAST_ACTIVE_TASK_STATUSES_PATH_ROOT, paths, taskStatusTranscoder);
  }

  public List getTasksOnSlave(List activeTaskIds, SingularitySlave slave) {
    List tasks = Lists.newArrayList();

    for (SingularityTaskId activeTaskId : activeTaskIds) {
      if (activeTaskId.getHost().equals(slave.getHost())) {
        Optional maybeTask = getTask(activeTaskId);
        if (maybeTask.isPresent() && slave.getId().equals(maybeTask.get().getOffer().getSlaveId().getValue())) {
          tasks.add(maybeTask.get());
        }
      }
    }

    return tasks;
  }

  public List getTaskHistoryUpdates(SingularityTaskId taskId) {
    List updates = getAsyncChildren(getUpdatesPath(taskId), taskHistoryUpdateTranscoder);
    Collections.sort(updates);
    return updates;
  }

  public Map> getTaskHistoryUpdates(Collection taskIds) {
    Map> map = Maps.newHashMapWithExpectedSize(taskIds.size());

    for (SingularityTaskId taskId : taskIds) {
      map.put(taskId, getTaskHistoryUpdates(taskId));
    }

    return map;
  }

  public List getHealthcheckResults(SingularityTaskId taskId) {
    List healthcheckResults = getAsyncChildren(getHealthcheckParentPath(taskId), healthcheckResultTranscoder);
    Collections.sort(healthcheckResults);
    return healthcheckResults;
  }

  public Optional getLastHealthcheck(SingularityTaskId taskId) {
    return getData(getLastHealthcheckPath(taskId), healthcheckResultTranscoder);
  }

  public Map getLastHealthcheck(Collection taskIds) {
    List paths = Lists.newArrayListWithCapacity(taskIds.size());
    for (SingularityTaskId taskId : taskIds) {
      paths.add(getLastHealthcheckPath(taskId));
    }

    List healthcheckResults = getAsync("healthchecks_by_ids", paths, healthcheckResultTranscoder);

    return Maps.uniqueIndex(healthcheckResults, SingularityTaskIdHolder.getTaskIdFunction());
  }

  public boolean taskHistoryUpdateExists(SingularityTaskHistoryUpdate taskHistoryUpdate) {
    return exists(getUpdatePath(taskHistoryUpdate.getTaskId(), taskHistoryUpdate.getTaskState()));
  }

  public SingularityCreateResult saveTaskHistoryUpdate(SingularityTaskHistoryUpdate taskHistoryUpdate) {
    webhookManager.enqueueTaskUpdate(taskHistoryUpdate);

    return create(getUpdatePath(taskHistoryUpdate.getTaskId(), taskHistoryUpdate.getTaskState()), taskHistoryUpdate, taskHistoryUpdateTranscoder);
  }

  public boolean isActiveTask(String taskId) {
    return exists(getActivePath(taskId));
  }

  public List getTaskIdsForRequest(String requestId) {
    return getChildrenAsIds(getRequestPath(requestId), taskIdTranscoder);
  }

  private enum TaskFilter {
    ACTIVE, INACTIVE;
  }

  public List getInactiveTaskIdsForRequest(String requestId) {
    return getTaskIdsForRequest(requestId, TaskFilter.INACTIVE);
  }

  public List getActiveTaskIdsForRequest(String requestId) {
    return getTaskIdsForRequest(requestId, TaskFilter.ACTIVE);
  }

  public List filterActiveTaskIds(List taskIds) {
    final List paths = Lists.newArrayListWithCapacity(taskIds.size());

    for (SingularityTaskId taskId : taskIds) {
      paths.add(getActivePath(taskId.getId()));
    }

    final List activeTaskIds = exists(ACTIVE_PATH_ROOT, paths, taskIdTranscoder);

    return activeTaskIds;
  }

  private List getTaskIdsForRequest(String requestId, TaskFilter taskFilter) {
    final List requestTaskIds = getTaskIdsForRequest(requestId);
    final List activeTaskIds = filterActiveTaskIds(requestTaskIds);

    if (taskFilter == TaskFilter.ACTIVE) {
      return activeTaskIds;
    }

    Iterables.removeAll(requestTaskIds, activeTaskIds);

    return requestTaskIds;
  }

  public Optional getTaskHistory(SingularityTaskId taskId) {
    final Optional task = getTask(taskId);

    if (!task.isPresent()) {
      return Optional.absent();
    }

    List taskUpdates = getTaskHistoryUpdates(taskId);
    Optional directory = getDirectory(taskId);
    List healthchecks = getHealthcheckResults(taskId);

    List loadBalancerUpdates = Lists.newArrayListWithCapacity(2);

    checkLoadBalancerHistory(loadBalancerUpdates, taskId, LoadBalancerRequestType.ADD);
    checkLoadBalancerHistory(loadBalancerUpdates, taskId, LoadBalancerRequestType.REMOVE);

    return Optional.of(new SingularityTaskHistory(taskUpdates, directory, healthchecks, task.get(), loadBalancerUpdates));
  }

  private void checkLoadBalancerHistory(List loadBalancerUpdates, SingularityTaskId taskId, LoadBalancerRequestType lbRequestType) {
    Optional lbHistory = getLoadBalancerState(taskId, lbRequestType);

    if (lbHistory.isPresent()) {
      loadBalancerUpdates.add(lbHistory.get());
    }
  }

  public boolean hasNotifiedOverdue(SingularityTaskId taskId) {
    return checkExists(getNotifiedOverduePath(taskId)).isPresent();
  }

  public void saveNotifiedOverdue(SingularityTaskId taskId) {
    save(getNotifiedOverduePath(taskId), Optional. absent());
  }

  public Optional getLoadBalancerState(SingularityTaskId taskId, LoadBalancerRequestType requestType) {
    return getData(getLoadBalancerStatePath(taskId, requestType), taskLoadBalancerUpdateTranscoder);
  }

  public SingularityPendingTask getPendingTask(SingularityPendingTaskId pendingTaskId) {
    return pendingTaskIdToPendingTaskFunction.apply(pendingTaskId);
  }

  public Optional getTask(SingularityTaskId taskId) {
    final String path = getTaskPath(taskId);

    return getData(path, taskTranscoder);
  }

  public List getPendingTaskIds() {
    return getChildrenAsIds(PENDING_PATH_ROOT, pendingTaskIdTranscoder);
  }

  public List getPendingTasks() {
    return Lists.transform(getPendingTaskIds(), pendingTaskIdToPendingTaskFunction);
  }

  public void createTaskAndDeletePendingTask(SingularityTask task) {
    try {
      createTaskAndDeletePendingTaskPrivate(task);
    } catch (Throwable t) {
      throw Throwables.propagate(t);
    }
  }

  public Map getTasks(Iterable taskIds) {
    final List paths = Lists.newArrayList();

    for (SingularityTaskId taskId : taskIds) {
      paths.add(getTaskPath(taskId));
    }

    return Maps.uniqueIndex(getAsync("tasks_by_ids", paths, taskTranscoder), SingularityTaskIdHolder.getTaskIdFunction());
  }

  private void createTaskAndDeletePendingTaskPrivate(SingularityTask task) throws Exception {
    delete(getPendingPath(task.getTaskRequest().getPendingTask().getPendingTaskId()));

    final long now = System.currentTimeMillis();

    saveTaskHistoryUpdate(new SingularityTaskHistoryUpdate(task.getTaskId(), now, ExtendedTaskState.TASK_LAUNCHED, Optional.absent()));
    saveLastActiveTaskStatus(new SingularityTaskStatusHolder(task.getTaskId(), Optional.absent(), now, serverId, Optional.of(task.getOffer().getSlaveId().getValue())));

    create(getTaskPath(task.getTaskId()), task, taskTranscoder);
    create(getActivePath(task.getTaskId().getId()));
  }

  public List getLBCleanupTasks() {
    return getChildrenAsIds(LB_CLEANUP_PATH_ROOT, taskIdTranscoder);
  }

  private String getLBCleanupPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(LB_CLEANUP_PATH_ROOT, taskId.getId());
  }

  private String getKilledPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(DRIVER_KILLED_PATH_ROOT, taskId.getId());
  }

  public SingularityDeleteResult deleteLBCleanupTask(SingularityTaskId taskId) {
    return delete(getLBCleanupPath(taskId));
  }

  public SingularityCreateResult createLBCleanupTask(SingularityTaskId taskId) {
    return create(getLBCleanupPath(taskId));
  }

  public SingularityCreateResult saveKilledRecord(SingularityKilledTaskIdRecord killedTaskIdRecord) {
    return save(getKilledPath(killedTaskIdRecord.getTaskId()), killedTaskIdRecord, killedTaskIdRecordTranscoder);
  }

  public List getKilledTaskIdRecords() {
    return getAsyncChildren(DRIVER_KILLED_PATH_ROOT, killedTaskIdRecordTranscoder);
  }

  public SingularityDeleteResult deleteKilledRecord(SingularityTaskId taskId) {
    return delete(getKilledPath(taskId));
  }

  public SingularityDeleteResult deleteLastActiveTaskStatus(SingularityTaskId taskId) {
    return delete(getLastActiveTaskStatusPath(taskId));
  }

  public SingularityCreateResult createCleanupTask(SingularityTaskCleanup cleanupTask) {
    StringBuilder msg = new StringBuilder(cleanupTask.getCleanupType().name());

    if (cleanupTask.getUser().isPresent()) {
      msg.append(" - ");
      msg.append(cleanupTask.getUser().get());
    }

    saveTaskHistoryUpdate(new SingularityTaskHistoryUpdate(cleanupTask.getTaskId(), cleanupTask.getTimestamp(), ExtendedTaskState.TASK_CLEANING, Optional.of(msg.toString())));

    return create(getCleanupPath(cleanupTask.getTaskId().getId()), cleanupTask, taskCleanupTranscoder);
  }

  public void deleteActiveTask(String taskId) {
    delete(getActivePath(taskId));
  }

  public void deletePendingTask(SingularityPendingTaskId pendingTaskId) {
    delete(getPendingPath(pendingTaskId));
  }

  public void deleteCleanupTask(String taskId) {
    delete(getCleanupPath(taskId));
  }

  public SingularityDeleteResult deleteTaskHistory(SingularityTaskId taskId) {
    return delete(getHistoryPath(taskId));
  }

}