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

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

import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.hubspot.singularity.ExtendedTaskState;
import com.hubspot.singularity.RequestUtilization;
import com.hubspot.singularity.SingularityAgent;
import com.hubspot.singularity.SingularityAgentUsageWithId;
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.SingularityTaskCleanup;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;
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.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@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 agents;
  private Map racks;
  private Set pendingTaskIdsToDelete;
  private Map requestUtilizations;
  private Map agentUsages;

  private volatile boolean active;

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

  public void activate() {
    active = true;
  }

  // Only for unit testing
  public void clear() {
    if (pendingTaskIdToPendingTask != null) {
      pendingTaskIdToPendingTask.clear();
    }
    if (activeTaskIds != null) {
      activeTaskIds.clear();
    }
    if (requests != null) {
      requests.clear();
    }
    if (cleanupTasks != null) {
      cleanupTasks.clear();
    }
    if (requestIdToDeployState != null) {
      requestIdToDeployState.clear();
    }
    if (killedTasks != null) {
      killedTasks.clear();
    }
    if (historyUpdates != null) {
      historyUpdates.clear();
    }
    if (agents != null) {
      agents.clear();
    }
    if (racks != null) {
      racks.clear();
    }
    if (pendingTaskIdsToDelete != null) {
      pendingTaskIdsToDelete.clear();
    }
    if (requestUtilizations != null) {
      requestUtilizations.clear();
    }
    if (agentUsages != null) {
      agentUsages.clear();
    }
  }

  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 cacheAgents(List slaves) {
    this.agents =
      slaves
        .stream()
        .collect(
          Collectors.toConcurrentMap(SingularityAgent::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 void cacheAgentUsages(Map slaveUsages) {
    this.agentUsages = new ConcurrentHashMap<>(slaveUsages);
  }

  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.ofNullable(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(SingularityTaskId taskId) {
    if (!active) {
      LOG.warn("deleteActiveTask {}, but not active", taskId);
      return;
    }

    activeTaskIds.remove(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 getNumActiveTasks(String requestId) {
    return (int) activeTaskIds
      .stream()
      .filter(t -> t.getRequestId().equals(requestId))
      .count();
  }

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

  public int getNumPendingTasks(String requestId) {
    return (int) pendingTaskIdToPendingTask
      .keySet()
      .stream()
      .filter(t -> t.getRequestId().equals(requestId))
      .count();
  }

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

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

    activeTaskIds.add(taskId);
  }

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

  public Optional getRequest(String requestId) {
    return Optional.ofNullable(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 int getNumCleaningTasks(String requestId) {
    return (int) cleanupTasks
      .keySet()
      .stream()
      .filter(t -> t.getRequestId().equals(requestId))
      .count();
  }

  public Optional getTaskCleanup(SingularityTaskId taskId) {
    return Optional.ofNullable(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.ofNullable(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 Optional getKilledTaskRecord(
    SingularityTaskId taskId
  ) {
    return Optional.ofNullable(killedTasks.get(taskId));
  }

  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.ofNullable(historyUpdates.get(taskId)).orElse(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 List getAgents() {
    return new ArrayList<>(agents.values());
  }

  public Optional getAgent(String agentId) {
    return Optional.ofNullable(agents.get(agentId));
  }

  public void putAgent(SingularityAgent agent) {
    if (!active) {
      LOG.warn("putSlave {}, but not active", agent);
    }

    agents.put(agent.getId(), agent);
  }

  public void removeAgent(String agentId) {
    if (!active) {
      LOG.warn("remove agent {}, but not active", agentId);
      return;
    }
    agents.remove(agentId);
  }

  public List getRacks() {
    return new ArrayList<>(racks.values());
  }

  public Optional getRack(String rackId) {
    return Optional.ofNullable(racks.get(rackId));
  }

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

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

  public void removeRack(String rackId) {
    if (!active) {
      LOG.warn("remove rack {}, but not active", rackId);
      return;
    }
    racks.remove(rackId);
  }

  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);
  }

  public void putAgentUsage(SingularityAgentUsageWithId agentUsage) {
    if (!active) {
      LOG.warn("putSlaveUsage {}, but not active", agentUsage);
    }

    agentUsages.put(agentUsage.getAgentId(), agentUsage);
  }

  public void removeAgentUsage(String agentId) {
    if (!active) {
      LOG.warn("removeSlaveUsage {}, but not active", agentId);
      return;
    }
    agentUsages.remove(agentId);
  }

  public Map getAgentUsages() {
    return new HashMap<>(agentUsages);
  }

  public Optional getAgentUsage(String agentId) {
    return Optional.ofNullable(agentUsages.get(agentId));
  }
}