From 865da9c525488b6bbb046a278567b4214d8ed7f5 Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Sat, 14 Dec 2024 00:43:11 -0300 Subject: [PATCH 1/4] fn/ContextGuard: test cancelling blocking context Make sure WgWait() doesn't block. --- fn/context_guard_test.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/fn/context_guard_test.go b/fn/context_guard_test.go index e13cba5dc0..76d63a50f5 100644 --- a/fn/context_guard_test.go +++ b/fn/context_guard_test.go @@ -298,6 +298,12 @@ func TestContextGuard(t *testing.T) { case <-time.After(time.Second): t.Fatalf("timeout") } + + // Cancel the context. + cancel() + + // Make sure wg's counter gets to 0 eventually. + g.WgWait() }) // Test that if we add the CustomTimeoutCGOpt option, then the context From 1750aec13dbaf0b498d760f373df1f596665c73e Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Sat, 14 Dec 2024 00:35:06 -0300 Subject: [PATCH 2/4] fn: remove uneeded argument of ctxBlocking Removed 'cancel' argument, because it is called only in case the context has already expired and the only action that cancel function did was cancelling the context. --- fn/context_guard.go | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/fn/context_guard.go b/fn/context_guard.go index 3868a3ff02..7c1d5f4832 100644 --- a/fn/context_guard.go +++ b/fn/context_guard.go @@ -149,7 +149,7 @@ func (g *ContextGuard) Create(ctx context.Context, } if opts.blocking { - g.ctxBlocking(ctx, cancel) + g.ctxBlocking(ctx) return ctx, cancel } @@ -196,14 +196,10 @@ func (g *ContextGuard) ctxQuitUnsafe(ctx context.Context, } // ctxBlocking spins off a goroutine that will block until the passed context -// is cancelled after which it will call the passed cancel function and -// decrement the wait group. -func (g *ContextGuard) ctxBlocking(ctx context.Context, - cancel context.CancelFunc) { - +// is cancelled after which it will decrement the wait group. +func (g *ContextGuard) ctxBlocking(ctx context.Context) { g.wg.Add(1) go func() { - defer cancel() defer g.wg.Done() select { From e9ab6037350ff6b820c19d250265265b5fdef5e5 Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Sat, 14 Dec 2024 15:15:40 -0300 Subject: [PATCH 3/4] fn/ContextGuard: clear store of cancel funcs If ContextGuard lives for some time after Quit method is called, the map won't be collected by GC. Optimization. --- fn/context_guard.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/fn/context_guard.go b/fn/context_guard.go index 7c1d5f4832..cfe5a32998 100644 --- a/fn/context_guard.go +++ b/fn/context_guard.go @@ -51,6 +51,10 @@ func (g *ContextGuard) Quit() { cancel() } + // Clear cancelFns. It is safe to use nil, because no write + // operations to it can happen after g.quit is closed. + g.cancelFns = nil + close(g.quit) }) } From 07c46680e9ba5b8c64956871c608b230168b4348 Mon Sep 17 00:00:00 2001 From: Boris Nagaev Date: Sat, 14 Dec 2024 14:51:40 -0300 Subject: [PATCH 4/4] 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 cfe5a32998..94b2d8ee8f 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 76d63a50f5..576ca5364f 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() +}