Skip to content

Commit d06ab0d

Browse files
authored
pickfirst: receive state updates via callback instead of UpdateSubConnState (#6495)
1 parent 7aceafc commit d06ab0d

File tree

3 files changed

+43
-24
lines changed

3 files changed

+43
-24
lines changed

clientconn_test.go

Lines changed: 16 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1143,18 +1143,10 @@ func testDefaultServiceConfigWhenResolverReturnInvalidServiceConfig(t *testing.T
11431143
}
11441144

11451145
type stateRecordingBalancer struct {
1146-
notifier chan<- connectivity.State
11471146
balancer.Balancer
11481147
}
11491148

1150-
func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
1151-
b.notifier <- s.ConnectivityState
1152-
b.Balancer.UpdateSubConnState(sc, s)
1153-
}
1154-
1155-
func (b *stateRecordingBalancer) ResetNotifier(r chan<- connectivity.State) {
1156-
b.notifier = r
1157-
}
1149+
func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {}
11581150

11591151
func (b *stateRecordingBalancer) Close() {
11601152
b.Balancer.Close()
@@ -1179,8 +1171,7 @@ func (b *stateRecordingBalancerBuilder) Build(cc balancer.ClientConn, opts balan
11791171
b.notifier = stateNotifications
11801172
b.mu.Unlock()
11811173
return &stateRecordingBalancer{
1182-
notifier: stateNotifications,
1183-
Balancer: balancer.Get("pick_first").Build(cc, opts),
1174+
Balancer: balancer.Get("pick_first").Build(&stateRecordingCCWrapper{cc, stateNotifications}, opts),
11841175
}
11851176
}
11861177

@@ -1192,6 +1183,20 @@ func (b *stateRecordingBalancerBuilder) nextStateNotifier() <-chan connectivity.
11921183
return ret
11931184
}
11941185

1186+
type stateRecordingCCWrapper struct {
1187+
balancer.ClientConn
1188+
notifier chan<- connectivity.State
1189+
}
1190+
1191+
func (ccw *stateRecordingCCWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
1192+
oldListener := opts.StateListener
1193+
opts.StateListener = func(s balancer.SubConnState) {
1194+
ccw.notifier <- s.ConnectivityState
1195+
oldListener(s)
1196+
}
1197+
return ccw.ClientConn.NewSubConn(addrs, opts)
1198+
}
1199+
11951200
// Keep reading until something causes the connection to die (EOF, server
11961201
// closed, etc). Useful as a tool for mindlessly keeping the connection
11971202
// healthy, since the client will error if things like client prefaces are not

pickfirst.go

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,12 @@ func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState
146146
return nil
147147
}
148148

149-
subConn, err := b.cc.NewSubConn(addrs, balancer.NewSubConnOptions{})
149+
var subConn balancer.SubConn
150+
subConn, err := b.cc.NewSubConn(addrs, balancer.NewSubConnOptions{
151+
StateListener: func(state balancer.SubConnState) {
152+
b.updateSubConnState(subConn, state)
153+
},
154+
})
150155
if err != nil {
151156
if b.logger.V(2) {
152157
b.logger.Infof("Failed to create new SubConn: %v", err)
@@ -168,7 +173,13 @@ func (b *pickfirstBalancer) UpdateClientConnState(state balancer.ClientConnState
168173
return nil
169174
}
170175

176+
// UpdateSubConnState is unused as a StateListener is always registered when
177+
// creating SubConns.
171178
func (b *pickfirstBalancer) UpdateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
179+
b.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", subConn, state)
180+
}
181+
182+
func (b *pickfirstBalancer) updateSubConnState(subConn balancer.SubConn, state balancer.SubConnState) {
172183
if b.logger.V(2) {
173184
b.logger.Infof("Received SubConn state update: %p, %+v", subConn, state)
174185
}

test/clientconn_state_transition_test.go

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -444,19 +444,9 @@ func (s) TestStateTransitions_MultipleAddrsEntersReady(t *testing.T) {
444444
}
445445

446446
type stateRecordingBalancer struct {
447-
notifier chan<- connectivity.State
448447
balancer.Balancer
449448
}
450449

451-
func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
452-
b.notifier <- s.ConnectivityState
453-
b.Balancer.UpdateSubConnState(sc, s)
454-
}
455-
456-
func (b *stateRecordingBalancer) ResetNotifier(r chan<- connectivity.State) {
457-
b.notifier = r
458-
}
459-
460450
func (b *stateRecordingBalancer) Close() {
461451
b.Balancer.Close()
462452
}
@@ -480,8 +470,7 @@ func (b *stateRecordingBalancerBuilder) Build(cc balancer.ClientConn, opts balan
480470
b.notifier = stateNotifications
481471
b.mu.Unlock()
482472
return &stateRecordingBalancer{
483-
notifier: stateNotifications,
484-
Balancer: balancer.Get("pick_first").Build(cc, opts),
473+
Balancer: balancer.Get("pick_first").Build(&stateRecordingCCWrapper{cc, stateNotifications}, opts),
485474
}
486475
}
487476

@@ -493,6 +482,20 @@ func (b *stateRecordingBalancerBuilder) nextStateNotifier() <-chan connectivity.
493482
return ret
494483
}
495484

485+
type stateRecordingCCWrapper struct {
486+
balancer.ClientConn
487+
notifier chan<- connectivity.State
488+
}
489+
490+
func (ccw *stateRecordingCCWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
491+
oldListener := opts.StateListener
492+
opts.StateListener = func(s balancer.SubConnState) {
493+
ccw.notifier <- s.ConnectivityState
494+
oldListener(s)
495+
}
496+
return ccw.ClientConn.NewSubConn(addrs, opts)
497+
}
498+
496499
// Keep reading until something causes the connection to die (EOF, server
497500
// closed, etc). Useful as a tool for mindlessly keeping the connection
498501
// healthy, since the client will error if things like client prefaces are not

0 commit comments

Comments
 (0)