io.trino.execution.resourcegroups.StochasticPriorityQueue Maven / Gradle / Ivy
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.trino.execution.resourcegroups;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ThreadLocalRandom;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
final class StochasticPriorityQueue
implements UpdateablePriorityQueue
{
private final Map> 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 root;
@Override
public boolean addOrUpdate(E element, long priority)
{
checkArgument(priority > 0, "priority must be positive");
if (root == null) {
root = new Node<>(Optional.empty(), element);
root.setTickets(priority);
index.put(element, root);
return true;
}
Node node = index.get(element);
if (node != null) {
node.setTickets(priority);
return false;
}
node = root.addNode(element, priority);
index.put(element, node);
return true;
}
@Override
public boolean contains(E element)
{
return index.containsKey(element);
}
@Override
public boolean remove(E element)
{
Node node = index.remove(element);
if (node == null) {
return false;
}
if (node.isLeaf() && node.equals(root)) {
root = null;
}
else if (node.isLeaf()) {
node.remove();
}
else {
// This is an intermediate node. Instead of removing it directly, remove a leaf
// and then replace the data in this node with the data in the leaf. This way
// we don't have to reorganize the tree structure.
Node leaf = root.findLeaf();
leaf.remove();
node.setTickets(leaf.getTickets());
node.setValue(leaf.getValue());
index.put(leaf.getValue(), node);
}
return true;
}
@Override
public E poll()
{
if (root == null) {
return null;
}
long winningTicket = ThreadLocalRandom.current().nextLong(root.getTotalTickets());
Node candidate = root;
while (!candidate.isLeaf()) {
long leftTickets = candidate.getLeft().map(Node::getTotalTickets).orElse(0L);
if (winningTicket < leftTickets) {
candidate = candidate.getLeft().get();
continue;
}
winningTicket -= leftTickets;
if (winningTicket < candidate.getTickets()) {
break;
}
winningTicket -= candidate.getTickets();
checkState(candidate.getRight().isPresent(), "Expected right node to contain the winner, but it does not exist");
candidate = candidate.getRight().get();
}
checkState(winningTicket < candidate.getTickets(), "Inconsistent winner");
E value = candidate.getValue();
remove(value);
return value;
}
@Override
public E peek()
{
throw new UnsupportedOperationException();
}
@Override
public int size()
{
return index.size();
}
@Override
public boolean isEmpty()
{
return index.isEmpty();
}
@Override
public Iterator iterator()
{
// Since poll() is not deterministic ordering is not required
return index.keySet().iterator();
}
private static final class Node
{
private Optional> parent;
private E value;
private Optional> left = Optional.empty();
private Optional> right = Optional.empty();
private long tickets;
private long totalTickets;
private int descendants;
private Node(Optional> parent, E value)
{
this.parent = parent;
this.value = value;
}
public E getValue()
{
return value;
}
public void setValue(E value)
{
this.value = value;
}
public Optional> getLeft()
{
return left;
}
public Optional> getRight()
{
return right;
}
public long getTotalTickets()
{
return totalTickets;
}
public long getTickets()
{
return tickets;
}
public void setTickets(long tickets)
{
checkArgument(tickets > 0, "tickets must be positive");
if (tickets == this.tickets) {
return;
}
long ticketDelta = tickets - this.tickets;
Node node = this;
// Update total tickets in this node and all ancestors
while (node != null) {
node.totalTickets += ticketDelta;
node = node.parent.orElse(null);
}
this.tickets = tickets;
}
public boolean isLeaf()
{
return left.isEmpty() && right.isEmpty();
}
public Node findLeaf()
{
int leftDecendants = left.map(node -> node.descendants).orElse(0);
int rightDecendants = right.map(node -> node.descendants).orElse(0);
if (leftDecendants == 0 && rightDecendants == 0) {
return left.orElse(right.orElse(this));
}
if (leftDecendants > rightDecendants) {
return left.get().findLeaf();
}
if (rightDecendants > leftDecendants) {
return right.get().findLeaf();
}
// For ties just go left
checkState(left.isPresent(), "Left child missing");
return left.get().findLeaf();
}
public void remove()
{
checkState(parent.isPresent(), "Cannot remove root node");
checkState(isLeaf(), "Can only remove leaf nodes");
Node parent = this.parent.get();
if (parent.getRight().map(node -> node.equals(this)).orElse(false)) {
parent.right = Optional.empty();
}
else {
checkState(parent.getLeft().map(node -> node.equals(this)).orElse(false), "Inconsistent parent pointer");
parent.left = Optional.empty();
}
while (parent != null) {
parent.descendants--;
parent.totalTickets -= tickets;
parent = parent.parent.orElse(null);
}
this.parent = Optional.empty();
}
public Node addNode(E value, long tickets)
{
// setTickets call in base case will update totalTickets
descendants++;
if (left.isPresent() && right.isPresent()) {
// Keep the tree balanced when inserting
if (left.get().descendants < right.get().descendants) {
return left.get().addNode(value, tickets);
}
return right.get().addNode(value, tickets);
}
Node child = new Node<>(Optional.of(this), value);
if (left.isPresent()) {
right = Optional.of(child);
}
else {
left = Optional.of(child);
}
child.setTickets(tickets);
return child;
}
}
}