無界阻塞佇列 PriorityBlockingQueue

介紹

PriorityBlockingQueue 是一個帶有優先順序的無界阻塞佇列，每次出隊返回的都是優先順序最高或者最低的元素。在內部是使用平衡二叉樹堆實現，所以遍歷元素不保證有序。

預設使用物件的 compareTo 方法進行比較，如果需要自定義比較規則可以自定義 comparators。

該類圖可以看到，PriorityBlockingQueue 內部有一個陣列 queue，用來存放佇列元素；size 用來存放元素個數；allocationSpinLock 是個自旋鎖，使用CAS操作來保證同時只有一個執行緒來進行擴容佇列，狀態只有 0 和 1，0表示當前沒有進行擴容，1表示正在擴容。由於是優先順序佇列，所以有一個比較器 comparator 用來比較大小，另外還有 lock 獨佔鎖，notEmpty 條件變數來實現 take 方法的阻塞，由於是無界佇列所以沒有 notFull 條件變數，所以 put 是非阻塞的。

//二叉樹最小堆的實現
private transient Object[] queue;
private transient int size;
private transient volatile int allocationSpinLock;
private transient Comparator<? super E> comparator;
private final ReentrantLock lock;
private final Condition notEmpty;

在建構函式中，預設佇列容量為11，預設比較器為 null，也就是預設使用元素的 compareTo 方法來確定優先順序，所以佇列元素必須實現 Comparable 介面。

private static final int DEFAULT_INITIAL_CAPACITY = 11;
public PriorityBlockingQueue() {
    this(DEFAULT_INITIAL_CAPACITY, null);
}

public PriorityBlockingQueue(int initialCapacity) {
    this(initialCapacity, null);
}

public PriorityBlockingQueue(int initialCapacity,
                             Comparator<? super E> comparator) {
    if (initialCapacity < 1)
        throw new IllegalArgumentException();
    this.lock = new ReentrantLock();
    this.notEmpty = lock.newCondition();
    this.comparator = comparator;
    this.queue = new Object[initialCapacity];
}

offer 操作

offer 操作的作用是在佇列中插入一個元素，由於是無界佇列，所以一直返回 true。

public boolean offer(E e) {
    if (e == null)
        throw new NullPointerException();
    final ReentrantLock lock = this.lock;
    lock.lock();
    int n, cap;
    Object[] array;
    //1. 如果當前元素個數 >= 佇列容量 則擴容
    while ((n = size) >= (cap = (array = queue).length))
        tryGrow(array, cap);
    try {
        //2. 預設比較器為null
        Comparator<? super E> cmp = comparator;
        if (cmp == null)
            siftUpComparable(n, e, array);
        else
            //3. 自定義比較器
            siftUpUsingComparator(n, e, array, cmp);
        //4. 佇列元素數量增加1，並喚醒notEmpty條件佇列中的一個阻塞執行緒
        size = n + 1;
        notEmpty.signal();
    } finally {
        lock.unlock();
    }
    return true;
}

如上程式碼並不複雜，我們主要看看如何進行擴容和在內部建堆。

我們先看擴容邏輯：

private void tryGrow(Object[] array, int oldCap) {
        lock.unlock(); // must release and then re-acquire main lock
        Object[] newArray = null;
        //1. CAS成功則擴容
        if (allocationSpinLock == 0 &&
            UNSAFE.compareAndSwapInt(this, allocationSpinLockOffset,
                                     0, 1)) {
            try {
                //oldCap<64則擴容執行oldCap+2，否則擴容50%，並且最大值為MAX_ARRAY_SIZE
                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;
            }
        }
       //2. 第一個執行緒CAS成功後，第二執行緒進入這段程式碼，然後第二個執行緒讓出CPU，儘量讓第一個執行緒獲取到鎖，但得不到保證
        if (newArray == null) // back off if another thread is allocating
            Thread.yield();
        lock.lock();
        if (newArray != null && queue == array) {
            queue = newArray;
            System.arraycopy(array, 0, newArray, 0, oldCap);
        }
}

tryGrow 的作用就是擴容，但是為什麼要在擴容前釋放鎖，然後使用 CAS 控制只有一個執行緒可以擴容成功？

其實不釋放鎖也是 ok 的，也就是在擴容期間一直持有該鎖，但是擴容需要時間，這段時間內佔用鎖的話那麼其他執行緒在這個時候就不能進行出隊和入隊操作，降低了併發性。所以為了提高效能，使用 CAS 來控制只有一個執行緒可以進行擴容，並且在擴容前釋放鎖，進而讓其他執行緒可以進行入隊和出隊操作。

擴容執行緒擴容完畢後會重置自旋鎖變數 allocationSpinLock 為 0，這裡並沒有使用 UNSAFE 方法的 CAS 進行設定是因為同時只可能有一個執行緒獲取到該鎖，並且 allocationSpinLock 被修飾為了 volatile 的。

我們接著看建堆演算法：

private static <T> void siftUpComparable(int k, T x, Object[] array) {
    Comparable<? super T> key = (Comparable<? super T>) x;
    // 佇列元素個數 > 0 則判斷插入位置，否則直接入隊
    while (k > 0) {
        int parent = (k - 1) >>> 1;
        Object e = array[parent];
        if (key.compareTo((T) e) >= 0)
            break;
        array[k] = e;
        k = parent;
    }
    array[k] = key;
}

