diff --git a/clientconn_test.go b/clientconn_test.go index c742e144a94c..e614045bee01 100644 --- a/clientconn_test.go +++ b/clientconn_test.go @@ -1146,7 +1146,9 @@ type stateRecordingBalancer struct { balancer.Balancer } -func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {} +func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) { + panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, s)) +} func (b *stateRecordingBalancer) Close() { b.Balancer.Close() diff --git a/test/balancer_switching_test.go b/test/balancer_switching_test.go index d18f89c98f5d..1a6a89d36294 100644 --- a/test/balancer_switching_test.go +++ b/test/balancer_switching_test.go @@ -483,10 +483,6 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) { }() return nil }, - UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) { - bal := bd.Data.(balancer.Balancer) - bal.UpdateSubConnState(sc, state) - }, }) // Push a resolver update with the service config specifying our stub diff --git a/test/balancer_test.go b/test/balancer_test.go index dd1fdc024843..df82f8b81115 100644 --- a/test/balancer_test.go +++ b/test/balancer_test.go @@ -87,6 +87,7 @@ func (b *testBalancer) UpdateClientConnState(state balancer.ClientConnState) err // Only create a subconn at the first time. if b.sc == nil { var err error + b.newSubConnOptions.StateListener = b.updateSubConnState b.sc, err = b.cc.NewSubConn(state.ResolverState.Addresses, b.newSubConnOptions) if err != nil { logger.Errorf("testBalancer: failed to NewSubConn: %v", err) @@ -99,21 +100,17 @@ func (b *testBalancer) UpdateClientConnState(state balancer.ClientConnState) err } func (b *testBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) { - logger.Infof("testBalancer: UpdateSubConnState: %p, %v", sc, s) - if b.sc != sc { - logger.Infof("testBalancer: ignored state change because sc is not recognized") - return - } - if s.ConnectivityState == connectivity.Shutdown { - b.sc = nil - return - } + panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, s)) +} + +func (b *testBalancer) updateSubConnState(s balancer.SubConnState) { + logger.Infof("testBalancer: updateSubConnState: %v", s) switch s.ConnectivityState { case connectivity.Ready: - b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{sc: sc, bal: b}}) + b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{bal: b}}) case connectivity.Idle: - b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{sc: sc, bal: b, idle: true}}) + b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{bal: b, idle: true}}) case connectivity.Connecting: b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{err: balancer.ErrNoSubConnAvailable, bal: b}}) case connectivity.TransientFailure: @@ -127,7 +124,6 @@ func (b *testBalancer) ExitIdle() {} type picker struct { err error - sc balancer.SubConn bal *testBalancer idle bool } @@ -137,14 +133,14 @@ func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) { return balancer.PickResult{}, p.err } if p.idle { - p.sc.Connect() + p.bal.sc.Connect() return balancer.PickResult{}, balancer.ErrNoSubConnAvailable } extraMD, _ := grpcutil.ExtraMetadata(info.Ctx) info.Ctx = nil // Do not validate context. p.bal.pickInfos = append(p.bal.pickInfos, info) p.bal.pickExtraMDs = append(p.bal.pickExtraMDs, extraMD) - return balancer.PickResult{SubConn: p.sc, Done: func(d balancer.DoneInfo) { p.bal.doneInfo = append(p.bal.doneInfo, d) }}, nil + return balancer.PickResult{SubConn: p.bal.sc, Done: func(d balancer.DoneInfo) { p.bal.doneInfo = append(p.bal.doneInfo, d) }}, nil } func (s) TestCredsBundleFromBalancer(t *testing.T) { @@ -397,16 +393,18 @@ func (s) TestAddressAttributesInNewSubConn(t *testing.T) { // Only use the first address. attr := attributes.New(testAttrKey, testAttrVal) addrs[0].Attributes = attr - sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{}) + var sc balancer.SubConn + sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{ + StateListener: func(state balancer.SubConnState) { + bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) + }, + }) if err != nil { return err } sc.Connect() return nil }, - UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) { - bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) - }, } stub.Register(attrBalancerName, bf) t.Logf("Registered balancer %s...", attrBalancerName) @@ -483,18 +481,20 @@ func (s) TestMetadataInAddressAttributes(t *testing.T) { return nil } // Only use the first address. + var sc balancer.SubConn sc, err := bd.ClientConn.NewSubConn([]resolver.Address{ imetadata.Set(addrs[0], metadata.Pairs(testMDKey, testMDValue)), - }, balancer.NewSubConnOptions{}) + }, balancer.NewSubConnOptions{ + StateListener: func(state balancer.SubConnState) { + bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) + }, + }) if err != nil { return err } sc.Connect() return nil }, - UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) { - bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) - }, } stub.Register(mdBalancerName, bf) t.Logf("Registered balancer %s...", mdBalancerName) @@ -717,16 +717,18 @@ func (s) TestAuthorityInBuildOptions(t *testing.T) { } // Only use the first address. - sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{}) + var sc balancer.SubConn + sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{ + StateListener: func(state balancer.SubConnState) { + bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) + }, + }) if err != nil { return err } sc.Connect() return nil }, - UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) { - bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}}) - }, } balancerName := "stub-balancer-" + test.name stub.Register(balancerName, bf) diff --git a/test/resolver_update_test.go b/test/resolver_update_test.go index 416f7175c53a..e9ae0df6ebd8 100644 --- a/test/resolver_update_test.go +++ b/test/resolver_update_test.go @@ -177,10 +177,6 @@ func (s) TestResolverUpdate_InvalidServiceConfigAfterGoodUpdate(t *testing.T) { ccs.BalancerConfig = nil return bal.UpdateClientConnState(ccs) }, - UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) { - bal := bd.Data.(balancer.Balancer) - bal.UpdateSubConnState(sc, state) - }, }) // Start a backend exposing the test service. diff --git a/test/subconn_test.go b/test/subconn_test.go index 2e61d62c9044..e8c8d936a9fb 100644 --- a/test/subconn_test.go +++ b/test/subconn_test.go @@ -59,7 +59,22 @@ func (s) TestSubConnEmpty(t *testing.T) { UpdateClientConnState: func(d *stub.BalancerData, ccs balancer.ClientConnState) error { if sc == nil { var err error - sc, err = d.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{}) + sc, err = d.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{ + StateListener: func(state balancer.SubConnState) { + switch state.ConnectivityState { + case connectivity.Ready: + d.ClientConn.UpdateState(balancer.State{ + ConnectivityState: connectivity.Ready, + Picker: &tsccPicker{sc: sc}, + }) + case connectivity.TransientFailure: + d.ClientConn.UpdateState(balancer.State{ + ConnectivityState: connectivity.TransientFailure, + Picker: base.NewErrPicker(fmt.Errorf("error connecting: %v", state.ConnectionError)), + }) + } + }, + }) if err != nil { t.Errorf("error creating initial subconn: %v", err) } @@ -81,20 +96,6 @@ func (s) TestSubConnEmpty(t *testing.T) { } return nil }, - UpdateSubConnState: func(d *stub.BalancerData, sc balancer.SubConn, scs balancer.SubConnState) { - switch scs.ConnectivityState { - case connectivity.Ready: - d.ClientConn.UpdateState(balancer.State{ - ConnectivityState: connectivity.Ready, - Picker: &tsccPicker{sc: sc}, - }) - case connectivity.TransientFailure: - d.ClientConn.UpdateState(balancer.State{ - ConnectivityState: connectivity.TransientFailure, - Picker: base.NewErrPicker(fmt.Errorf("error connecting: %v", scs.ConnectionError)), - }) - } - }, } stub.Register("tscc", bal)