Repository: grpc/grpc-go

Issue:
outlierdetection, pickfirst: Remove the balancer attribute to skip ejection using the connectivity state listener
As part of https://github.com/grpc/grpc-go/pull/7908, a workaround was added to make Outlier Detection aware that the health listener is to be used for ejection updates. This workaround should be removed once all petiole policies delegate to pickfirst and ejection happens only using the health listener.

Hints from issue comments:

Solution patch:
diff --git a/balancer/pickfirst/pickfirstleaf/pickfirstleaf.go b/balancer/pickfirst/pickfirstleaf/pickfirstleaf.go
index 494314f2..e6204725 100644
--- a/balancer/pickfirst/pickfirstleaf/pickfirstleaf.go
+++ b/balancer/pickfirst/pickfirstleaf/pickfirstleaf.go
@@ -54,18 +54,9 @@ func init() {
 	balancer.Register(pickfirstBuilder{})
 }
 
-type (
-	// enableHealthListenerKeyType is a unique key type used in resolver
-	// attributes to indicate whether the health listener usage is enabled.
-	enableHealthListenerKeyType struct{}
-	// managedByPickfirstKeyType is an attribute key type to inform Outlier
-	// Detection that the generic health listener is being used.
-	// TODO: https://github.com/grpc/grpc-go/issues/7915 - Remove this when
-	// implementing the dualstack design. This is a hack. Once Dualstack is
-	// completed, outlier detection will stop sending ejection updates through
-	// the connectivity listener.
-	managedByPickfirstKeyType struct{}
-)
+// enableHealthListenerKeyType is a unique key type used in resolver
+// attributes to indicate whether the health listener usage is enabled.
+type enableHealthListenerKeyType struct{}
 
 var (
 	logger = grpclog.Component("pick-first-leaf-lb")
@@ -149,17 +140,6 @@ func EnableHealthListener(state resolver.State) resolver.State {
 	return state
 }
 
-// IsManagedByPickfirst returns whether an address belongs to a SubConn
-// managed by the pickfirst LB policy.
-// TODO: https://github.com/grpc/grpc-go/issues/7915 - This is a hack to disable
-// outlier_detection via the with connectivity listener when using pick_first.
-// Once Dualstack changes are complete, all SubConns will be created by
-// pick_first and outlier detection will only use the health listener for
-// ejection. This hack can then be removed.
-func IsManagedByPickfirst(addr resolver.Address) bool {
-	return addr.BalancerAttributes.Value(managedByPickfirstKeyType{}) != nil
-}
-
 type pfConfig struct {
 	serviceconfig.LoadBalancingConfig `json:"-"`
 
@@ -186,7 +166,6 @@ type scData struct {
 }
 
 func (b *pickfirstBalancer) newSCData(addr resolver.Address) (*scData, error) {
-	addr.BalancerAttributes = addr.BalancerAttributes.WithValue(managedByPickfirstKeyType{}, true)
 	sd := &scData{
 		rawConnectivityState: connectivity.Idle,
 		effectiveState:       connectivity.Idle,
diff --git a/internal/balancer/stub/stub.go b/internal/balancer/stub/stub.go
index 71e3cb61..aec5551f 100644
--- a/internal/balancer/stub/stub.go
+++ b/internal/balancer/stub/stub.go
@@ -49,8 +49,8 @@ type BalancerData struct {
 	ClientConn balancer.ClientConn
 	// BuildOptions is set by the builder.
 	BuildOptions balancer.BuildOptions
-	// Data may be used to store arbitrary user data.
-	Data any
+	// ChildBalancer holds a child balancer.
+	ChildBalancer balancer.Balancer
 }
 
 type bal struct {
diff --git a/xds/internal/balancer/outlierdetection/balancer.go b/xds/internal/balancer/outlierdetection/balancer.go
index bc5cea40..1b0149e3 100644
--- a/xds/internal/balancer/outlierdetection/balancer.go
+++ b/xds/internal/balancer/outlierdetection/balancer.go
@@ -32,7 +32,6 @@ import (
 	"time"
 
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/internal/buffer"
@@ -467,12 +466,10 @@ func (b *outlierDetectionBalancer) UpdateState(s balancer.State) {
 func (b *outlierDetectionBalancer) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
 	oldListener := opts.StateListener
 	scw := &subConnWrapper{
-		addresses:                  addrs,
-		scUpdateCh:                 b.scUpdateCh,
-		listener:                   oldListener,
-		latestRawConnectivityState: balancer.SubConnState{ConnectivityState: connectivity.Idle},
-		latestHealthState:          balancer.SubConnState{ConnectivityState: connectivity.Connecting},
-		healthListenerEnabled:      len(addrs) == 1 && pickfirstleaf.IsManagedByPickfirst(addrs[0]),
+		addresses:         addrs,
+		scUpdateCh:        b.scUpdateCh,
+		listener:          oldListener,
+		latestHealthState: balancer.SubConnState{ConnectivityState: connectivity.Connecting},
 	}
 	opts.StateListener = func(state balancer.SubConnState) { b.updateSubConnState(scw, state) }
 	b.mu.Lock()
diff --git a/xds/internal/balancer/outlierdetection/subconn_wrapper.go b/xds/internal/balancer/outlierdetection/subconn_wrapper.go
index 2c390d23..b893ea68 100644
--- a/xds/internal/balancer/outlierdetection/subconn_wrapper.go
+++ b/xds/internal/balancer/outlierdetection/subconn_wrapper.go
@@ -40,16 +40,6 @@ type subConnWrapper struct {
 	// that.
 
 	listener func(balancer.SubConnState)
-	// healthListenerEnabled indicates whether the leaf LB policy is using a
-	// generic health listener. When enabled, ejection updates are sent via the
-	// health listener instead of the connectivity listener. Once Dualstack
-	// changes are complete, all SubConns will be created by pickfirst which
-	// uses the health listener.
-	// TODO: https://github.com/grpc/grpc-go/issues/7915 - Once Dualstack
-	// changes are complete, all SubConns will be created by pick_first and
-	// outlier detection will only use the health listener for ejection and
-	// this field can be removed.
-	healthListenerEnabled bool
 
 	scUpdateCh *buffer.Unbounded
 
@@ -67,9 +57,6 @@ type subConnWrapper struct {
 	// latestHealthState is tracked to update the child policy during
 	// unejection.
 	latestHealthState balancer.SubConnState
-	// latestRawConnectivityState is tracked to update the child policy during
-	// unejection.
-	latestRawConnectivityState balancer.SubConnState
 
 	// Access to the following fields are protected by a mutex. These fields
 	// should not be accessed from outside this file, instead use methods
@@ -114,12 +101,7 @@ func (scw *subConnWrapper) RegisterHealthListener(listener func(balancer.SubConn
 	// pick_first the universal leaf policy (see A61), both these mechanisms
 	// started using the new health listener to make health signal visible to
 	// the petiole policies without affecting the underlying connectivity
-	// management of the pick_first policy
-	if !scw.healthListenerEnabled {
-		logger.Errorf("Health listener unexpectedly registered on SubConn %v.", scw)
-		return
-	}
-
+	// management of the pick_first policy.
 	scw.mu.Lock()
 	defer scw.mu.Unlock()
 
@@ -157,12 +139,6 @@ func (scw *subConnWrapper) updateSubConnHealthState(scs balancer.SubConnState) {
 // updateSubConnConnectivityState stores the latest connectivity state for
 // unejection and updates the raw connectivity listener.
 func (scw *subConnWrapper) updateSubConnConnectivityState(scs balancer.SubConnState) {
-	scw.latestRawConnectivityState = scs
-	// If the raw connectivity listener is used for ejection, and the SubConn is
-	// ejected, don't send the update.
-	if scw.ejected && !scw.healthListenerEnabled {
-		return
-	}
 	if scw.listener != nil {
 		scw.listener(scs)
 	}
@@ -176,12 +152,6 @@ func (scw *subConnWrapper) clearHealthListener() {
 
 func (scw *subConnWrapper) handleUnejection() {
 	scw.ejected = false
-	if !scw.healthListenerEnabled {
-		// If scw.latestRawConnectivityState has never been written to will
-		// default to connectivity IDLE, which is fine.
-		scw.updateSubConnConnectivityState(scw.latestRawConnectivityState)
-		return
-	}
 	// If scw.latestHealthState has never been written to will use the health
 	// state CONNECTING set during object creation.
 	scw.updateSubConnHealthState(scw.latestHealthState)
@@ -192,12 +162,6 @@ func (scw *subConnWrapper) handleEjection() {
 	stateToUpdate := balancer.SubConnState{
 		ConnectivityState: connectivity.TransientFailure,
 	}
-	if !scw.healthListenerEnabled {
-		if scw.listener != nil {
-			scw.listener(stateToUpdate)
-		}
-		return
-	}
 	scw.mu.Lock()
 	defer scw.mu.Unlock()
 	if scw.healthListener != nil {



Test patch:
diff --git a/balancer/endpointsharding/endpointsharding_test.go b/balancer/endpointsharding/endpointsharding_test.go
index 0b995521..3a71503c 100644
--- a/balancer/endpointsharding/endpointsharding_test.go
+++ b/balancer/endpointsharding/endpointsharding_test.go
@@ -222,13 +222,13 @@ func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
-			bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
+			bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 	stub.Register(name, bf)
@@ -303,16 +303,16 @@ func (s) TestEndpointShardingExitIdle(t *testing.T) {
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
-			bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
+			bd.ChildBalancer = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, epOpts)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).ExitIdle()
+			bd.ChildBalancer.ExitIdle()
 		},
 	}
 	stub.Register(name, bf)
diff --git a/balancer/lazy/lazy_ext_test.go b/balancer/lazy/lazy_ext_test.go
index 1b1aef6a..53ca6574 100644
--- a/balancer/lazy/lazy_ext_test.go
+++ b/balancer/lazy/lazy_ext_test.go
@@ -79,19 +79,19 @@ func (s) TestExitIdle(t *testing.T) {
 
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).ExitIdle()
+			bd.ChildBalancer.ExitIdle()
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
-			bd.Data.(balancer.Balancer).ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 	stub.Register(t.Name(), bf)
@@ -144,16 +144,16 @@ func (s) TestPicker(t *testing.T) {
 
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
 		},
 		ExitIdle: func(*stub.BalancerData) {
 			t.Log("Ignoring call to ExitIdle, calling the picker should make the lazy balancer exit IDLE state.")
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 
@@ -201,24 +201,24 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
 
 	childBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			if resolverErrorReceived.HasFired() {
 				t.Error("Received resolver error before resolver state.")
 			}
 			resolverStateReceived = true
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
 			if !resolverStateReceived {
 				t.Error("Received resolver error before resolver state.")
 			}
 			resolverErrorReceived.Fire()
-			bd.Data.(balancer.Balancer).ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 
@@ -227,19 +227,19 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
 
 	topLevelBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
 		},
 		ExitIdle: func(*stub.BalancerData) {
 			t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
-			bd.Data.(balancer.Balancer).ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 
@@ -306,17 +306,17 @@ func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
 
 	childBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
 			t.Error("Received unexpected resolver error.")
-			bd.Data.(balancer.Balancer).ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 
@@ -325,16 +325,16 @@ func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
 
 	topLevelBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
 		},
 		ExitIdle: func(*stub.BalancerData) {
 			t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 
@@ -407,19 +407,19 @@ func (s) TestExitIdlePassthrough(t *testing.T) {
 
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
+			bd.ChildBalancer = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).ExitIdle()
+			bd.ChildBalancer.ExitIdle()
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
-			bd.Data.(balancer.Balancer).ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 	}
 	stub.Register(t.Name(), bf)
diff --git a/balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go b/balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go
index 25bb2944..df8c06f4 100644
--- a/balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go
+++ b/balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go
@@ -1241,14 +1241,14 @@ func (s) TestPickFirstLeaf_HealthListenerEnabled(t *testing.T) {
 	defer cancel()
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
@@ -1289,15 +1289,15 @@ func (s) TestPickFirstLeaf_HealthListenerNotEnabled(t *testing.T) {
 				healthListenerCh: healthListenerCh,
 				subConnStateCh:   make(chan balancer.SubConnState, 5),
 			}
-			bd.Data = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			// Functions like a non-petiole policy by not configuring the use
 			// of health listeners.
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
@@ -1345,14 +1345,14 @@ func (s) TestPickFirstLeaf_HealthUpdates(t *testing.T) {
 				healthListenerCh: healthListenerCh,
 				subConnStateCh:   scConnectivityStateCh,
 			}
-			bd.Data = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirstleaf.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
diff --git a/balancer/rls/balancer_test.go b/balancer/rls/balancer_test.go
index 6d13ccca..fd53b731 100644
--- a/balancer/rls/balancer_test.go
+++ b/balancer/rls/balancer_test.go
@@ -689,19 +689,17 @@ func (s) TestPickerUpdateOnDataCacheSizeDecrease(t *testing.T) {
 	stub.Register(topLevelBalancerName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			ccWrapper = &testCCWrapper{ClientConn: bd.ClientConn}
-			bd.Data = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
 		},
 		ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			parser := balancer.Get(Name).(balancer.ConfigParser)
 			return parser.ParseConfig(sc)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 
@@ -1070,19 +1068,17 @@ func (s) TestUpdateStatePauses(t *testing.T) {
 	stub.Register(topLevelBalancerName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			ccWrapper = &testCCWrapper{ClientConn: bd.ClientConn}
-			bd.Data = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(Name).Build(ccWrapper, bd.BuildOptions)
 		},
 		ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			parser := balancer.Get(Name).(balancer.ConfigParser)
 			return parser.ParseConfig(sc)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 
@@ -1098,10 +1094,10 @@ func (s) TestUpdateStatePauses(t *testing.T) {
 	}
 	stub.Register(childPolicyName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			cfg := &childPolicyConfig{}
@@ -1111,7 +1107,7 @@ func (s) TestUpdateStatePauses(t *testing.T) {
 			return cfg, nil
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
+			bal := bd.ChildBalancer
 			bd.ClientConn.UpdateState(balancer.State{ConnectivityState: connectivity.Idle, Picker: &testutils.TestConstPicker{Err: balancer.ErrNoSubConnAvailable}})
 			bd.ClientConn.UpdateState(balancer.State{ConnectivityState: connectivity.Connecting, Picker: &testutils.TestConstPicker{Err: balancer.ErrNoSubConnAvailable}})
 
diff --git a/internal/balancergroup/balancergroup_test.go b/internal/balancergroup/balancergroup_test.go
index 65e5af63..9dfad1a7 100644
--- a/internal/balancergroup/balancergroup_test.go
+++ b/internal/balancergroup/balancergroup_test.go
@@ -562,15 +562,14 @@ func (s) TestBalancerGracefulSwitch(t *testing.T) {
 	childPolicyName := t.Name()
 	stub.Register(childPolicyName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.ResolverState.Endpoints = ccs.ResolverState.Endpoints[1:]
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	})
 	cfgJSON := json.RawMessage(fmt.Sprintf(`[{%q: {}}]`, t.Name()))
diff --git a/internal/idle/idle_e2e_test.go b/internal/idle/idle_e2e_test.go
index 2cdf340a..3b606a9d 100644
--- a/internal/idle/idle_e2e_test.go
+++ b/internal/idle/idle_e2e_test.go
@@ -98,19 +98,17 @@ func registerWrappedRoundRobinPolicy(t *testing.T) chan struct{} {
 	closeCh := make(chan struct{}, 1)
 	stub.Register(roundrobin.Name, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = rrBuilder.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			select {
 			case closeCh <- struct{}{}:
 			default:
 			}
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 	t.Cleanup(func() { balancer.Register(rrBuilder) })
diff --git a/test/balancer_switching_test.go b/test/balancer_switching_test.go
index 2c4672f3..db543e63 100644
--- a/test/balancer_switching_test.go
+++ b/test/balancer_switching_test.go
@@ -488,17 +488,16 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) {
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			pf := balancer.Get(pickfirst.Name)
-			bd.Data = pf.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = pf.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
 			close(ccUpdateCh)
 			go func() {
 				<-waitToProceed
-				bal.UpdateClientConnState(ccs)
+				bd.ChildBalancer.UpdateClientConnState(ccs)
 			}()
 			return nil
 		},
diff --git a/test/balancer_test.go b/test/balancer_test.go
index 049aeb7f..7e3280d3 100644
--- a/test/balancer_test.go
+++ b/test/balancer_test.go
@@ -859,14 +859,13 @@ func (s) TestMetadataInPickResult(t *testing.T) {
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			cc := &testCCWrapper{ClientConn: bd.ClientConn}
-			bd.Data = balancer.Get(pickfirst.Name).Build(cc, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(cc, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	})
 
@@ -1025,16 +1024,16 @@ func (s) TestSubConn_RegisterHealthListener(t *testing.T) {
 				ClientConn: cc,
 				scChan:     scChan,
 			}
-			bd.Data = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).ExitIdle()
+			bd.ChildBalancer.ExitIdle()
 		},
 	}
 
@@ -1136,13 +1135,13 @@ func (s) TestSubConn_RegisterHealthListener_RegisterTwice(t *testing.T) {
 					}
 				},
 			}
-			bd.Data = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
@@ -1232,13 +1231,13 @@ func (s) TestSubConn_RegisterHealthListener_NilListener(t *testing.T) {
 					}
 				},
 			}
-			bd.Data = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
diff --git a/test/healthcheck_test.go b/test/healthcheck_test.go
index 0f7ec54a..18d9685c 100644
--- a/test/healthcheck_test.go
+++ b/test/healthcheck_test.go
@@ -1193,13 +1193,13 @@ func (s) TestHealthCheckUnregisterHealthListener(t *testing.T) {
 					close(readyUpdateReceivedCh)
 				},
 			}
-			bd.Data = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(ccw, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
diff --git a/test/resolver_update_test.go b/test/resolver_update_test.go
index 141cd553..e62d915a 100644
--- a/test/resolver_update_test.go
+++ b/test/resolver_update_test.go
@@ -161,10 +161,10 @@ func (s) TestResolverUpdate_InvalidServiceConfigAfterGoodUpdate(t *testing.T) {
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			pf := balancer.Get(pickfirst.Name)
-			bd.Data = pf.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = pf.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			cfg := &wrappingBalancerConfig{}
@@ -177,10 +177,9 @@ func (s) TestResolverUpdate_InvalidServiceConfigAfterGoodUpdate(t *testing.T) {
 			if _, ok := ccs.BalancerConfig.(*wrappingBalancerConfig); !ok {
 				return fmt.Errorf("received balancer config of unsupported type %T", ccs.BalancerConfig)
 			}
-			bal := bd.Data.(balancer.Balancer)
 			ccUpdateCh.Send(ccs)
 			ccs.BalancerConfig = nil
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	})
 
diff --git a/xds/internal/balancer/cdsbalancer/cdsbalancer_security_test.go b/xds/internal/balancer/cdsbalancer/cdsbalancer_security_test.go
index 7f647bfe..2df4fbf5 100644
--- a/xds/internal/balancer/cdsbalancer/cdsbalancer_security_test.go
+++ b/xds/internal/balancer/cdsbalancer/cdsbalancer_security_test.go
@@ -109,18 +109,16 @@ func registerWrappedCDSPolicyWithNewSubConnOverride(t *testing.T, ch chan *xdscr
 				ClientConn:      bd.ClientConn,
 				handshakeInfoCh: ch,
 			}
-			bd.Data = cdsBuilder.Build(ccWrapper, bd.BuildOptions)
+			bd.ChildBalancer = cdsBuilder.Build(ccWrapper, bd.BuildOptions)
 		},
 		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			return cdsBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 	t.Cleanup(func() { balancer.Register(cdsBuilder) })
diff --git a/xds/internal/balancer/cdsbalancer/cdsbalancer_test.go b/xds/internal/balancer/cdsbalancer/cdsbalancer_test.go
index 2cebfe52..4501a2a5 100644
--- a/xds/internal/balancer/cdsbalancer/cdsbalancer_test.go
+++ b/xds/internal/balancer/cdsbalancer/cdsbalancer_test.go
@@ -116,7 +116,7 @@ func registerWrappedClusterResolverPolicy(t *testing.T) (chan serviceconfig.Load
 
 	stub.Register(clusterresolver.Name, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = clusterresolverBuilder.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = clusterresolverBuilder.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			return clusterresolverBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
@@ -126,24 +126,21 @@ func registerWrappedClusterResolverPolicy(t *testing.T) (chan serviceconfig.Load
 			case lbCfgCh <- ccs.BalancerConfig:
 			default:
 			}
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		ResolverError: func(bd *stub.BalancerData, err error) {
 			select {
 			case resolverErrCh <- err:
 			default:
 			}
-			bal := bd.Data.(balancer.Balancer)
-			bal.ResolverError(err)
+			bd.ChildBalancer.ResolverError(err)
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).ExitIdle()
+			bd.ChildBalancer.ExitIdle()
 			close(exitIdleCh)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 			close(closeCh)
 		},
 	})
@@ -164,19 +161,17 @@ func registerWrappedCDSPolicy(t *testing.T) chan balancer.Balancer {
 	stub.Register(cdsBuilder.Name(), stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			bal := cdsBuilder.Build(bd.ClientConn, bd.BuildOptions)
-			bd.Data = bal
+			bd.ChildBalancer = bal
 			cdsBalancerCh <- bal
 		},
 		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			return cdsBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 	t.Cleanup(func() { balancer.Register(cdsBuilder) })
diff --git a/xds/internal/balancer/clustermanager/clustermanager_test.go b/xds/internal/balancer/clustermanager/clustermanager_test.go
index b996d88e..91166c17 100644
--- a/xds/internal/balancer/clustermanager/clustermanager_test.go
+++ b/xds/internal/balancer/clustermanager/clustermanager_test.go
@@ -651,14 +651,13 @@ func TestClusterGracefulSwitch(t *testing.T) {
 	childPolicyName := t.Name()
 	stub.Register(childPolicyName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	})
 	// Same cluster, different balancer type.
diff --git a/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go b/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go
index ef28c7b3..379783f9 100644
--- a/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go
+++ b/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go
@@ -207,19 +207,17 @@ func (s) TestConfigUpdate_ChildPolicyChange(t *testing.T) {
 			return pfBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
 		},
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = pfBuilder.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = pfBuilder.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			select {
 			case lbCfgCh <- ccs.BalancerConfig:
 			default:
 			}
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 
diff --git a/xds/internal/balancer/clusterresolver/e2e_test/balancer_test.go b/xds/internal/balancer/clusterresolver/e2e_test/balancer_test.go
index 517915dd..8e7f3d1f 100644
--- a/xds/internal/balancer/clusterresolver/e2e_test/balancer_test.go
+++ b/xds/internal/balancer/clusterresolver/e2e_test/balancer_test.go
@@ -342,7 +342,7 @@ func (s) TestOutlierDetectionConfigPropagationToChildPolicy(t *testing.T) {
 	lbCfgCh := make(chan serviceconfig.LoadBalancingConfig, 1)
 	stub.Register(priority.Name, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = priorityBuilder.Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = priorityBuilder.Build(bd.ClientConn, bd.BuildOptions)
 		},
 		ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			return priorityBuilder.(balancer.ConfigParser).ParseConfig(lbCfg)
@@ -352,12 +352,10 @@ func (s) TestOutlierDetectionConfigPropagationToChildPolicy(t *testing.T) {
 			case lbCfgCh <- ccs.BalancerConfig:
 			default:
 			}
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.Close()
+			bd.ChildBalancer.Close()
 		},
 	})
 	defer balancer.Register(priorityBuilder)
diff --git a/xds/internal/balancer/clusterresolver/e2e_test/eds_impl_test.go b/xds/internal/balancer/clusterresolver/e2e_test/eds_impl_test.go
index c127e0db..3123bfab 100644
--- a/xds/internal/balancer/clusterresolver/e2e_test/eds_impl_test.go
+++ b/xds/internal/balancer/clusterresolver/e2e_test/eds_impl_test.go
@@ -1256,14 +1256,14 @@ func (s) TestEDS_EndpointWithMultipleAddresses(t *testing.T) {
 			resolverState.Store(&resolver.State{})
 			stub.Register(roundrobin.Name, stub.BalancerFuncs{
 				Init: func(bd *stub.BalancerData) {
-					bd.Data = originalRRBuilder.Build(bd.ClientConn, bd.BuildOptions)
+					bd.ChildBalancer = originalRRBuilder.Build(bd.ClientConn, bd.BuildOptions)
 				},
 				Close: func(bd *stub.BalancerData) {
-					bd.Data.(balancer.Balancer).Close()
+					bd.ChildBalancer.Close()
 				},
 				UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 					resolverState.Store(&ccs.ResolverState)
-					return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+					return bd.ChildBalancer.UpdateClientConnState(ccs)
 				},
 			})
 
diff --git a/xds/internal/balancer/outlierdetection/balancer_test.go b/xds/internal/balancer/outlierdetection/balancer_test.go
index 061c1880..25b6a221 100644
--- a/xds/internal/balancer/outlierdetection/balancer_test.go
+++ b/xds/internal/balancer/outlierdetection/balancer_test.go
@@ -698,20 +698,27 @@ func (s) TestUpdateAddresses(t *testing.T) {
 	scsCh := testutils.NewChannel()
 	var scw1, scw2 balancer.SubConn
 	var err error
+	connectivityCh := make(chan struct{})
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
 			scw1, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw1, state: state}) },
+				StateListener: func(balancer.SubConnState) {},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw1.Connect()
 			scw2, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address2"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw2, state: state}) },
+				StateListener: func(state balancer.SubConnState) {
+					if state.ConnectivityState == connectivity.Ready {
+						close(connectivityCh)
+					}
+				},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw2.Connect()
 			bd.ClientConn.UpdateState(balancer.State{
 				ConnectivityState: connectivity.Ready,
 				Picker: &rrPicker{
@@ -753,6 +760,33 @@ func (s) TestUpdateAddresses(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 
+	// Transition SubConns to READY so that they can register a health listener.
+	for range 2 {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("Timed out waiting for creation of new SubConn.")
+		case sc := <-tcc.NewSubConnCh:
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		}
+	}
+
+	// Register health listeners after all the connectivity updates are
+	// processed to avoid data races while accessing the health listener within
+	// the TestClientConn.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for all SubConns to become READY.")
+	case <-connectivityCh:
+	}
+
+	scw1.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw1, state: healthState})
+	})
+	scw2.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw2, state: healthState})
+	})
+
 	// Setup the system to where one address is ejected and one address
 	// isn't.
 	select {
@@ -829,7 +863,7 @@ func (s) TestUpdateAddresses(t *testing.T) {
 	}
 	if err = scwsEqual(gotSCWS.(subConnWithState), subConnWithState{
 		sc:    scw1,
-		state: balancer.SubConnState{ConnectivityState: connectivity.Idle},
+		state: balancer.SubConnState{ConnectivityState: connectivity.Connecting},
 	}); err != nil {
 		t.Fatalf("Error in Sub Conn update: %v", err)
 	}
@@ -1005,26 +1039,34 @@ func (s) TestEjectUnejectSuccessRate(t *testing.T) {
 	scsCh := testutils.NewChannel()
 	var scw1, scw2, scw3 balancer.SubConn
 	var err error
+	connectivityCh := make(chan struct{})
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
 			scw1, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw1, state: state}) },
+				StateListener: func(balancer.SubConnState) {},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw1.Connect()
 			scw2, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address2"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw2, state: state}) },
+				StateListener: func(balancer.SubConnState) {},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw2.Connect()
 			scw3, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address3"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw3, state: state}) },
+				StateListener: func(state balancer.SubConnState) {
+					if state.ConnectivityState == connectivity.Ready {
+						close(connectivityCh)
+					}
+				},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw3.Connect()
 			bd.ClientConn.UpdateState(balancer.State{
 				ConnectivityState: connectivity.Ready,
 				Picker: &rrPicker{
@@ -1069,6 +1111,37 @@ func (s) TestEjectUnejectSuccessRate(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 
+	// Transition the SubConns to READY so that they can register health
+	// listeners.
+	for range 3 {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("Timed out waiting for creation of new SubConn.")
+		case sc := <-tcc.NewSubConnCh:
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		}
+	}
+
+	// Register health listeners after all the connectivity updates are
+	// processed to avoid data races while accessing the health listener within
+	// the TestClientConn.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for all SubConns to become READY.")
+	case <-connectivityCh:
+	}
+
+	scw1.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw1, state: healthState})
+	})
+	scw2.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw2, state: healthState})
+	})
+	scw3.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw3, state: healthState})
+	})
+
 	select {
 	case <-ctx.Done():
 		t.Fatalf("timeout while waiting for a UpdateState call on the ClientConn")
@@ -1098,9 +1171,12 @@ func (s) TestEjectUnejectSuccessRate(t *testing.T) {
 
 		// Since no addresses are ejected, a SubConn update should forward down
 		// to the child.
-		od.updateSubConnState(scw1.(*subConnWrapper), balancer.SubConnState{
-			ConnectivityState: connectivity.Connecting,
-		})
+		od.scUpdateCh.Put(&scHealthUpdate{
+			scw: scw1.(*subConnWrapper),
+			state: balancer.SubConnState{
+				ConnectivityState: connectivity.Connecting,
+			}},
+		)
 
 		gotSCWS, err := scsCh.Receive(ctx)
 		if err != nil {
@@ -1163,8 +1239,9 @@ func (s) TestEjectUnejectSuccessRate(t *testing.T) {
 		// that address should not be forwarded downward. These SubConn updates
 		// will be cached to update the child sometime in the future when the
 		// address gets unejected.
-		od.updateSubConnState(scw3.(*subConnWrapper), balancer.SubConnState{
-			ConnectivityState: connectivity.Connecting,
+		od.scUpdateCh.Put(&scHealthUpdate{
+			scw:   scw3.(*subConnWrapper),
+			state: balancer.SubConnState{ConnectivityState: connectivity.Connecting},
 		})
 		sCtx, cancel = context.WithTimeout(context.Background(), defaultTestShortTimeout)
 		defer cancel()
@@ -1215,29 +1292,37 @@ func (s) TestEjectFailureRate(t *testing.T) {
 	scsCh := testutils.NewChannel()
 	var scw1, scw2, scw3 balancer.SubConn
 	var err error
+	connectivityCh := make(chan struct{})
 	stub.Register(t.Name(), stub.BalancerFuncs{
 		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
 			if scw1 != nil { // UpdateClientConnState was already called, no need to recreate SubConns.
 				return nil
 			}
 			scw1, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw1, state: state}) },
+				StateListener: func(balancer.SubConnState) {},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw1.Connect()
 			scw2, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address2"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw2, state: state}) },
+				StateListener: func(balancer.SubConnState) {},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw2.Connect()
 			scw3, err = bd.ClientConn.NewSubConn([]resolver.Address{{Addr: "address3"}}, balancer.NewSubConnOptions{
-				StateListener: func(state balancer.SubConnState) { scsCh.Send(subConnWithState{sc: scw3, state: state}) },
+				StateListener: func(scs balancer.SubConnState) {
+					if scs.ConnectivityState == connectivity.Ready {
+						close(connectivityCh)
+					}
+				},
 			})
 			if err != nil {
 				t.Errorf("error in od.NewSubConn call: %v", err)
 			}
+			scw3.Connect()
 			return nil
 		},
 	})
@@ -1283,6 +1368,36 @@ func (s) TestEjectFailureRate(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 
+	// Transition the SubConns to READY so that they can register health
+	// listeners.
+	for range 3 {
+		select {
+		case <-ctx.Done():
+			t.Fatal("Timed out waiting for creation of new SubConn.")
+		case sc := <-tcc.NewSubConnCh:
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		}
+	}
+	// Register health listeners after all the connectivity updates are
+	// processed to avoid data races while accessing the health listener within
+	// the TestClientConn.
+	select {
+	case <-ctx.Done():
+		t.Fatal("Context timed out waiting for all SubConns to become READY.")
+	case <-connectivityCh:
+	}
+
+	scw1.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw1, state: healthState})
+	})
+	scw2.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw2, state: healthState})
+	})
+	scw3.RegisterHealthListener(func(healthState balancer.SubConnState) {
+		scsCh.Send(subConnWithState{sc: scw3, state: healthState})
+	})
+
 	select {
 	case <-ctx.Done():
 		t.Fatalf("timeout while waiting for a UpdateState call on the ClientConn")
@@ -1378,7 +1493,7 @@ func (s) TestEjectFailureRate(t *testing.T) {
 		}
 		if err = scwsEqual(gotSCWS.(subConnWithState), subConnWithState{
 			sc:    scw3,
-			state: balancer.SubConnState{ConnectivityState: connectivity.Idle},
+			state: balancer.SubConnState{ConnectivityState: connectivity.Connecting},
 		}); err != nil {
 			t.Fatalf("Error in Sub Conn update: %v", err)
 		}
@@ -1626,15 +1741,15 @@ func (s) TestPickFirstHealthListenerDisabled(t *testing.T) {
 	lbChan := make(chan *outlierDetectionBalancer, 1)
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
-			lbChan <- bd.Data.(*outlierDetectionBalancer)
+			bd.ChildBalancer = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
+			lbChan <- bd.ChildBalancer.(*outlierDetectionBalancer)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.BalancerConfig = odCfg
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
@@ -1729,15 +1844,15 @@ func (s) TestMultipleAddressesPerEndpoint(t *testing.T) {
 	lbChan := make(chan *outlierDetectionBalancer, 1)
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
-			lbChan <- bd.Data.(*outlierDetectionBalancer)
+			bd.ChildBalancer = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
+			lbChan <- bd.ChildBalancer.(*outlierDetectionBalancer)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.BalancerConfig = odCfg
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
@@ -1905,15 +2020,15 @@ func (s) TestEjectionStateResetsWhenEndpointAddressesChange(t *testing.T) {
 	lbChan := make(chan *outlierDetectionBalancer, 1)
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
-			lbChan <- bd.Data.(*outlierDetectionBalancer)
+			bd.ChildBalancer = balancer.Get(Name).Build(bd.ClientConn, bd.BuildOptions)
+			lbChan <- bd.ChildBalancer.(*outlierDetectionBalancer)
 		},
 		Close: func(bd *stub.BalancerData) {
-			bd.Data.(balancer.Balancer).Close()
+			bd.ChildBalancer.Close()
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			ccs.BalancerConfig = odCfg
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	}
 
diff --git a/xds/internal/balancer/priority/balancer_test.go b/xds/internal/balancer/priority/balancer_test.go
index dee06006..3aac85ce 100644
--- a/xds/internal/balancer/priority/balancer_test.go
+++ b/xds/internal/balancer/priority/balancer_test.go
@@ -1718,11 +1718,10 @@ func (s) TestPriority_IgnoreReresolutionRequestTwoChildren(t *testing.T) {
 	stub.Register(childPolicyName, stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
 			ccCh.Send(bd.ClientConn)
-			bd.Data = balancer.Get(roundrobin.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.ChildBalancer = balancer.Get(roundrobin.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			bal := bd.Data.(balancer.Balancer)
-			return bal.UpdateClientConnState(ccs)
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
 		},
 	})