From ba494ed57f51075b14177bca071bb72542e3b890 Mon Sep 17 00:00:00 2001 From: Arjan Bal Date: Sun, 25 Aug 2024 18:01:02 +0530 Subject: [PATCH] Wait for go routines to finish before cancelling the context --- picker_wrapper_test.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/picker_wrapper_test.go b/picker_wrapper_test.go index 33eb86f94b6f..c1f5c1699809 100644 --- a/picker_wrapper_test.go +++ b/picker_wrapper_test.go @@ -21,6 +21,7 @@ package grpc import ( "context" "fmt" + "sync" "sync/atomic" "testing" "time" @@ -147,12 +148,15 @@ func (s) TestBlockingPickSCNotReady(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout) defer cancel() // All goroutines should block because subConn is not ready. + wg := sync.WaitGroup{} + wg.Add(goroutineCount) for i := goroutineCount; i > 0; i-- { go func() { if tr, _, err := bp.pick(ctx, true, balancer.PickInfo{}); err != nil || tr != testT { t.Errorf("bp.pick returned non-nil error: %v", err) } atomic.AddUint64(&finishedCount, 1) + wg.Done() }() } time.Sleep(time.Millisecond) @@ -160,4 +164,6 @@ func (s) TestBlockingPickSCNotReady(t *testing.T) { t.Errorf("finished goroutines count: %v, want 0", c) } bp.updatePicker(&testingPicker{sc: testSC, maxCalled: goroutineCount}) + // Wait for all pickers to finish before the context is cancelled. + wg.Wait() }