com.hubspot.singularity.data.history.BlendedHistoryHelper Maven / Gradle / Ivy

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

import com.google.common.collect.Lists;
import com.hubspot.singularity.SingularityTask;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;
import com.hubspot.singularity.SingularityTaskIdHistory;
import com.hubspot.singularity.data.TaskManager;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.SortedMap;
import java.util.TreeMap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public abstract class BlendedHistoryHelper {
  private static final Logger LOG = LoggerFactory.getLogger(BlendedHistoryHelper.class);

  protected final boolean sqlEnabled;

  public BlendedHistoryHelper(boolean sqlEnabled) {
    this.sqlEnabled = sqlEnabled;
  }

  protected abstract List getFromZk(Q id);

  protected abstract List getFromHistory(Q id, int historyStart, int numFromHistory);

  protected abstract Optional getTotalCount(Q id, boolean canSkipZk);

  public List getTaskHistoriesFor(
    TaskManager taskManager,
    Collection taskIds
  ) {
    Map tasks = taskManager.getTasks(taskIds);
    Map> map = taskManager.getTaskHistoryUpdates(
      taskIds
    );

    List histories = Lists.newArrayListWithCapacity(
      taskIds.size()
    );

    for (SingularityTaskId taskId : taskIds) {
      List historyUpdates = map.get(taskId);
      SingularityTask task = tasks.get(taskId);
      if (task != null) {
        histories.add(
          SingularityTaskIdHistory.fromTaskIdAndTaskAndUpdates(
            taskId,
            task,
            historyUpdates
          )
        );
      }
    }

    Collections.sort(histories);
    return histories;
  }

  protected boolean queryUsesZkFirst(Q id) {
    return true;
  }

  protected Comparator getComparator(Q id) {
    throw new IllegalStateException(
      "Comparator requested for query which doesn't implement it"
    );
  }

  public Optional getBlendedHistoryCount(Q id, boolean canSkipZk) {
    return getTotalCount(id, canSkipZk);
  }

  public List getBlendedHistory(Q id, Integer limitStart, Integer limitCount) {
    return getBlendedHistory(id, limitStart, limitCount, false);
  }

  public List getBlendedHistory(
    Q id,
    Integer limitStart,
    Integer limitCount,
    boolean canSkipZk
  ) {
    final List fromZk;
    if (sqlEnabled && canSkipZk) {
      fromZk = new ArrayList<>();
    } else {
      fromZk = getFromZk(id);
    }

    List returned = null;

    if (queryUsesZkFirst(id)) {
      final int numFromZk = Math.max(0, Math.min(limitCount, fromZk.size() - limitStart));

      final Integer numFromHistory = limitCount - numFromZk;
      final Integer historyStart = Math.max(0, limitStart - fromZk.size());

      returned = Lists.newArrayListWithCapacity(limitCount);

      if (numFromZk > 0) {
        returned.addAll(fromZk.subList(limitStart, limitStart + numFromZk));
      }

      if (numFromHistory > 0) {
        returned.addAll(getFromHistory(id, historyStart, numFromHistory));
      }
    } else {
      returned = getOrderedFromHistory(id, limitStart, limitCount, fromZk);
    }

    return returned;
  }

  private List getOrderedFromHistory(
    Q id,
    Integer limitStart,
    Integer limitCount,
    List fromZk
  ) {
    SortedMap returnedMap = new TreeMap<>(getComparator(id));
    for (T item : fromZk) {
      returnedMap.put(item, false);
    }

    int historyLimitStart = 0;
    List fromHistory = getFromHistory(id, historyLimitStart, limitCount);
    for (T item : fromHistory) {
      returnedMap.put(item, true);
    }

    int currentStartIndex = 0;
    while (
      !foundAllFromHistoryAndTrimResults(
        returnedMap,
        currentStartIndex,
        getLastRelevantHistoryItemIndex(returnedMap, currentStartIndex),
        limitStart,
        limitCount,
        fromHistory.size()
      )
    ) {
      if (returnedMap.isEmpty()) {
        return getFromHistory(id, limitStart - currentStartIndex, limitCount);
      } else {
        historyLimitStart += limitCount;
        fromHistory = getFromHistory(id, historyLimitStart, limitCount);
        for (T item : fromHistory) {
          returnedMap.put(item, true);
        }
      }
    }
    return new ArrayList<>(returnedMap.keySet());
  }

  private int getLastRelevantHistoryItemIndex(
    SortedMap returnedMap,
    Integer currentStartIndex
  ) {
    int highestHistoryItemIndex = 0;
    int index = 0;
    for (Map.Entry entry : returnedMap.entrySet()) {
      if (entry.getValue()) {
        highestHistoryItemIndex = index;
      }
      index++;
    }
    return currentStartIndex + highestHistoryItemIndex;
  }

  private boolean foundAllFromHistoryAndTrimResults(
    SortedMap returnedMap,
    Integer currentStartIndex,
    Integer lastRelevantHistoryItemIndex,
    Integer limitStart,
    Integer limitCount,
    int numFromHistory
  ) {
    boolean foundAllFromHistory = false;
    List toRemove = new ArrayList<>();
    if (numFromHistory == 0 || lastRelevantHistoryItemIndex > limitStart + limitCount) {
      List current = new ArrayList<>(returnedMap.keySet());
      toRemove.addAll(
        current.subList(0, Math.min(limitStart - currentStartIndex, current.size()))
      );
      toRemove.addAll(
        current.subList(
          Math.min(limitStart - currentStartIndex + limitCount, current.size()),
          current.size()
        )
      );
      foundAllFromHistory = true;
    } else {
      toRemove =
        toRemove.subList(
          0,
          Math.min(
            Math.min(lastRelevantHistoryItemIndex, limitStart - currentStartIndex),
            toRemove.size()
          )
        );
      currentStartIndex += toRemove.size();
    }
    for (T item : toRemove) {
      returnedMap.remove(item);
    }

    if (!foundAllFromHistory) {
      LOG.trace(
        "Current start index is {}, querying for more history",
        currentStartIndex
      );
    }

    return foundAllFromHistory;
  }
}