Merge commit 'b2e0a54eb08f0b392ceb12146b06b8846646dd51' into release-v0.1.18

Author: cosnicolaou

v23/context/context.go

@@ -303,8 +303,12 @@ func WithRootCancel(parent *T) (*T, CancelFunc) {
 		// Forward the cancelation from the root context to the newly
 		// created context.
 		go func() {
-			<-rootCtx.Done()
-			cancel()
+			select {
+			case <-rootCtx.Done():
+				cancel()
+			case <-ctx.Done():
+				cancel()
+			}
 		}()
 	} else if atomic.AddInt32(&nRootCancelWarning, 1) < 3 {
 		vlog.Errorf("context.WithRootCancel: context %+v is not derived from root v23 context.\n", parent)

v23/context/context_test.go

@@ -9,6 +9,8 @@ import (
 	gocontext "context"
 	"fmt"
 	"os"
+	"runtime"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -320,6 +322,36 @@ func TestRootCancelChain(t *testing.T) {
 	}
 }
 
+func TestRootCancelGoroutineLeak(t *testing.T) {
+	rootCtx, rootcancel := context.RootContext()
+	const iterations = 1024
+	for i := 0; i != iterations; i++ {
+		_, cancel := context.WithRootCancel(rootCtx)
+		cancel()
+	}
+
+	// Arbitrary threshold to wait for the goroutines in the created contexts
+	// above to exit. This threshold was arbitrarily created after running
+	// `go test -count=10000 -run TestRootCancelGoroutineLeak$` and verifying
+	// that the tests did not fail flakily.
+	const waitThreshold = 8*time.Millisecond
+	time.Sleep(waitThreshold)
+
+	// Verify that goroutines no longer exist in the runtime stack.
+	buf := make([]byte, 2<<20)
+	buf = buf[:runtime.Stack(buf, true)]
+	count := 0
+	for _, g := range strings.Split(string(buf), "\n\n") {
+		if strings.Contains(g, "v.io/v23/context.WithRootCancel.func1") {
+			count++
+		}
+	}
+	if count != 0 {
+		t.Errorf("expected 0 but got %d: goroutine leaking in WithRootCancel", count)
+	}
+	rootcancel()
+}
+
 func TestRootCancel_GoContext(t *testing.T) {
 	root, rootcancel := context.RootContext()