【JUC系列】CountDownLatch实现原理

简单示例

public class Main {
    private static final int NUM = 3;

    public static void main(String[] args) throws InterruptedException {
        CountDownLatch latch = new CountDownLatch(NUM);
        for (int i = 0; i < NUM; i++) {
            new Thread(() -> {
                try {
                    Thread.sleep(2000);
                    System.out.println(Thread.currentThread().getName() + "运行完毕");
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    latch.countDown();
                }
            }).start();
        }
        latch.await();
        System.out.println("主线程运行完毕");
    }
}

输出结果如下：

Thread-2运行完毕
Thread-0运行完毕
Thread-1运行完毕
主线程运行完毕

主线程会等到3个子线程执行完毕才会执行。

源码解析

CountDownLatch里面有一个继承AQS的内部类Sync，其实是AQS来支持CountDownLatch的各项操作的。

CountDownLatch(int count)

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

    private static final class Sync extends AbstractQueuedSynchronizer {

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

countDown()

countDown()方法会对计数器进行减1的操作，当计数器值为0时，将会唤醒在阻塞队列中等待的所有线程。其内部调用了Sync的releaseShared(1)方法

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

AbstractQueuedSynchronizer#releaseShared。

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

CountDownLatch.Sync#tryReleaseShared，通过CAS将state减1，当state=0时，方法返回true。

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

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

AbstractQueuedSynchronizer#unparkSuccessor。唤醒线程。

private void unparkSuccessor(Node node) {
    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;
    }
    if (s != null)
        LockSupport.unpark(s.thread);
}

await()

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

AbstractQueuedSynchronizer#acquireSharedInterruptibly。

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

CountDownLatch.Sync#tryAcquireShared。判断当前状态是否为0。

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

CountDownLatch.Sync#tryAcquireShared方法返回一个负值时，将会调用AbstractQueuedSynchronizer#doAcquireSharedInterruptibly()方法，将调用await()方法的线程加入到阻塞队列中，并将此线程挂起。

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

AbstractQueuedSynchronizer#setHeadAndPropagate。

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

AbstractQueuedSynchronizer#shouldParkAfterFailedAcquire。设置pred的状态为可唤醒，如果前面一个节点的状态就是Node.SIGNAL，则返回true。

    private static boolean shouldParkAfterFailedAcquire(Node pred, Node node) {
        int ws = pred.waitStatus;
        if (ws == Node.SIGNAL)
            /*
             * This node has already set status asking a release
             * to signal it, so it can safely park.
             */
            return true;
        if (ws > 0) {
            /*
             * Predecessor was cancelled. Skip over predecessors and
             * indicate retry.
             */
            do {
                node.prev = pred = pred.prev;
            } while (pred.waitStatus > 0);
            pred.next = node;
        } else {
            /*
             * waitStatus must be 0 or PROPAGATE.  Indicate that we
             * need a signal, but don't park yet.  Caller will need to
             * retry to make sure it cannot acquire before parking.
             */
            compareAndSetWaitStatus(pred, ws, Node.SIGNAL);
        }
        return false;
    }

AbstractQueuedSynchronizer#parkAndCheckInterrupt，进行阻塞。

private final boolean parkAndCheckInterrupt() {
    LockSupport.park(this);
    return Thread.interrupted();
}

总结概括

1. 初始化，设置state为固定值。
2. countDown，每次都在进行state--。当state=0的时候，进行线程的释放。
3. await，会创建节点线程，进行阻塞。