CountDownLatch原始碼閱讀

酒冽發表於2021-12-25

簡介

CountDownLatch是JUC提供的一個執行緒同步工具，主要功能就是協調多個執行緒之間的同步，或者說實現執行緒之間的通訊

CountDown，數數字，只能往下數。Latch，門閂。光看名字就能明白這個CountDownLatch是如何使用的了哈哈。CountDownLatch就相當於一個計數器，計數器的初值通過構造方法的引數來設定。呼叫CountDownLatch例項的await方法的執行緒，會等待計數器變為0才會被喚醒，繼續向下執行。那麼計數器如何變為0呢？
如果其他執行緒呼叫該CountDownLatch例項的countDown方法，會將計數值減1。當減為0時，會讓那些因呼叫await方法而阻塞等待的執行緒繼續執行。這樣就實現了這些執行緒之間的同步功能

作者：酒冽出處：https://www.cnblogs.com/frankiedyz/p/15730573.html
版權：本文版權歸作者和部落格園共有
轉載：歡迎轉載，但未經作者同意，必須保留此段宣告；必須在文章中給出原文連線；否則必究法律責任

使用場景

直接介紹或許太抽象，對於初學者來說很難理解，最好的方式就是通過一個實際場景來引入

有一種通用的場景：主執行緒開啟多個子執行緒去並行執行多個子任務，等待所有子執行緒執行完畢，主執行緒收集子執行緒的執行結果並統計

場景示例

比如，主執行緒等A和B給它轉賬，等收齊所有錢再一併放入銀行賺利息。示例程式碼如下：

public class TestCountDownLatch {

    // 這裡必須是原子類，要保證對money的修改是原子性操作，才能保證執行緒安全
    // 僅把money設定為volatile int是不行的哦
    private static final AtomicInteger money = new AtomicInteger(0);

