@@ -22,7 +22,7 @@ import (
22
22
"context"
23
23
"fmt"
24
24
"io"
25
- "sync"
25
+ "sync/atomic "
26
26
27
27
"google.golang.org/grpc/balancer"
28
28
"google.golang.org/grpc/codes"
@@ -33,35 +33,43 @@ import (
33
33
"google.golang.org/grpc/status"
34
34
)
35
35
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
+
36
47
// pickerWrapper is a wrapper of balancer.Picker. It blocks on certain pick
37
48
// actions and unblock when there's a picker update.
38
49
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 ]
43
52
statsHandlers []stats.Handler // to record blocking picker calls
44
53
}
45
54
46
55
func newPickerWrapper (statsHandlers []stats.Handler ) * pickerWrapper {
47
- return & pickerWrapper {
48
- blockingCh : make (chan struct {}),
56
+ pw := & pickerWrapper {
49
57
statsHandlers : statsHandlers ,
50
58
}
59
+ pw .pickerGen .Store (& pickerGeneration {
60
+ blockingCh : make (chan struct {}),
61
+ })
62
+ return pw
51
63
}
52
64
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.
54
67
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 )
65
73
}
66
74
67
75
// doneChannelzWrapper performs the following:
@@ -98,20 +106,17 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
98
106
var lastPickErr error
99
107
100
108
for {
101
- pw .mu .Lock ()
102
- if pw .done {
103
- pw .mu .Unlock ()
109
+ pg := pw .pickerGen .Load ()
110
+ if pg == nil {
104
111
return nil , balancer.PickResult {}, ErrClientConnClosing
105
112
}
106
-
107
- if pw .picker == nil {
108
- ch = pw .blockingCh
113
+ if pg .picker == nil {
114
+ ch = pg .blockingCh
109
115
}
110
- if ch == pw .blockingCh {
116
+ if ch == pg .blockingCh {
111
117
// This could happen when either:
112
118
// - 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.
115
120
select {
116
121
case <- ctx .Done ():
117
122
var errStr string
@@ -145,9 +150,8 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
145
150
}
146
151
}
147
152
148
- ch = pw .blockingCh
149
- p := pw .picker
150
- pw .mu .Unlock ()
153
+ ch = pg .blockingCh
154
+ p := pg .picker
151
155
152
156
pickResult , err := p .Pick (info )
153
157
if err != nil {
@@ -197,24 +201,15 @@ func (pw *pickerWrapper) pick(ctx context.Context, failfast bool, info balancer.
197
201
}
198
202
199
203
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 )
207
206
}
208
207
209
208
// reset clears the pickerWrapper and prepares it for being used again when idle
210
209
// mode is exited.
211
210
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 )
218
213
}
219
214
220
215
// dropError is a wrapper error that indicates the LB policy wishes to drop the
0 commit comments