CountDownLatch源码分析

CountDownLatch的作用是可以让主线程等待副线程执行任务，当全部副线程执行完任务后主线程再执行自己的任务。如下是主线程与副线程的生命周期：

案例

public class CountDownLatchDemo {

    static class CountDownLatchTask implements Runnable{
        CountDownLatch countDownLatch;
        public CountDownLatchTask(CountDownLatch countDownLatch){
            this.countDownLatch = countDownLatch;
        }

        @Override
        public void run() {
            try {
                Thread.sleep(new Random().nextInt(10) * 1000);
                System.out.println(Thread.currentThread().getName()+" 已经执行完任务");
              //将计数器减1  
              countDownLatch.countDown();
            }catch (Exception e){
                e.printStackTrace();
            }
        }
    }


    public static void main(String[] args) throws Exception {
        int size = 3;
        CountDownLatch count = new CountDownLatch(size);
        CountDownLatchTask countDownLatchTask = new CountDownLatchTask(count);

        for(int i = 1; i <= size; i ++){
            Thread thread = new Thread(countDownLatchTask,"thread"+i);
            thread.start();
        }
				//主线程进入阻塞状态
        count.await();
        System.out.println("副线程任务已经全部执行完了，开始执行主线程任务");
        Thread.sleep(2000);
        System.out.println("主线程任务执行完了");
    }
}

执行结果：从打印的结果看出，主线程等副线程执行完之后才会被唤醒执行自己的任务。

thread2 已经执行完任务
thread3 已经执行完任务
thread1 已经执行完任务
副线程任务已经全部执行完了，开始执行主线程任务
主线程任务执行完了

源码分析

构造方法

构造方法带有一个int型参数，用来指定主线程等待的副线程数量。方法体中会创建一个Sync内部类，Sync继承自AQS。

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

Sync是一个共享锁，它实现了AQS的tryAcquireShared：尝试获取锁，和 tryReleaseShared：释放锁。

private static final class Sync extends AbstractQueuedSynchronizer {
    private static final long serialVersionUID = 4982264981922014374L;
	//1. 指定主线程需要等待的线程数量
    Sync(int count) {
        setState(count);
    }

    int getCount() {
        return getState();
    }
		
  	//尝试获取共享锁，当state =0时，说明副线程已经全部释放了锁
    protected int tryAcquireShared(int acquires) {
        return (getState() == 0) ? 1 : -1;
    }
		
  	//自旋释放共享锁
  	//一般的情况下，释放锁都是state +releases，但是这里是-1
    protected boolean tryReleaseShared(int releases) {
       	
        for (;;) {
            int c = getState();
            if (c == 0)
                return false;
            int nextc = c-1;
            if (compareAndSetState(c, nextc))
                return nextc == 0;
        }
    }
}

await 阻塞主线程。

当主线程启动完所有的副线程后，主线程将调用await方法进入阻塞状态，并等待所有的副线程完成任务并释放锁。

public void await() throws InterruptedException {
  	//调用的是AQS的acquireSharedInterruptibly方法
    sync.acquireSharedInterruptibly(1);
}

 public final void acquireSharedInterruptibly(int arg)
            throws InterruptedException {
  //判断线程释放被中断过 
 	if (Thread.interrupted())
            throw new InterruptedException();
  // 1.调用Sync的tryAcquireShared尝试获取锁
  if (tryAcquireShared(arg) < 0)
    	//2.自旋
      doAcquireSharedInterruptibly(arg);
 }

1. tryAcquireShared 尝试获取共享锁

主线程调用tryAcquireShared来获取共享锁，当state=0时，说明副线程已经全部执行了释放锁的操作即countDown，主线程不用进入阻塞状态执。

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

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

countDown 释放共享锁

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

countDown会调用AQS的releaseShared方法，在releaseShared方法中会先调用CountDownLatch的内部类的tryReleaseShared方法释放线程，如果调用成功，会尝试唤醒阻塞的线程。

public final boolean releaseShared(int arg) {
  	//1.释放共享锁
    if (tryReleaseShared(arg)) {
      	// 2.唤醒阻塞的线程
        doReleaseShared();
        return true;
    }
    return false;
}

1. 自旋释放锁

protected boolean tryReleaseShared(int releases) {
       	
        for (;;) {
            int c = getState();
            if (c == 0)
                return false;
            int nextc = c-1;
          	// cas更新
            if (compareAndSetState(c, nextc))
                return nextc == 0;
        }
}

2.唤醒阻塞的线程

线程成功释放锁后，将唤醒队列中的阻塞线程。

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;            // 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是一个倒计时栅栏，通过构造函数指定某个调用await的线程需要等待多少其他的线程完成任务并调用CountDownLatch的countDown释放锁后才能运行。