PriorityBlockingQueue源码解析

前言

	/*
     * The implementation uses an array-based binary heap, with public
     * operations protected with a single lock. However, allocation
     * during resizing uses a simple spinlock (used only while not
     * holding main lock) in order to allow takes to operate
     * concurrently with allocation.  This avoids repeated
     * postponement of waiting consumers and consequent element
     * build-up. The need to back away from lock during allocation
     * makes it impossible to simply wrap delegated
     * java.util.PriorityQueue operations within a lock, as was done
     * in a previous version of this class. To maintain
     * interoperability, a plain PriorityQueue is still used during
     * serialization, which maintains compatibility at the expense of
     * transiently doubling overhead.
     */

以上为源码对PriorityBlockingQueue的解释，从中我们可以得到几点关键信息：

1. PriorityBlockingQueue是一个通过数组实现的堆(完全二叉树)
2. 公有方法通过Lock实现同步
3. 在进行扩容操作时，使用spinlock(自旋锁)

何为堆

在进行源码阅读前，先补充一下基本数据结构的概念，这样对阅读也能起到很大的帮助。堆是一颗完全二叉树，所以堆可以通过数组实现，也不用担心内存浪费的问题。

特点

1. 堆中某个结点的值总是不大于或不小于其父结点的值；
2. 堆总是一棵完全二叉树。

其中根结点最大的堆叫做最大堆，根结点最小的堆叫做最小堆

建堆

以最小堆为例，数据【3，1，6，2，5，7，4】如何建堆

首先从树的底部找到第一个非叶子结点，即6，然后将其和左子结点和右子结点比较，并进行交换操作：

紧接着选择下一个非叶子结点，即1，进行同样的操作：(此时1已经是最小，不用进行交换)

最后处理根节点3，与1交换后，仍要继续和2交换，确保3下沉以后，子树仍然是最小堆

至此一个最小堆建立完成。

总结下来流程为：从树的尾部开始，依次对非叶子结点进行下沉操作。

插入

向上述例子插入值为0的结点，操作流程为首先将待插入的结点添加至堆尾部，然后对此结点进行上浮操作。

删除

对于堆的删除操作，首先将堆顶元素弹出(置空)，然后将最后一个结点放入根结点，最后对根结点进行下沉操作。

源码分析

初始化

// 默认队列大小
private static final int DEFAULT_INITIAL_CAPACITY = 11;

//最大容量
private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;

//数组，堆
private transient Object[] queue;

//比较器，如果不传，则对象应该实现Comparable接口
private transient Comparator<? super E> comparator;

//实现同步使用的锁
private final ReentrantLock lock = new ReentrantLock();

//扩容时用的自旋锁
private transient volatile int allocationSpinLock;

以上为PriorityBlockingQueue较为重要的属性，接下来我们看构造方法

// 普通构造函数最终执行到此
public PriorityBlockingQueue(int initialCapacity, Comparator<? super E> comparator) {
    if (initialCapacity < 1)
        throw new IllegalArgumentException();
    this.comparator = comparator;
    this.queue = new Object[Math.max(1, initialCapacity)];
}

//如果通过其他集合构造，则走到这里
public PriorityBlockingQueue(Collection<? extends E> c) {
    boolean heapify = true; // true if not known to be in heap order
    boolean screen = true;  // true if must screen for nulls
    if (c instanceof SortedSet<?>) {
        SortedSet<? extends E> ss = (SortedSet<? extends E>) c;
        this.comparator = (Comparator<? super E>) ss.comparator();
        heapify = false;
    }
    else if (c instanceof PriorityBlockingQueue<?>) {
        PriorityBlockingQueue<? extends E> pq =
            (PriorityBlockingQueue<? extends E>) c;
        this.comparator = (Comparator<? super E>) pq.comparator();
        screen = false;
        if (pq.getClass() == PriorityBlockingQueue.class) // exact match
            heapify = false;
    }
    Object[] es = c.toArray();
    int n = es.length;
    // Android-changed: Defend against c.toArray (incorrectly) not returning Object[]
    //                  (see b/204397945)
    // if (c.getClass() != java.util.ArrayList.class)
    if (es.getClass() != Object[].class)
        es = Arrays.copyOf(es, n, Object[].class);
    if (screen && (n == 1 || this.comparator != null)) {
        for (Object e : es)
            if (e == null)
                throw new NullPointerException();
    }
    this.queue = ensureNonEmpty(es);
    this.size = n;
    if (heapify)
        heapify();
}

