JUC之CountDownLatch源码分析

简介

CountDownLatch一般用作多线程倒计时计数器，通过 await()方法等待一组线程执行完成。

示例

public class CountDownDemo {

    public static void main(String[] args) throws Exception {

        CountDownLatch latch = new CountDownLatch(2);
        for (int i = 0; i < 2; i++) {
            new Thread(() -> {
                try {
                    System.out.println(new Date(System.currentTimeMillis()) + ": " + Thread.currentThread().getName());
                    Thread.sleep(5000L);
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            }, "Thread-" + i).start();
        }

        latch.await();

        System.out.println(new Date(System.currentTimeMillis()) + ": " + "main thread end");

    }

}

源码分析

通过示例分析，CountDownLatch通过三个步骤实现多线程倒计时计数器

• 创建CountDownLatch实例并设置初始计数器数值

• 子线程调用countDown()方法将计数器数值减1并在剩余计数器数值为0时调用LockSupport.unpark()方法唤醒通过await()方法阻塞的线程

• 主线程调用await()方法阻塞（内部通过LockSupport.part()方法阻塞）

接下来就通过分析源码看看时如何通过着三个步骤实现多线程倒计时计数器的

创建CountDownLatch实例的源码分析

public CountDownLatch(int count) {
    if (count < 0) throw new IllegalArgumentException("count < 0");
    this.sync = new Sync(count); // 创建Sync实例，并设置status初始值
}

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

    Sync(int count) {
        setState(count);
    }
}
public abstract class AbstractQueuedSynchronizer
    extends AbstractOwnableSynchronizer
    implements java.io.Serializable {
  
    protected final void setState(int newState) {
        state = newState;
    }
}

这个过程比较简单，只是创建了一个Sync实例，其继承至AQS，并设置计数器(status)的初始值，后续都是通过对该字段的操作来实现效果的。

countDown()方法源码分析

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

/**
* 计数器数值减1，如果剩余计数器数值为0，则唤醒阻塞线程
**/
public final boolean releaseShared(int arg) {
    if (tryReleaseShared(arg)) {
        doReleaseShared();
        return true;
    }
    return false;
}

/**
* 自旋锁，通过CAS操作，将计数器数值减1，直接操作成功
* 可能存在多个线程同时更新status内容，通过CAS更新有可能失败
* 只有status等于0时，此时需要唤醒阻塞线程，否则不需要唤醒阻塞线程
**/
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; // 计数器数值减1
        if (compareAndSetState(c, nextc)) // 通过CAS更新最新的计数器数值
            return nextc == 0; // 如果更新后的计数器数值为0，此时所有调用await方法阻塞的线程均可以唤醒
    }
}

分析 doReleaseShared()方法的源码时，先看下AQS队列的结构，这样理解起来比较容易，左侧的图标志目前没有线程阻塞，右侧的图表示当前有一个线程处于阻塞状态。

/**
* 唤醒阻塞的线程（该方法时AQS类中的公共方法，Semaphore等JUC的工具类都是通过该方法唤醒线程 ）
**/
private void doReleaseShared() {
    for (;;) {
        Node h = head; // 指向Node节点
        if (h != null && h != tail) { // 表示AQS队列属于上图右侧这种状态，存在需要被唤醒的线程
            int ws = h.waitStatus;
            if (ws == Node.SIGNAL) { // Node.SIGNAL表明当前节点的后继节点需要唤醒
                // 0表示后继节点不需要唤醒，当时需要唤醒
                if (!compareAndSetWaitStatus(h, Node.SIGNAL, 0))
                    continue;            // loop to recheck cases
                unparkSuccessor(h); // 唤醒当前节点的后继节点，即Node1
            }
            else if (ws == 0 &&
                     !compareAndSetWaitStatus(h, 0, Node.PROPAGATE))
                continue;                // loop on failed CAS
        }
        if (h == head)                   // 如果head节点被更改，说明有线程节点唤醒，继续循环，检查是否有节点可以唤醒
            break;
    }
}

/**
* 唤醒node节点的后继节点
*/
private void unparkSuccessor(Node node) {
    int ws = node.waitStatus;
    if (ws < 0)
        compareAndSetWaitStatus(node, ws, 0);

    Node s = node.next; // s即Node1节点

    // 如果Node1节点为空或着已经被取消，就尝试获取下一个节点，即Node2，依次类推
    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;
    }
    // 唤醒AQS队列中从head节点开始第一个需要被唤醒的线程
    // 如果Node1符合条件就唤醒Node节点线程，否则就唤醒Nod2，依次类推
    if (s != null)
        LockSupport.unpark(s.thread);
}

