io.prestosql.execution.resourcegroups.IndexedPriorityQueue Maven / Gradle / Ivy

Go to download
/*
 * 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.prestosql.execution.resourcegroups;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Iterators.transform;
import static java.util.Objects.requireNonNull;

/**
 * A priority queue with constant time contains(E) and log time remove(E)
 * Ties are broken by insertion order
 */
public final class IndexedPriorityQueue
        implements UpdateablePriorityQueue
{
    private final Map> index = new HashMap<>();
    private final Set> 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;

    @Override
    public boolean addOrUpdate(E element, long priority)
    {
        Entry entry = index.get(element);
        if (entry != null) {
            queue.remove(entry);
            Entry newEntry = new Entry<>(element, priority, entry.getGeneration());
            queue.add(newEntry);
            index.put(element, newEntry);
            return false;
        }
        Entry newEntry = new Entry<>(element, priority, generation);
        generation++;
        queue.add(newEntry);
        index.put(element, newEntry);
        return true;
    }

    @Override
    public boolean contains(E element)
    {
        return index.containsKey(element);
    }

    @Override
    public boolean remove(E element)
    {
        Entry entry = index.remove(element);
        if (entry != null) {
            queue.remove(entry);
            return true;
        }
        return false;
    }

    @Override
    public E poll()
    {
        Iterator> iterator = queue.iterator();
        if (!iterator.hasNext()) {
            return null;
        }
        Entry entry = iterator.next();
        iterator.remove();
        checkState(index.remove(entry.getValue()) != null, "Failed to remove entry from index");
        return entry.getValue();
    }

    @Override
    public E peek()
    {
        Iterator> iterator = queue.iterator();
        if (!iterator.hasNext()) {
            return null;
        }
        Entry entry = iterator.next();
        return entry.getValue();
    }

    @Override
    public int size()
    {
        return queue.size();
    }

    @Override
    public boolean isEmpty()
    {
        return queue.isEmpty();
    }

    @Override
    public Iterator iterator()
    {
        return transform(queue.iterator(), Entry::getValue);
    }

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

        private Entry(E value, long priority, long generation)
        {
            this.value = requireNonNull(value, "value is null");
            this.priority = priority;
            this.generation = generation;
        }

        public E getValue()
        {
            return value;
        }

        public long getPriority()
        {
            return priority;
        }

        public long getGeneration()
        {
            return generation;
        }
    }
}