task pool: avoid race condition on shutdown/close

When we close a task pool, we use a condition variable to wait for
workers to exit, if any workers are still running. The workers,
in turn, use the condition variable to notify that they've exited.
Unfortunately, the workers aren't using the lock (the rationale was
that it's all atomic ops so they don't need to), which means that
it's possible to have the following sequence:

	Close(): obtain current live count
	worker: decrement live count
	worker: send broadcast to condition variable
	Close(): wait on condition variable

To resolve this, we make the worker update also request the lock.

We add a simple reproducer for this. Note that simple doesn't mean
it fires completely reliably; on my laptop, the test causes a test
timeout about 60% of the time without the fix. If you add a short
delay between sampling the live count and waiting on the condition
variable, the deadlocks move from "60% chance of hitting it in
a million trials" to "nearly always".

task/pool.go

@@ -85,7 +85,12 @@ func (p *Pool) Stats() (live, unblocked, target int) {
 
 // Close is a Shutdown followed by waiting for all jobs to exit.
 func (p *Pool) Close() {
+	// important to note: p.cond.Wait() is actually releasing this lock,
+	// then reacquiring it when the wait succeeds. This means that
+	// nothing which uses the lock can trigger between our read of
+	// live, and our wait on the condition variable...
 	p.mu.Lock()
+	defer p.mu.Unlock()
 	p.Shutdown()
 	live := atomic.LoadInt32(&p.live)
 	for live > 0 {
@@ -116,6 +121,22 @@ func (p *Pool) addWorker() {
 // too many unblocked goroutines, otherwise it exits.
 func (p *Pool) work() {
 	defer func() {
+		// The lock prevents our modification of p.live from
+		// happening between the read of p.live and the wait on
+		// the condition variable in p.Close. Otherwise, it's
+		// possible for these to interleave as:
+		//
+		// p.Close        this function
+		// -------        -------------
+		// read p.live
+		//                modify p.live
+		//                broadcast to p.cond
+		// p.Cond.Wait
+		//
+		// and the wait never terminates because the broadcast
+		// happened before that.
+		p.mu.Lock()
+		defer p.mu.Unlock()
 		live := atomic.AddInt32(&p.live, -1)
 		if p.stats != nil {
 			p.stats.PoolSize(int(live))

task/pool_test.go

@@ -430,3 +430,29 @@ func TestPoolStartup(t *testing.T) {
 		t.Fatalf("expected no more adds, got %d including previous 3", v)
 	}
 }
+
+// There was a race condition in Pool.Close(), where it was
+// possible to have a worker thread broadcast to the condition
+// variable *after* the Close had checked the current value of
+// p.live, but *before* it had gotten to waiting. This is a very
+// narrow window. If you're chasing this down, consider adding
+// a short time.Sleep before the `p.Cond.Wait` call in `Pool.Close`,
+// with which this test would typically deadlock within the first
+// few iterations. The 1M repetitions produced about 65% failures
+// on my laptop, with successes taking under two seconds to run.
+//
+// The key interaction is that the individual worker pool `work`
+// calls are exiting almost immediately after `p.targetN` gets set
+// to zero; if they take any time to exit, the Close will
+// be waiting on the condition variable before they get there.
+func TestPoolShutdown(t *testing.T) {
+	for i := 0; i < 1000000; i++ {
+		ch := make(chan struct{})
+		doSomething := func() {
+			<-ch
+		}
+		p := NewPool(3, doSomething, nil)
+		close(ch)
+		p.Close()
+	}
+}