countDown()方法每次都会将计数器数值减1，并判断减1后的数值是否为0，为0表示需要唤醒调用await()方法处于阻塞状态的线程（唤醒AQS队列中第一个节点所属线程，即Node1）。

await()方法源码分析

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

/**
* 1.判断计数器是否为0，为0就直接结束，不阻塞
* 2.如果计数器数值为0，阻塞线程，等待唤醒
*/
public final void acquireSharedInterruptibly(int arg)
        throws InterruptedException {
    if (Thread.interrupted())
        throw new InterruptedException();
    if (tryAcquireShared(arg) < 0)
        doAcquireSharedInterruptibly(arg);
}

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

/**
* 这是AQS类中公共方法，主要做了两件事
* 1.将当前线程的节点加入到AQS队列
* 2.自旋，首先尝试获取锁，如果获取到锁，就尝试通知后继节点尝试
*/ 
private void doAcquireSharedInterruptibly(int arg)
    throws InterruptedException {
    final Node node = addWaiter(Node.SHARED); // 该方法执行之后的AQS队列结果即为上面图的右侧，一个双端队列
    boolean failed = true;
    try {
        for (;;) {
            final Node p = node.predecessor();
            // 此时就是判断当前线程所书的节点的前继节点是否为head节点，如果为head节点，即判断当前线程所属节点是不是AQS队列中的第一个节点
            if (p == head) {
                // 当前线程所属节点时AQS队列的第一个节点，则再次判断剩余计数器是否为0，如果为0，就不需要阻塞
                int r = tryAcquireShared(arg);
                if (r >= 0) {
                    setHeadAndPropagate(node, r); // 表明后继节点可能也可以被唤醒
                    p.next = null; // help GC
                    failed = false;
                    return;
                }
            }
            // 如果当前线程节点所属节点不是AQS队列中的第一个节点，根据公平原则，需要阻塞
            if (shouldParkAfterFailedAcquire(p, node) &&
                parkAndCheckInterrupt())
                throw new InterruptedException();
        }
    } finally {
        if (failed)
            cancelAcquire(node);
    }
}

private void setHeadAndPropagate(Node node, int propagate) {
    Node h = head; // setHead会重置head节点，临时变量保持旧的head节点
    setHead(node); //  设置head节点为node节点

    // propagate > 0 说明此时有剩余许可
    // h.waitStatus < 0 表示后继节点为等待状态
    // (h = head) == null 表明node节点也为空
    // h.waitStatus < 0 表示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();
    }
}

/**
* 判断当前线程是否需要阻塞
*/
private static boolean shouldParkAfterFailedAcquire(Node pred, Node node) {
    int ws = pred.waitStatus;
    if (ws == Node.SIGNAL)
        return true;
    if (ws > 0) {
        // 此时表明ws节点被取消，需要缓存AQS队列，将已取消的节点从AQS队列移除
        do {
            node.prev = pred = pred.prev;
        } while (pred.waitStatus > 0);
        pred.next = node;
    } else {
        compareAndSetWaitStatus(pred, ws, Node.SIGNAL);
    }
    return false;
}

/**
* 线程在此处阻塞
*/
private final boolean parkAndCheckInterrupt() {
    LockSupport.park(this);
    return Thread.interrupted();
}

附录

AQS同步队列中的节点状态说明

状态	说明
SIGNAL	值为-1，后继节点的线程处于等待状态，而当前节点的线程如果释放了同步状态或者被取消，那么就会通知后继节点，让后继节点的线程能够运行
CONDITION	值为-2，节点在等待队列中，节点线程等待在Condition上，不过当其他的线程对Condition调用了signal()方法后，该节点就会从等待队列转移到同步队列中，然后开始尝试对同步状态的获取
PROPAGATE	值为-3，表示下一次的共享式同步状态获取将会无条件的被传播下去
CANCELLED	值为1，由于超时或中断，该节点被取消。 节点进入该状态将不再变化。特别是具有取消节点的线程永远不会再次阻塞
INITIAL	值为0，初始状态