v23/context: fix goroutine leak in WithRootCancel (#214)

Whenever a child context is created via `WithRootCancel`, a goroutine
spawns to gracefully handle closing the child context whenever the root
context gets canceled.

However, the current implementation leaks goroutines whenever the child
context exits before the root context exits. This pull request looks to
fix the problem by exiting the goroutine whenever the child context is
done.

See `TestRootCancelGoroutineLeak` for the reproduction use case and test
that demonstrates the leak. This is also easily visible via `pprof`
using the `goroutine` module.

Co-authored-by: Richard Huang <rhuang@grailbio.com>

Author: richhx

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()