熟悉二叉堆的話，該段程式碼並不複雜，我們看下圖具體結構：

首先我們看parent = (k - 1) >>> 1，首先 k - 1 就是拿到當前真正的下標的位置，隨後 >>> 1拿到父節點的位置，該圖我們得知，k = 7，執行(k - 1) >>> 1之後得到的parent = 3，根據下標我們知道是元素 6。

PriorityQueue 是一個完全二叉樹，且不允許出現 null 節點，其父節點都比葉子節點小，這個是堆排序中的最小堆。二叉樹存入陣列的方式很簡單，就是從上到下，從左到右。完全二叉樹可以和陣列中的位置一一對應：

• 左葉子節點 = 父節點下標 * 2 + 1
• 右葉子節點 = 父節點下標 * 2 + 2
• 父節點 = (葉子節點 - 1) / 2

實際上就是將要插入的元素 x 和它的父節點元素 6 做對比，如果比父節點大就一直向上移動。

poll 操作

poll 操作的作用是獲取佇列內部堆樹的根節點元素，如果佇列為空，則返回 null。

public E poll() {
    final ReentrantLock lock = this.lock;
    //獲取獨佔鎖
    lock.lock();
    try {
        return dequeue();
    } finally {
        //釋放獨佔鎖
        lock.unlock();
    }
}

我們主要看一下 dequeue 方法。

private E dequeue() {
    int n = size - 1;
    //佇列為空，返回null
    if (n < 0)
        return null;
    else {
        Object[] array = queue;
        //1.獲取頭部元素
        E result = (E) array[0];
        //2. 獲取隊尾元素，並賦值為null
        E x = (E) array[n];
        array[n] = null;
        Comparator<? super E> cmp = comparator;
        if (cmp == null)//3.
            siftDownComparable(0, x, array, n);
        else
            siftDownUsingComparator(0, x, array, n, cmp);
        size = n; //4.
        return result;
    }
}

該方法如果佇列為空則直接返回 null，否則執行程式碼（1）獲取陣列第一個元素作為返回值存放到變數 Result 中，這裡需要注意，陣列裡面的第一個元素是優先順序最小或者最大的元素，出隊操作就是返回這個元素。然後程式碼（2）獲取佇列尾部元素並存放到變數 x 中，且置空尾部節點，然後執行程式碼（3）將變數 x 插入到陣列下標為 0 的位置，之後重新調整堆為最大或者最小堆，然後返回。這裡重要的是，去掉堆的根節點後，如何使用剩下的節點重新調整一個最大或者最小堆。下面我們看下 siftDownComparable 的實現。

private static <T> void siftDownComparable(int k, T x, Object[] array,
                                           int n) {
    if (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 = array[child];
            int right = child + 1;
            if (right < n &&
                ((Comparable<? super T>) c).compareTo((T) array[right]) > 0)
                c = array[child = right];
            if (key.compareTo((T) c) <= 0)
                break;
            array[k] = c;
            k = child;
        }
        array[k] = key;
    }
}

由於佇列陣列第 0 個元素為根，因此出隊時要移除它。這時陣列就不再是最小的堆了，所以需要調整堆。具體是從被移除的樹根的左右子樹中找一個最小的值來當樹根，左右子樹又會找自己左右子樹裡面那個最小值，這是一個遞迴過程，直到葉子節點結束遞迴。

假設目前佇列內容如下圖：

上圖中樹根的 leftChildVal = 4; rightChildVal = 6；由於4 < 6，所以c = 4。然後由於11 > 4，也就是key > c，所以使用元素 4 覆蓋樹根節點的值。

然後樹根的左子樹樹根的左右孩子節點中的 leftChildVal = 8; rightChildVal = 10；由於8 < 10，所以c = 8。然後由於11 > 8，也就是 key > c，所以元素 8 作為樹根左子樹的根節點，現在樹的形狀如下圖第三步所示。這時候判斷是否k < half，結果為 false，所以退出迴圈。然後把x = 11的元素設定到陣列下標為3的地方，這時候堆樹如下圖第四步所示，至此調整堆完畢。

put 操作

put 操作內部呼叫的是 offer 操作，由於是無界佇列，所以不需要阻塞。

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

take 操作

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;
}

size 操作

獲取佇列元素個數。如下程式碼在返回 size 前加了鎖，以保證在呼叫 size 方法時不會有其他執行緒進行入隊和出隊操作。另外，由於 size 變數沒有被修飾為 volatie 的，所以這裡加鎖也保證了在多執行緒下 size 變數的記憶體可見性。

public int size() {
    final ReentrantLock lock = this.lock;
    lock.lock();
    try {
        return size;
    } finally {
        lock.unlock();
    }
}

總結

PriorityBlockingQueue 類似於 ArrayBlockingQueue，在內部使用一個獨佔鎖來控制同時只有一個執行緒可以進行入隊和出隊操作。另外，PriorityBlockingQueue 只使用了一個 notEmpty 條件變數而沒有使用 notFull，因為是無界佇列，執行 put 操作時永遠不會處於 await 狀態，所以也不需要被喚醒。而 take 方法是阻塞方法，並且是可被中斷的。當需要存放有優先順序的元素時該佇列比較有用。