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.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.api.transaction.CuratorTransactionFinal;
import org.apache.curator.utils.ZKPaths;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.annotation.Timed;
import com.google.common.base.Function;
import com.google.common.base.Optional;
import com.google.common.base.Preconditions;
import com.google.common.base.Predicate;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
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.mesos.JavaUtils;
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.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.SingularityTaskMetadata;
import com.hubspot.singularity.SingularityTaskShellCommandHistory;
import com.hubspot.singularity.SingularityTaskShellCommandRequest;
import com.hubspot.singularity.SingularityTaskShellCommandRequestId;
import com.hubspot.singularity.SingularityTaskShellCommandUpdate;
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;
import com.hubspot.singularity.event.SingularityEventListener;
import com.hubspot.singularity.scheduler.SingularityLeaderCache;

@Singleton
public class TaskManager extends CuratorAsyncManager {

  private static final Logger LOG = LoggerFactory.getLogger(CuratorAsyncManager.class);

  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";
  private 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 FINISHED_TASK_MAIL_QUEUE = TASKS_ROOT + "/mailqueue";
  private static final String SHELL_REQUESTS_QUEUE_PATH_ROOT = TASKS_ROOT + "/shellqueue";
  private static final String PENDING_TASKS_TO_DELETE_PATH_ROOT = TASKS_ROOT + "/pendingdeletes";
  private static final String UNHEALTHY_KILL_PATH_ROOT = TASKS_ROOT + "/unhealthyKill";

  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 CONTAINER_ID_KEY = "CONTAINER_ID";
  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 SHELLS_PATH = "/shells";
  private static final String SHELL_REQUEST_KEY = "REQUEST";
  private static final String SHELL_UPDATES_PATH = "/updates";

  private static final String HEALTHCHECKS_PATH = "/healthchecks";
  private static final String HEALTHCHECKS_FINISHED_PATH = "/healthchecks-finished";
  private static final String STARTUP_HEALTHCHECK_PATH_SUFFIX = "-NOT_STARTED";

  private static final String METADATA_PATH = "/metadata";
  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 Transcoder pendingTaskTranscoder;
  private final Transcoder taskShellCommandRequestTranscoder;
  private final Transcoder taskShellCommandUpdateTranscoder;
  private final Transcoder taskMetadataTranscoder;

  private final IdTranscoder pendingTaskIdTranscoder;
  private final IdTranscoder taskIdTranscoder;

  private final ZkCache taskCache;
  private final SingularityWebCache webCache;
  private final SingularityLeaderCache leaderCache;

  private final SingularityEventListener singularityEventListener;
  private final String serverId;

  @Inject
  public TaskManager(CuratorFramework curator, SingularityConfiguration configuration, MetricRegistry metricRegistry, SingularityEventListener singularityEventListener,
      IdTranscoder pendingTaskIdTranscoder, IdTranscoder taskIdTranscoder, Transcoder taskLoadBalancerHistoryUpdateTranscoder,
      Transcoder taskStatusTranscoder, Transcoder healthcheckResultTranscoder, Transcoder taskTranscoder,
      Transcoder taskCleanupTranscoder, Transcoder taskHistoryUpdateTranscoder, Transcoder pendingTaskTranscoder,
      Transcoder killedTaskIdRecordTranscoder, Transcoder taskShellCommandRequestTranscoder,
      Transcoder taskShellCommandUpdateTranscoder,  Transcoder taskMetadataTranscoder,
      ZkCache taskCache, SingularityWebCache webCache, SingularityLeaderCache leaderCache,
      @Named(SingularityMainModule.SERVER_ID_PROPERTY) String serverId) {
    super(curator, configuration, metricRegistry);

    this.healthcheckResultTranscoder = healthcheckResultTranscoder;
    this.taskTranscoder = taskTranscoder;
    this.taskStatusTranscoder = taskStatusTranscoder;
    this.killedTaskIdRecordTranscoder = killedTaskIdRecordTranscoder;
    this.taskCleanupTranscoder = taskCleanupTranscoder;
    this.taskHistoryUpdateTranscoder = taskHistoryUpdateTranscoder;
    this.taskIdTranscoder = taskIdTranscoder;
    this.pendingTaskTranscoder = pendingTaskTranscoder;
    this.taskShellCommandRequestTranscoder = taskShellCommandRequestTranscoder;
    this.pendingTaskIdTranscoder = pendingTaskIdTranscoder;
    this.taskLoadBalancerUpdateTranscoder = taskLoadBalancerHistoryUpdateTranscoder;
    this.singularityEventListener = singularityEventListener;
    this.taskCache = taskCache;
    this.taskShellCommandUpdateTranscoder = taskShellCommandUpdateTranscoder;
    this.taskMetadataTranscoder = taskMetadataTranscoder;

    this.webCache = webCache;
    this.leaderCache = leaderCache;
    this.serverId = serverId;
  }

