com.hubspot.singularity.scheduler.SingularityLeaderCache Maven / Gradle / Ivy

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

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Optional;
import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.hubspot.singularity.ExtendedTaskState;
import com.hubspot.singularity.RequestUtilization;
import com.hubspot.singularity.SingularityKilledTaskIdRecord;
import com.hubspot.singularity.SingularityPendingTask;
import com.hubspot.singularity.SingularityPendingTaskId;
import com.hubspot.singularity.SingularityRack;
import com.hubspot.singularity.SingularityRequestDeployState;
import com.hubspot.singularity.SingularityRequestWithState;
import com.hubspot.singularity.SingularitySlave;
import com.hubspot.singularity.SingularityTask;
import com.hubspot.singularity.SingularityTaskCleanup;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;

@Singleton
public class SingularityLeaderCache {

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

  private Map pendingTaskIdToPendingTask;
  private Set activeTaskIds;
  private Map requests;
  private Map cleanupTasks;
  private Map requestIdToDeployState;
  private Map killedTasks;
  private Map> historyUpdates;
  private Map slaves;
  private Map racks;
  private Set pendingTaskIdsToDelete;
  private ConcurrentMap requestUtilizations;

  private volatile boolean active;

  @Inject
  public SingularityLeaderCache() {
    this.active = false;
  }

  public void activate() {
    active = true;
  }

  public void cachePendingTasks(List pendingTasks) {
    this.pendingTaskIdToPendingTask = new ConcurrentHashMap<>(pendingTasks.size());
    pendingTasks.forEach((t) -> pendingTaskIdToPendingTask.put(t.getPendingTaskId(), t));
  }

  public void cachePendingTasksToDelete(List pendingTaskIds) {
    this.pendingTaskIdsToDelete = new HashSet<>(pendingTaskIds.size());
    pendingTaskIdsToDelete.addAll(pendingTaskIds);
  }

  public void cacheActiveTaskIds(List activeTaskIds) {
    this.activeTaskIds = Collections.synchronizedSet(new HashSet(activeTaskIds.size()));
    activeTaskIds.forEach(this.activeTaskIds::add);
  }

  public void cacheRequests(List requestsWithState) {
    this.requests = new ConcurrentHashMap<>(requestsWithState.size());
    requestsWithState.forEach((r) -> requests.put(r.getRequest().getId(), r));
  }

  public void cacheCleanupTasks(List cleanups) {
    this.cleanupTasks = new ConcurrentHashMap<>(cleanups.size());
    cleanups.forEach((c) -> cleanupTasks.put(c.getTaskId(), c));
  }

  public void cacheRequestDeployStates(Map requestDeployStates) {
    this.requestIdToDeployState = new ConcurrentHashMap<>(requestDeployStates.size());
    requestIdToDeployState.putAll(requestDeployStates);
  }

  public void cacheKilledTasks(List killedTasks) {
    this.killedTasks = new ConcurrentHashMap<>(killedTasks.size());
    killedTasks.forEach((k) -> this.killedTasks.put(k.getTaskId(), k));
  }

  public void cacheTaskHistoryUpdates(Map> historyUpdates) {
    this.historyUpdates = new ConcurrentHashMap<>(historyUpdates.size());
    historyUpdates.entrySet().stream().forEach((e) ->
        this.historyUpdates.put(
            e.getKey(),
            e.getValue().stream()
                .collect(Collectors.toMap((u) -> u.getTaskState(), (u) -> u)))
    );
  }

  public void cacheSlaves(List slaves) {
    this.slaves = slaves.stream().collect(Collectors.toConcurrentMap(SingularitySlave::getId, Function.identity()));
  }

  public void cacheRacks(List racks) {
    this.racks = racks.stream().collect(Collectors.toConcurrentMap(SingularityRack::getId, Function.identity()));
  }
  public void stop() {
    active = false;
  }

  public void cacheRequestUtilizations(Map requestUtilizations) {
    this.requestUtilizations = new ConcurrentHashMap<>(requestUtilizations);
  }

  public boolean active() {
    return active;
  }

  public List getPendingTasks() {
    return new ArrayList<>(pendingTaskIdToPendingTask.values());
  }

  public List getPendingTaskIds() {
    return new ArrayList<>(pendingTaskIdToPendingTask.keySet());
  }

  public List getPendingTaskIdsForRequest(String requestId) {
    Set allPendingTaskIds = new HashSet<>(pendingTaskIdToPendingTask.keySet());
    return allPendingTaskIds.stream()
        .filter(t -> t.getRequestId().equals(requestId))
        .collect(Collectors.toList());
  }

  public List getPendingTaskIdsToDelete() {
    synchronized (pendingTaskIdsToDelete) {
      return new ArrayList<>(pendingTaskIdsToDelete);
    }
  }

  public void markPendingTaskForDeletion(SingularityPendingTaskId taskId) {
    pendingTaskIdsToDelete.add(taskId);
  }

