并发编程之CountDownLatch源码解析

简介

CountDownLatch允许一个或多个线程等待其它线程的操作执行完毕后再执行后续的操作。是通过一个计数器来实现的，计数器的初始值是线程的数量。每当一个线程执行完毕后，计数器的值就-1，当计数器的值为0时，表示所有线程都执行完毕，然后等待的线程就可恢复工作。
CountDownLatch的通常用法和Thread.join()有点类似，等待其它线程都完成后再执行主任务。

简单使用

public class CountDownLatchTest {
    public static void main(String[] args) {
        final CountDownLatch latch = new CountDownLatch(2);
        System.out.println("主线程开始执行…… ……");
        //第一个子线程执行
        ExecutorService es1 = Executors.newSingleThreadExecutor();
        es1.execute(() -> {
            try {
                Thread.sleep(2000);
                System.out.println("子线程："+Thread.currentThread().getName()+"执行");
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            latch.countDown();
        });
        es1.shutdown();

        //第二个子线程执行
        ExecutorService es2 = Executors.newSingleThreadExecutor();
        es2.execute(() -> {
            try {
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
            System.out.println("子线程："+Thread.currentThread().getName()+"执行");
            latch.countDown();
        });
        es2.shutdown();
        System.out.println("等待两个线程执行完毕…… ……");
        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        System.out.println("两个子线程都执行完毕，继续执行主线程");
    }
}

运行结果：

主线程开始执行…… ……
等待两个线程执行完毕…… ……
子线程：pool-1-thread-1执行
子线程：pool-2-thread-1执行
两个子线程都执行完毕，继续执行主线程

可以看出，main线程执行到latch.await()处被阻塞了，一直等到两个子线程执行完调用latch.countDown()后，main线程继续执行。

源码解析

Sync内部类

private static final class Sync extends AbstractQueuedSynchronizer {
        private static final long serialVersionUID = 4982264981922014374L;
		//CountDownLatch初始化调用Sync初始化
		//传入count，表示初始次数
        Sync(int count) {
            setState(count);//设置state为count
        }
		//获取state
        int getCount() {
            return getState();
        }
		//尝试获取共享锁
        protected int tryAcquireShared(int acquires) {
        	// state等于0返回1，count减为0的时候获取总是成功
        	// state不等于0返回-1，count不为0的时候总是要排队
            return (getState() == 0) ? 1 : -1;
        }
		// 尝试释放锁
        protected boolean tryReleaseShared(int releases) {
            //自旋
            for (;;) {
            	//获取state值
                int c = getState();
                //等于0则无法再释放
                if (c == 0)
                    return false;
                //state-1
                int nextc = c-1;
                //cas写入
                if (compareAndSetState(c, nextc))
                	//减去1后state为0则表示无等待线程，会唤醒后面排队的线程
                    return nextc == 0;
            }
        }
    }

CountDownLatch构造方法

	//传入count表示初始次数，即调用countDown次数
    public CountDownLatch(int count) {
    	//count校验，不可小于0
        if (count < 0) throw new IllegalArgumentException("count < 0");
        //调用Sync类初始化
        this.sync = new Sync(count);
    }

await方法

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

调用aqs的acquireSharedInterruptibly方法

    public final void acquireSharedInterruptibly(int arg)
            throws InterruptedException {
         //线程是否中止，
        if (Thread.interrupted())
            throw new InterruptedException();
         //tryAcquireShared调用内部类Sync的tryAcquireShared方法
         //state不等于0则返回-1，表示有线程任在执行
         //调用doAcquireSharedInterruptibly阻塞当前线程
        if (tryAcquireShared(arg) < 0)
            doAcquireSharedInterruptibly(arg);
    }

继续扒doAcquireSharedInterruptibly

private void doAcquireSharedInterruptibly(int arg)
        throws InterruptedException {
        //添加Node至aqs同步队列，方法详情见前文 AQS原理
        final Node node = addWaiter(Node.SHARED);
        boolean failed = true;
        try {
        	//自旋
            for (;;) {
            	//node上一节点
                final Node p = node.predecessor();
                //上一节点为队首，表示排同步队列第一个，可尝试获取锁
                if (p == head) {
                	//state等于0返回1
                    int r = tryAcquireShared(arg);
                    //等待线程都结束进入if，等待队列设置null垃圾回收，返回
                    if (r >= 0) {
                        setHeadAndPropagate(node, r);
                        p.next = null;
                        failed = false;
                        return;
                    }
                }
                //shouldParkAfterFailedAcquire
                //parkAndCheckInterrupt方法见前文 AQS原理
                //parkAndCheckInterrupt最终调用 LockSupport.park(this)阻塞线程
                if (shouldParkAfterFailedAcquire(p, node) &&
                    parkAndCheckInterrupt())
                    throw new InterruptedException();
            }
        } finally {
            if (failed)
                cancelAcquire(node);
        }
    }

await()方法是等待其它线程完成的方法，它会先尝试获取一下共享锁，如果失败则进入AQS的队列中排队等待被唤醒。

shutdown方法

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

调用aqs的releaseShared方法

    public final boolean releaseShared(int arg) {
    	//内部类Sync方法，尝试释放锁
    	//tryReleaseShared返回true表示state为0
    	//所有等待线程执行完成，可唤醒线程
        if (tryReleaseShared(arg)) {
            doReleaseShared();
            return true;
        }
        return false;
    }

doReleaseShared唤醒线程操作

 private void doReleaseShared() {
 		//自旋
        for (;;) {
        	//同步队列首节点
            Node h = head;
            if (h != null && h != tail) {
                int ws = h.waitStatus;
                // 如果头节点状态为SIGNAL，说明要唤醒下一个节点
                if (ws == Node.SIGNAL) {
                    if (!compareAndSetWaitStatus(h, Node.SIGNAL, 0))
                        continue;           
                    // 唤醒下一个节点
                    unparkSuccessor(h);
                }
                // 把头节点的状态改为PROPAGATE成功才会跳到下面的if
                else if (ws == 0 &&
                         !compareAndSetWaitStatus(h, 0, Node.PROPAGATE))
                    continue;                
            }
            // 如果唤醒后head没变，则跳出循环
            if (h == head)                   
                break;
        }
    }

注意：doReleaseShared是个死循环，一直唤醒所有状态为SIGNAL的线程，直到h == head时表示全部唤醒后退出循环（head为同步队列头，每唤醒线程head头会变化，相当于队列左移），这也是AQS的共享锁模式。
为何CountDownLatch需要用共享锁模式?因为CountDownLatch的await()多个线程可以调用多次，当调用多次的时候这些线程都要进入AQS队列中排队，当count次数减为0的时候，它们都需要被唤醒，继续执行任务，如果使用互斥锁则不行，互斥锁在多个线程之间是互斥的，一次只能唤醒一个，不能保证当count减为0的时候这些调用了await()方法等待的线程都被唤醒。

countDown()方法，会释放一个共享锁，也就是count/state的次数会减1。
tryReleaseShared()每次会把count/state的次数减1，当其减为0的时候返回true，这时候才会唤醒等待的线程。

总结

1、CountDownLatch使用AQS的共享锁机制实现；
2、每次调用await()方法的时候会尝试获取锁，这里的获取锁其实是检查AQS的state变量的值是否为0；
3、当count/state的值减为0的时候会唤醒排队着的线程；