这里主要关注一下当我们想通过其他Collection集合构造PriorityBlockingQueue时的行为。

1. 判断是否为SortedSet 和 PriorityBlockingQueue，如果是的话，直接获取想要的属性然后Arrays.copyOf完事
2. 如果不是，则需要通过heapify();进行堆化，也就是用这些数据重新建堆

heapify()

private void heapify() {
    final Object[] es = queue;
    int n = size, i = (n >>> 1) - 1;
    final Comparator<? super E> cmp;
    if ((cmp = comparator) == null)
        for (; i >= 0; i--)
            siftDownComparable(i, (E) es[i], es, n);
    else
        for (; i >= 0; i--)
            siftDownUsingComparator(i, (E) es[i], es, n, cmp);
}

这个方法从尾部的第一个非叶子结点开始向根结点开始遍历，并判断是否初始化comparator选择下沉函数，这里我们随便选择一个进行深入。

siftDownComparable

private static <T> void siftDownComparable(int k, T x, Object[] es, int n) {
    // assert n > 0;
    Comparable<? super T> key = (Comparable<? super T>)x;
    int half = n >>> 1;           // loop while a non-leaf
    while (k < half) {
        int child = (k << 1) + 1; // assume left child is least
        Object c = es[child];
        int right = child + 1;
        if (right < n &&
            ((Comparable<? super T>) c).compareTo((T) es[right]) > 0)
            c = es[child = right];
        if (key.compareTo((T) c) <= 0)
            break;
        es[k] = c;
        k = child;
    }
    es[k] = key;
}

这个方法就是下沉操作的具体实现，整个流程和我们上面介绍的建堆流程一样，大家可以参照着阅读一下。

入队

PriorityBlockingQueue的插入包括put、add、offer

put

public void put(E e) {
   	offer(e); // never need to block
}

put方法实则调用了offer。

大家注意这里还有一句注释never need to block（不需要阻塞）这是为什么呢？

这是因为put方法是来自接口BlockingQueue，在这个接口上put的定义是阻塞插入，如果队列满了，则会等待。
然而PriorityBlockingQueue底层是通过无界数组实现的，能够扩容所以这里不会阻塞

add

public boolean add(E e) {
    return offer(e);
}

add方法也是调用了offer， 不同就是加上了返回值。但是同样也是因为底层通过无界数组实现的原因，目前我的水平还没看出这个返回值的意义。

offer

终于到了关键代码部分了

public boolean offer(E e) {
    if (e == null)
        throw new NullPointerException();
    final ReentrantLock lock = this.lock;
    // 1. 先上锁
    lock.lock();
    int n, cap;
    Object[] es;
    while ((n = size) >= (cap = (es = queue).length))
    	//2. 扩容
        tryGrow(es, cap);
    try {
        final Comparator<? super E> cmp;
        // 3. 上浮操作
        if ((cmp = comparator) == null)
            siftUpComparable(n, e, es);
        else
            siftUpUsingComparator(n, e, es, cmp);
        size = n + 1;
        notEmpty.signal();
    } finally {
        lock.unlock();
    }
    return true;
}

三个关键点已经在注释中标注出来

tryGrow

private void tryGrow(Object[] array, int oldCap) {
	//1. 释放锁，切换allocationSpinLock自旋锁
    lock.unlock(); // must release and then re-acquire main lock
    Object[] newArray = null;
    if (allocationSpinLock == 0 &&
        ALLOCATIONSPINLOCK.compareAndSet(this, 0, 1)) {
        try {
        	//2. 扩容规则
            int newCap = oldCap + ((oldCap < 64) ?
                                   (oldCap + 2) : // grow faster if small
                                   (oldCap >> 1));
            if (newCap - MAX_ARRAY_SIZE > 0) {    // possible overflow
                int minCap = oldCap + 1;
                if (minCap < 0 || minCap > MAX_ARRAY_SIZE)
                    throw new OutOfMemoryError();
                newCap = MAX_ARRAY_SIZE;
            }
            if (newCap > oldCap && queue == array)
                newArray = new Object[newCap];
        } finally {
            allocationSpinLock = 0;
        }
    }
    //3. 如果其他线程进行插入操作，则释放yield
    if (newArray == null) // back off if another thread is allocating
        Thread.yield();
        
    //4. System.arraycopy时还是要重新上锁
    lock.lock();
    if (newArray != null && queue == array) {
        queue = newArray;
        System.arraycopy(array, 0, newArray, 0, oldCap);
    }
}

