Fix goroutine leak in WatchSet.watchMany

watch.go

@@ -91,10 +91,14 @@ func (w WatchSet) WatchCtx(ctx context.Context) error {
 
 // watchMany is used if there are many watchers.
 func (w WatchSet) watchMany(ctx context.Context) error {
+	// Cancel all watcher goroutines when return.
+	watcherCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	// Set up a goroutine for each watcher.
 	triggerCh := make(chan struct{}, 1)
 	watcher := func(chunk []<-chan struct{}) {
-		if err := watchFew(ctx, chunk); err == nil {
+		if err := watchFew(watcherCtx, chunk); err == nil {
 			select {
 			case triggerCh <- struct{}{}:
 			default:

watch_test.go

@@ -1,8 +1,11 @@
 package memdb
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"runtime/pprof"
+	"strings"
 	"testing"
 	"time"
 )
@@ -240,6 +243,62 @@ func TestWatch_AddWithLimit(t *testing.T) {
 	}
 }
 
+func TestWatchCtxLeak(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// We add a large number of channels to a WatchSet then
+	// call WatchCtx. If one of those channels fires, we
+	// expect to see all the goroutines spawned by WatchCtx
+	// cleaned up.
+	pprof.Do(ctx, pprof.Labels("foo", "bar"), func(ctx context.Context) {
+		ws := NewWatchSet()
+		fireCh := make(chan struct{})
+		ws.Add(fireCh)
+		for i := 0; i < 10000; i++ {
+			watchCh := make(chan struct{})
+			ws.Add(watchCh)
+		}
+		result := make(chan error)
+		go func() {
+			result <- ws.WatchCtx(ctx)
+		}()
+
+		fireCh <- struct{}{}
+
+		if err := <-result; err != nil {
+			t.Fatalf("expected no err got: %v", err)
+		}
+	})
+
+	numRetries := 3
+	var gced bool
+	for i := 0; i < numRetries; i++ {
+		var pb bytes.Buffer
+		profiler := pprof.Lookup("goroutine")
+		if profiler == nil {
+			t.Fatal("unable to find profile")
+		}
+		err := profiler.WriteTo(&pb, 1)
+		if err != nil {
+			t.Fatalf("unable to read profile: %v", err)
+		}
+		// If the debug profile dump contains the string "foo",
+		// it means one of the goroutines spawned in pprof.Do above
+		// still appears in the capture.
+		if !strings.Contains(pb.String(), "foo") {
+			gced = true
+			break
+		} else {
+			t.Log("retrying")
+			time.Sleep(1 * time.Second)
+		}
+	}
+	if !gced {
+		t.Errorf("goroutines were not garbage collected after %d retries", numRetries)
+	}
+}
+
 func BenchmarkWatch(b *testing.B) {
 	ws := NewWatchSet()
 	for i := 0; i < 1024; i++ {