Skip to content

Commit e22436a

Browse files
authored
pickerwrapper: use atomic instead of locks (#7214)
1 parent 0020ccf commit e22436a

File tree

2 files changed

+42
-43
lines changed

2 files changed

+42
-43
lines changed

balancer_wrapper.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,10 @@ func (ccb *ccBalancerWrapper) UpdateAddresses(sc balancer.SubConn, addrs []resol
198198
func (ccb *ccBalancerWrapper) UpdateState(s balancer.State) {
199199
ccb.cc.mu.Lock()
200200
defer ccb.cc.mu.Unlock()
201+
if ccb.cc.conns == nil {
202+
// The CC has been closed; ignore this update.
203+
return
204+
}
201205

202206
ccb.mu.Lock()
203207
if ccb.closed {

picker_wrapper.go

Lines changed: 38 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import (
2222
"context"
2323
"fmt"
2424
"io"
25-
"sync"
25+
"sync/atomic"
2626

2727
"google.golang.org/grpc/balancer"
2828
"google.golang.org/grpc/codes"
@@ -33,35 +33,43 @@ import (
3333
"google.golang.org/grpc/status"
3434
)
3535

36+
// pickerGeneration stores a picker and a channel used to signal that a picker
37+
// newer than this one is available.
38+
type pickerGeneration struct {
39+
// picker is the picker produced by the LB policy. May be nil if a picker
40+
// has never been produced.
41+
picker balancer.Picker
42+
// blockingCh is closed when the picker has been invalidated because there
43+
// is a new one available.
44+
blockingCh chan struct{}
45+
}
46+
3647
// pickerWrapper is a wrapper of balancer.Picker. It blocks on certain pick
3748
// actions and unblock when there's a picker update.
3849
type pickerWrapper struct {
39-
mu sync.Mutex
40-
done bool
41-
blockingCh chan struct{}
42-
picker balancer.Picker
50+
// If pickerGen holds a nil pointer, the pickerWrapper is closed.
51+
pickerGen atomic.Pointer[pickerGeneration]
4352
statsHandlers []stats.Handler // to record blocking picker calls
4453
}
4554

4655
func newPickerWrapper(statsHandlers []stats.Handler) *pickerWrapper {
47-
return &pickerWrapper{
48-
blockingCh: make(chan struct{}),
56+
pw := &pickerWrapper{
4957
statsHandlers: statsHandlers,
5058
}
59+
pw.pickerGen.Store(&pickerGeneration{
60+
blockingCh: make(chan struct{}),
61+
})
62+
return pw
5163
}
5264

53-
// updatePicker is called by UpdateBalancerState. It unblocks all blocked pick.
65+
// updatePicker is called by UpdateState calls from the LB policy. It
66+
// unblocks all blocked pick.
5467
func (pw *pickerWrapper) updatePicker(p balancer.Picker) {
55-
pw.mu.Lock()
56-
if pw.done {
57-
pw.mu.Unlock()
58-
return
59-
}
60-
pw.picker = p
61-
// pw.blockingCh should never be nil.
62-
close(pw.blockingCh)
63-
pw.blockingCh = make(chan struct{})
64-
pw.mu.Unlock()
68+
old := pw.pickerGen.Swap(&pickerGeneration{
69+
picker: p,
70+
blockingCh: make(chan struct{}),
71+
})
72+
close(old.blockingCh)
6573
}
6674

6775
// doneChannelzWrapper performs the following:
@@ -98,20 +106,17 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
98106
var lastPickErr error
99107

100108
for {
101-
pw.mu.Lock()
102-
if pw.done {
103-
pw.mu.Unlock()
109+
pg := pw.pickerGen.Load()
110+
if pg == nil {
104111
return nil, balancer.PickResult{}, ErrClientConnClosing
105112
}
106-
107-
if pw.picker == nil {
108-
ch = pw.blockingCh
113+
if pg.picker == nil {
114+
ch = pg.blockingCh
109115
}
110-
if ch == pw.blockingCh {
116+
if ch == pg.blockingCh {
111117
// This could happen when either:
112118
// - pw.picker is nil (the previous if condition), or
113-
// - has called pick on the current picker.
114-
pw.mu.Unlock()
119+
// - we have already called pick on the current picker.
115120
select {
116121
case <-ctx.Done():
117122
var errStr string
@@ -145,9 +150,8 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
145150
}
146151
}
147152

148-
ch = pw.blockingCh
149-
p := pw.picker
150-
pw.mu.Unlock()
153+
ch = pg.blockingCh
154+
p := pg.picker
151155

152156
pickResult, err := p.Pick(info)
153157
if err != nil {
@@ -197,24 +201,15 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
197201
}
198202

199203
func (pw *pickerWrapper) close() {
200-
pw.mu.Lock()
201-
defer pw.mu.Unlock()
202-
if pw.done {
203-
return
204-
}
205-
pw.done = true
206-
close(pw.blockingCh)
204+
old := pw.pickerGen.Swap(nil)
205+
close(old.blockingCh)
207206
}
208207

209208
// reset clears the pickerWrapper and prepares it for being used again when idle
210209
// mode is exited.
211210
func (pw *pickerWrapper) reset() {
212-
pw.mu.Lock()
213-
defer pw.mu.Unlock()
214-
if pw.done {
215-
return
216-
}
217-
pw.blockingCh = make(chan struct{})
211+
old := pw.pickerGen.Swap(&pickerGeneration{blockingCh: make(chan struct{})})
212+
close(old.blockingCh)
218213
}
219214

220215
// dropError is a wrapper error that indicates the LB policy wishes to drop the

0 commit comments

Comments
 (0)