    public static void main(String[] args) {
        CountDownLatch latch = new CountDownLatch(2);

        Thread threadA = new Thread(() -> {
            try {
                System.out.println("A轉賬30元給我");
                Thread.sleep(1000);
                
                // 執行緒A和執行緒B對money必須要使用CAS修改，否則可能會出錯
                int origin = money.get();
                while (!money.compareAndSet(origin, origin + 30)) {
                    origin = money.get();
                    continue;
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            latch.countDown();
            System.out.println("A轉賬完成");
        });

        Thread threadB = new Thread(() -> {
            try {
                System.out.println("B轉賬70元給我");
                Thread.sleep(1000);
                int origin = money.get();
                while (!money.compareAndSet(origin, origin + 30)) {
                    origin = money.get();
                    continue;
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            latch.countDown();
            System.out.println("B轉賬完成");
        });

        System.out.println("等A和B轉賬給我...");
        threadA.start();
        threadB.start();
        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        // System.out.println(latch.getCount());
        System.out.println("轉賬完成，將錢存入銀行...");

        // 這裡就不用CAS修改money，因為只有main執行緒對money做修改
        int origin = money.get();
        money.set(origin * 2);
        System.out.println("將錢去除，本金加利潤一共" + money.get() + "元");
    }
}

命令列輸出如下：

點選檢視程式碼

等A和B轉賬給我...
A轉賬30元給我
B轉賬70元給我
A轉賬完成
B轉賬完成
轉賬完成，將錢存入銀行...
將錢去除，本金加利潤一共120元

Process finished with exit code 0

得益於CountDownLatch的同步功能，當上述程式碼執行結束時，money的值必定是120，而不會是0、60或140。因為主執行緒會一致await直到執行緒A和執行緒B都執行完latch.countDown()才會繼續往下執行

實現原理

CountDownLatch的實現原理其實就是AbstractQueuedSynchronizer（AQS）

CountDownLatch有一個內部類Sync，它實現了AQS類定義的部分鉤子方法，CountDownLatch通過Sync類例項sync實現了所有功能，呼叫CountDownLatch的方法都會委託給sync域來執行

// 所有CountDown的功能都是委託給這個Sync類物件來完成
private final Sync sync;

因此，要搞懂CountDownLatch，必須搞懂AQS以及Sync類。接下來就跟我一起來剖析一下原始碼，看看這個Sync到底幹了些啥

原始碼剖析

構造方法

CountDownLatch的建構函式中可以傳入count引數，表明必須呼叫count次countDown方法，才能讓呼叫await的執行緒繼續向下執行。其原始碼如下：

public CountDownLatch(int count) {
    if (count < 0) throw new IllegalArgumentException("count < 0");
    this.sync = new Sync(count);
}

實際上是初始化了一個Sync類物件，並注入到CountDownLatch的sync域中。Sync構造方法如下：

Sync(int count) {
    setState(count);
}

Sync構造方法實際上就是設定了AQS中的state，將其設定為初始計數值為count

重要結論：AQS的state就表示CountDownLatch當前的計數值

await

await方法是一個例項方法，呼叫它的執行緒一直阻塞等待，直到CountDownLatch物件的計數值降為0，才能被喚醒。如果呼叫await時計數值就已經是0，就不會被阻塞

await方法是響應中斷的：

如果一個執行緒在呼叫await方法之前就已經被中斷，那麼呼叫時會直接丟擲中斷異常
如果一個執行緒呼叫await方法阻塞等待過程中，收到中斷訊號，就會丟擲中斷異常

await方法是響應中斷的：

如果一個執行緒在呼叫await方法之前就已經被中斷，那麼呼叫時會直接丟擲中斷異常
如果一個執行緒呼叫await方法阻塞等待過程中，收到中斷訊號，就會丟擲中斷異常

說了這麼多，還是來看看await的原始碼吧：

public void await() throws InterruptedException {
    sync.acquireSharedInterruptibly(1);
}

可以看到，這個方法實際上是委託給了Sync類物件sync來執行，這裡的acquireSharedInterruptibly已經由Sync類的父類AQS提供了實現，原始碼如下：

public final void acquireSharedInterruptibly(int arg)
    throws InterruptedException {
    if (Thread.interrupted())
        throw new InterruptedException();
    if (tryAcquireShared(arg) < 0)
        doAcquireSharedInterruptibly(arg);
}

這段程式碼在全網最詳細的AbstractQueuedSynchronizer(AQS)原始碼剖析系列中有詳細介紹，不過那裡並沒有涉及到具體的應用類（如CountDownLatch這種），只是高屋建瓴地分析過，這裡正好藉助CountDownLatch來更好地理解它

從acquireSharedInterruptibly原始碼中可以看到，如果執行緒在呼叫await之前就已經被設定了中斷狀態，那麼會直接丟擲InterruptedException異常

該方法接下來會呼叫鉤子方法tryAcquireShared，Sync類為該方法提供了具體實現，原始碼如下：

protected int tryAcquireShared(int acquires) {
    return (getState() == 0) ? 1 : -1;
}

AQS雖然沒有為tryAcquireShared提供具體實現，但是規定了返回值的含義：

負數：表明獲取失敗，該執行緒需要被加入同步佇列阻塞等待
0：表明獲取共享資源成功，但是後續獲取共享資源一定不會成功
正數：表明獲取共享資源成功，而且後續的獲取也可能成功

讓我們來分析一下Sync類實現的tryAcquireShared方法：

如果state不為0，即CountDownLatch物件的計數值還沒減到0，則返回-1，會繼續執行doAcquireSharedInterruptibly方法，將呼叫await的執行緒加入同步佇列阻塞等待
如果state為0，即CountDownLatch物件的計數值已經減為0，則返回1，呼叫await會直接返回，不會被阻塞

注：doAcquireSharedInterruptibly的具體分析見全網最詳細的AbstractQueuedSynchronizer(AQS)原始碼剖析系列

countDown

countDown方法也是例項方法，呼叫它會將CountDownLatch物件的計數值減1。如果正好減為0，那麼會將所有因呼叫await而被阻塞的執行緒都喚醒。其原始碼如下：

public void countDown() {
    sync.releaseShared(1);
}

該方法實際上委託給sync來執行，這裡的releaseShared方法在Sync類的父類AQS中提供了具體實現，其原始碼如下：

public final boolean releaseShared(int arg) {
    if (tryReleaseShared(arg)) {
        doReleaseShared();
        return true;
    }
    return false;
}

該方法首先會呼叫鉤子方法tryReleaseShared，該方法在Sync類中提供了具體實現，原始碼如下：

protected boolean tryReleaseShared(int releases) {
    // Decrement count; signal when transition to zero
    for (;;) {
        int c = getState();
        if (c == 0)
            return false;
        int nextc = c-1;
        if (compareAndSetState(c, nextc))
            return nextc == 0;
    }
}

AQS雖然沒有為tryReleaseShared提供具體實現，但是規定了返回值的含義：

true：此次釋放資源的行為可能會讓一個阻塞等待中的執行緒被喚醒
false：otherwise

讓我們來分析一下Sync類實現的tryReleaseShared方法：
該方法所有程式碼都包含在一個for迴圈中，這是為了應對CAS失敗的情況。迴圈體內CAS修改state，即將CountDownLatch的計數值減1
如果CountDownLatch的計數值減1後變成0，則返回true。那麼releaseShared方法會繼續呼叫doReleaseShared方法，喚醒同步佇列中的後續執行緒
如果不為0，則返回false，無事發生~

注：doReleaseShared方法的作用是喚醒隊首執行緒，並確保狀態傳播，該方法的詳細解釋見全網最詳細的AbstractQueuedSynchronizer(AQS)原始碼剖析系列

getCount

getCount方法就是返回CountDownLatch物件當前的計數值，原始碼如下：

public long getCount() {
    return sync.getCount();
}

實際上委託給了sync物件的getCount方法來執行，其原始碼如下：

int getCount() {
    return getState();
}

其實就是呼叫AQS的getState方法，返回當前的state，即CountDownLatch的計數值，很簡單哦~

CountDownLatch與join方法的區別

當然，在上述場景中，也可以使用Thread的物件方法join來實現這一點，在主執行緒中呼叫所有子執行緒的join方法，再執行結果收集和統計任務。上面例子如果使用join方法來實現，程式碼如下：

public class TestJoin {
    private static final AtomicInteger money = new AtomicInteger(0);

    public static void main(String[] args) {
        Thread threadA = new Thread(() -> {
            try {
                System.out.println("A轉賬30元給我");
                Thread.sleep(1000);

                int origin = money.get();
                while (!money.compareAndSet(origin, origin + 30)) {
                    origin = money.get();
                    continue;
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            System.out.println("A轉賬完成");
        });

        Thread threadB = new Thread(() -> {
            try {
                System.out.println("B轉賬70元給我");
                Thread.sleep(1000);
                int origin = money.get();
                while (!money.compareAndSet(origin, origin + 30)) {
                    origin = money.get();
                    continue;
                }
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            System.out.println("B轉賬完成");
        });

        System.out.println("等A和B轉賬給我...");
        threadA.start();
        threadB.start();
        try {
            threadA.join();
            threadB.join();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        System.out.println("轉賬完成，將錢存入銀行...");

        int origin = money.get();
        money.set(origin * 2);
        System.out.println("將錢去除，本金加利潤一共" + money.get() + "元");
    }
}