CountDownLatch源码解析

简介

CountDownLatch是并发包中提供的一个可用于控制多个线程同时开始某动作的类，可以看做是一个计数器，计数器操作是院子操作，同时只能有一个线程去操作这个计数器。可以向CountDownLatch对象设置一个初始的数字作为计数值，任何调用这个对象上的await()方法都会阻塞，直到这个计数器的计数值被其他的线程减为0为止。

CountDownLatch是基于AQS的共享锁实现的，对AQS不了解的可以看我的另一篇：AQS深度解析

简单例子

先来个例子热热身。

private static CountDownLatch latch = new CountDownLatch(20);
public static void main(String[] args) {
    new Thread(new MyRunnable(),"thread-1").start();
    new Thread(new MyRunnable(),"thread-2").start();
    while (latch.getCount()!=0){
        latch.countDown();
        System.out.println("latch = "+latch.getCount());
    }
}
static class MyRunnable implements  Runnable{
    public void run() {
        try {
            System.out.println(Thread.currentThread().getName()+":latch执行await()。。。。。");
            latch.await();
            System.out.println(Thread.currentThread().getName()+":latch执行结束。。。。。");
        }catch (Exception e){
            e.printStackTrace();
        }
    }
}

执行结果：

latch = 19
thread-1:latch执行await()。。。。。
thread-2:latch执行await()。。。。。
latch = 18
.......
latch = 2
latch = 1
latch = 0
thread-1:latch执行结束。。。。。
thread-2:latch执行结束。。。。。

实现原理

CountDownLatch的实现原理就想上图一样简单，在new一个CountDownLatch对象的时候要设置state的值，该值设置成功后不能改变，也就是CountDownLatch不能复用。当一个线程调用countDown()方法时，state值-1，直到当state=0时，就会恢复同步队列中被挂起的线程。当state!=0时，有线程调用了await()方法，该线程就会加入到同步队列，并被挂起。知道state=0时，被恢复。

源码解析

state

下面是CountDownLatch初始化时设置state值的代码，很简单

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

countDown()

下面是这个方法的源码：

public void countDown() {
    sync.releaseShared(1);
}
//AQS
public final boolean releaseShared(int arg) {
    if (tryReleaseShared(arg)) {
        doReleaseShared();
        return true;
    }
    return false;
}

主要的逻辑在releaseShared()方法中。
tryReleaseShared()做释放逻辑，doReleaseShared()做真正的释放动作。
下面首先看下tryReleaseShared()

protected boolean tryReleaseShared(int releases) {
    // 通过自旋减量：也就是 --state;
    for (;;) {
        int c = getState();
        //如果state等于0，说明已经不能再做减操作了
        if (c == 0)
            return false;
        //使用cas自减state的值，减后state值等于0就返回true
        int nextc = c-1;
        if (compareAndSetState(c, nextc))
            return nextc == 0;
    }
}

1. 通过自旋保证state设置失败后的重试操作。
2. 获取state的值，如果state=0，则直接返回false
3. state!=0,则state–,并cas修改state的值，失败则自旋重新设置。
4. 成功判断state-1==0，false,挂起的线程举行等待。
5. true,调用doReleaseShared()恢复队列中挂起的线程。

下面就是通过自旋恢复被挂起的线程，很简单不细说。

private void doReleaseShared() {
    for (;;) {
        //head头节点不是一个线程节点，只作为队列的前驱
        Node h = head;
        if (h != null && h != tail) {
            int ws = h.waitStatus;
            //如果线程被挂起，先利用cas修改waitStatus的值，然后再恢复被挂起的线程
            if (ws == Node.SIGNAL) {
                if (!compareAndSetWaitStatus(h, Node.SIGNAL, 0))
                    continue;            // loop to recheck cases
                unparkSuccessor(h);
            }
            //ws==0说明，state的值已经为0这时候的线程可以直接继续执行，并且设置为以后的线程节点都传播此行为。
            else if (ws == 0 &&
                     !compareAndSetWaitStatus(h, 0, Node.PROPAGATE))
                continue;                // loop on failed CAS
        }
        if (h == head)                   // loop if head changed
            break;
    }
}
恢复被挂起的线程
private void unparkSuccessor(Node node) {
   //----可忽略start------
    int ws = node.waitStatus;
    if (ws < 0)
        compareAndSetWaitStatus(node, ws, 0);
    Node s = node.next;
    if (s == null || s.waitStatus > 0) {
        s = null;
        //从后向前遍历链表，
        for (Node t = tail; t != null && t != node; t = t.prev)
            if (t.waitStatus <= 0)
                s = t;
    }
    //------end-------
    //恢复队列中被挂起的线程
    if (s != null)
        LockSupport.unpark(s.thread);
}

