Skip to content

Commit 5143e0d

Browse files
vinothkumarr227vinothkumarr227
committed
Add ChildBalancer field to BalancerData for direct access and update all references
1 parent d2e8366 commit 5143e0d

File tree

19 files changed

+121
-141
lines changed

19 files changed

+121
-141
lines changed

‎balancer/endpointsharding/endpointsharding_test.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -222,13 +222,13 @@ func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
222222
bf := stub.BalancerFuncs{
223223
Init: func(bd *stub.BalancerData) {
224224
epOpts := endpointsharding.Options{DisableAutoReconnect: true}
225-
bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
225+
bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
226226
},
227227
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
228-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
228+
return bd.ChildBalancer.UpdateClientConnState(ccs)
229229
},
230230
Close: func(bd *stub.BalancerData) {
231-
bd.Data.(balancer.Balancer).Close()
231+
bd.ChildBalancer.Close()
232232
},
233233
}
234234
stub.Register(name, bf)
@@ -303,16 +303,16 @@ func (s) TestEndpointShardingExitIdle(t *testing.T) {
303303
bf := stub.BalancerFuncs{
304304
Init: func(bd *stub.BalancerData) {
305305
epOpts := endpointsharding.Options{DisableAutoReconnect: true}
306-
bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
306+
bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
307307
},
308308
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
309-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
309+
return bd.ChildBalancer.UpdateClientConnState(ccs)
310310
},
311311
Close: func(bd *stub.BalancerData) {
312-
bd.Data.(balancer.Balancer).Close()
312+
bd.ChildBalancer.Close()
313313
},
314314
ExitIdle: func(bd *stub.BalancerData) {
315-
bd.Data.(balancer.Balancer).ExitIdle()
315+
bd.ChildBalancer.ExitIdle()
316316
},
317317
}
318318
stub.Register(name, bf)

‎balancer/lazy/lazy_ext_test.go‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -79,19 +79,19 @@ func (s) TestExitIdle(t *testing.T) {
7979

8080
bf := stub.BalancerFuncs{
8181
Init: func(bd *stub.BalancerData) {
82-
bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
82+
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
8383
},
8484
ExitIdle: func(bd *stub.BalancerData) {
85-
bd.Data.(balancer.Balancer).ExitIdle()
85+
bd.ChildBalancer.ExitIdle()
8686
},
8787
ResolverError: func(bd *stub.BalancerData, err error) {
88-
bd.Data.(balancer.Balancer).ResolverError(err)
88+
bd.ChildBalancer.ResolverError(err)
8989
},
9090
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
91-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
91+
return bd.ChildBalancer.UpdateClientConnState(ccs)
9292
},
9393
Close: func(bd *stub.BalancerData) {
94-
bd.Data.(balancer.Balancer).Close()
94+
bd.ChildBalancer.Close()
9595
},
9696
}
9797
stub.Register(t.Name(), bf)
@@ -144,16 +144,16 @@ func (s) TestPicker(t *testing.T) {
144144

145145
bf := stub.BalancerFuncs{
146146
Init: func(bd *stub.BalancerData) {
147-
bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
147+
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
148148
},
149149
ExitIdle: func(*stub.BalancerData) {
150150
t.Log("Ignoring call to ExitIdle, calling the picker should make the lazy balancer exit IDLE state.")
151151
},
152152
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
153-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
153+
return bd.ChildBalancer.UpdateClientConnState(ccs)
154154
},
155155
Close: func(bd *stub.BalancerData) {
156-
bd.Data.(balancer.Balancer).Close()
156+
bd.ChildBalancer.Close()
157157
},
158158
}
159159

@@ -201,24 +201,24 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
201201

202202
childBF := stub.BalancerFuncs{
203203
Init: func(bd *stub.BalancerData) {
204-
bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
204+
bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
205205
},
206206
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
207207
if resolverErrorReceived.HasFired() {
208208
t.Error("Received resolver error before resolver state.")
209209
}
210210
resolverStateReceived = true
211-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
211+
return bd.ChildBalancer.UpdateClientConnState(ccs)
212212
},
213213
ResolverError: func(bd *stub.BalancerData, err error) {
214214
if !resolverStateReceived {
215215
t.Error("Received resolver error before resolver state.")
216216
}
217217
resolverErrorReceived.Fire()
218-
bd.Data.(balancer.Balancer).ResolverError(err)
218+
bd.ChildBalancer.ResolverError(err)
219219
},
220220
Close: func(bd *stub.BalancerData) {
221-
bd.Data.(balancer.Balancer).Close()
221+
bd.ChildBalancer.Close()
222222
},
223223
}
224224

