From 07c46680e9ba5b8c64956871c608b230168b4348 Mon Sep 17 00:00:00 2001
From: Boris Nagaev <bnagaev@gmail.com>
Date: Sat, 14 Dec 2024 14:51:40 -0300
Subject: [PATCH] fn/ContextGuard: use context.AfterFunc to wait

Simplifies context cancellation handling by using context.AfterFunc instead of a
goroutine to wait for context cancellation. This approach avoids the overhead of
a goroutine during the waiting period.

For ctxQuitUnsafe, since g.quit is closed only in the Quit method (which also
cancels all associated contexts), waiting on context cancellation ensures the
same behavior without unnecessary dependency on g.quit.

Added a test to ensure that the Create method does not launch any goroutines.
---
 fn/context_guard.go      | 37 +++++++++++++++++--------------------
 fn/context_guard_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 20 deletions(-)

diff --git a/fn/context_guard.go b/fn/context_guard.go
index cfe5a3299..94b2d8ee8 100644
--- a/fn/context_guard.go
+++ b/fn/context_guard.go
@@ -173,9 +173,10 @@ func (g *ContextGuard) Create(ctx context.Context,
 	return ctx, cancel
 }
 
-// ctxQuitUnsafe spins off a goroutine that will block until the passed context
-// is cancelled or until the quit channel has been signaled after which it will
-// call the passed cancel function and decrement the wait group.
+// ctxQuitUnsafe increases the wait group counter, waits until the context is
+// cancelled and decreases the wait group counter. It stores the passed cancel
+// function and returns a wrapped version, which removed the stored one and
+// calls it. The Quit method calls all the stored cancel functions.
 //
 // NOTE: the caller must hold the ContextGuard's mutex before calling this
 // function.
@@ -185,31 +186,27 @@ func (g *ContextGuard) ctxQuitUnsafe(ctx context.Context,
 	cancel = g.addCancelFnUnsafe(cancel)
 
 	g.wg.Add(1)
-	go func() {
-		defer cancel()
-		defer g.wg.Done()
 
-		select {
-		case <-g.quit:
-
-		case <-ctx.Done():
-		}
-	}()
+	// We don't have to wait on g.quit here: g.quit can be closed only in
+	// the Quit method, which also closes the context we are waiting for.
+	context.AfterFunc(ctx, func() {
+		g.wg.Done()
+	})
 
 	return cancel
 }
 
-// ctxBlocking spins off a goroutine that will block until the passed context
-// is cancelled after which it will decrement the wait group.
+// ctxBlocking increases the wait group counter, waits until the context is
+// cancelled and decreases the wait group counter.
+//
+// NOTE: the caller must hold the ContextGuard's mutex before calling this
+// function.
 func (g *ContextGuard) ctxBlocking(ctx context.Context) {
 	g.wg.Add(1)
-	go func() {
-		defer g.wg.Done()
 
-		select {
-		case <-ctx.Done():
-		}
-	}()
+	context.AfterFunc(ctx, func() {
+		g.wg.Done()
+	})
 }
 
 // addCancelFnUnsafe adds a context cancel function to the manager and returns a
diff --git a/fn/context_guard_test.go b/fn/context_guard_test.go
index 76d63a50f..576ca5364 100644
--- a/fn/context_guard_test.go
+++ b/fn/context_guard_test.go
@@ -2,8 +2,11 @@ package fn
 
 import (
 	"context"
+	"runtime"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 // TestContextGuard tests the behaviour of the ContextGuard.
@@ -439,3 +442,36 @@ func TestContextGuard(t *testing.T) {
 		}
 	})
 }
+
+// TestContextGuardCountGoroutines makes sure that ContextGuard doesn't create
+// any goroutines while waiting for contexts.
+func TestContextGuardCountGoroutines(t *testing.T) {
+	// NOTE: t.Parallel() is not called in this test because it relies on an
+	// accurate count of active goroutines. Running other tests in parallel
+	// would introduce additional goroutines, leading to unreliable results.
+
+	g := NewContextGuard()
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// Count goroutines before contexts are created.
+	count1 := runtime.NumGoroutine()
+
+	// Create 1000 contexts of each type.
+	for i := 0; i < 1000; i++ {
+		_, _ = g.Create(ctx)
+		_, _ = g.Create(ctx, WithBlockingCG())
+		_, _ = g.Create(ctx, WithTimeoutCG())
+		_, _ = g.Create(ctx, WithBlockingCG(), WithTimeoutCG())
+	}
+
+	// Make sure no new goroutine was launched.
+	count2 := runtime.NumGoroutine()
+	require.LessOrEqual(t, count2, count1)
+
+	// Cancel root context.
+	cancel()
+
+	// Make sure wg's counter gets to 0 eventually.
+	g.WgWait()
+}