CountDownLatch 基于 AQS 共享锁模式实现原理分析

class Driver2 { // ...
    void main() throws InterruptedException {
        CountDownLatch doneSignal = new CountDownLatch(N);
        Executor e = ...
        for (int i = 0; i < N; ++i) {// create and start threads
            e.execute(new WorkerRunnable(doneSignal, i));
        }
        doneSignal.await();           // wait for all to finish
    }
 }

class WorkerRunnable implements Runnable {
    private final CountDownLatch doneSignal;
    private final int i;
    WorkerRunnable(CountDownLatch doneSignal, int i) {
        this.doneSignal = doneSignal;
        this.i = i;
    }
    public void run() {
        try {
            doWork(i);
            doneSignal.countDown();
        } catch (InterruptedException ex) {} // return;
    }
    void doWork() { ... }
}

protected boolean tryAcquire(int arg) {
    throw new UnsupportedOperationException();
}

protected boolean tryRelease(int arg) {
    throw new UnsupportedOperationException();
}

protected int tryAcquireShared(int arg) {
    throw new UnsupportedOperationException();
}

protected boolean tryReleaseShared(int arg) {
    throw new UnsupportedOperationException();
}

protected boolean isHeldExclusively() {
    throw new UnsupportedOperationException();
}

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

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

/**
 * The synchronization state.
 */
private volatile int state;

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

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

/**
 * Acquires in shared interruptible mode.
 * @param arg the acquire argument
 */
private void doAcquireSharedInterruptibly(int arg)
        throws InterruptedException {
    //加入等待队列
    final Node node = addWaiter(Node.SHARED);
    boolean failed = true;
    try {
        //进入 CAS 循环
        for (;;) {
            //当一个节点(关联一个线程)进入等待队列后， 获取此节点的 prev 节点 
            final Node p = node.predecessor();
            // 如果获取到的 prev 是 head，也就是队列中第一个等待线程
            if (p == head) {
                // 再次尝试申请 反应到 CountDownLatch 就是查看是否还有线程需要等待(state是否为0)
                int r = tryAcquireShared(arg);
                // 如果 r >=0 说明 没有线程需要等待了 state==0
                if (r >= 0) {
                    //尝试将第一个线程关联的节点设置为 head 
                    setHeadAndPropagate(node, r);
                    p.next = null; // help GC
                    failed = false;
                    return;
                }
            }
            //经过自旋tryAcquireShared后，state还不为0，就会到这里，第一次的时候，
            //waitStatus是0，那么node的waitStatus就会被置为SIGNAL，第二次再走到这里，
            //就会用LockSupport的park方法把当前线程阻塞住
            if (shouldParkAfterFailedAcquire(p, node) &&
                    parkAndCheckInterrupt())
                throw new InterruptedException();
        }
    } finally {
        if (failed)
            cancelAcquire(node);
    }
}

private Node addWaiter(Node mode) {
    Node node = new Node(Thread.currentThread(), mode);
    // Try the fast path of enq; backup to full enq on failure
    // 尝试快速入队操作，因为大多数时候尾节点不为 null
    Node pred = tail;
    if (pred != null) {
        node.prev = pred;
        if (compareAndSetTail(pred, node)) {
            pred.next = node;
            return node;
        }
    }
    //如果尾节点为空(也就是队列为空) 或者尝试CAS入队失败(由于并发原因)，进入enq方法
    enq(node);
    return node;
}

/**
 * Inserts node into queue, initializing if necessary. See picture above.
 * @param node the node to insert
 * @return node's predecessor
 */
private Node enq(final Node node) {
    // 死循环 + CAS 保证所有节点都入队
    for (;;) {
        Node t = tail;
        // 如果队列为空 设置一个空节点作为 head
        if (t == null) { // Must initialize
            if (compareAndSetHead(new Node()))
                tail = head;
        } else {
            // 加入队尾
            node.prev = t;
            if (compareAndSetTail(t, node)) {
                t.next = node;
                return t;
            }
        }
    }
}

/**
 * Sets head of queue, and checks if successor may be waiting
 * in shared mode, if so propagating if either propagate > 0 or
 * PROPAGATE status was set.
 *
 * @param node the node
 * @param propagate the return value from a tryAcquireShared
 */
private void setHeadAndPropagate(Node node, int propagate) {
    // 备份现在的 head
    Node h = head; // Record old head for check below
    // 抢到锁的线程被唤醒 将这个节点设置为head
    setHead(node);
    /*
     * Try to signal next queued node if:
     *   Propagation was indicated by caller,
     *     or was recorded (as h.waitStatus either before
     *     or after setHead) by a previous operation
     *     (note: this uses sign-check of waitStatus because
     *      PROPAGATE status may transition to SIGNAL.)
     * and
     *   The next node is waiting in shared mode,
     *     or we don't know, because it appears null
     *
     * The conservatism in both of these checks may cause
     * unnecessary wake-ups, but only when there are multiple
     * racing acquires/releases, so most need signals now or soon
     * anyway.
     */
    // propagate 一般都会大于 0 或者存在可被唤醒的线程
    if (propagate > 0 || h == null || h.waitStatus < 0 ||
            (h = head) == null || h.waitStatus < 0) {
        Node s = node.next;
        // 只有一个节点 或者是共享模式，释放所有等待线程，各自尝试抢占锁
        if (s == null || s.isShared())
            doReleaseShared();
    }
}

//线程已被 cancelled ，这种状态的节点将会被忽略，并移出队列
static final int CANCELLED =  1;
// 表示当前线程已被挂起，并且后继节点可以尝试抢占锁
static final int SIGNAL    = -1;
//线程正在等待某些条件
static final int CONDITION = -2;
//共享模式下 无条件所有等待线程尝试抢占锁
static final int PROPAGATE = -3;

// CountDownLatch 类 countDown() 方法
public void countDown() {
    sync.releaseShared(1);
}

// AQS 类
public final boolean releaseShared(int arg) {
    // arg 为固定值 1
    // 如果计数器state 为 0 返回true，前提是调用 countDown() 之前能已经为 0
    if (tryReleaseShared(arg)) {
        // 唤醒等待队列的线程
        doReleaseShared();
        return true;
    }
    return false;
}

// CountDownLatch 类重写 AQS 方法
protected boolean tryReleaseShared(int releases) {
    // Decrement count; signal when transition to zero
    // 依然是循环 + CAS 配合 实现计数器减 1
    for (;;) {
        int c = getState();
        if (c == 0)
            return false;
        int nextc = c-1;
        if (compareAndSetState(c, nextc))
            return nextc == 0;
    }
}

// AQS 类
private void doReleaseShared() {
    for (;;) {
        Node h = head;
        if (h != null && h != tail) {
            int ws = h.waitStatus;
            // 如果节点状态为 SIGNAL，则他的 next 节点也可以尝试被唤醒
            if (ws == Node.SIGNAL) {
                if (!compareAndSetWaitStatus(h, Node.SIGNAL, 0))
                    continue;            // loop to recheck cases
                unparkSuccessor(h);
            }
            // 将节点状态设置为 PROPAGATE，表示要向下传播，依次唤醒
            else if (ws == 0 &&
                        !compareAndSetWaitStatus(h, 0, Node.PROPAGATE))
                continue;                // loop on failed CAS
        }
        if (h == head)                   // loop if head changed
            break;
    }
}