【Presto源码学习】ResourceGroups调度策略

private Queue<InternalResourceGroup> eligibleSubGroups = new FifoQueue<>();
private TieredQueue<ManagedQueryExecution> queuedQueries = new TieredQueue<>(FifoQueue::new);

submit(InternalResourceGroupManager.java):151
-createGroupIfNecessary(InternalResourceGroupManager.java)
submit(InternalResourceGroupManager.java):152
-run(InternalResourceGroup.java):339
--enqueueQuery(InternalResourceGroup.java):339

final class TieredQueue<E> implements UpdateablePriorityQueue<E> {
  private final UpdateablePriorityQueue<E> highPriorityQueue;
  private final UpdateablePriorityQueue<E> lowPriorityQueue;
  
  public boolean addOrUpdate(E element, long priority) {
    return lowPriorityQueue.addOrUpdate(element, priority);
  }
  
  public boolean prioritize(E element, long priority) {
    return highPriorityQueue.addOrUpdate(element, priority);
  }

//InternalResourceGroup.enqueueQuery()
int priority = getQueryPriority(query.getSession());
if (query.isRetry()) {
  queuedQueries.prioritize(query, priority);
} else {
  queuedQueries.addOrUpdate(query, priority);
}

queryFinished(InternalResourceGroup.java):797
-updateEligibility(InternalResourceGroup.java):735
--isEligibleToStartNext(InternalResourceGroup.java)
-updateEligibility(InternalResourceGroup.java):736
--addOrUpdateSubGroup(InternalResourceGroup.java):898
---addOrUpdateSubGroup(InternalResourceGroup.java)

private boolean isEligibleToStartNext() {
    synchronized (root) {
      if (!canRunMore()) {
        return false;
      }
      return !queuedQueries.isEmpty() || !eligibleSubGroups.isEmpty();
    }
}

//sub group
{
  "name": "pipeline_${USER}",
  "softMemoryLimit": "50%",
  "hardConcurrencyLimit": 5,
  "maxQueued": 100
}
//selector
{
  "source": ".*pipeline.*",
  "group": "global.pipeline.pipeline_${USER}"
}

submit(InternalResourceGroupManager.java):151
-createGroupIfNecessary(InternalResourceGroupManager.java):311
--createGroupIfNecessary(InternalResourceGroupManager.java)

//InternalResourceGroupManager.createGroupIfNecessary()
private synchronized void createGroupIfNecessary(SelectionContext<C> context, Executor executor) {
  ResourceGroupId id = context.getResourceGroupId();
  if (!groups.containsKey(id)) {
    InternalResourceGroup group;
    if (id.getParent().isPresent()) {
      createGroupIfNecessary(new SelectionContext<>(id.getParent().get(), context.getContext()), executor);
      //省略其余代码

submit(InternalResourceGroupManager.java):151
-createGroupIfNecessary(InternalResourceGroupManager.java):339
--configure(FileResourceGroupConfigurationManager.java):108
---configureGroup(AbstractResourceConfigurationManager.java):209
----setSchedulingPolicy(InternalResourceGroup.java):545

//InternalResourceGroup.setSchedulingPolicy()
Queue<InternalResourceGroup> queue;
TieredQueue<ManagedQueryExecution> queryQueue;
switch (policy) {
  case FAIR:
    queue = new FifoQueue<>();
    queryQueue = new TieredQueue<>(FifoQueue::new);
    break;
  case WEIGHTED:
    queue = new StochasticPriorityQueue<>();
    queryQueue = new TieredQueue<>(StochasticPriorityQueue::new);
    break;
  case WEIGHTED_FAIR:
    queue = new WeightedFairQueue<>();
    queryQueue = new TieredQueue<>(IndexedPriorityQueue::new);
    break;
  case QUERY_PRIORITY:
    // Sub groups must use query priority to ensure ordering
    for (InternalResourceGroup group : subGroups.values()) {
      group.setSchedulingPolicy(QUERY_PRIORITY);
    }
    queue = new IndexedPriorityQueue<>();
    queryQueue = new TieredQueue<>(IndexedPriorityQueue::new);
    break;
}

final class FifoQueue<E> implements UpdateablePriorityQueue<E> {
  private final Set<E> delegate = new LinkedHashSet<>();

  public boolean addOrUpdate(E element, long priority) {
    return delegate.add(element);
  }

final class StochasticPriorityQueue<E> implements UpdateablePriorityQueue<E> {
  private final Map<E, Node<E>> index = new HashMap<>();

  // This is a Fenwick tree, where each node has weight equal to the sum of its weight
  // and all its children's weights
  private Node<E> root;

private static final class Node<E> {
  private Optional<Node<E>> parent;
  private E value;
  private Optional<Node<E>> left = Optional.empty();
  private Optional<Node<E>> right = Optional.empty();
  private long tickets;
  private long totalTickets;
  private int descendants;

weighted: queued queries are selected stochastically in proportion to their priority, specified via the query_priority session property. Sub groups are selected to start new queries in proportion to their schedulingWeight.

final class IndexedPriorityQueue<E> implements UpdateablePriorityQueue<E> {
  private final Map<E, Entry<E>> index = new HashMap<>();
  private final Set<Entry<E>> queue = new TreeSet<>((entry1, entry2) -> {
    int priorityComparison = Long.compare(entry2.getPriority(), entry1.getPriority());
    if (priorityComparison != 0) {
      return priorityComparison;
    }
    return Long.compare(entry1.getGeneration(), entry2.getGeneration());
  });

  private long generation;
  
private static final class Entry<E> {
  private final E value;
  private final long priority;
  private final long generation;

private void addOrUpdateSubGroup(Queue<InternalResourceGroup> queue, InternalResourceGroup group) {
  if (schedulingPolicy == WEIGHTED_FAIR) {
    ((WeightedFairQueue<InternalResourceGroup>) queue).addOrUpdate(group, new Usage(group.getSchedulingWeight(), group.getRunningQueries()));
  }

//WeightedFairQueue.java
public static class Usage {
  private final int share;
  private final int utilization;

final class WeightedFairQueue<E> implements Queue<E> {
  private final Map<E, Node<E>> index = new LinkedHashMap<>();
  private long currentLogicalTime;
  
private static final class Node<E> implements Comparable<Node<E>> {
  private final E value;
  private final long logicalCreateTime;
  private Usage usage;

weighted_fair: sub-groups are selected based on their schedulingWeight and the number of queries they are already running concurrently. The expected share of running queries for a sub-group is computed based on the weights for all currently eligible sub-groups. The sub-group with the least concurrency relative to its share is selected to start the next query.

set session query_priority=10;

String url = "jdbc:presto://ip:host/hive/tpch_parquet";
Properties properties = new Properties();
properties.setProperty("user", "presto");
Connection connection = DriverManager.getConnection(url, properties);
PrestoConnection pc = connection.unwrap(PrestoConnection.class);
//设置client tags和source
pc.setClientInfo("ClientTags", "test,dev");
pc.setClientInfo("ApplicationName", "app");
//设置session property，这里就是设置了priority
pc.setSessionProperty("query_priority", "5");

private void addOrUpdateSubGroup(Queue<InternalResourceGroup> queue, InternalResourceGroup group) {
  if (schedulingPolicy == WEIGHTED_FAIR) {
    //省略部分代码
  } else {
    ((UpdateablePriorityQueue<InternalResourceGroup>) queue).addOrUpdate(group, getSubGroupSchedulingPriority(schedulingPolicy, group));
  }
}

private static long getSubGroupSchedulingPriority(SchedulingPolicy policy, InternalResourceGroup group) {
  if (policy == QUERY_PRIORITY) {
    return group.getHighestQueryPriority();
  } else {
    return group.computeSchedulingWeight();
  }
}

private long computeSchedulingWeight() {
  if (runningQueries.size() + descendantRunningQueries >= softConcurrencyLimit) {
    return schedulingWeight;
  }

  return (long) Integer.MAX_VALUE * schedulingWeight;
}