tryGrow对数组进行扩容，当容量小于64时，一次扩容大小为2，当大于64时直接扩容2倍。

并且在申请内存时，会释放锁，改用ALLOCATIONSPINLOCK，其作用就和开头的注释写的一样，为了出队和申请内存能同时进行。

在申请内存时，如果有其他线程也想要插入元素，则会走到Thread.yield()释放当前cpu使用权重新竞争。很大程度（也不一定）能让原先申请内存的线程再申请完成后再次获取到cpu使用权，完成剩余的插入操作。

siftUpComparable

private static <T> void siftUpComparable(int k, T x, Object[] es) {
    Comparable<? super T> key = (Comparable<? super T>) x;
    while (k > 0) {
        int parent = (k - 1) >>> 1;
        Object e = es[parent];
        if (key.compareTo((T) e) >= 0)
            break;
        es[k] = e;
        k = parent;
    }
    es[k] = key;
}

这里就是上浮操作，大家参照我们上面讲的插入流程阅读即可

出队

出队操作包括：poll、peek、remove、take

这几个操作最终都是调用的dequeue()不同的是出队的表现

poll

public E poll() {
    final ReentrantLock lock = this.lock;
    lock.lock();
    try {
        return dequeue();
    } finally {
        lock.unlock();
    }
}

poll方法通过lock加锁后直接返回dequeue

peek

public E peek() {
    final ReentrantLock lock = this.lock;
    lock.lock();
    try {
        return (E) queue[0];
    } finally {
        lock.unlock();
    }
}

peek返回队首元素，但是不出队，如果队列为空则返回null

remove

public boolean remove(Object o) {
    final ReentrantLock lock = this.lock;
    lock.lock();
    try {
        int i = indexOf(o);
        if (i == -1)
            return false;
        removeAt(i);
        return true;
    } finally {
        lock.unlock();
    }
}

remove先通过indexOf找到对应object所在下标，再通过removeAt移除。indexOf通过equals方法做比较，所以想通过remove进行出队，应该重写equals方法

private void removeAt(int i) {
    final Object[] es = queue;
    final int n = size - 1;
    if (n == i) // removed last element
        es[i] = null;
    else {
        E moved = (E) es[n];
        es[n] = null;
        final Comparator<? super E> cmp;
        if ((cmp = comparator) == null)
            siftDownComparable(i, moved, es, n);
        else
            siftDownUsingComparator(i, moved, es, n, cmp);
        if (es[i] == moved) {
            if (cmp == null)
                siftUpComparable(i, moved, es);
            else
                siftUpUsingComparator(i, moved, es, cmp);
        }
    }
    size = n;
}

这里首先先将最后一个结点A覆盖到第i个结点B，然后进行下沉操作，当下沉完毕以后。

如果此时A无需下沉， 也就是(es[i] == moved)，则需要对A进行上浮操作。

就比如一个最小堆，如果最后一个元素是0（最小），此时他和中间某个值进行覆盖操作，在下沉操作完成后，它实际上应该是最上层根节点，所以还要进行上浮操作。

take

public E take() throws InterruptedException {
    final ReentrantLock lock = this.lock;
    lock.lockInterruptibly();
    E result;
    try {
        while ( (result = dequeue()) == null)
            notEmpty.await();
    } finally {
        lock.unlock();
    }
    return result;
}

take操作进行阻塞出队操作，如果队列为空，则一直等待。

dequeue()

private E dequeue() {
    // assert lock.isHeldByCurrentThread();
    final Object[] es;
    final E result;

    if ((result = (E) ((es = queue)[0])) != null) {
        final int n;
        //取出最后一个元素
        final E x = (E) es[(n = --size)];
        //置空最后一个元素
        es[n] = null;
        if (n > 0) {
            final Comparator<? super E> cmp;
            if ((cmp = comparator) == null)
                siftDownComparable(0, x, es, n);
            else
                siftDownUsingComparator(0, x, es, n, cmp);
        }
    }
    return result;
}

终于到了具体的出队操作，这里也是参照上面讲的出队流程阅读即可了，就是一个下沉操作

总结

PriorityBlockingQueue源码并不是很复杂，但是其中涉及到的知识点却非常多，我觉得很值得仔细阅读一番。

最后它的名字说明了一切：优先级阻塞队列