CountDownLatch源码浅析

CountDownLatch也是juc包下的工具类，用于实现线程同步。

CountDownLacth 和juc（java.util.concurrent）包下的其他工具类一样，存在一个继承了AQS（AbstractQueuedSynchronizer）的内部类Sync。通过Sync来操作同步状态state来实现线程的同步。
相较于ReentrantLock，CountDownLacth 的源码更简单，大部分的功能都是在AQS中完成的。不同的是ReentrantLock使用的是独占锁，CountDownLacth 使用的共享锁

CountDownLatch

热身

CountDownLatch的使用很简单。

1. 创建一个CountDownLatch 对象，设置计数器（state）
2. 调用 countDown()方法，计数器就会减一
3. 调用await()方法，线程阻塞。直到countDown 的计数器为0。

第一种场景：需要大量的运算，分配至多个线程处理。子线程完成之后把结果交给主线程处理。

public class CountDownLatchDemo {
    public static void main(String[] args) {
        final CountDownLatch latch = new CountDownLatch(4);
        new Thread(()->{
        	//doSomething 
            latch.countDown();
            System.out.println(Thread.currentThread().getName()+"  over");
        }).start();
        new Thread(()->{
      		//doSomething 
            latch.countDown();
            System.out.println(Thread.currentThread().getName()+"  over");
        }).start();
        new Thread(()->{
       		//doSomething 
            latch.countDown();
            System.out.println(Thread.currentThread().getName()+"  over");
        }).start();
        new Thread(()->{
        	//doSomething 
            latch.countDown();
            System.out.println(Thread.currentThread().getName()+"  over");
        }).start();
        
        try {
       		 //线程阻塞等待子线程全部返回数据
            latch.await();
            System.out.println(Thread.currentThread().getName()+"  over");
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }
}

Thread-0 over
Thread-1 over
Thread-2 over
Thread-3 over
main over

还有一种常用的用法是资源初始化。当主线线程资源初始化完后之后，其他线程才开始执行任务。

public class CountDownLatchDemo {
    public static void main(String[] args) {
        final CountDownLatch latch = new CountDownLatch(1);

        new Thread(()->{
            try {
                latch.await();
                System.out.println(Thread.currentThread().getName()+"  over");
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }).start();
        new Thread(()->{
            try {
                latch.await();
                System.out.println(Thread.currentThread().getName()+"  over");
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }).start();


        try{
            //doSomething
            Thread.sleep(1000);
            System.out.println("资源初始化完成。。。。");
        }catch (InterruptedException e){
            e.printStackTrace();
        }
        latch.countDown();
    }
}

资源初始化完成。。。。
Thread-0 over
Thread-1 over

源码解析

CountDownLatch 的功能都是通过内部类Sync 实现的。 所以源码解析的重点是在AQS中。
下面的代码解析分三部分：

• 创建对象
• 调用await方法
• 调用countDown方法

创建CountDownLatch对象

我们创建CountDownLatch对象时需要传入一个计数器值。

 	/**
     * Constructs a {@code CountDownLatch} initialized with the given count.
     */
    public CountDownLatch(int count) {
        if (count < 0) throw new IllegalArgumentException("count < 0");
        //创建一个Sync对象
        this.sync = new Sync(count);
    }

可以看到CountDownLatch构造函数并没有什么多余的操作，仅仅是创建了一个Sync对象。

Sync 是CountDownLatch中的内部静态类继承了AQS，对Sync的操作最终会作用于AQS中。

 private static final class Sync extends AbstractQueuedSynchronizer

Sync.java

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

AbstractQueuedSynchronizer.java

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

Sync的构造函数调用 setState方法，该方法的作用是设置AQS中state的值。在AQS中state表示重入锁（计数器），在CountDownLatch中表示计数器。

await

线程被挂起直到计数器归零线程被唤醒或线程被中断。

CountDownLatch.java

	/**
	* Causes the current thread to wait until the latch has counted down to zero, unless the thread is	interrupted.
	* If the current count is zero then this method returns immediately.
	* If the current count is greater than zero then the current thread becomes disabled for thread scheduling purposes and lies dormant until one of two things happen:
	*      The count reaches zero due to invocations of the countDown method; or
	*      Some other thread interrupts the current thread.
	* If the current thread:
	*     has its interrupted status set on entry to this method; or
	*     is interrupted while waiting,
	* then InterruptedException is thrown and the current thread's interrupted status is cleared.
	*/
    public void await() throws InterruptedException {
        sync.acquireSharedInterruptibly(1);
    }

AbstractQueuedSynchronizer.java

	/**
	*Acquires in shared mode, aborting if interrupted. Implemented by first checking interrupt status, then invoking at least once tryAcquireShared, returning on success. Otherwise the thread is queued, possibly repeatedly blocking and unblocking, invoking tryAcquireShared until success or the thread is interrupted.
	*/
   public final void acquireSharedInterruptibly(int arg)
            throws InterruptedException {
        // 线程中断直接抛出异常
        if (Thread.interrupted())
            throw new InterruptedException();
        
        if (tryAcquireShared(arg) < 0)
        	//state!=0，线程阻塞并加入等待队列
            doAcquireSharedInterruptibly(arg);
    }

---

tryAcquireShared 用来获取共享锁，默认实现就是抛出异常，获取共享锁的逻辑需要在子类中指定。返回结果>=0表示获取成功，负数表示获取失败，需要将线程添加到队列中。

Sync.java

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

如果计数器为0时将可以获取共享锁（state==0），否则加入等待队列中。getState 方法返回值就是 创建CountDownLatch对象时传递的值。

共享锁获取失败则将当前线程封装为Node节点加入到等待队列中。如果有看过ReentrantLock会发现，该代码逻辑和AQS中方法acquireQueued逻辑大致相同。

【AbstractQueuedSynchronizer.java】

	 /**
     * 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 {
            for (;;) {
                final Node p = node.predecessor();
				//只有头结点的下一节点才有机会获取共享锁。
                if (p == head) {
                	//当 state==0时 ，tryAcquireShared()>1 
                    int r = tryAcquireShared(arg);
                    if (r >= 0) {
                    	//将当前节点设置为头结点，并唤醒下一等待节点
                        setHeadAndPropagate(node, r);
                        p.next = null; // help GC
                        failed = false;
                        return;
                    }
                }
                //此处代码和acquireQueued 方法完全一致
                //shouldParkAfterFailedAcquire 修改节点状态
                //parkAndCheckInterrupt 将线程挂起，如果线程在阻塞过程中调用了interrupt方法，则线程被唤醒后相应中断请求
                if (shouldParkAfterFailedAcquire(p, node) &&
                    parkAndCheckInterrupt())
                    throw new InterruptedException();
            }
        } finally {
        	//异常处理，取消节点的执行
            if (failed)
                cancelAcquire(node);
        }
    }

【AbstractQueuedSynchronizer.java】

    /**
     * Creates and enqueues node for current thread and given mode.
     *
     * @param mode Node.EXCLUSIVE for exclusive, Node.SHARED for shared
     * @return the new 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
        Node pred = tail;
        if (pred != null) {
            node.prev = pred;
            if (compareAndSetTail(pred, node)) {
                pred.next = node;
                return node;
            }
        }
        enq(node);
        return node;
    }

这样节点就添加成功了

countDown

调用countDown 方法计数器会减一，当state=0时，会唤醒等待队列中的线程

CountDownLatch.java

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

【AbstractQueuedSynchronizer.java】

 /**
     * Releases in shared mode.  Implemented by unblocking one or more
     * threads if {@link #tryReleaseShared} returns true.
     *
     * @param arg the release argument.  This value is conveyed to
     *        {@link #tryReleaseShared} but is otherwise uninterpreted
     *        and can represent anything you like.
     * @return the value returned from {@link #tryReleaseShared}
     */
    public final boolean releaseShared(int arg) {
		//释放共享锁
        if (tryReleaseShared(arg)) {
        	// 唤醒阻塞线程
            doReleaseShared();
            return true;
        }
        return false;
    }

当共享锁已经全部释放时，唤醒因获取共享锁而被阻塞的线程。

【Sync.java】

 protected boolean tryReleaseShared(int releases) {
            // Decrement count; signal when transition to zero
            for (;;) {
                int c = getState();
                // c==0， 说明共享锁已经全部释放
                if (c == 0)
                    return false;
                int nextc = c-1;
                if (compareAndSetState(c, nextc))
                	//说明该操作之后，共享锁已经全部释放。
                    return nextc == 0;
            }
        }

判断当前是否还有共享锁未释放。

【AbstractQueuedSynchronizer.java】

	 /**
     * Release action for shared mode -- signals successor and ensures
     * propagation. (Note: For exclusive mode, release just amounts
     * to calling unparkSuccessor of head if it needs signal.)
     */
    private void doReleaseShared() {
        /*
         * Ensure that a release propagates, even if there are other
         * in-progress acquires/releases.  This proceeds in the usual
         * way of trying to unparkSuccessor of head if it needs
         * signal. But if it does not, status is set to PROPAGATE to
         * ensure that upon release, propagation continues.
         * Additionally, we must loop in case a new node is added
         * while we are doing this. Also, unlike other uses of
         * unparkSuccessor, we need to know if CAS to reset status
         * fails, if so rechecking.
         */
        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;
        }
    }

最后唤醒之后就是如下结构