  public void deletePendingTask(SingularityPendingTaskId pendingTaskId) {
    if (!active) {
      LOG.warn("deletePendingTask {}, but not active", pendingTaskId);
      return;
    }
    if (pendingTaskIdsToDelete.contains(pendingTaskId)) {
      pendingTaskIdsToDelete.remove(pendingTaskId);
    }
    pendingTaskIdToPendingTask.remove(pendingTaskId);
  }

  public Optional getPendingTask(SingularityPendingTaskId pendingTaskId) {
    return Optional.fromNullable(pendingTaskIdToPendingTask.get(pendingTaskId));
  }

  public void savePendingTask(SingularityPendingTask pendingTask) {
    if (!active) {
      LOG.warn("savePendingTask {}, but not active", pendingTask);
      return;
    }

    pendingTaskIdToPendingTask.put(pendingTask.getPendingTaskId(), pendingTask);
  }

  public void deleteActiveTaskId(String taskId) {
    if (!active) {
      LOG.warn("deleteActiveTask {}, but not active", taskId);
      return;
    }

    activeTaskIds.remove(SingularityTaskId.valueOf(taskId));
  }

  public List exists(List taskIds) {
    List activeTaskIds = new ArrayList<>(taskIds.size());
    for (SingularityTaskId taskId : taskIds) {
      if (this.activeTaskIds.contains(taskId)) {
        activeTaskIds.add(taskId);
      }
    }
    return activeTaskIds;
  }

  public List getActiveTaskIds() {
    return new ArrayList<>(activeTaskIds);
  }

  public List getActiveTaskIdsForRequest(String requestId) {
    Set allActiveTaskIds;
    synchronized (activeTaskIds) {
      allActiveTaskIds = new HashSet<>(activeTaskIds);
    }
    return allActiveTaskIds.stream()
        .filter(t -> t.getRequestId().equals(requestId))
        .collect(Collectors.toList());
  }

  public List getActiveTaskIdsAsStrings() {
    List localActiveTaskIds = getActiveTaskIds();
    List strings = new ArrayList<>(localActiveTaskIds.size());
    for (SingularityTaskId taskId : localActiveTaskIds) {
      strings.add(taskId.getId());
    }
    return strings;
  }

  public List getInactiveTaskIds(List taskIds) {
    List inactiveTaskIds = new ArrayList<>(taskIds.size());
    for (SingularityTaskId taskId : taskIds) {
      if (!activeTaskIds.contains(taskId)) {
        inactiveTaskIds.add(taskId);
      }
    }
    return inactiveTaskIds;
  }

  public int getNumActiveTasks() {
    return activeTaskIds.size();
  }

  public int getNumPendingTasks() {
    return pendingTaskIdToPendingTask.size();
  }

  public boolean isActiveTask(String taskId) {
    return activeTaskIds.contains(SingularityTaskId.valueOf(taskId));
  }

  public void putActiveTask(SingularityTask task) {
    if (!active) {
      LOG.warn("putActiveTask {}, but not active", task.getTaskId());
      return;
    }

    activeTaskIds.add(task.getTaskId());
  }

  public List getRequests() {
    return new ArrayList<>(requests.values());
  }

  public Optional getRequest(String requestId) {
    return Optional.fromNullable(requests.get(requestId));
  }

  public void putRequest(SingularityRequestWithState requestWithState) {
    if (!active) {
      LOG.warn("putRequest {}, but not active", requestWithState.getRequest().getId());
      return;
    }

    requests.put(requestWithState.getRequest().getId(), requestWithState);
  }

  public void deleteRequest(String reqeustId) {
    if (!active) {
      LOG.warn("deleteRequest {}, but not active", reqeustId);
      return;
    }

    requests.remove(reqeustId);
  }

  public List getCleanupTasks() {
    return new ArrayList<>(cleanupTasks.values());
  }

  public List getCleanupTaskIds() {
    return new ArrayList<>(cleanupTasks.keySet());
  }

  public Optional getTaskCleanup(SingularityTaskId taskId) {
    return Optional.fromNullable(cleanupTasks.get(taskId));
  }

  public void deleteTaskCleanup(SingularityTaskId taskId) {
    if (!active) {
      LOG.warn("deleteTaskCleanup {}, but not active", taskId);
      return;
    }

    cleanupTasks.remove(taskId);
  }

  public void saveTaskCleanup(SingularityTaskCleanup cleanup) {
    if (!active) {
      LOG.warn("saveTaskCleanup {}, but not active", cleanup);
      return;
    }

    cleanupTasks.put(cleanup.getTaskId(), cleanup);
  }

  public void createTaskCleanupIfNotExists(SingularityTaskCleanup cleanup) {
    if (!active) {
      LOG.warn("createTaskCleanupIfNotExists {}, but not active", cleanup);
      return;
    }

    cleanupTasks.putIfAbsent(cleanup.getTaskId(), cleanup);
  }

