pickerwrapper: use atomic instead of locks

balancer_wrapper.go

@@ -198,6 +198,10 @@ func (ccb *ccBalancerWrapper) UpdateAddresses(sc balancer.SubConn, addrs []resol
 func (ccb *ccBalancerWrapper) UpdateState(s balancer.State) {
 	ccb.cc.mu.Lock()
 	defer ccb.cc.mu.Unlock()
+	if ccb.cc.conns == nil {
+		// The CC has been closed; ignore this update.
+		return
+	}
 
 	ccb.mu.Lock()
 	if ccb.closed {

picker_wrapper.go

@@ -22,7 +22,7 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"sync"
+	"sync/atomic"
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/codes"
@@ -33,35 +33,42 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+// pickerGeneration stores a picker and a channel used to signal that a picker
+// newer than this one is available.
+type pickerGeneration struct {
+	// picker is the picker produced by the LB policy.  May be nil if a picker
+	// has never been produced.
+	picker balancer.Picker
+	// blockingCh is closed when the picker has been invalidated because there
+	// is a new one available.
+	blockingCh chan struct{}
+}
+
 // pickerWrapper is a wrapper of balancer.Picker. It blocks on certain pick
 // actions and unblock when there's a picker update.
 type pickerWrapper struct {
-	mu            sync.Mutex
-	done          bool
-	blockingCh    chan struct{}
-	picker        balancer.Picker
+	// If pickerGen holds a nil pointer, the pickerWrapper is closed.
+	pickerGen     atomic.Pointer[pickerGeneration]
 	statsHandlers []stats.Handler // to record blocking picker calls
 }
 
 func newPickerWrapper(statsHandlers []stats.Handler) *pickerWrapper {
-	return &pickerWrapper{
-		blockingCh:    make(chan struct{}),
+	pw := &pickerWrapper{
 		statsHandlers: statsHandlers,
 	}
+	pw.pickerGen.Store(&pickerGeneration{
+		blockingCh: make(chan struct{}),
+	})
+	return pw
 }
 
 // updatePicker is called by UpdateBalancerState. It unblocks all blocked pick.
 func (pw *pickerWrapper) updatePicker(p balancer.Picker) {
-	pw.mu.Lock()
-	if pw.done {
-		pw.mu.Unlock()
-		return
-	}
-	pw.picker = p
-	// pw.blockingCh should never be nil.
-	close(pw.blockingCh)
-	pw.blockingCh = make(chan struct{})
-	pw.mu.Unlock()
+	old := pw.pickerGen.Swap(&pickerGeneration{
+		picker:     p,
+		blockingCh: make(chan struct{}),
+	})
+	close(old.blockingCh)
 }
 
 // doneChannelzWrapper performs the following:
@@ -98,20 +105,17 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 	var lastPickErr error
 
 	for {
-		pw.mu.Lock()
-		if pw.done {
-			pw.mu.Unlock()
+		pg := pw.pickerGen.Load()
+		if pg == nil {
 			return nil, balancer.PickResult{}, ErrClientConnClosing
 		}
-
-		if pw.picker == nil {
-			ch = pw.blockingCh
+		if pg.picker == nil {
+			ch = pg.blockingCh
 		}
-		if ch == pw.blockingCh {
+		if ch == pg.blockingCh {
 			// This could happen when either:
 			// - pw.picker is nil (the previous if condition), or
-			// - has called pick on the current picker.
-			pw.mu.Unlock()
+			// - we have already called pick on the current picker.
 			select {
 			case <-ctx.Done():
 				var errStr string
@@ -145,9 +149,8 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 			}
 		}
 
-		ch = pw.blockingCh
-		p := pw.picker
-		pw.mu.Unlock()
+		ch = pg.blockingCh
+		p := pg.picker
 
 		pickResult, err := p.Pick(info)
 		if err != nil {
@@ -197,24 +200,15 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
 }
 
 func (pw *pickerWrapper) close() {
-	pw.mu.Lock()
-	defer pw.mu.Unlock()
-	if pw.done {
-		return
-	}
-	pw.done = true
-	close(pw.blockingCh)
+	old := pw.pickerGen.Swap(nil)
+	close(old.blockingCh)
 }
 
 // reset clears the pickerWrapper and prepares it for being used again when idle
 // mode is exited.
 func (pw *pickerWrapper) reset() {
-	pw.mu.Lock()
-	defer pw.mu.Unlock()
-	if pw.done {
-		return
-	}
-	pw.blockingCh = make(chan struct{})
+	old := pw.pickerGen.Swap(&pickerGeneration{blockingCh: make(chan struct{})})
+	close(old.blockingCh)
 }
 
 // dropError is a wrapper error that indicates the LB policy wishes to drop the