From 118c31bcdd478f6f98e1fe8cffb930813147520f Mon Sep 17 00:00:00 2001 From: Tom Wieczorek Date: Fri, 5 Dec 2025 12:01:14 +0100 Subject: [PATCH] gracefulswitch: Wait for all goroutines on close Goroutines spawned during balancer swaps could outlive the call to Balancer.Close(). Monitor these via a wait group and wait for them to finish before returning from Close(). This prevents any noticeable side effects that could otherwise occur after Close() returns. RELEASE NOTES: - Closing a graceful switch balancer will now block until all pending goroutines complete. Signed-off-by: Tom Wieczorek --- internal/balancer/gracefulswitch/gracefulswitch.go | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/internal/balancer/gracefulswitch/gracefulswitch.go b/internal/balancer/gracefulswitch/gracefulswitch.go index ba25b8988718..f38de74a4933 100644 --- a/internal/balancer/gracefulswitch/gracefulswitch.go +++ b/internal/balancer/gracefulswitch/gracefulswitch.go @@ -67,6 +67,10 @@ type Balancer struct { // balancerCurrent before the UpdateSubConnState is called on the // balancerCurrent. currentMu sync.Mutex + + // activeGoroutines tracks all the goroutines that this balancer has started + // and that should be waited on when the balancer closes. + activeGoroutines sync.WaitGroup } // swap swaps out the current lb with the pending lb and updates the ClientConn. @@ -76,7 +80,9 @@ func (gsb *Balancer) swap() { cur := gsb.balancerCurrent gsb.balancerCurrent = gsb.balancerPending gsb.balancerPending = nil + gsb.activeGoroutines.Add(1) go func() { + defer gsb.activeGoroutines.Done() gsb.currentMu.Lock() defer gsb.currentMu.Unlock() cur.Close() @@ -274,6 +280,7 @@ func (gsb *Balancer) Close() { currentBalancerToClose.Close() pendingBalancerToClose.Close() + gsb.activeGoroutines.Wait() } // balancerWrapper wraps a balancer.Balancer, and overrides some Balancer @@ -324,7 +331,12 @@ func (bw *balancerWrapper) UpdateState(state balancer.State) { defer bw.gsb.mu.Unlock() bw.lastState = state + // If Close() acquires the mutex before UpdateState(), the balancer + // will already have been removed from the current or pending state when + // reaching this point. if !bw.gsb.balancerCurrentOrPending(bw) { + // Returning here ensures that (*Balancer).swap() is not invoked after + // (*Balancer).Close() and therefore prevents "use after close". return }