@@ -227,19 +227,19 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
227227

228228
topLevelBF := stub.BalancerFuncs{
229229
Init: func(bd *stub.BalancerData) {
230-
bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
230+
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
231231
},
232232
ExitIdle: func(*stub.BalancerData) {
233233
t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
234234
},
235235
ResolverError: func(bd *stub.BalancerData, err error) {
236-
bd.Data.(balancer.Balancer).ResolverError(err)
236+
bd.ChildBalancer.ResolverError(err)
237237
},
238238
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
239-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
239+
return bd.ChildBalancer.UpdateClientConnState(ccs)
240240
},
241241
Close: func(bd *stub.BalancerData) {
242-
bd.Data.(balancer.Balancer).Close()
242+
bd.ChildBalancer.Close()
243243
},
244244
}
245245

@@ -306,17 +306,17 @@ func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
306306

307307
childBF := stub.BalancerFuncs{
308308
Init: func(bd *stub.BalancerData) {
309-
bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
309+
bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
310310
},
311311
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
312-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
312+
return bd.ChildBalancer.UpdateClientConnState(ccs)
313313
},
314314
ResolverError: func(bd *stub.BalancerData, err error) {
315315
t.Error("Received unexpected resolver error.")
316-
bd.Data.(balancer.Balancer).ResolverError(err)
316+
bd.ChildBalancer.ResolverError(err)
317317
},
318318
Close: func(bd *stub.BalancerData) {
319-
bd.Data.(balancer.Balancer).Close()
319+
bd.ChildBalancer.Close()
320320
},
321321
}
322322

@@ -325,16 +325,16 @@ func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
325325

326326
topLevelBF := stub.BalancerFuncs{
327327
Init: func(bd *stub.BalancerData) {
328-
bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
328+
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
329329
},
330330
ExitIdle: func(*stub.BalancerData) {
331331
t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
332332
},
333333
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
334-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
334+
return bd.ChildBalancer.UpdateClientConnState(ccs)
335335
},
336336
Close: func(bd *stub.BalancerData) {
337-
bd.Data.(balancer.Balancer).Close()
337+
bd.ChildBalancer.Close()
338338
},
339339
}
340340

@@ -407,19 +407,19 @@ func (s) TestExitIdlePassthrough(t *testing.T) {
407407

408408
bf := stub.BalancerFuncs{
409409
Init: func(bd *stub.BalancerData) {
410-
bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
410+
bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
411411
},
412412
ExitIdle: func(bd *stub.BalancerData) {
413-
bd.Data.(balancer.Balancer).ExitIdle()
413+
bd.ChildBalancer.ExitIdle()
414414
},
415415
ResolverError: func(bd *stub.BalancerData, err error) {
416-
bd.Data.(balancer.Balancer).ResolverError(err)
416+
bd.ChildBalancer.ResolverError(err)
417417
},
418418
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
419-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
419+
return bd.ChildBalancer.UpdateClientConnState(ccs)
420420
},
421421
Close: func(bd *stub.BalancerData) {
422-
bd.Data.(balancer.Balancer).Close()
422+
bd.ChildBalancer.Close()
423423
},
424424
}
425425
stub.Register(t.Name(), bf)

‎balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1241,14 +1241,14 @@ func (s) TestPickFirstLeaf_HealthListenerEnabled(t *testing.T) {
12411241
defer cancel()
12421242
bf := stub.BalancerFuncs{
12431243
Init: func(bd *stub.BalancerData) {
1244-
bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
1244+
bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
12451245
},
12461246
Close: func(bd *stub.BalancerData) {
1247-
bd.Data.(balancer.Balancer).Close()
1247+
bd.ChildBalancer.Close()
12481248
},
12491249
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
12501250
ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
1251-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
1251+
return bd.ChildBalancer.UpdateClientConnState(ccs)
12521252
},
12531253
}
12541254

