Repository: grpc/grpc-go

Issue:
[Flaky test]: Test/DNSResolver_ExponentialBackoff
https://github.com/grpc/grpc-go/actions/runs/11409299053/job/31749223628

Discussion:
@purnesh42H: @aranjans as oncall to measure the flakiness@aranjans: FAILED in 9 out of 100000.
http://fusion2/c2049abb-ebc7-407e-8804-0a755971ffa2@pvsravani: https://github.com/grpc/grpc-go/actions/runs/12645537324/job/35234840419?pr=7951@aranjans: https://github.com/grpc/grpc-go/actions/runs/13005060741/job/36270375062

Solution patch:
diff --git a/balancer/endpointsharding/endpointsharding.go b/balancer/endpointsharding/endpointsharding.go
index f6c4e2c9..e812f1fd 100644
--- a/balancer/endpointsharding/endpointsharding.go
+++ b/balancer/endpointsharding/endpointsharding.go
@@ -26,37 +26,17 @@
 package endpointsharding
 
 import (
-	"encoding/json"
 	"errors"
-	"fmt"
 	rand "math/rand/v2"
 	"sync"
 	"sync/atomic"
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/base"
-	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
-	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/resolver"
-	"google.golang.org/grpc/serviceconfig"
 )
 
-var (
-	// PickFirstConfig is a pick first config without shuffling enabled.
-	PickFirstConfig serviceconfig.LoadBalancingConfig
-	logger          = grpclog.Component("endpoint-sharding")
-)
-
-func init() {
-	var err error
-	PickFirstConfig, err = ParseConfig(json.RawMessage(fmt.Sprintf("[{%q: {}}]", pickfirstleaf.Name)))
-	if err != nil {
-		logger.Fatal(err)
-	}
-}
-
 // ChildState is the balancer state of a child along with the endpoint which
 // identifies the child balancer.
 type ChildState struct {
@@ -68,26 +48,30 @@ type ChildState struct {
 	Balancer balancer.ExitIdler
 }
 
-// NewBalancer returns a load balancing policy that manages homogeneous child
-// policies each owning a single endpoint. The balancer will automatically call
-// ExitIdle on its children if they report IDLE connectivity state.
-func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	return newBlanacer(cc, opts, true)
+// Options are the options to configure the behaviour of the
+// endpointsharding balancer.
+type Options struct {
+	// DisableAutoReconnect allows the balancer to keep child balancer in the
+	// IDLE state until they are explicitly triggered to exit using the
+	// ChildState obtained from the endpointsharding picker. When set to false,
+	// the endpointsharding balancer will automatically call ExitIdle on child
+	// connections that report IDLE.
+	DisableAutoReconnect bool
 }
 
-// NewBalancerWithoutAutoReconnect returns a load balancing policy that manages
-// homogeneous child policies each owning a single endpoint. The balancer will
-// allow children to remain in IDLE state until triggered to exit idle state
-// using the ChildState obtained using the endpointsharding picker.
-func NewBalancerWithoutAutoReconnect(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	return newBlanacer(cc, opts, false)
-}
+// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
+// type as the balancer.Builder.Build method.
+type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
 
-func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, autoReconnect bool) balancer.Balancer {
+// NewBalancer returns a load balancing policy that manages homogeneous child
+// policies each owning a single endpoint. The endpointsharding balancer
+// forwards the LoadBalancingConfig in ClientConn state updates to its children.
+func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions, childBuilder ChildBuilderFunc, esOpts Options) balancer.Balancer {
 	es := &endpointSharding{
-		cc:                  cc,
-		bOpts:               opts,
-		enableAutoReconnect: autoReconnect,
+		cc:           cc,
+		bOpts:        opts,
+		esOpts:       esOpts,
+		childBuilder: childBuilder,
 	}
 	es.children.Store(resolver.NewEndpointMap())
 	return es
@@ -97,9 +81,10 @@ func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, autoReconne
 // balancer with child config for every unique Endpoint received. It updates the
 // child states on any update from parent or child.
 type endpointSharding struct {
-	cc                  balancer.ClientConn
-	bOpts               balancer.BuildOptions
-	enableAutoReconnect bool
+	cc           balancer.ClientConn
+	bOpts        balancer.BuildOptions
+	esOpts       Options
+	childBuilder ChildBuilderFunc
 
 	// childMu synchronizes calls to any single child. It must be held for all
 	// calls into a child. To avoid deadlocks, do not acquire childMu while
@@ -160,7 +145,7 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 				es:         es,
 			}
 			childBalancer.childState.Balancer = childBalancer
-			childBalancer.child = gracefulswitch.NewBalancer(childBalancer, es.bOpts)
+			childBalancer.child = es.childBuilder(childBalancer, es.bOpts)
 		}
 		newChildren.Set(endpoint, childBalancer)
 		if err := childBalancer.updateClientConnStateLocked(balancer.ClientConnState{
@@ -334,7 +319,7 @@ func (bw *balancerWrapper) UpdateState(state balancer.State) {
 	bw.es.mu.Lock()
 	bw.childState.State = state
 	bw.es.mu.Unlock()
-	if state.ConnectivityState == connectivity.Idle && bw.es.enableAutoReconnect {
+	if state.ConnectivityState == connectivity.Idle && !bw.es.esOpts.DisableAutoReconnect {
 		bw.ExitIdle()
 	}
 	bw.es.updateState()
@@ -370,12 +355,3 @@ func (bw *balancerWrapper) closeLocked() {
 	bw.child.Close()
 	bw.isClosed = true
 }
-
-// ParseConfig parses a child config list and returns an LB config to use with
-// the endpointsharding balancer.
-//
-// cfg is expected to be a JSON array of LB policy names + configs as the
-// format of the loadBalancingConfig field in ServiceConfig.
-func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
-	return gracefulswitch.ParseConfig(cfg)
-}
diff --git a/balancer/lazy/lazy.go b/balancer/lazy/lazy.go
index c680b9bb..7368d8f4 100644
--- a/balancer/lazy/lazy.go
+++ b/balancer/lazy/lazy.go
@@ -26,44 +26,35 @@
 package lazy
 
 import (
-	"encoding/json"
 	"fmt"
 	"sync"
 
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/resolver"
-	"google.golang.org/grpc/serviceconfig"
 
 	internalgrpclog "google.golang.org/grpc/internal/grpclog"
 )
 
-func init() {
-	balancer.Register(builder{})
-}
-
 var (
-	// PickfirstConfig is the LB policy config json for a pick_first load
-	// balancer that is lazily initialized.
-	PickfirstConfig = fmt.Sprintf("{\"childPolicy\": [{%q: {}}]}", pickfirstleaf.Name)
-	logger          = grpclog.Component("lazy-lb")
+	logger = grpclog.Component("lazy-lb")
 )
 
 const (
-	// Name is the name of the lazy balancer.
-	Name      = "lazy"
 	logPrefix = "[lazy-lb %p] "
 )
 
-type builder struct{}
+// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
+// type as the balancer.Builder.Build method.
+type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
 
-func (builder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+// NewBalancer is the constructor for the lazy balancer.
+func NewBalancer(cc balancer.ClientConn, bOpts balancer.BuildOptions, childBuilder ChildBuilderFunc) balancer.Balancer {
 	b := &lazyBalancer{
 		cc:           cc,
 		buildOptions: bOpts,
+		childBuilder: childBuilder,
 	}
 	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
 	cc.UpdateState(balancer.State{
@@ -77,16 +68,13 @@ func (builder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balanc
 	return b
 }
 
-func (builder) Name() string {
-	return Name
-}
-
 type lazyBalancer struct {
 	// The following fields are initialized at build time and read-only after
 	// that and therefore do not need to be guarded by a mutex.
 	cc           balancer.ClientConn
 	buildOptions balancer.BuildOptions
 	logger       *internalgrpclog.PrefixLogger
+	childBuilder ChildBuilderFunc
 
 	// The following fields are accessed while handling calls to the idlePicker
 	// and when handling ClientConn state updates. They are guarded by a mutex.
@@ -119,12 +107,6 @@ func (lb *lazyBalancer) ResolverError(err error) {
 func (lb *lazyBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
 	lb.mu.Lock()
 	defer lb.mu.Unlock()
-	childLBCfg, ok := ccs.BalancerConfig.(lbCfg)
-	if !ok {
-		lb.logger.Errorf("Got LB config of unexpected type: %v", ccs.BalancerConfig)
-		return balancer.ErrBadResolverState
-	}
-	ccs.BalancerConfig = childLBCfg.childLBCfg
 	if lb.delegate != nil {
 		return lb.delegate.UpdateClientConnState(ccs)
 	}
@@ -148,7 +130,7 @@ func (lb *lazyBalancer) ExitIdle() {
 		}
 		return
 	}
-	lb.delegate = gracefulswitch.NewBalancer(lb.cc, lb.buildOptions)
+	lb.delegate = lb.childBuilder(lb.cc, lb.buildOptions)
 	if lb.latestClientConnState != nil {
 		if err := lb.delegate.UpdateClientConnState(*lb.latestClientConnState); err != nil {
 			if err == balancer.ErrBadResolverState {
@@ -165,25 +147,6 @@ func (lb *lazyBalancer) ExitIdle() {
 	}
 }
 
-type lbCfg struct {
-	serviceconfig.LoadBalancingConfig
-	childLBCfg serviceconfig.LoadBalancingConfig
-}
-
-func (b builder) ParseConfig(lbConfig json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
-	jsonReprsentation := &struct {
-		ChildPolicy json.RawMessage
-	}{}
-	if err := json.Unmarshal(lbConfig, jsonReprsentation); err != nil {
-		return nil, err
-	}
-	childCfg, err := gracefulswitch.ParseConfig(jsonReprsentation.ChildPolicy)
-	if err != nil {
-		return nil, err
-	}
-	return lbCfg{childLBCfg: childCfg}, nil
-}
-
 // idlePicker is used when the SubConn is IDLE and kicks the SubConn into
 // CONNECTING when Pick is called.
 type idlePicker struct {
diff --git a/balancer/leastrequest/leastrequest.go b/balancer/leastrequest/leastrequest.go
index 0157d7ee..d25f9178 100644
--- a/balancer/leastrequest/leastrequest.go
+++ b/balancer/leastrequest/leastrequest.go
@@ -41,9 +41,8 @@ const Name = "least_request_experimental"
 
 var (
 	// randuint32 is a global to stub out in tests.
-	randuint32               = rand.Uint32
-	endpointShardingLBConfig = endpointsharding.PickFirstConfig
-	logger                   = grpclog.Component("least-request")
+	randuint32 = rand.Uint32
+	logger     = grpclog.Component("least-request")
 )
 
 func init() {
@@ -91,7 +90,7 @@ func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Ba
 		ClientConn:        cc,
 		endpointRPCCounts: resolver.NewEndpointMap(),
 	}
-	b.child = endpointsharding.NewBalancer(b, bOpts)
+	b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirstleaf.Name).Build, endpointsharding.Options{})
 	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", b))
 	b.logger.Infof("Created")
 	return b
@@ -129,11 +128,11 @@ func (lrb *leastRequestBalancer) UpdateClientConnState(ccs balancer.ClientConnSt
 	lrb.mu.Lock()
 	lrb.choiceCount = lrCfg.ChoiceCount
 	lrb.mu.Unlock()
-	// Enable the health listener in pickfirst children for client side health
-	// checks and outlier detection, if configured.
-	ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
-	ccs.BalancerConfig = endpointShardingLBConfig
-	return lrb.child.UpdateClientConnState(ccs)
+	return lrb.child.UpdateClientConnState(balancer.ClientConnState{
+		// Enable the health listener in pickfirst children for client side health
+		// checks and outlier detection, if configured.
+		ResolverState: pickfirstleaf.EnableHealthListener(ccs.ResolverState),
+	})
 }
 
 type endpointState struct {
diff --git a/balancer/weightedroundrobin/balancer.go b/balancer/weightedroundrobin/balancer.go
index fab3c5d7..f8b6914c 100644
--- a/balancer/weightedroundrobin/balancer.go
+++ b/balancer/weightedroundrobin/balancer.go
@@ -84,9 +84,6 @@ var (
 	})
 )
 
-// endpointSharding which specifies pick first children.
-var endpointShardingLBConfig = endpointsharding.PickFirstConfig
-
 func init() {
 	balancer.Register(bb{})
 }
@@ -103,7 +100,7 @@ func (bb) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Ba
 		scToWeight:       make(map[balancer.SubConn]*endpointWeight),
 	}
 
-	b.child = endpointsharding.NewBalancer(b, bOpts)
+	b.child = endpointsharding.NewBalancer(b, bOpts, balancer.Get(pickfirstleaf.Name).Build, endpointsharding.Options{})
 	b.logger = prefixLogger(b)
 	b.logger.Infof("Created")
 	return b
@@ -230,14 +227,12 @@ func (b *wrrBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error
 	b.updateEndpointsLocked(ccs.ResolverState.Endpoints)
 	b.mu.Unlock()
 
-	// Make pickfirst children use health listeners for outlier detection to
-	// work.
-	ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
 	// This causes child to update picker inline and will thus cause inline
 	// picker update.
 	return b.child.UpdateClientConnState(balancer.ClientConnState{
-		BalancerConfig: endpointShardingLBConfig,
-		ResolverState:  ccs.ResolverState,
+		// Make pickfirst children use health listeners for outlier detection to
+		// work.
+		ResolverState: pickfirstleaf.EnableHealthListener(ccs.ResolverState),
 	})
 }
 
diff --git a/examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go b/examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go
index 61a50c01..f53919ee 100644
--- a/examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go
+++ b/examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go
@@ -27,12 +27,11 @@ import (
 	_ "google.golang.org/grpc" // to register pick_first
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/serviceconfig"
 )
 
-var gracefulSwitchPickFirst = endpointsharding.PickFirstConfig
-
 func init() {
 	balancer.Register(customRoundRobinBuilder{})
 }
@@ -69,7 +68,7 @@ func (customRoundRobinBuilder) Build(cc balancer.ClientConn, bOpts balancer.Buil
 		ClientConn: cc,
 		bOpts:      bOpts,
 	}
-	crr.Balancer = endpointsharding.NewBalancer(crr, bOpts)
+	crr.Balancer = endpointsharding.NewBalancer(crr, bOpts, balancer.Get(pickfirstleaf.Name).Build, endpointsharding.Options{})
 	return crr
 }
 
@@ -100,8 +99,7 @@ func (crr *customRoundRobin) UpdateClientConnState(state balancer.ClientConnStat
 	// is guaranteed to happen since the aggregator will always call
 	// UpdateChildState in its UpdateClientConnState.
 	return crr.Balancer.UpdateClientConnState(balancer.ClientConnState{
-		BalancerConfig: gracefulSwitchPickFirst,
-		ResolverState:  state.ResolverState,
+		ResolverState: state.ResolverState,
 	})
 }
 



Test patch:
diff --git a/balancer/endpointsharding/endpointsharding_test.go b/balancer/endpointsharding/endpointsharding_test.go
index 2c1fe3b6..86691db3 100644
--- a/balancer/endpointsharding/endpointsharding_test.go
+++ b/balancer/endpointsharding/endpointsharding_test.go
@@ -30,6 +30,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
@@ -80,7 +81,7 @@ func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptio
 		ClientConn: cc,
 		bOpts:      opts,
 	}
-	fp.Balancer = endpointsharding.NewBalancer(fp, opts)
+	fp.Balancer = endpointsharding.NewBalancer(fp, opts, balancer.Get(pickfirstleaf.Name).Build, endpointsharding.Options{})
 	return fp
 }
 
@@ -103,10 +104,7 @@ func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) err
 		return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
 	}
 
-	return fp.Balancer.UpdateClientConnState(balancer.ClientConnState{
-		BalancerConfig: endpointsharding.PickFirstConfig,
-		ResolverState:  state.ResolverState,
-	})
+	return fp.Balancer.UpdateClientConnState(state)
 }
 
 func (fp *fakePetiole) UpdateState(state balancer.State) {
@@ -182,13 +180,11 @@ func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
 	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = endpointsharding.NewBalancerWithoutAutoReconnect(bd.ClientConn, bd.BuildOptions)
+			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
+			bd.Data = 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(balancer.ClientConnState{
-				BalancerConfig: endpointsharding.PickFirstConfig,
-				ResolverState:  ccs.ResolverState,
-			})
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			bd.Data.(balancer.Balancer).Close()
diff --git a/balancer/lazy/lazy_ext_test.go b/balancer/lazy/lazy_ext_test.go
index a822b884..bff01489 100644
--- a/balancer/lazy/lazy_ext_test.go
+++ b/balancer/lazy/lazy_ext_test.go
@@ -20,7 +20,6 @@ package lazy_test
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"strings"
@@ -77,7 +76,25 @@ func (s) TestExitIdle(t *testing.T) {
 		},
 	})
 
-	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": %s}]}`, lazy.Name, lazy.PickfirstConfig)
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
+		},
+		ExitIdle: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.ExitIdler).ExitIdle()
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			bd.Data.(balancer.Balancer).ResolverError(err)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.Balancer).Close()
+		},
+	}
+	stub.Register(t.Name(), bf)
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
 	opts := []grpc.DialOption{
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithDefaultServiceConfig(json),
@@ -86,7 +103,6 @@ func (s) TestExitIdle(t *testing.T) {
 	cc, err := grpc.NewClient(mr.Scheme()+":///", opts...)
 	if err != nil {
 		t.Fatalf("grpc.NewClient(_) failed: %v", err)
-
 	}
 	defer cc.Close()
 
@@ -125,23 +141,15 @@ func (s) TestPicker(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
 	defer cancel()
 
-	lazyCfg, err := balancer.Get(lazy.Name).(balancer.ConfigParser).ParseConfig(json.RawMessage(lazy.PickfirstConfig))
-	if err != nil {
-		t.Fatalf("Failed to parse service config: %v", err)
-	}
-
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(lazy.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
 		},
 		ExitIdle: func(bd *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(balancer.ClientConnState{
-				BalancerConfig: lazyCfg,
-				ResolverState:  ccs.ResolverState,
-			})
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			bd.Data.(balancer.Balancer).Close()
@@ -216,14 +224,9 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
 	childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
 	stub.Register(childBalName, childBF)
 
-	lazyCfgJSON := fmt.Sprintf("{\"childPolicy\": [{%q: {}}]}", childBalName)
-	lazyCfg, err := balancer.Get(lazy.Name).(balancer.ConfigParser).ParseConfig(json.RawMessage(lazyCfgJSON))
-	if err != nil {
-		t.Fatalf("Failed to parse service config: %v", err)
-	}
 	topLevelBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(lazy.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
 		},
 		ExitIdle: func(bd *stub.BalancerData) {
 			t.Log("Ignoring call to ExitIdle to delay lazy child creation until RPC time.")
@@ -232,10 +235,7 @@ func (s) TestGoodUpdateThenResolverError(t *testing.T) {
 			bd.Data.(balancer.Balancer).ResolverError(err)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(balancer.ClientConnState{
-				BalancerConfig: lazyCfg,
-				ResolverState:  ccs.ResolverState,
-			})
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			bd.Data.(balancer.Balancer).Close()
@@ -320,23 +320,15 @@ func (s) TestResolverErrorThenGoodUpdate(t *testing.T) {
 	childBalName := strings.ReplaceAll(strings.ToLower(t.Name())+"_child", "/", "")
 	stub.Register(childBalName, childBF)
 
-	lazyCfgJSON := fmt.Sprintf("{\"childPolicy\": [{%q: {}}]}", childBalName)
-	lazyCfg, err := balancer.Get(lazy.Name).(balancer.ConfigParser).ParseConfig(json.RawMessage(lazyCfgJSON))
-	if err != nil {
-		t.Fatalf("Failed to parse service config: %v", err)
-	}
 	topLevelBF := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = balancer.Get(lazy.Name).Build(bd.ClientConn, bd.BuildOptions)
+			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(childBalName).Build)
 		},
 		ExitIdle: func(bd *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(balancer.ClientConnState{
-				BalancerConfig: lazyCfg,
-				ResolverState:  ccs.ResolverState,
-			})
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			bd.Data.(balancer.Balancer).Close()
@@ -410,7 +402,25 @@ func (s) TestExitIdlePassthrough(t *testing.T) {
 		},
 	})
 
-	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": %s}]}`, lazy.Name, lazy.PickfirstConfig)
+	bf := stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.Data = lazy.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build)
+		},
+		ExitIdle: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.ExitIdler).ExitIdle()
+		},
+		ResolverError: func(bd *stub.BalancerData, err error) {
+			bd.Data.(balancer.Balancer).ResolverError(err)
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.Balancer).Close()
+		},
+	}
+	stub.Register(t.Name(), bf)
+	json := fmt.Sprintf(`{"loadBalancingConfig": [{"%s": {}}]}`, t.Name())
 	opts := []grpc.DialOption{
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithDefaultServiceConfig(json),
diff --git a/credentials/alts/alts_test.go b/credentials/alts/alts_test.go
index 6c083a2f..48f871a0 100644
--- a/credentials/alts/alts_test.go
+++ b/credentials/alts/alts_test.go
@@ -36,6 +36,7 @@ import (
 	altspb "google.golang.org/grpc/credentials/alts/internal/proto/grpc_gcp"
 	"google.golang.org/grpc/credentials/alts/internal/testutil"
 	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils"
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
 	testpb "google.golang.org/grpc/interop/grpc_testing"
@@ -323,7 +324,7 @@ func (s) TestFullHandshake(t *testing.T) {
 	defer wait.Wait()
 	stopHandshaker, handshakerAddress := startFakeHandshakerService(t, &wait)
 	defer stopHandshaker()
-	stopServer, serverAddress := startServer(t, handshakerAddress, &wait)
+	stopServer, serverAddress := startServer(t, handshakerAddress)
 	defer stopServer()
 
 	// Ping the server, authenticating with ALTS.
@@ -349,7 +350,7 @@ func (s) TestConcurrentHandshakes(t *testing.T) {
 	defer wait.Wait()
 	stopHandshaker, handshakerAddress := startFakeHandshakerService(t, &wait)
 	defer stopHandshaker()
-	stopServer, serverAddress := startServer(t, handshakerAddress, &wait)
+	stopServer, serverAddress := startServer(t, handshakerAddress)
 	defer stopServer()
 
 	// Ping the server, authenticating with ALTS.
@@ -444,31 +445,22 @@ func startFakeHandshakerService(t *testing.T, wait *sync.WaitGroup) (stop func()
 	return func() { s.Stop() }, listener.Addr().String()
 }
 
-func startServer(t *testing.T, handshakerServiceAddress string, wait *sync.WaitGroup) (stop func(), address string) {
+func startServer(t *testing.T, handshakerServiceAddress string) (stop func(), address string) {
 	listener, err := testutils.LocalTCPListener()
 	if err != nil {
 		t.Fatalf("LocalTCPListener() failed: %v", err)
 	}
 	serverOpts := &ServerOptions{HandshakerServiceAddress: handshakerServiceAddress}
 	creds := NewServerCreds(serverOpts)
-	s := grpc.NewServer(grpc.Creds(creds))
-	testgrpc.RegisterTestServiceServer(s, &testServer{})
-	wait.Add(1)
-	go func() {
-		defer wait.Done()
-		if err := s.Serve(listener); err != nil {
-			t.Errorf("s.Serve(%v) failed: %v", listener, err)
-		}
-	}()
-	return func() { s.Stop() }, listener.Addr().String()
-}
-
-type testServer struct {
-	testgrpc.UnimplementedTestServiceServer
-}
-
-func (s *testServer) UnaryCall(_ context.Context, _ *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
-	return &testpb.SimpleResponse{
-		Payload: &testpb.Payload{},
-	}, nil
+	stub := &stubserver.StubServer{
+		Listener: listener,
+		UnaryCallF: func(context.Context, *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			return &testpb.SimpleResponse{
+				Payload: &testpb.Payload{},
+			}, nil
+		},
+		S: grpc.NewServer(grpc.Creds(creds)),
+	}
+	stubserver.StartTestService(t, stub)
+	return func() { stub.S.Stop() }, listener.Addr().String()
 }
diff --git a/internal/resolver/dns/dns_resolver_test.go b/internal/resolver/dns/dns_resolver_test.go
index 47285052..8b01ce04 100644
--- a/internal/resolver/dns/dns_resolver_test.go
+++ b/internal/resolver/dns/dns_resolver_test.go
@@ -626,20 +626,17 @@ func (s) TestDNSResolver_ExponentialBackoff(t *testing.T) {
 					}
 				}
 
+				if i == retries-1 {
+					// Update resolver.ClientConn to not return an error
+					// anymore before last resolution retry to ensure that
+					// last resolution attempt doesn't back off.
+					returnNilErr.Store(true)
+				}
+
 				// Unblock the DNS resolver's backoff by pushing the current time.
 				timeChan <- time.Now()
 			}
 
-			// Update resolver.ClientConn to not return an error anymore.
-			returnNilErr.Store(true)
-
-			// Unblock the DNS resolver's backoff, if ongoing, while we set the
-			// test clientConn to not return an error anymore.
-			select {
-			case timeChan <- time.Now():
-			default:
-			}
-
 			// Verify that the DNS resolver does not backoff anymore.
 			sCtx, sCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
 			defer sCancel()
diff --git a/test/stats_test.go b/test/stats_test.go
index 0c36b6f1..5ee00bb3 100644
--- a/test/stats_test.go
+++ b/test/stats_test.go
@@ -27,8 +27,10 @@ import (
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/interop"
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/stats"
 )
@@ -59,18 +61,15 @@ func (s) TestPeerForClientStatsHandler(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	s := grpc.NewServer()
-	testgrpc.RegisterTestServiceServer(s, interop.NewTestServer())
-	errCh := make(chan error)
-	go func() {
-		errCh <- s.Serve(l)
-	}()
-	defer func() {
-		s.Stop()
-		if err := <-errCh; err != nil {
-			t.Error(err)
-		}
-	}()
+	ss := &stubserver.StubServer{
+		Listener: l,
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+		S: grpc.NewServer(),
+	}
+	stubserver.StartTestService(t, ss)
+	defer ss.S.Stop()
 
 	// Create client with stats handler and do some calls.
 	cc, err := grpc.NewClient(
diff --git a/xds/internal/server/listener_wrapper_test.go b/xds/internal/server/listener_wrapper_test.go
index 981d65e4..e181a7b8 100644
--- a/xds/internal/server/listener_wrapper_test.go
+++ b/xds/internal/server/listener_wrapper_test.go
@@ -23,7 +23,6 @@ import (
 	"fmt"
 	"net"
 	"strconv"
-	"sync"
 	"testing"
 	"time"
 
@@ -31,6 +30,7 @@ import (
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/internal/testutils/xds/e2e"
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
@@ -159,14 +159,6 @@ func (s) TestListenerWrapper(t *testing.T) {
 	}
 }
 
-type testService struct {
-	testgrpc.TestServiceServer
-}
-
-func (*testService) EmptyCall(context.Context, *testpb.Empty) (*testpb.Empty, error) {
-	return &testpb.Empty{}, nil
-}
-
 // TestConnsCleanup tests that the listener wrapper clears it's connection
 // references when connections close. It sets up a listener wrapper and gRPC
 // Server, and connects to the server 100 times and makes an RPC each time, and
@@ -220,14 +212,15 @@ func (s) TestConnsCleanup(t *testing.T) {
 		}
 	}
 
-	server := grpc.NewServer(grpc.Creds(insecure.NewCredentials()))
-	testgrpc.RegisterTestServiceServer(server, &testService{})
-	wg := sync.WaitGroup{}
-	go func() {
-		if err := server.Serve(lw); err != nil {
-			t.Errorf("failed to serve: %v", err)
-		}
-	}()
+	ss := &stubserver.StubServer{
+		Listener: lis,
+		EmptyCallF: func(context.Context, *testpb.Empty) (*testpb.Empty, error) {
+			return &testpb.Empty{}, nil
+		},
+		S: grpc.NewServer(grpc.Creds(insecure.NewCredentials())),
+	}
+	stubserver.StartTestService(t, ss)
+	defer ss.S.Stop()
 
 	// Make 100 connections to the server, and make an RPC on each one.
 	for i := 0; i < 100; i++ {
@@ -255,6 +248,4 @@ func (s) TestConnsCleanup(t *testing.T) {
 		t.Fatalf("timeout waiting for lis wrapper conns to clear, size: %v", lenConns)
 	}
 
-	server.Stop()
-	wg.Wait()
 }