CountDownLatch源码剖析

在并发多线程之AQS源码分析（上）和并发多线程之AQS源码分析（下）中分析了AQS的独占锁下的方法。当然AQS也提供了共享锁的实现。接下来就先简单分析一下AQS中关于共享锁的实现原理流程及AQS源码，然后再剖析一下基于AQS实现的共享锁CountDownLatch。

AQS中共享锁的实现

在AQS对线程的调用阻塞时通过同步队列实现的，将阻塞的线程添加至同步队列中，当满足条件时则从队列中唤醒。对于独占模式，从队列中每次唤醒一个线程，但是对于共享模式则需要将所有的共享模式的线程均唤醒，那么在锁释放时，获取在线程获取到锁时均需要判断后续节点是否是共享锁，是共享锁则依次全部唤醒。
其原理如下：

如图所示：
在同步队列中的第一个线程获取到锁资源后，后续节点是共享节点的话，则会挨个的依次唤醒节点。

AQS中共享锁相关的方法

接下来来剖析一下AQS中与共享锁相关的方法。

1. 共享锁释放。

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

   该方法的主要流程就是通过不停的循环从头节点向后遍历依次唤醒后续节点。这里可能会有个疑问，为什么需要在循环中各种CAS操作判断呢？
   1>第一个判断就是代表头节点为SIGNAL,证明后续有节点并且可能被阻塞，所以需要将头节点状态修改为0，修改状态失败时则需要继续循环唤醒。修改成功则会唤醒后续节点。
   2>第二个判断是头节点状态为0，则证明可能后续可能没有节点，但是在将节点状态修由0修改为PROPAGATE失败时，则证明头节点不为0了，那证明有其他线程修改了状态，可能是新增加节点，所以需要继续循环唤醒新追加的节点。
   3>最后的h==head判断，证明后续节点没有继续执行，或者说已经到了同步队列尾部，此时说明不需要再唤醒后续节点了则结束循环。为什么说没有后续节点需要唤醒了呢？这是因为后续节点被唤醒并且获取锁则会修改头节点为当前线程节点，那么head就会变更，所以h==head证明没有后续节点在执行了。

2. 获取锁后设置头节点并传播唤醒。

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

   该方法首先是将当前获取锁的线程设置为头节点，然后判断后续节点是否是共享节点，是共享节点的话则调用doReleaseShared()进行传播唤醒后续节点。单看这个方法可能会有一个疑问，就是多个线程在获取锁时，并且三个线程为ABC，在同步队列中的顺序也同样是这么排序，当在获取锁的过程中，AB挂起阻塞，C还未挂起，此时条件已经满足了（也就是在所谓的自旋转的过程中获取了锁），那么会直接从C开始执行并唤醒后续节点，而AB则被丢弃了。显然这种问题是不会存在的，可以看第三步说明的方法。

3. 共享锁抢占资源。

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

该方法的逻辑与独占锁的逻辑没有太大的区别，就是在获取锁是将setHead方法该为了setHeadAndPropagate进行唤醒的传播了，在步骤2有一个疑问，这里可以看获取资源的前一步判断 if (p == head)，这是是只有前置节点为头节点的线程才可以抢占，所以就不会出现步骤2的疑问了。关于共享锁相关的AQS方法已经介绍完毕，接下来来剖析一下CountDownLatch是如何使用AQS的共享锁框架的。

CountDownLatch原理浅析

这里从一个CountDownLatch实例开始：

public class CountDownLatchDemo {

    public static void main(String[] args) throws InterruptedException {
        CountDownLatch countDownLatch = new CountDownLatch(3);

        Thread t1 = new Thread(new Task(countDownLatch));
        Thread t2 = new Thread(new Task(countDownLatch));
        Thread t3 = new Thread(new Task(countDownLatch));


        t1.start();
        t2.start();
        t3.start();

        countDownLatch.await();
        System.out.println("3个任务完成");
    }

    static class Task implements Runnable {

        private final CountDownLatch countDownLatch;

        Task(CountDownLatch countDownLatch) {
            this.countDownLatch = countDownLatch;
        }


        @Override
        public void run() {
            System.out.println("a task finished");
            countDownLatch.countDown();
        }
    }
}

结果如下：

a task finished
a task finished
a task finished
3个任务完成

CountDownLatch提供了两个方法await和countDown两个方法，在初始化时会指定资源的数量，每调用一次countDown方法则将资源减一，在这个过程中所有调用await方法的线程均会阻塞，当通过调用countDown方法将资源减至0时，则唤醒所有通过await方法阻塞的线程执行。

CountDownLatch源码剖析

接下来来剖析一下CountDownLatch的源码，在并发多线程之Lock及ReentrantLock源码剖析中了解到基于AQS实现特定功能的均会实现有一个静态内部类Sycn，该类继承了AbstractQueuedSynchronizer类，对于共享锁则实现一下两个方法：

protected int tryAcquireShared(int arg)
protected int tryAcquireShared(int arg)

那么就来看看在CountDownLatch类中Sync的两个重写方法：
共享获取锁：

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

可以看到获取资源的方法是state状态为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;
                if (compareAndSetState(c, nextc))
                    return nextc == 0;
            }
        }

通过该方法可以知道，将state值减一。当然state已经为0则什么也不做，只是返回falese。
在来看一下Sync的构造方法：

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

调用setState方法将state设置为指定count。
了解了静态内部类Sync的实现后，来看一下CountDownLatch的相关方法。首先来看一下构造方法：

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

调用了Sync的构造方法初始化指定的资源数量。

countDown方法

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

调用了AQS的共享释放，将state减一，同时判断state是否释放完毕（state==0）,释放完成则唤醒等待的线程，这里就不再进入方法剖析了，相关方法已经梳理过。

await方法

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

方法主要逻辑就是判断资源是否释放完毕（state==0），释放完成则获取锁并唤醒后续节点执行，没有释放完毕则进入同步队列并阻塞直到被唤醒。