JUC之CountDownLatch源码分析

CountDownLatch是AbstractQueuedSynchronizer中共享锁模式的一个的实现，是一个同步工具类，用来协调多个线程之间的同步。CountDownLatch能够使一个或多个线程在等待另外一些线程完成各自工作之后，再继续执行。CountDownLatch内部使用一个计数器进行实现线程通知条件，计数器初始值为进行通知线程的数量。当每一个通知线程完成自己任务后，计数器的值就会减一。当计数器的值为0时，表示所有的通知线程都已经完成一些任务，然后在CountDownLatch上所有等待的线程就可以恢复执行接下来的任务。基本上CountDownLatch的原理就是这样，下面我们一起去看看源码。

public class CountDownLatch {

    private static final class Sync extends AbstractQueuedSynchronizer {
        private static final long serialVersionUID = 4982264981922014374L;

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

    private final Sync sync;

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

从上面简略的源码可以看出，CountDownLatch和ReentrantLock一样，在内部声明了一个继承AbstractQueuedSynchronizer的Sync内部类（重写了父类的tryAcquireShared(int acquires)和tryReleaseShared(int releases)），并在声明了一个sync属性。CountDownLatch只有一个有参构造器，参数count就是上面说的的进行通知的线程数目，说白点也就是countDown()方法需要被调用的次数。

CountDownLatch的主要方法是wait()和countDown()，我们先看wait()方法源码。

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

和ReentrantLock一样，CountDownLatch依然算是一件外衣，实际还是靠sync进行操作。我们接着看AQS的acquireSharedInterruptibly(int arg)方法（实际上参数在CountDownLatch里没什么用）

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

看到先判断当前线程是否是中断状态，然后调用子类重写的tryAcquireShared(int acquires)方法去判断是否需要进行阻塞（也即是尝试获取锁），如果返回值小于0 ，就调用doAcquireSharedInterruptibly(int acquires)方法进行线程阻塞。先看tryAcquireShared(int acquires)方法

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

源码很简单，就是看下state是否等于0，等于0，就返回1，代表不需要线程阻塞，不等于0（实际上state只会大于或者等于0），就返回-1，表示需要进行线程阻塞。这里有个伏笔就是如果CountDownLatch的计数器state被减至0时，后续再有线程调用CountDownLatch的wait()方法时，会直接往下执行调用者方法的代码，不会造成线程阻塞。

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) {
                int r = tryAcquireShared(arg);
                if (r >= 0) {
                    setHeadAndPropagate(node, r);
                    p.next = null; // help GC
                    failed = false;
                    return;
                }
            }
            if (shouldParkAfterFailedAcquire(p, node) &&
                parkAndCheckInterrupt())
                throw new InterruptedException();
        }
    } finally {
        if (failed)
            cancelAcquire(node);
    }
}

在doAcquireSharedInterruptibly(int acquires)方法中进行线程阻塞的步骤依然是先调用addWaiter(Node mode)方法将该线程封装到AQS内部的CLH队列的Node.SHARE（共享）模式的Node节点，并放入到队尾，然后在循环中去尝试持有锁和进行线程阻塞。在循环体内，先获取到前任队尾，然后判断前任队尾是否是队首head，如果是就调用tryAcquireShared(int acquires)尝试获取锁，如果返回1表示获取到了锁，就调用setHeadAndPropagate(Node node, int propagate)方法将节点设置head然后再往下传播，解除后续节点的线程阻塞状态（这就是共享锁的核心）。如果返回-1，表示没有获取到锁，就调用shouldParkAfterFailedAcquire(Node pre, Node node)进行pre节点的waitStatus赋值为Node.SIGNAL，然后在墨迹一次循环，调用parkAndCheckInterrupt()方法进行线程阻塞。我们先看setHeadAndPropagate(Node node, int propagate)方法源码

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

在setHeadAndPropagate(Node node, int propagate)方法中，直接将node设置从head，因为参数propagate为始终为1（到这一步就表示获取了共享锁，state等于0，在tryAcquireShared(int acquires)方法中就只会返回1），所以也就是后面直接获取head的next节点，如果head的next节点存在，并且是共享模式，就调用doReleaseShared()方法去释放CLH中head的next节点。

