CountDownLatch 源码分析

本文是建立在了解 AbstractQueuedSynchronizer 的基础上的。

构造时可通过new CountDownLatch(n) 传入 n 作为计数值，调用 countDown() 方法可以递减计数器，在计数器归零之前其他线程调用 await() 方法时会被阻塞，计数变为 0 时会唤醒所有阻塞的线程。若在计数器为 0 时调用 await() 方法会直接返回，不会阻塞。

CountDownLatch 是基于 AbstractQueuedSynchronizer 实现的，简称 AQS，AQS 中包括 acquire 和 release 操作，acquire 操作在失败时会被阻塞，release 操作不会发生阻塞，所以 await() 方法对应 acquire，countDown() 方法对应 release。

AQS 这是两种模式，独占（exclusive）和共享（shared）模式。

当以独占模式 acquire 成功时，其他线程的 acquire 都不会成功，会被阻塞。在共享模式下，多个线程的 acquire 可能会成功，也可能不会。

独占模式需要实现 tryAcquire、tryRelease，可选实现 isHeldExclusively。

共享模式需要实现 tryAcquireShared、tryReleaseShared。

当 CountDownLatch 计数器归零时，多个阻塞线程会被同时唤醒，而不是只唤醒一个，说明其工作在共享模式下，所以需要实现 tryAcquireShared、tryReleaseShared。

int tryAcquireShared(int arg)
- 失败时返回负值；
- 如果共享模式的 acquire 成功但后续的共享模式 acquire 无法成功，则返回零；
- 如果共享模式的 acquire 成功且后续的共享模式 acquire 也可能成功，则返回正值，此时后续等待的线程必须检查资源的可用性。
- （支持三种不同的返回值，使得该方法可用于某些获取操作有时仅独占的情况。）一旦成功，表示该对象已被获取。
boolean tryReleaseShared(int arg)
- 如果此共享模式的释放可能使等待的线程 acquire（共享或独占）成功，则返回 true；否则返回 false

有了上面的推测之后，再进行源码分析就简单多了。

CountDownLatch 的各个方法都是简单的委托给了内部类 Sync。

countDown() 委托给了 sync.releaseShared(1)，await() 委托给了 sync.acquireSharedInterruptibly(1)。

java 复制代码

public class CountDownLatch {

    private final Sync sync;

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

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

    public boolean await(long timeout, TimeUnit unit)
        throws InterruptedException {
        return sync.tryAcquireSharedNanos(1, unit.toNanos(timeout));
    }

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

    public long getCount() {
        return sync.getCount();
    }

    public String toString() {
        return super.toString() + "[Count = " + sync.getCount() + "]";
    }
}

Sync 继承自 AbstractQueuedSynchronizer，以 state 存储计数器的值。

java 复制代码

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

    Sync(int count) {
        // 设置计数器的值
        setState(count);
    }

    int getCount() {
        // 获取计数器的值
        return getState();
    }

    protected int tryAcquireShared(int acquires) {
        // 计数器为 0 则返回 1，acquire 成功且后续的 acquire 也会成功
        // 计数器不为 0 则返回 -1， acquire 失败且后续的 acquire 也会失败
        return (getState() == 0) ? 1 : -1;
    }

    protected boolean tryReleaseShared(int releases) {
        // Decrement count; signal when transition to zero
        for (;;) {
            // 获取当前计时器的值
            int c = getState();
            // 已经为零则 release 失败，且不会唤醒阻塞的线程，因为阻塞的线程已经被为 0 的那一次 release 唤醒了。
            // 计数器为 0 后再调用 await() 不会再阻塞，也就不再需要唤醒
            if (c == 0)
                return false;
            // 计数器减一
            int nextc = c-1;
            // cas 设置新计数器的值，如果失败则循序重试
            if (compareAndSetState(c, nextc))
                // cas 成功后，如果 nextc 不为 0 则返回 false。
                // 如果 nextc 为 0 则返回 true，返回 true 时 AQS 会自动唤醒阻塞的线程。
                return nextc == 0;
        }
    }
}