ctxwatch: ContextWatcher goroutine leak

pgconn/ctxwatch/context_watcher.go

@@ -8,21 +8,21 @@ import (
 // ContextWatcher watches a context and performs an action when the context is canceled. It can watch one context at a
 // time.
 type ContextWatcher struct {
-	handler     Handler
-	unwatchChan chan struct{}
+	handler Handler
 
-	lock              sync.Mutex
-	watchInProgress   bool
-	onCancelWasCalled bool
+	// Lock protects the members below.
+	lock sync.Mutex
+	// Stop is the handle for an "after func". See [context.AfterFunc].
+	stop func() bool
+	done chan struct{}
 }
 
 // NewContextWatcher returns a ContextWatcher. onCancel will be called when a watched context is canceled.
 // OnUnwatchAfterCancel will be called when Unwatch is called and the watched context had already been canceled and
 // onCancel called.
 func NewContextWatcher(handler Handler) *ContextWatcher {
 	cw := &ContextWatcher{
-		handler:     handler,
-		unwatchChan: make(chan struct{}),
+		handler: handler,
 	}
 
 	return cw
@@ -33,25 +33,16 @@ func (cw *ContextWatcher) Watch(ctx context.Context) {
 	cw.lock.Lock()
 	defer cw.lock.Unlock()
 
-	if cw.watchInProgress {
-		panic("Watch already in progress")
+	if cw.stop != nil {
+		panic("watch already in progress")
 	}
 
-	cw.onCancelWasCalled = false
-
 	if ctx.Done() != nil {
-		cw.watchInProgress = true
-		go func() {
-			select {
-			case <-ctx.Done():
-				cw.handler.HandleCancel(ctx)
-				cw.onCancelWasCalled = true
-				<-cw.unwatchChan
-			case <-cw.unwatchChan:
-			}
-		}()
-	} else {
-		cw.watchInProgress = false
+		cw.done = make(chan struct{})
+		cw.stop = context.AfterFunc(ctx, func() {
+			cw.handler.HandleCancel(ctx)
+			close(cw.done)
+		})
 	}
 }
 
@@ -61,12 +52,12 @@ func (cw *ContextWatcher) Unwatch() {
 	cw.lock.Lock()
 	defer cw.lock.Unlock()
 
-	if cw.watchInProgress {
-		cw.unwatchChan <- struct{}{}
-		if cw.onCancelWasCalled {
+	if cw.stop != nil {
+		if !cw.stop() {
+			<-cw.done
 			cw.handler.HandleUnwatchAfterCancel()
 		}
-		cw.watchInProgress = false
+		cw.stop = nil
 	}
 }
 

pgconn/ctxwatch/synctest_test.go

@@ -0,0 +1,62 @@
+//go:build go1.26
+
+package ctxwatch_test
+
+import (
+	"context"
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"testing/synctest"
+
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
+)
+
+func TestContextWatchGoroutineBuildup(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		var cancelFuncCalls int64
+		var cleanupFuncCalls int64
+		h := &testHandler{
+			handleCancel: func(context.Context) {
+				atomic.AddInt64(&cancelFuncCalls, 1)
+			},
+			handleUnwatchAfterCancel: func() {
+				atomic.AddInt64(&cleanupFuncCalls, 1)
+			},
+		}
+		ctx, done := context.WithCancel(t.Context())
+		defer done()
+		floor := runtime.NumGoroutine()
+
+		for range 10 {
+			cw := ctxwatch.NewContextWatcher(h)
+			cw.Watch(ctx)
+			defer cw.Unwatch()
+		}
+		synctest.Wait()
+		done()
+		for range 10 {
+			cw := ctxwatch.NewContextWatcher(h)
+			cw.Watch(ctx)
+			cw.Unwatch()
+		}
+
+		synctest.Wait()
+		outstanding := runtime.NumGoroutine() - floor
+		t.Log("outstanding goroutines:", outstanding)
+		if outstanding != 0 {
+			t.Fail()
+		}
+
+		actualCancelFuncCalls := atomic.LoadInt64(&cancelFuncCalls)
+		t.Log("cancel:", actualCancelFuncCalls)
+		if actualCancelFuncCalls != 20 {
+			t.Fail()
+		}
+		actualCleanupFuncCalls := atomic.LoadInt64(&cleanupFuncCalls)
+		t.Log("cleanup:", actualCleanupFuncCalls)
+		if actualCleanupFuncCalls != 10 {
+			t.Fail()
+		}
+	})
+}