@@ -1289,15 +1289,15 @@ func (s) TestPickFirstLeaf_HealthListenerNotEnabled(t *testing.T) {
12891289
healthListenerCh: healthListenerCh,
12901290
subConnStateCh: make(chan balancer.SubConnState, 5),
12911291
}
1292-
bd.Data = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
1292+
bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
12931293
},
12941294
Close: func(bd *stub.BalancerData) {
1295-
bd.Data.(balancer.Balancer).Close()
1295+
bd.ChildBalancer.Close()
12961296
},
12971297
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
12981298
// Functions like a non-petiole policy by not configuring the use
12991299
// of health listeners.
1300-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
1300+
return bd.ChildBalancer.UpdateClientConnState(ccs)
13011301
},
13021302
}
13031303

@@ -1345,14 +1345,14 @@ func (s) TestPickFirstLeaf_HealthUpdates(t *testing.T) {
13451345
healthListenerCh: healthListenerCh,
13461346
subConnStateCh: scConnectivityStateCh,
13471347
}
1348-
bd.Data = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
1348+
bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
13491349
},
13501350
Close: func(bd *stub.BalancerData) {
1351-
bd.Data.(balancer.Balancer).Close()
1351+
bd.ChildBalancer.Close()
13521352
},
13531353
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
13541354
ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
1355-
return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
1355+
return bd.ChildBalancer.UpdateClientConnState(ccs)
13561356
},
13571357
}
13581358

‎balancer/rls/balancer_test.go‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -689,19 +689,17 @@ func (s) TestPickerUpdateOnDataCacheSizeDecrease(t *testing.T) {
689689
stub.Register(topLevelBalancerName, stub.BalancerFuncs{
690690
Init: func(bd *stub.BalancerData) {
691691
ccWrapper = &testCCWrapper{ClientConn: bd.ClientConn}
692-
bd.Data = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
692+
bd.ChildBalancer = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
693693
},
694694
ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
695695
parser := balancer.Get(Name).(balancer.ConfigParser)
696696
return parser.ParseConfig(sc)
697697
},
698698
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
699-
bal := bd.Data.(balancer.Balancer)
700-
return bal.UpdateClientConnState(ccs)
699+
return bd.ChildBalancer.UpdateClientConnState(ccs)
701700
},
702701
Close: func(bd *stub.BalancerData) {
703-
bal := bd.Data.(balancer.Balancer)
704-
bal.Close()
702+
bd.ChildBalancer.Close()
705703
},
706704
})
707705

@@ -1070,19 +1068,17 @@ func (s) TestUpdateStatePauses(t *testing.T) {
10701068
stub.Register(topLevelBalancerName, stub.BalancerFuncs{
10711069
Init: func(bd *stub.BalancerData) {
10721070
ccWrapper = &testCCWrapper{ClientConn: bd.ClientConn}
1073-
bd.Data = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
1071+
bd.ChildBalancer = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
10741072
},
10751073
ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
10761074
parser := balancer.Get(Name).(balancer.ConfigParser)
10771075
return parser.ParseConfig(sc)
10781076
},
10791077
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
1080-
bal := bd.Data.(balancer.Balancer)
1081-
return bal.UpdateClientConnState(ccs)
1078+
return bd.ChildBalancer.UpdateClientConnState(ccs)
10821079
},
10831080
Close: func(bd *stub.BalancerData) {
1084-
bal := bd.Data.(balancer.Balancer)
1085-
bal.Close()
1081+
bd.ChildBalancer.Close()
10861082
},
10871083
})
10881084