恢复线程的说完了，下面来看下挂起线程的方法。

await()

下面是await()的源码：

CountDownLatch
public void await() throws InterruptedException {
    sync.acquireSharedInterruptibly(1);
}
AQS：尝试获取共享锁：失败则中断线程
public final void acquireSharedInterruptibly(int arg)
        throws InterruptedException {
    if (Thread.interrupted())
        throw new InterruptedException();
    //判断state的值是否等于0：(getState() == 0) ? 1 : -1;
    //state！=0说明获取共享锁失败,则要挂起
    if (tryAcquireShared(arg) < 0)
        //到这里说明state不等于0
        doAcquireSharedInterruptibly(arg);
}

CountDownLatch.await()中，并没有什么逻辑，主要的都在acquireSharedInterruptibly()方法中：

1. 线程中断就抛出异常

2. 使用tryAcquireShared(arg)判断state是否等于0

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

3. state !=0,调用doAcquireSharedInterruptibly(arg)方法将线程加入到队列中，并挂起线程。

下面看下doAcquireSharedInterruptibly()是怎么加入到队列并挂起的。

private void doAcquireSharedInterruptibly(int arg)
    throws InterruptedException {
    //向队列中添加一个共享节点
    final Node node = addWaiter(Node.SHARED);
    boolean failed = true;
    try {
        for (;;) {
            final Node p = node.predecessor();
            if (p == head) {
            // 如果新建节点的前一个节点，就是 Head，说明当前节点是 AQS 队列中等待获取锁的第一个节点，按照 FIFO 的原则，可以直接尝试获取锁。
                int r = tryAcquireShared(arg);
                if (r >= 0) {
                // 获取成功，需要将当前节点设置为 AQS 队列中的第一个节点，这是 AQS 的规则 // 队列的头节点表示已经获取锁的节点，所以会吧next，thread设为null，gcroot回收 
                    setHeadAndPropagate(node, r);//唯一的退出条件，也就是await()方法返回的条件很重要！！
                    p.next = null; // help GC
                    failed = false;
                    return;//到这里返回
                }
            }
            if (shouldParkAfterFailedAcquire(p, node) &&
                parkAndCheckInterrupt())// // 检查下是否需要将当前节点挂起 
                throw new InterruptedException();
        }
    } finally {
        if (failed)//如果失败或出现异常，失败 取消该节点，以便唤醒后续节点
            cancelAcquire(node);
    }
}

这个方法有点长，但是主要方法就四个：

//1. 向队列中添加一个共享节点
final Node node = addWaiter(Node.SHARED);
//2. 设置当前节点为head，并设置node节点的传播属性
setHeadAndPropagate(node, r);
//3. 挂起线程
parkAndCheckInterrupt();
//4. 失败或异常，取消线程
cancelAcquire(node);

其中1、3、4都很简单，并且在之前分析ReentrantLock时都说过，这里就不详细说了，只有第二个setHeadAndPropagate(node, r);没见到过，下面来看看它：

//设置当前节点为头节点，并设置传播状态
private void setHeadAndPropagate(Node node, int propagate) {
    Node h = head; // Record old head for check below
    setHead(node);
    if (propagate > 0 || h == null || h.waitStatus < 0 ||
        (h = head) == null || h.waitStatus < 0) {
        Node s = node.next;
        if (s == null || s.isShared())
            doReleaseShared();
    }
}

注意，执行到这个方法说明state的值已经为0，可以恢复被挂起的线程了
上面代码中setHead(node)用于设置head节点，如果是共享节点就调用doReleaseShared()恢复被挂起的节点。

总结

最后总结一下，从源码中能看到，AQS中具体实现了共享模式下获取锁acquireSharedInterruptibly()和释放锁releaseShared()，定义了tryAcquireShared()和tryReleaseShared()方法让子类去实现具体的获取和释放的逻辑，其实主要的就是利用cas对state值的操作。

以上只是我粗略的理解，如果有理解不对、表达不到的地方，期待大家留意。

如果觉得能让你理解一二，请看我写的另一篇文章，可以对AQS有更好的理解：
ReentrantLock源码解析