shouldParkAfterFailedAcquire(Node pre, Node node)和parkAndCheckInterrupt()两个方法在JUC之ReentrantLock源码分析一篇博客中已经讲过了，这里不再赘述了。

doReleaseShared()这个方法其实也是countDown()方法的核心实现，这里先卖个关子，后面会讲到。我们接着看doReleaseShared()方法。

private void doReleaseShared() {
    for (;;) {
        Node h = head;
        if (h != null && h != tail) {
            int ws = h.waitStatus;
            if (ws == Node.SIGNAL) {
                if (!compareAndSetWaitStatus(h, Node.SIGNAL, 0))
                    continue;            // loop to recheck cases
                unparkSuccessor(h);
            }
            else if (ws == 0 &&
                     !compareAndSetWaitStatus(h, 0, Node.PROPAGATE))
                continue;                // loop on failed CAS
        }
        if (h == head)                   // loop if head changed
            break;
    }
}

当调用wait()方法进行线程阻塞等待被别的线程解除阻塞后，对于AQS中共享锁最核心的代码就是doReleaseShared()这个方法。先获取head节点，如果head节点存在并且有后续节点，就先判断head节点的状态，如果状态是Node.SIGNAL（表示后续节点需要锁释放通知），将head节点状态改为0，然后解除下一节点的线程阻塞状态，然后判断下之前获取的head和现在的head是否还是同一个，如果是就结束循环，如果不是，那就接着循环。其实就算是存在线程在执行完unparkSuccessor(Node node)方法后失去了CPU的执行权，一直到被解除线程阻塞的next节点坐上了head节点后才有机会获取到CPU执行权这种情况，无非就是之前获取到head和现在的head不相同了，大不了再循环一次，也算是多线程去解除当前head节点的next节点线程阻塞，影响不大；如果状态是0，尝试将状态有0改为Node.PROPAGATE，然后重复循环，head节点是0的这种状态在CountDownLatch中应该是不会出现的，可能是AQS对别的类的兼容，也可能是我眼拙没看出来。如果head为null或者head与tail相同，就结束循环。

到这里CountDownLatch的wait()方法就分析完成了，这里总结下wait()方法流程：
　　1、如果state是0，直接往下执行调用者的代码，不会进行线程等待阻塞
　　2、将当前线程封装到共享模式的Node节点中，然后放入到CLH队列的队尾
　　3、将前任队尾的waitStatus改变为Node.SIGNAL，然后调用LockSupport.park()阻塞当前线程，等待前节点唤醒
　　4、被前节点唤醒后，把自己设为head，然后去唤醒next节点

我们看完了wait()方法，现在在来看下countDown()方法的源码

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

一如既往的简单，直接调用AQS的releaseShared(int arg)方法，我们直接去看AQS的releaseShared(int arg)方法

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

CountDownLatch方法
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的releaseShared(int arg)中先去调用一定要被子类重写的tryReleaseShared(int releases)方法，返回值表示是否需要进行唤醒操作，如果返回true，那就调用doReleaseShared()方法去解除head节点的next节点线程阻塞状态。是的，你没看错，就是我们前面分析的doReleaseShared()方法，他们复用了同一个方法。而在CountDownLatch的tryReleaseShared(int releases)方法实现也非常简单，就是判断下当前state是否是0，如果是0，表示已经减完了，不需要再减了，等待线程已经在被依次唤醒了，返回false表示不需要去唤醒后续节点了。最后再看看减完后的state是否是等于0，等于0，表示需要去解除后续节点的阻塞状态；不等于0（其实是大于0），表示调用countDown()方法去减state的次数还不够，暂时还不能解除后续节点的阻塞状态。

countDown()方法比较简单，我们也总结下countDown()流程：
　　1、如果state为0，表示已经有线程在解除CLH队列节点的阻塞状态了，这里直接结束
　　2、如果state减去1后不等于0，表示还要等有线程再次调用countDown()方法进行state减1，这里直接结束
　　3、如果state减去1后等于0，表示已经线程调用countDown()方法的次数已经达到最初设定的次数，然后去解除CLH队列上节点的阻塞状态