  public Optional getRequestDeployState(String requestId) {
    return Optional.fromNullable(requestIdToDeployState.get(requestId));
  }

  public Map getRequestDeployStateByRequestId() {
    return new HashMap<>(requestIdToDeployState);
  }

  public Map getRequestDeployStateByRequestId(Collection requestIds) {
    Map allDeployStates = new HashMap<>(requestIdToDeployState);
    return allDeployStates.entrySet().stream()
        .filter((e) -> requestIds.contains(e.getKey()))
        .collect(Collectors.toMap((e) -> e.getKey(), (e) -> e.getValue())
    );
  }

  public void deleteRequestDeployState(String requestId) {
    if (!active) {
      LOG.warn("deleteRequestDeployState {}, but not active", requestId);
      return;
    }

    requestIdToDeployState.remove(requestId);
  }

  public void putRequestDeployState(SingularityRequestDeployState requestDeployState) {
    if (!active) {
      LOG.warn("putRequestDeployState {}, but not active", requestDeployState.getRequestId());
      return;
    }

    requestIdToDeployState.put(requestDeployState.getRequestId(), requestDeployState);
  }

  public List getKilledTasks() {
    return new ArrayList<>(killedTasks.values());
  }

  public void addKilledTask(SingularityKilledTaskIdRecord killedTask) {
    if (!active) {
      LOG.warn("addKilledTask {}, but not active", killedTask.getTaskId().getId());
      return;
    }
    killedTasks.put(killedTask.getTaskId(), killedTask);
  }

  public void deleteKilledTask(SingularityTaskId killedTaskId) {
    if (!active) {
      LOG.warn("deleteKilledTask {}, but not active", killedTaskId.getId());
      return;
    }
    killedTasks.remove(killedTaskId);
  }

  public List getTaskHistoryUpdates(SingularityTaskId taskId) {
    List updates = new ArrayList<>(Optional.fromNullable(historyUpdates.get(taskId)).or(new HashMap<>()).values());
    Collections.sort(updates);
    return updates;
  }

  public Map> getTaskHistoryUpdates(Collection taskIds) {
    Map> allHistoryUpdates = new HashMap<>(historyUpdates);
    return allHistoryUpdates.entrySet().stream()
        .filter((e) -> taskIds.contains(e.getKey()))
        .collect(Collectors.toMap(Map.Entry::getKey, (e) -> new ArrayList<>(e.getValue().values()))
    );
  }

  public void saveTaskHistoryUpdate(SingularityTaskHistoryUpdate taskHistoryUpdate, boolean overwrite) {
    if (!active) {
      LOG.warn("saveTaskHistoryUpdate {}, but not active", taskHistoryUpdate);
      return;
    }
    historyUpdates.putIfAbsent(taskHistoryUpdate.getTaskId(), new ConcurrentHashMap<>());
    if (overwrite) {
      historyUpdates.get(taskHistoryUpdate.getTaskId()).put(taskHistoryUpdate.getTaskState(), taskHistoryUpdate);
    } else {
      historyUpdates.get(taskHistoryUpdate.getTaskId()).putIfAbsent(taskHistoryUpdate.getTaskState(), taskHistoryUpdate);
    }
  }

  public void deleteTaskHistoryUpdate(SingularityTaskId taskId, ExtendedTaskState state) {
    if (!active) {
      LOG.warn("deleteTaskHistoryUpdate {}, but not active", taskId);
      return;
    }
    historyUpdates.getOrDefault(taskId, new HashMap<>()).remove(state);
  }

  public void deleteTaskHistory(SingularityTaskId taskId) {
    if (!active) {
      LOG.warn("deleteTaskHistory {}, but not active", taskId);
      return;
    }
    historyUpdates.remove(taskId);
  }

  public Collection getSlaves() {
    return slaves.values();
  }

  public Optional getSlave(String slaveId) {
    return Optional.fromNullable(slaves.get(slaveId));
  }

  public void putSlave(SingularitySlave slave) {
    if (!active) {
      LOG.warn("putSlave {}, but not active", slave);
    }

    slaves.put(slave.getId(), slave);
  }

  public Collection getRacks() {
    return racks.values();
  }

  public Optional getRack(String rackName) {
    return Optional.fromNullable(racks.get(rackName));
  }

  public void putRack(SingularityRack rack) {
    if (!active) {
      LOG.warn("putSlave {}, but not active", rack);
    }

    racks.put(rack.getId(), rack);
  }

  public void putRequestUtilization(RequestUtilization requestUtilization) {
    if (!active) {
      LOG.warn("putRequestUtilization {}, but not active", requestUtilization);
    }

    requestUtilizations.put(requestUtilization.getRequestId(), requestUtilization);
  }

  public void removeRequestUtilization(String requestId) {
    if (!active) {
      LOG.warn("removeRequestUtilization {}, but not active", requestId);
      return;
    }
    requestUtilizations.remove(requestId);
  }

  public Map getRequestUtilizations() {
    return new HashMap<>(requestUtilizations);
  }
}