@@ -1098,10 +1094,10 @@ func (s) TestUpdateStatePauses(t *testing.T) {
10981094
}
10991095
stub.Register(childPolicyName, stub.BalancerFuncs{
11001096
Init: func(bd *stub.BalancerData) {
1101-
bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
1097+
bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
11021098
},
11031099
Close: func(bd *stub.BalancerData) {
1104-
bd.Data.(balancer.Balancer).Close()
1100+
bd.ChildBalancer.Close()
11051101
},
11061102
ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
11071103
cfg := &childPolicyConfig{}
@@ -1111,7 +1107,7 @@ func (s) TestUpdateStatePauses(t *testing.T) {
11111107
return cfg, nil
11121108
},
11131109
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
1114-
bal := bd.Data.(balancer.Balancer)
1110+
bal := bd.ChildBalancer
11151111
bd.ClientConn.UpdateState(balancer.State{ConnectivityState: connectivity.Idle, Picker: &testutils.TestConstPicker{Err: balancer.ErrNoSubConnAvailable}})
11161112
bd.ClientConn.UpdateState(balancer.State{ConnectivityState: connectivity.Connecting, Picker: &testutils.TestConstPicker{Err: balancer.ErrNoSubConnAvailable}})
11171113

‎internal/balancer/stub/stub.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,9 @@ type BalancerData struct {
5151
BuildOptions balancer.BuildOptions
5252
// Data may be used to store arbitrary user data.
5353
Data any
54+
// ChildBalancer holds the child balancer directly, avoiding the need
55+
// to type-assert from Data.
56+
ChildBalancer balancer.Balancer
5457
}
5558

5659
type bal struct {

‎internal/balancergroup/balancergroup_test.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -562,15 +562,14 @@ func (s) TestBalancerGracefulSwitch(t *testing.T) {
562562
childPolicyName := t.Name()
563563
stub.Register(childPolicyName, stub.BalancerFuncs{
564564
Init: func(bd *stub.BalancerData) {
565-
bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
565+
bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
566566
},
567567
Close: func(bd *stub.BalancerData) {
568-
bd.Data.(balancer.Balancer).Close()
568+
bd.ChildBalancer.Close()
569569
},
570570
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
571571
ccs.ResolverState.Endpoints = ccs.ResolverState.Endpoints[1:]
572-
bal := bd.Data.(balancer.Balancer)
573-
return bal.UpdateClientConnState(ccs)
572+
return bd.ChildBalancer.UpdateClientConnState(ccs)
574573
},
575574
})
576575
cfgJSON := json.RawMessage(fmt.Sprintf(`[{%q: {}}]`, t.Name()))

‎internal/idle/idle_e2e_test.go‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -98,19 +98,17 @@ func registerWrappedRoundRobinPolicy(t *testing.T) chan struct{} {
9898
closeCh := make(chan struct{}, 1)
9999
stub.Register(roundrobin.Name, stub.BalancerFuncs{
100100
Init: func(bd *stub.BalancerData) {
101-
bd.Data = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
101+
bd.ChildBalancer = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
102102
},
103103
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
104-
bal := bd.Data.(balancer.Balancer)
105-
return bal.UpdateClientConnState(ccs)
104+
return bd.ChildBalancer.UpdateClientConnState(ccs)
106105
},
107106
Close: func(bd *stub.BalancerData) {
108107
select {
109108
case closeCh <- struct{}{}:
110109
default:
111110
}
112-
bal := bd.Data.(balancer.Balancer)
113-
bal.Close()
111+
bd.ChildBalancer.Close()
114112
},
115113
})
116114
t.Cleanup(func() { balancer.Register(rrBuilder) })

‎test/balancer_switching_test.go‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -488,17 +488,16 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) {
488488
stub.Register(t.Name(), stub.BalancerFuncs{
489489
Init: func(bd *stub.BalancerData) {
490490
pf := balancer.Get(pickfirst.Name)
491-
bd.Data = pf.Build(bd.ClientConn, bd.BuildOptions)
491+
bd.ChildBalancer = pf.Build(bd.ClientConn, bd.BuildOptions)
492492
},
493493
Close: func(bd *stub.BalancerData) {
494-
bd.Data.(balancer.Balancer).Close()
494+
bd.ChildBalancer.Close()
495495
},
496496
UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
497-
bal := bd.Data.(balancer.Balancer)
498497
close(ccUpdateCh)
499498
go func() {
500499
<-waitToProceed
501-
bal.UpdateClientConnState(ccs)
500+
bd.ChildBalancer.UpdateClientConnState(ccs)
502501
}()
503502
return nil
504503
},

0 commit comments

Comments
 (0)