1/*
2 *
3 * Copyright 2017 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19// Package pickfirst contains the pick_first load balancing policy.
20package pickfirst
21
22import (
23 "encoding/json"
24 "errors"
25 "fmt"
26 rand "math/rand/v2"
27
28 "google.golang.org/grpc/balancer"
29 "google.golang.org/grpc/balancer/pickfirst/internal"
30 "google.golang.org/grpc/connectivity"
31 "google.golang.org/grpc/grpclog"
32 "google.golang.org/grpc/internal/envconfig"
33 internalgrpclog "google.golang.org/grpc/internal/grpclog"
34 "google.golang.org/grpc/internal/pretty"
35 "google.golang.org/grpc/resolver"
36 "google.golang.org/grpc/serviceconfig"
37
38 _ "google.golang.org/grpc/balancer/pickfirst/pickfirstleaf" // For automatically registering the new pickfirst if required.
39)
40
41func init() {
42 if envconfig.NewPickFirstEnabled {
43 return
44 }
45 balancer.Register(pickfirstBuilder{})
46}
47
48var logger = grpclog.Component("pick-first-lb")
49
50const (
51 // Name is the name of the pick_first balancer.
52 Name = "pick_first"
53 logPrefix = "[pick-first-lb %p] "
54)
55
56type pickfirstBuilder struct{}
57
58func (pickfirstBuilder) Build(cc balancer.ClientConn, _ balancer.BuildOptions) balancer.Balancer {
59 b := &pickfirstBalancer{cc: cc}
60 b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
61 return b
62}
63
64func (pickfirstBuilder) Name() string {
65 return Name
66}
67
68type pfConfig struct {
69 serviceconfig.LoadBalancingConfig `json:"-"`
70
71 // If set to true, instructs the LB policy to shuffle the order of the list
72 // of endpoints received from the name resolver before attempting to
73 // connect to them.
74 ShuffleAddressList bool `json:"shuffleAddressList"`
75}
76
77func (pickfirstBuilder) ParseConfig(js json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
78 var cfg pfConfig
79 if err := json.Unmarshal(js, &cfg); err != nil {
80 return nil, fmt.Errorf("pickfirst: unable to unmarshal LB policy config: %s, error: %v", string(js), err)
81 }
82 return cfg, nil
83}
84
85type pickfirstBalancer struct {
86 logger *internalgrpclog.PrefixLogger
87 state connectivity.State
88 cc balancer.ClientConn
89 subConn balancer.SubConn
90}
91
92func (b *pickfirstBalancer) ResolverError(err error) {
93 if b.logger.V(2) {
94 b.logger.Infof("Received error from the name resolver: %v", err)
95 }
96 if b.subConn == nil {
97 b.state = connectivity.TransientFailure
98 }
99
100 if b.state != connectivity.TransientFailure {
101 // The picker will not change since the balancer does not currently
102 // report an error.
103 return
104 }
105 b.cc.UpdateState(balancer.State{
106 ConnectivityState: connectivity.TransientFailure,
107 Picker: &picker{err: fmt.Errorf("name resolver error: %v", err)},
108 })
109}
110
111// Shuffler is an interface for shuffling an address list.
112type Shuffler interface {
113 ShuffleAddressListForTesting(n int, swap func(i, j int))
114}
115
116// ShuffleAddressListForTesting pseudo-randomizes the order of addresses. n
117// is the number of elements. swap swaps the elements with indexes i and j.
118func ShuffleAddressListForTesting(n int, swap func(i, j int)) { rand.Shuffle(n, swap) }
119
120func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState) error {
121 if len(state.ResolverState.Addresses) == 0 && len(state.ResolverState.Endpoints) == 0 {
122 // The resolver reported an empty address list. Treat it like an error by
123 // calling b.ResolverError.
124 if b.subConn != nil {
125 // Shut down the old subConn. All addresses were removed, so it is
126 // no longer valid.
127 b.subConn.Shutdown()
128 b.subConn = nil
129 }
130 b.ResolverError(errors.New("produced zero addresses"))
131 return balancer.ErrBadResolverState
132 }
133 // We don't have to guard this block with the env var because ParseConfig
134 // already does so.
135 cfg, ok := state.BalancerConfig.(pfConfig)
136 if state.BalancerConfig != nil && !ok {
137 return fmt.Errorf("pickfirst: received illegal BalancerConfig (type %T): %v", state.BalancerConfig, state.BalancerConfig)
138 }
139
140 if b.logger.V(2) {
141 b.logger.Infof("Received new config %s, resolver state %s", pretty.ToJSON(cfg), pretty.ToJSON(state.ResolverState))
142 }
143
144 var addrs []resolver.Address
145 if endpoints := state.ResolverState.Endpoints; len(endpoints) != 0 {
146 // Perform the optional shuffling described in gRFC A62. The shuffling will
147 // change the order of endpoints but not touch the order of the addresses
148 // within each endpoint. - A61
149 if cfg.ShuffleAddressList {
150 endpoints = append([]resolver.Endpoint{}, endpoints...)
151 internal.RandShuffle(len(endpoints), func(i, j int) { endpoints[i], endpoints[j] = endpoints[j], endpoints[i] })
152 }
153
154 // "Flatten the list by concatenating the ordered list of addresses for each
155 // of the endpoints, in order." - A61
156 for _, endpoint := range endpoints {
157 // "In the flattened list, interleave addresses from the two address
158 // families, as per RFC-8304 section 4." - A61
159 // TODO: support the above language.
160 addrs = append(addrs, endpoint.Addresses...)
161 }
162 } else {
163 // Endpoints not set, process addresses until we migrate resolver
164 // emissions fully to Endpoints. The top channel does wrap emitted
165 // addresses with endpoints, however some balancers such as weighted
166 // target do not forward the corresponding correct endpoints down/split
167 // endpoints properly. Once all balancers correctly forward endpoints
168 // down, can delete this else conditional.
169 addrs = state.ResolverState.Addresses
170 if cfg.ShuffleAddressList {
171 addrs = append([]resolver.Address{}, addrs...)
172 rand.Shuffle(len(addrs), func(i, j int) { addrs[i], addrs[j] = addrs[j], addrs[i] })
173 }
174 }
175
176 if b.subConn != nil {
177 b.cc.UpdateAddresses(b.subConn, addrs)
178 return nil
179 }
180
181 var subConn balancer.SubConn
182 subConn, err := b.cc.NewSubConn(addrs, balancer.NewSubConnOptions{
183 StateListener: func(state balancer.SubConnState) {
184 b.updateSubConnState(subConn, state)
185 },
186 })
187 if err != nil {
188 if b.logger.V(2) {
189 b.logger.Infof("Failed to create new SubConn: %v", err)
190 }
191 b.state = connectivity.TransientFailure
192 b.cc.UpdateState(balancer.State{
193 ConnectivityState: connectivity.TransientFailure,
194 Picker: &picker{err: fmt.Errorf("error creating connection: %v", err)},
195 })
196 return balancer.ErrBadResolverState
197 }
198 b.subConn = subConn
199 b.state = connectivity.Idle
200 b.cc.UpdateState(balancer.State{
201 ConnectivityState: connectivity.Connecting,
202 Picker: &picker{err: balancer.ErrNoSubConnAvailable},
203 })
204 b.subConn.Connect()
205 return nil
206}
207
208// UpdateSubConnState is unused as a StateListener is always registered when
209// creating SubConns.
210func (b *pickfirstBalancer) UpdateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
211 b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", subConn, state)
212}
213
214func (b *pickfirstBalancer) updateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
215 if b.logger.V(2) {
216 b.logger.Infof("Received SubConn state update: %p, %+v", subConn, state)
217 }
218 if b.subConn != subConn {
219 if b.logger.V(2) {
220 b.logger.Infof("Ignored state change because subConn is not recognized")
221 }
222 return
223 }
224 if state.ConnectivityState == connectivity.Shutdown {
225 b.subConn = nil
226 return
227 }
228
229 switch state.ConnectivityState {
230 case connectivity.Ready:
231 b.cc.UpdateState(balancer.State{
232 ConnectivityState: state.ConnectivityState,
233 Picker: &picker{result: balancer.PickResult{SubConn: subConn}},
234 })
235 case connectivity.Connecting:
236 if b.state == connectivity.TransientFailure {
237 // We stay in TransientFailure until we are Ready. See A62.
238 return
239 }
240 b.cc.UpdateState(balancer.State{
241 ConnectivityState: state.ConnectivityState,
242 Picker: &picker{err: balancer.ErrNoSubConnAvailable},
243 })
244 case connectivity.Idle:
245 if b.state == connectivity.TransientFailure {
246 // We stay in TransientFailure until we are Ready. Also kick the
247 // subConn out of Idle into Connecting. See A62.
248 b.subConn.Connect()
249 return
250 }
251 b.cc.UpdateState(balancer.State{
252 ConnectivityState: state.ConnectivityState,
253 Picker: &idlePicker{subConn: subConn},
254 })
255 case connectivity.TransientFailure:
256 b.cc.UpdateState(balancer.State{
257 ConnectivityState: state.ConnectivityState,
258 Picker: &picker{err: state.ConnectionError},
259 })
260 }
261 b.state = state.ConnectivityState
262}
263
264func (b *pickfirstBalancer) Close() {
265}
266
267func (b *pickfirstBalancer) ExitIdle() {
268 if b.subConn != nil && b.state == connectivity.Idle {
269 b.subConn.Connect()
270 }
271}
272
273type picker struct {
274 result balancer.PickResult
275 err error
276}
277
278func (p *picker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
279 return p.result, p.err
280}
281
282// idlePicker is used when the SubConn is IDLE and kicks the SubConn into
283// CONNECTING when Pick is called.
284type idlePicker struct {
285 subConn balancer.SubConn
286}
287
288func (i *idlePicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
289 i.subConn.Connect()
290 return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
291}