BlendedHistoryHelper.java

package com.hubspot.singularity.data.history;

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.SortedMap;
import java.util.TreeMap;

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

import com.google.common.base.Optional;
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;

public abstract class BlendedHistoryHelper<T, Q> {

  private static final Logger LOG = LoggerFactory.getLogger(BlendedHistoryHelper.class);
  protected abstract List<T> getFromZk(Q id);
  protected abstract List<T> getFromHistory(Q id, int historyStart, int numFromHistory);

  protected abstract Optional<Integer> getTotalCount(Q id);

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

    List<SingularityTaskIdHistory> histories = Lists.newArrayListWithCapacity(taskIds.size());

    for (SingularityTaskId taskId : taskIds) {
      List<SingularityTaskHistoryUpdate> 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<T> getComparator(Q id) {
    throw new IllegalStateException("Comparator requested for query which doesn't implement it");
  }

  public Optional<Integer> getBlendedHistoryCount(Q id) {
    return getTotalCount(id);
  }

  public List<T> getBlendedHistory(Q id, Integer limitStart, Integer limitCount) {
    final List<T> fromZk = getFromZk(id);

    List<T> 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<T> getOrderedFromHistory(Q id, Integer limitStart, Integer limitCount, List<T> fromZk) {
    SortedMap<T, Boolean> returnedMap = new TreeMap<>(getComparator(id));
    for (T item : fromZk) {
      returnedMap.put(item, false);
    }

    int historyLimitStart = 0;
    List<T> 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<T, Boolean> returnedMap, Integer currentStartIndex) {
    int highestHistoryItemIndex = 0;
    int index = 0;
    for (Map.Entry<T, Boolean> entry : returnedMap.entrySet()) {
      if (entry.getValue()) {
        highestHistoryItemIndex = index;
      }
      index ++;
    }
    return currentStartIndex + highestHistoryItemIndex;
  }

  private boolean foundAllFromHistoryAndTrimResults(SortedMap<T, Boolean> returnedMap, Integer currentStartIndex, Integer lastRelevantHistoryItemIndex, Integer limitStart, Integer limitCount, int numFromHistory) {
    boolean foundAllFromHistory = false;
    List<T> toRemove = new ArrayList<>();
    if (numFromHistory == 0 || lastRelevantHistoryItemIndex > limitStart + limitCount) {
      List<T> 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;
  }

}