CountDownLatch源码解析

文章只是简单记述下CountDownLatch的原理，看过AQS、ReentrantLock和ReentrantReadWriteLock的文章后，可以大致清楚了解和AQS相关工具的使用。

CountDownLatch本身也是通过一个内部类Sync实现AQS，所以其构造方法内部也是实现如何去初始化这个Sync的过程。

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

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

    protected final void setState(int newState) {
        state = newState;
    }

内部初始化方法只是简单设置state的初始化状态值。最主要的时await和countDown两个方法的使用。

1、await方法解析

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

内部的await时直接调用sync的方法，可以通过方法名知道这个方法时响应中断的，并且默认传递了参数1。除此之外await还有其他响应超时的方法实现。

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

下面来看doAcquireSharedInterruptibly的实现

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

看过AQS的是实现应该知道addWaiter本身的作用就是初始化队列(如果队列不存在)，并且让当前节点入队。所以这里可以看到CountDownLatch并不像ReentrantLock那样尝试获取锁，而是直接就加入队列中并且等待唤醒。所以这里在挂起之前，或者时从parkAndCheckInterrupt唤醒之前，都是会先检查当前状态（就是初始化设置的state）是否为0，具体实现是在tryAcquireShared中

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

而且我们也可以看到一点，一旦r>=0成功，表示资源为0可以释放await阻塞的方法，会调用setHeadAndPropagate，也表示允许多个线程同时调用同一个CountDownLatch.await方法，并且当CountDownLatch.state的值同时为0时一起唤醒。所以我这里有一个疑问，会不会有多个线程调用CountDownLatch.countDown()导致state为负数呢？

上面的方法要注意一个点，就是CountDownLatch是如何同时唤醒多个节点的，准确说应该是唤醒队列的其他节点。这里我们来看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();
        }
    }

CountDownLatch这里能走到setHendAndPropagte，说明tryAcquireShared返回的结果必然大于或者等于0，其实在CountDownLatch这里能进来只能返回结果1，所以里说当然的这里会进入到if条件判断里面。在CountDownLatch中所有节点的属性都是Node.SHARED的，所以自然也可以走到doReleaseShared这个方法里面。并通过调用方法unparkSucessor持续唤醒下一个节点。这里设置头节点为PROPAGATE的意义不大，可以忽略。

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

这里有一个疑问，如果CountDownLatch.countDown已经使得节点状态为0，并且队列中的节点已经全部被唤醒，此时如果再加入一个节点会怎样？

这里不难思考，在没加入一个新节点之前，队列的头尾节点都指向一个，所以在doReleaseShared这里应该是因为head = tail而直接退出了方法，所以新加入的节点应该直接看

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

新加入的节点因为头节点是head，所以会进入tryAcquireShared，并且返回结果1，后续也就和上面一样。总结来说新加入的节点并不会阻塞，而是直接通过。这样说明了CountDownLatch本身其实不支持重复使用，如果需要重复使用的话需要使用Cyclicbarrier。

2、countDown方法解析

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

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

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

countDown方法比较简单，就是判断当前的state是否为0，是的话返回false，否则将当前state值减1之后CAS。这里其实然我有点意外的是为什么两次判断state为0的时候返回值会不一样。

我大概是这么猜测的，tryReleaseShared第一次进来的时候，拿到状态为0，说明CountDownLatch已经被别的线程设置为0了，就是说有别的线程去唤醒队列中的等待队列。所以直接返回false可以避免同时有多个线程去唤醒队列中的节点。退一步讲如果真的有多个线程去同时唤醒，doReleaseShared中第一个if判断就已经把头节点设置为0，阻止其他线程唤醒。

而第二个nextc == 0，本身就是CAS成功，说明当前线程是最后一个把state设置为0的，所以需要它去唤醒队列中的节点，才让其返回true，进入方法doReleaseShared方法。而doReleaseShared会唤醒第一个头节点，头节点之后再acquireQueue中又会通过setHeadAndPropage中再次调用doReleaseShared方法，依次唤醒队列中的节点。

3、CountDownLatch使用注意事项

A、countDown这个方法一定要放在finally块里面，如果少了一个countDown，那些处在await的方法必然会一直阻塞；

B、如果使用带过期时间的await时，如果达到timeout时间，被await阻塞的主线程会继续往下执行，此时进行countdown的任务线程可能仍在执行中。通常我们使用countDownLatch时，场景都是主线程依赖子线程的运行结果，在上述情况下，极有可能发生线程安全问题，需要我们判断await方法的返回值，返回false表明是存在超时，需要谨慎处理。（比如传给任务线程一个HashMap，任务线程会对该Map进行增删，处理完后会交给主线程遍历，在上面这种情况下就会发生ConcurrentModificationException）。

解决办法：

1. await返回false,主线程主动抛出异常，终止接下来的操作。
2. 任务中加一个任务完成标志位（需要volatile修饰），countdown完成之后将标志位置为完成状态，存在超时则判断标志位，只处理标志位完成的任务
3. 把任务封装成Callable，直接使用ExecutorService接口带超时参数的invokeAll方法（任务超时后会被cancel掉），Future接口带了isCancelled方法可以得到任务运行是否被取消。

public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks,long timeout, TimeUnit unit)

以上解决方法是网上提出的，我个人觉得如果真的await超时了，需要结合业务逻辑判断执行countDown的方法逻辑是否会影响到执行await的方法逻辑，是否需要再await超时返回的时候中端掉，如果是的话最好把执行countDown的方法加上中断标志位检查，而不是放在future接口中交给系统去操作，最后视情况确定是否要结束执行await的方法。