  // since we can't use creatingParentsIfNeeded in transactions
  public void createRequiredParents() {
    create(HISTORY_PATH_ROOT);
    create(ACTIVE_PATH_ROOT);
  }

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

  private String getMetadataParentPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), METADATA_PATH);
  }

  private String getTaskMetadataPath(SingularityTaskMetadata taskMetadata) {
    return ZKPaths.makePath(getMetadataParentPath(taskMetadata.getTaskId()), String.format("%s-%s", taskMetadata.getTimestamp(), taskMetadata.getType()));
  }

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

  private String getHealthchecksFinishedPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), HEALTHCHECKS_FINISHED_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()), String.format("%s%s", Long.toString(healthcheck.getTimestamp()), healthcheck.isStartup() ? STARTUP_HEALTHCHECK_PATH_SUFFIX : ""));
  }

  private String getShellRequestQueuePath(SingularityTaskShellCommandRequest shellRequest) {
    return ZKPaths.makePath(SHELL_REQUESTS_QUEUE_PATH_ROOT, shellRequest.getId().getId());
  }

  private String getFinishedTaskMailQueuePath(SingularityTaskId taskId) {
    return ZKPaths.makePath(FINISHED_TASK_MAIL_QUEUE, taskId.getId());
  }

  private String getShellsParentPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), SHELLS_PATH);
  }

  private String getShellHistoryParentPath(SingularityTaskShellCommandRequestId shellRequestId) {
    return ZKPaths.makePath(getShellsParentPath(shellRequestId.getTaskId()), shellRequestId.getSubIdForTaskHistory());
  }

  private String getShellHistoryRequestPath(SingularityTaskShellCommandRequestId shellRequestId) {
    return ZKPaths.makePath(getShellHistoryParentPath(shellRequestId), SHELL_REQUEST_KEY);
  }

  private String getShellHistoryUpdateParentPath(SingularityTaskShellCommandRequestId shellRequestId) {
    return ZKPaths.makePath(getShellHistoryParentPath(shellRequestId), SHELL_UPDATES_PATH);
  }

  private String getShellHistoryUpdatePath(SingularityTaskShellCommandUpdate shellUpdate) {
    return ZKPaths.makePath(getShellHistoryUpdateParentPath(shellUpdate.getShellRequestId()), shellUpdate.getUpdateType().name());
  }

  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 getContainerIdPath(SingularityTaskId taskId) {
    return ZKPaths.makePath(getHistoryPath(taskId), CONTAINER_ID_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 getPendingTasksToDeletePath(SingularityPendingTaskId pendingTaskId) { return ZKPaths.makePath(PENDING_TASKS_TO_DELETE_PATH_ROOT, pendingTaskId.getId()); }

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

  private String getUnhealthyKillParentPath(String requestId) {
    return ZKPaths.makePath(UNHEALTHY_KILL_PATH_ROOT, requestId);
  }

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

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

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

  public int getNumActiveTasks() {
    if (leaderCache.active()) {
      return leaderCache.getNumActiveTasks();
    }
    return getNumChildren(ACTIVE_PATH_ROOT);
  }

  public int getNumScheduledTasks() {
    if (leaderCache.active()) {
      return leaderCache.getNumPendingTasks();
    }
    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 saveContainerId(SingularityTaskId taskId, String containerId) {
    save(getContainerIdPath(taskId), Optional.of(containerId.getBytes(UTF_8)));
  }

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

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

  public Optional getContainerId(SingularityTaskId taskId) {
    return getData(getContainerIdPath(taskId), StringTranscoder.INSTANCE);
  }

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

      save(getHealthcheckPath(healthcheckResult), bytes);
      save(getLastHealthcheckPath(healthcheckResult.getTaskId()), bytes);
    } else {
      LOG.warn("Healthchecks have finished, could not save new result {}", healthcheckResult);
    }
  }

  private boolean canSaveNewHealthcheck(SingularityTaskHealthcheckResult healthcheckResult) {
    return !exists(getHealthchecksFinishedPath(healthcheckResult.getTaskId()));
  }

  public void markHealthchecksFinished(SingularityTaskId taskId) {
    create(getHealthchecksFinishedPath(taskId));
  }

  public SingularityCreateResult savePendingTask(SingularityPendingTask task) {
    final String pendingPath = getPendingPath(task.getPendingTaskId());

    leaderCache.savePendingTask(task);

    return save(pendingPath, task, pendingTaskTranscoder);
  }

  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("getAllTaskIds", paths, taskIdTranscoder);
  }

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

  public List getActiveTaskIdsAsStrings() {
    if (leaderCache.active()) {
      return leaderCache.getActiveTaskIdsAsStrings();
    }
    return getChildren(ACTIVE_PATH_ROOT);
  }

  public List getActiveTaskIds() {
    return getActiveTaskIds(false);
  }

  public List getActiveTaskIds(boolean useWebCache) {
    if (leaderCache.active()) {
      return leaderCache.getActiveTaskIds();
    }

    if (useWebCache && webCache.useCachedActiveTasks()) {
      return webCache.getActiveTaskIds();
    }

    return getTaskIds(ACTIVE_PATH_ROOT);
  }

  public List getCleanupTaskIds() {
    if (leaderCache.active()) {
      return leaderCache.getCleanupTaskIds();
    }

    return getChildrenAsIds(CLEANUP_PATH_ROOT, taskIdTranscoder);
  }

  public List getCleanupTasks(boolean useWebCache) {
    if (leaderCache.active()) {
      return leaderCache.getCleanupTasks();
    }

    if (useWebCache && webCache.useCachedCleanupTasks()) {
      return webCache.getCleanupTasks();
    }

    List cleanupTasks = fetchCleanupTasks();

    if (useWebCache) {
      webCache.cacheTaskCleanup(cleanupTasks);
    }

    return cleanupTasks;
  }

  public List getCleanupTasks() {
    return getCleanupTasks(false);
  }

  public Optional getTaskCleanup(String taskId) {
    if (leaderCache.active()) {
      return leaderCache.getTaskCleanup(SingularityTaskId.valueOf(taskId));
    }
    return getData(getCleanupPath(taskId), taskCleanupTranscoder);
  }

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

  public List getActiveTasks() {
    return getActiveTasks(false);
  }

  public List getActiveTasks(boolean useWebCache) {
    if (useWebCache && webCache.useCachedActiveTasks()) {
      return webCache.getActiveTasks();
    }

    List children = Lists.transform(getActiveTaskIds(), new Function() {

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

    });

    List activeTasks = getAsync("getActiveTasks", children, taskTranscoder, taskCache);

    if (useWebCache) {
      webCache.cacheActiveTasks(activeTasks);
    }

    return activeTasks;
  }

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

  @Timed
  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("getLastActiveTaskStatusesFor", paths, taskStatusTranscoder);
  }

  public List getTasksOnSlave(Collection activeTaskIds, SingularitySlave slave) {
    final List tasks = Lists.newArrayList();
    final String sanitizedHost = JavaUtils.getReplaceHyphensWithUnderscores(slave.getHost());

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

    return tasks;
  }

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

  public Optional getTaskHistoryUpdate(SingularityTaskId taskId, ExtendedTaskState taskState) {
    return getData(getUpdatePath(taskId, taskState), taskHistoryUpdateTranscoder);
  }

  public Map> getTaskHistoryUpdates(Collection taskIds) {
    if (leaderCache.active()) {
      return leaderCache.getTaskHistoryUpdates(taskIds);
    }
    Map pathsMap = Maps.newHashMap();
    for (SingularityTaskId taskId : taskIds) {
      pathsMap.put(getHistoryPath(taskId), taskId);
    }

    return getAsyncNestedChildDataAsMap("getTaskHistoryUpdates", pathsMap, UPDATES_PATH, taskHistoryUpdateTranscoder);
  }

  public Map> getAllTaskHistoryUpdates() {
    return getTaskHistoryUpdates(getAllTaskIds());
  }

  public int getNumHealthchecks(SingularityTaskId taskId) {
    return getNumChildren(getHealthcheckParentPath(taskId));
  }

  public int getNumNonstartupHealthchecks(SingularityTaskId taskId) {
    int numChecks = 0;
    List checks = getChildren(getHealthcheckParentPath(taskId));
    for (String check : checks) {
      if (!check.endsWith(STARTUP_HEALTHCHECK_PATH_SUFFIX)) {
        numChecks++;
      }
    }
    return numChecks;
  }

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

  public void clearStartupHealthchecks(SingularityTaskId taskId) {
    Optional maybeLastHealthcheck = getLastHealthcheck(taskId);
    String parentPath = getHealthcheckParentPath(taskId);
    for (String healthcheckPath : getChildren(parentPath)) {
      String fullPath = ZKPaths.makePath(parentPath, healthcheckPath);
      if (healthcheckPath.endsWith(STARTUP_HEALTHCHECK_PATH_SUFFIX) && (!maybeLastHealthcheck.isPresent() || !getHealthcheckPath(maybeLastHealthcheck.get()).equals(fullPath))) {
        delete(fullPath);
      }
    }
  }

  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("getLastHealthcheck", paths, healthcheckResultTranscoder);

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

  public SingularityCreateResult saveTaskHistoryUpdate(SingularityTaskHistoryUpdate taskHistoryUpdate) {
    return saveTaskHistoryUpdate(taskHistoryUpdate, false);
  }

  @Timed
  public SingularityCreateResult saveTaskHistoryUpdate(SingularityTaskHistoryUpdate taskHistoryUpdate, boolean overwriteExisting) {
    singularityEventListener.taskHistoryUpdateEvent(taskHistoryUpdate);

    if (overwriteExisting) {
      Optional maybeExisting = getTaskHistoryUpdate(taskHistoryUpdate.getTaskId(), taskHistoryUpdate.getTaskState());
      LOG.info("Found existing history {}", maybeExisting);
      SingularityTaskHistoryUpdate updateWithPrevious;
      if (maybeExisting.isPresent()) {
        updateWithPrevious = taskHistoryUpdate.withPrevious(maybeExisting.get());
        LOG.info("Will save new update {}", updateWithPrevious);

      } else {
        updateWithPrevious = taskHistoryUpdate;
      }

      if (leaderCache.active()) {
        leaderCache.saveTaskHistoryUpdate(updateWithPrevious, overwriteExisting);
      }

      return save(getUpdatePath(taskHistoryUpdate.getTaskId(), taskHistoryUpdate.getTaskState()), updateWithPrevious, taskHistoryUpdateTranscoder);
    } else {
      if (leaderCache.active()) {
        leaderCache.saveTaskHistoryUpdate(taskHistoryUpdate, overwriteExisting);
      }
      return create(getUpdatePath(taskHistoryUpdate.getTaskId(), taskHistoryUpdate.getTaskState()), taskHistoryUpdate, taskHistoryUpdateTranscoder);
    }
  }

  public SingularityDeleteResult deleteTaskHistoryUpdate(SingularityTaskId taskId, ExtendedTaskState state, Optional previousStateUpdate) {
    if (previousStateUpdate.isPresent()) {
      singularityEventListener.taskHistoryUpdateEvent(previousStateUpdate.get());
    }
    if (leaderCache.active()) {
      leaderCache.deleteTaskHistoryUpdate(taskId, state);
    }
    return delete(getUpdatePath(taskId, state));
  }

  public boolean isActiveTask(String taskId) {
    if (leaderCache.active()) {
      return leaderCache.isActiveTask(taskId);
    }

    return exists(getActivePath(taskId));
  }

  public SingularityCreateResult markUnhealthyKill(SingularityTaskId taskId) {
    return create(getUnhealthyKillPath(taskId));
  }

  public int getNumUnhealthyKills(String requestId) {
    return getNumChildren(getUnhealthyKillParentPath(requestId));
  }

  public SingularityDeleteResult clearUnhealthyKills(String requestId) {
    return delete(getUnhealthyKillParentPath(requestId));
  }

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

  public Optional getTaskByRunId(String requestId, String runId) {
    Map activeTasks = getTasks(getActiveTaskIdsForRequest(requestId));
    for (Map.Entry entry : activeTasks.entrySet()) {
      if (entry.getValue().getTaskRequest().getPendingTask().getRunId().isPresent() && entry.getValue().getTaskRequest().getPendingTask().getRunId().get().equals(runId)) {
        return Optional.of(entry.getKey());
      }
    }
    Map inactiveTasks = getTasks(getInactiveTaskIdsForRequest(requestId));
    for (Map.Entry entry : inactiveTasks.entrySet()) {
      if (entry.getValue().getTaskRequest().getPendingTask().getRunId().isPresent() && entry.getValue().getTaskRequest().getPendingTask().getRunId().get().equals(runId)) {
        return Optional.of(entry.getKey());
      }
    }
    return Optional.absent();
  }

  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) {
    if (leaderCache.active()) {
      return leaderCache.exists(taskIds);
    }

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

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

    return exists("filterActiveTaskIds", paths, taskIdTranscoder);
  }

  public int getNumLaunchingTasks() {
    List activeTaskIds = getActiveTaskIds();

    final Map paths = Maps.newHashMapWithExpectedSize(activeTaskIds.size());

    for (SingularityTaskId taskId : activeTaskIds) {
      paths.put(getUpdatePath(taskId, ExtendedTaskState.TASK_RUNNING), taskId);
    }

    return notExists("getNumLaunchingTasks", paths).size();
  }

  public List filterInactiveTaskIds(List taskIds) {
    if (leaderCache.active()) {
      return leaderCache.getInactiveTaskIds(taskIds);
    }

    final Map pathsMap = Maps.newHashMap();

    for (SingularityTaskId taskId : taskIds) {
      pathsMap.put(getActivePath(taskId.getId()), taskId);
    }

    return notExists("filterInactiveTaskIds", pathsMap);
  }

  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 List getTaskIdsForDeploy(String requestId, final String deployId, TaskFilter taskFilter) {
    List requestTaskIds = getTaskIdsForRequest(requestId, taskFilter);
    final Iterable deployTaskIds = Iterables.filter(requestTaskIds, new Predicate() {
      @Override
      public boolean apply(SingularityTaskId input) {
        return input.getDeployId().equals(deployId);
      }
    });
    return ImmutableList.copyOf(deployTaskIds);
  }

  public List getActiveTaskIdsForDeploy(String requestId, final String deployId) {
    return getTaskIdsForDeploy(requestId, deployId, TaskFilter.ACTIVE);
  }

  public List getInactiveTaskIdsForDeploy(String requestId, final String deployId) {
    return getTaskIdsForDeploy(requestId, deployId, TaskFilter.INACTIVE);
  }

  public List getInactiveTaskIds(List requestIds) {
    List paths = Lists.newArrayListWithCapacity(requestIds.size());
    for (String requestId : requestIds) {
      paths.add(getRequestPath(requestId));
    }

    List taskIds = getChildrenAsIdsForParents("getInactiveTaskIds", paths, taskIdTranscoder);

    return filterInactiveTaskIds(taskIds);
  }

  public Optional getTaskHistory(SingularityTaskId taskId) {
    final Optional task = getTaskCheckCache(taskId, true);

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

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

    List loadBalancerUpdates = Lists.newArrayListWithCapacity(2);

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

    List shellCommandHistory = getTaskShellCommandHistory(taskId);

    List taskMetadata = getTaskMetadata(taskId);

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

  public List getTaskShellCommandHistory(SingularityTaskId taskId) {
    List shellRequests = getTaskShellCommandRequestsForTask(taskId);
    List shellCommandHistory = new ArrayList<>(shellRequests.size());

    for (SingularityTaskShellCommandRequest shellRequest : shellRequests) {
      shellCommandHistory.add(new SingularityTaskShellCommandHistory(shellRequest, getTaskShellCommandUpdates(shellRequest.getId())));
    }

    return shellCommandHistory;
  }

  private List getTaskMetadata(SingularityTaskId taskId) {
    List taskMetadata = getAsyncChildren(getMetadataParentPath(taskId), taskMetadataTranscoder);
    Collections.sort(taskMetadata);
    return taskMetadata;
  }

  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 Optional getPendingTask(SingularityPendingTaskId pendingTaskId) {
    if (leaderCache.active()) {
      return leaderCache.getPendingTask(pendingTaskId);
    }
    return getData(getPendingPath(pendingTaskId), pendingTaskTranscoder);
  }

  private Optional getTaskCheckCache(SingularityTaskId taskId, boolean shouldCheckExists) {
    final String path = getTaskPath(taskId);

    return getData(path, taskTranscoder, taskCache, shouldCheckExists);
  }

  @Timed
  public Optional getTask(SingularityTaskId taskId) {
    return getTaskCheckCache(taskId, false);
  }

  public boolean taskExistsInZk(SingularityTaskId taskId) {
    return checkExists(getTaskPath(taskId)).isPresent();
  }

  public void activateLeaderCache() {
    leaderCache.cachePendingTasks(fetchPendingTasks());
    leaderCache.cachePendingTasksToDelete(getPendingTasksMarkedForDeletion());
    leaderCache.cacheActiveTaskIds(getTaskIds(ACTIVE_PATH_ROOT));
    leaderCache.cacheCleanupTasks(fetchCleanupTasks());
    leaderCache.cacheKilledTasks(fetchKilledTaskIdRecords());
    leaderCache.cacheTaskHistoryUpdates(getAllTaskHistoryUpdates());
  }

  private List fetchPendingTasks() {
    return getAsyncChildren(PENDING_PATH_ROOT, pendingTaskTranscoder);
  }

  public List getPendingTaskIds() {
    return getPendingTaskIds(false);
  }

  public List getPendingTaskIds(boolean useWebCache) {
    if (leaderCache.active()) {
      return leaderCache.getPendingTaskIds();
    }

    if (useWebCache && webCache.useCachedPendingTasks()) {
      return webCache.getPendingTaskIds();
    }

    return getChildrenAsIds(PENDING_PATH_ROOT, pendingTaskIdTranscoder);
  }

  public List getPendingTaskIdsForRequest(final String requestId) {
    List pendingTaskIds = getPendingTaskIds();
    return pendingTaskIds.stream()
        .filter(pendingTaskId -> pendingTaskId.getRequestId().equals(requestId))
        .collect(Collectors.collectingAndThen(Collectors.toList(), ImmutableList::copyOf));
  }

  public List getPendingTasksForRequest(final String requestId) {
    return getAsync(
        PENDING_PATH_ROOT,
        getPendingTaskIdsForRequest(requestId).stream().map(this::getPendingPath).collect(Collectors.toList()),
        pendingTaskTranscoder
    );
  }

  public List getPendingTasks() {
    return getPendingTasks(false);
  }

  public List getPendingTasks(boolean useWebCache) {
    if (leaderCache.active()) {
      return leaderCache.getPendingTasks();
    }

    if (useWebCache && webCache.useCachedPendingTasks()) {
      return webCache.getPendingTasks();
    }

    List pendingTasks = fetchPendingTasks();

    if (useWebCache) {
      webCache.cachePendingTasks(pendingTasks);
    }

    return pendingTasks;
  }

  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("getTasks", paths, taskTranscoder, taskCache), SingularityTaskIdHolder.getTaskIdFunction());
  }

  private void createTaskAndDeletePendingTaskPrivate(SingularityTask task) throws Exception {
    // TODO: Should more of the below be done within a transaction?
    deletePendingTask(task.getTaskRequest().getPendingTask().getPendingTaskId());

    final long now = System.currentTimeMillis();

    String msg = String.format("Task launched because of %s", task.getTaskRequest().getPendingTask().getPendingTaskId().getPendingType().name());

    if (task.getTaskRequest().getPendingTask().getUser().isPresent()) {
      msg = String.format("%s by %s", msg, task.getTaskRequest().getPendingTask().getUser().get());
    }

    if (task.getTaskRequest().getPendingTask().getMessage().isPresent()) {
      msg = String.format("%s (%s)", msg, task.getTaskRequest().getPendingTask().getMessage().get());
    }

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

    try {
      final String path = getTaskPath(task.getTaskId());

      CuratorTransactionFinal transaction = curator.inTransaction().create().forPath(path, taskTranscoder.toBytes(task)).and();

      transaction.create().forPath(getActivePath(task.getTaskId().getId())).and().commit();

      leaderCache.putActiveTask(task);
      taskCache.set(path, task);
    } catch (KeeperException.NodeExistsException nee) {
      LOG.error("Task or active path already existed for {}", task.getTaskId());
    }
  }

  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) {
    if (leaderCache.active()) {
      leaderCache.addKilledTask(killedTaskIdRecord);
    }
    return save(getKilledPath(killedTaskIdRecord.getTaskId()), killedTaskIdRecord, killedTaskIdRecordTranscoder);
  }

  public List getKilledTaskIdRecords() {
    if (leaderCache.active()) {
      return leaderCache.getKilledTasks();
    }
    return fetchKilledTaskIdRecords();
  }

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

  public SingularityDeleteResult deleteKilledRecord(SingularityTaskId taskId) {
    if (leaderCache.active()) {
      leaderCache.deleteKilledTask(taskId);
    }
    return delete(getKilledPath(taskId));
  }

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

  public SingularityCreateResult saveTaskMetadata(SingularityTaskMetadata taskMetadata) {
    return save(getTaskMetadataPath(taskMetadata), taskMetadata, taskMetadataTranscoder);
  }

  public SingularityCreateResult saveTaskShellCommandRequestToQueue(SingularityTaskShellCommandRequest shellRequest) {
    return save(getShellRequestQueuePath(shellRequest), shellRequest, taskShellCommandRequestTranscoder);
  }

  public SingularityCreateResult saveTaskFinishedInMailQueue(SingularityTaskId taskId) {
    return save(getFinishedTaskMailQueuePath(taskId), Optional.absent());
  }

  public List getTaskFinishedMailQueue() {
    return getChildrenAsIds(FINISHED_TASK_MAIL_QUEUE, taskIdTranscoder);
  }

  public SingularityDeleteResult deleteFinishedTaskMailQueue(SingularityTaskId taskId) {
    return delete(getFinishedTaskMailQueuePath(taskId));
  }

  public SingularityCreateResult saveTaskShellCommandRequestToTask(SingularityTaskShellCommandRequest shellRequest) {
    return save(getShellHistoryRequestPath(shellRequest.getId()), shellRequest, taskShellCommandRequestTranscoder);
  }

  public SingularityCreateResult saveTaskShellCommandUpdate(SingularityTaskShellCommandUpdate shellUpdate) {
    return save(getShellHistoryUpdatePath(shellUpdate), shellUpdate, taskShellCommandUpdateTranscoder);
  }

  public List getTaskShellCommandUpdates(SingularityTaskShellCommandRequestId shellRequestId) {
    return getAsyncChildren(getShellHistoryUpdateParentPath(shellRequestId), taskShellCommandUpdateTranscoder);
  }

  public List getTaskShellCommandRequestsForTask(SingularityTaskId taskId) {
    final String parentPath = getShellsParentPath(taskId);
    List children = getChildren(parentPath);
    List paths = Lists.newArrayListWithCapacity(children.size());

    for (String child : children) {
      paths.add(ZKPaths.makePath(parentPath, ZKPaths.makePath(child, SHELL_REQUEST_KEY)));
    }

    List shellRequests = getAsync("getTaskShellCommandRequestsForTask", paths, taskShellCommandRequestTranscoder);

    Collections.sort(shellRequests);

    return shellRequests;
  }

  public List getAllQueuedTaskShellCommandRequests() {
    return getAsyncChildren(SHELL_REQUESTS_QUEUE_PATH_ROOT, taskShellCommandRequestTranscoder);
  }

  public SingularityDeleteResult deleteTaskShellCommandRequestFromQueue(SingularityTaskShellCommandRequest shellRequest) {
    return delete(getShellRequestQueuePath(shellRequest));
  }

  public SingularityCreateResult saveTaskCleanup(SingularityTaskCleanup cleanup) {
    saveTaskHistoryUpdate(cleanup);
    leaderCache.saveTaskCleanup(cleanup);
    return save(getCleanupPath(cleanup.getTaskId().getId()), cleanup, taskCleanupTranscoder);
  }

  private void saveTaskHistoryUpdate(SingularityTaskCleanup cleanup) {
    StringBuilder msg = new StringBuilder(cleanup.getCleanupType().name());

    if (cleanup.getUser().isPresent()) {
      msg.append(" by ");
      msg.append(cleanup.getUser().get());
    }

    if (cleanup.getMessage().isPresent()) {
      msg.append(" - ");
      msg.append(cleanup.getMessage().get());
    }

    saveTaskHistoryUpdate(new SingularityTaskHistoryUpdate(cleanup.getTaskId(), cleanup.getTimestamp(), ExtendedTaskState.TASK_CLEANING, Optional.of(msg.toString()), Optional.absent()), true);
  }

  public SingularityCreateResult createTaskCleanup(SingularityTaskCleanup cleanup) {
    leaderCache.createTaskCleanupIfNotExists(cleanup);
    final SingularityCreateResult result = create(getCleanupPath(cleanup.getTaskId().getId()), cleanup, taskCleanupTranscoder);

    if (result == SingularityCreateResult.CREATED) {
      saveTaskHistoryUpdate(cleanup);
    }

    return result;
  }

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

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

  public void markPendingTaskForDeletion(SingularityPendingTaskId pendingTaskId) {
    leaderCache.markPendingTaskForDeletion(pendingTaskId);
    create(getPendingTasksToDeletePath(pendingTaskId));
  }

  public List getPendingTasksMarkedForDeletion() {
    if (leaderCache.active()) {
      return leaderCache.getPendingTaskIdsToDelete();
    }

    return getChildrenAsIds(PENDING_TASKS_TO_DELETE_PATH_ROOT, pendingTaskIdTranscoder);
  }

  public void deleteCleanupTask(String taskId) {
    leaderCache.deleteTaskCleanup(SingularityTaskId.valueOf(taskId));
    delete(getCleanupPath(taskId));
  }

  public SingularityDeleteResult deleteTaskHistory(SingularityTaskId taskId) {
    taskCache.delete(getTaskPath(taskId));
    if (leaderCache.active()) {
      leaderCache.deleteTaskHistory(taskId);
    }
    return delete(getHistoryPath(taskId));
  }

  public void purgeStaleRequests(List activeRequestIds, long deleteBeforeTime) {
    final List requestIds = getChildren(HISTORY_PATH_ROOT);
    for (String requestId : requestIds) {
      if (!activeRequestIds.contains(requestId)) {
        String path = getRequestPath(requestId);
        Optional maybeStat = checkExists(path);
        if (maybeStat.isPresent() && maybeStat.get().getMtime() < deleteBeforeTime && getChildren(path).size() == 0) {
          delete(path);
        }
      }
    }
  }

  public SingularityDeleteResult deleteRequestId(String requestId) {
    return delete(getRequestPath(requestId));
  }
}