snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/balancergroup/balancergroup.go

@@ -0,0 +1,607 @@
+/*
+ * Copyright 2019 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Package balancergroup implements a utility struct to bind multiple balancers
+// into one balancer.
+package balancergroup
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal/balancer/gracefulswitch"
+	"google.golang.org/grpc/internal/cache"
+	"google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// subBalancerWrapper is used to keep the configurations that will be used to start
+// the underlying balancer. It can be called to start/stop the underlying
+// balancer.
+//
+// When the config changes, it will pass the update to the underlying balancer
+// if it exists.
+//
+// TODO: move to a separate file?
+type subBalancerWrapper struct {
+	// subBalancerWrapper is passed to the sub-balancer as a ClientConn
+	// wrapper, only to keep the state and picker.  When sub-balancer is
+	// restarted while in cache, the picker needs to be resent.
+	//
+	// It also contains the sub-balancer ID, so the parent balancer group can
+	// keep track of SubConn/pickers and the sub-balancers they belong to. Some
+	// of the actions are forwarded to the parent ClientConn with no change.
+	// Some are forward to balancer group with the sub-balancer ID.
+	balancer.ClientConn
+	id    string
+	group *BalancerGroup
+
+	mu    sync.Mutex
+	state balancer.State
+
+	// The static part of sub-balancer. Keeps balancerBuilders and addresses.
+	// To be used when restarting sub-balancer.
+	builder balancer.Builder
+	// Options to be passed to sub-balancer at the time of creation.
+	buildOpts balancer.BuildOptions
+	// ccState is a cache of the addresses/balancer config, so when the balancer
+	// is restarted after close, it will get the previous update. It's a pointer
+	// and is set to nil at init, so when the balancer is built for the first
+	// time (not a restart), it won't receive an empty update. Note that this
+	// isn't reset to nil when the underlying balancer is closed.
+	ccState *balancer.ClientConnState
+	// The dynamic part of sub-balancer. Only used when balancer group is
+	// started. Gets cleared when sub-balancer is closed.
+	balancer *gracefulswitch.Balancer
+}
+
+// UpdateState overrides balancer.ClientConn, to keep state and picker.
+func (sbc *subBalancerWrapper) UpdateState(state balancer.State) {
+	sbc.mu.Lock()
+	sbc.state = state
+	sbc.group.updateBalancerState(sbc.id, state)
+	sbc.mu.Unlock()
+}
+
+// NewSubConn overrides balancer.ClientConn, so balancer group can keep track of
+// the relation between subconns and sub-balancers.
+func (sbc *subBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	return sbc.group.newSubConn(sbc, addrs, opts)
+}
+
+func (sbc *subBalancerWrapper) updateBalancerStateWithCachedPicker() {
+	sbc.mu.Lock()
+	if sbc.state.Picker != nil {
+		sbc.group.updateBalancerState(sbc.id, sbc.state)
+	}
+	sbc.mu.Unlock()
+}
+
+func (sbc *subBalancerWrapper) startBalancer() {
+	if sbc.balancer == nil {
+		sbc.balancer = gracefulswitch.NewBalancer(sbc, sbc.buildOpts)
+	}
+	sbc.group.logger.Infof("Creating child policy of type %q for child %q", sbc.builder.Name(), sbc.id)
+	sbc.balancer.SwitchTo(sbc.builder)
+	if sbc.ccState != nil {
+		sbc.balancer.UpdateClientConnState(*sbc.ccState)
+	}
+}
+
+// exitIdle invokes the ExitIdle method on the sub-balancer, a gracefulswitch
+// balancer.
+func (sbc *subBalancerWrapper) exitIdle() {
+	b := sbc.balancer
+	if b == nil {
+		return
+	}
+	b.ExitIdle()
+}
+
+func (sbc *subBalancerWrapper) updateClientConnState(s balancer.ClientConnState) error {
+	sbc.ccState = &s
+	b := sbc.balancer
+	if b == nil {
+		// A sub-balancer is closed when it is removed from the group or the
+		// group is closed as a whole, and is not expected to receive updates
+		// after that. But when used with the priority LB policy a sub-balancer
+		// (and the whole balancer group) could be closed because it's the lower
+		// priority, but it can still get address updates.
+		return nil
+	}
+	return b.UpdateClientConnState(s)
+}
+
+func (sbc *subBalancerWrapper) resolverError(err error) {
+	b := sbc.balancer
+	if b == nil {
+		// A sub-balancer is closed when it is removed from the group or the
+		// group is closed as a whole, and is not expected to receive updates
+		// after that. But when used with the priority LB policy a sub-balancer
+		// (and the whole balancer group) could be closed because it's the lower
+		// priority, but it can still get address updates.
+		return
+	}
+	b.ResolverError(err)
+}
+
+func (sbc *subBalancerWrapper) stopBalancer() {
+	if sbc.balancer == nil {
+		return
+	}
+	sbc.balancer.Close()
+	sbc.balancer = nil
+}
+
+// BalancerGroup takes a list of balancers, each behind a gracefulswitch
+// balancer, and make them into one balancer.
+//
+// Note that this struct doesn't implement balancer.Balancer, because it's not
+// intended to be used directly as a balancer. It's expected to be used as a
+// sub-balancer manager by a high level balancer.
+//
+//	Updates from ClientConn are forwarded to sub-balancers
+//	- service config update
+//	- address update
+//	- subConn state change
+//	  - find the corresponding balancer and forward
+//
+//	Actions from sub-balances are forwarded to parent ClientConn
+//	- new/remove SubConn
+//	- picker update and health states change
+//	  - sub-pickers are sent to an aggregator provided by the parent, which
+//	    will group them into a group-picker. The aggregated connectivity state is
+//	    also handled by the aggregator.
+//	- resolveNow
+//
+// Sub-balancers are only built when the balancer group is started. If the
+// balancer group is closed, the sub-balancers are also closed. And it's
+// guaranteed that no updates will be sent to parent ClientConn from a closed
+// balancer group.
+type BalancerGroup struct {
+	cc        balancer.ClientConn
+	buildOpts balancer.BuildOptions
+	logger    *grpclog.PrefixLogger
+
+	// stateAggregator is where the state/picker updates will be sent to. It's
+	// provided by the parent balancer, to build a picker with all the
+	// sub-pickers.
+	stateAggregator BalancerStateAggregator
+
+	// outgoingMu guards all operations in the direction:
+	// ClientConn-->Sub-balancer. Including start, stop, resolver updates and
+	// SubConn state changes.
+	//
+	// The corresponding boolean outgoingStarted is used to stop further updates
+	// to sub-balancers after they are closed.
+	outgoingMu         sync.Mutex
+	outgoingClosed     bool
+	idToBalancerConfig map[string]*subBalancerWrapper
+	// Cache for sub-balancers when they are removed. This is `nil` if caching
+	// is disabled by passing `0` for Options.SubBalancerCloseTimeout`.
+	deletedBalancerCache *cache.TimeoutCache
+
+	// incomingMu is to make sure this balancer group doesn't send updates to cc
+	// after it's closed.
+	//
+	// We don't share the mutex to avoid deadlocks (e.g. a call to sub-balancer
+	// may call back to balancer group inline. It causes deadlock if they
+	// require the same mutex).
+	//
+	// We should never need to hold multiple locks at the same time in this
+	// struct. The case where two locks are held can only happen when the
+	// underlying balancer calls back into balancer group inline. So there's an
+	// implicit lock acquisition order that outgoingMu is locked before
+	// incomingMu.
+
+	// incomingMu guards all operations in the direction:
+	// Sub-balancer-->ClientConn. Including NewSubConn, RemoveSubConn. It also
+	// guards the map from SubConn to balancer ID, so updateSubConnState needs
+	// to hold it shortly to potentially delete from the map.
+	//
+	// UpdateState is called by the balancer state aggregator, and it will
+	// decide when and whether to call.
+	//
+	// The corresponding boolean incomingStarted is used to stop further updates
+	// from sub-balancers after they are closed.
+	incomingMu      sync.Mutex
+	incomingClosed  bool // This boolean only guards calls back to ClientConn.
+	scToSubBalancer map[balancer.SubConn]*subBalancerWrapper
+}
+
+// Options wraps the arguments to be passed to the BalancerGroup ctor.
+type Options struct {
+	// CC is a reference to the parent balancer.ClientConn.
+	CC balancer.ClientConn
+	// BuildOpts contains build options to be used when creating sub-balancers.
+	BuildOpts balancer.BuildOptions
+	// StateAggregator is an implementation of the BalancerStateAggregator
+	// interface to aggregate picker and connectivity states from sub-balancers.
+	StateAggregator BalancerStateAggregator
+	// Logger is a group specific prefix logger.
+	Logger *grpclog.PrefixLogger
+	// SubBalancerCloseTimeout is the amount of time deleted sub-balancers spend
+	// in the idle cache. A value of zero here disables caching of deleted
+	// sub-balancers.
+	SubBalancerCloseTimeout time.Duration
+}
+
+// New creates a new BalancerGroup. Note that the BalancerGroup
+// needs to be started to work.
+func New(opts Options) *BalancerGroup {
+	var bc *cache.TimeoutCache
+	if opts.SubBalancerCloseTimeout != time.Duration(0) {
+		bc = cache.NewTimeoutCache(opts.SubBalancerCloseTimeout)
+	}
+
+	return &BalancerGroup{
+		cc:              opts.CC,
+		buildOpts:       opts.BuildOpts,
+		stateAggregator: opts.StateAggregator,
+		logger:          opts.Logger,
+
+		deletedBalancerCache: bc,
+		idToBalancerConfig:   make(map[string]*subBalancerWrapper),
+		scToSubBalancer:      make(map[balancer.SubConn]*subBalancerWrapper),
+	}
+}
+
+// AddWithClientConn adds a balancer with the given id to the group. The
+// balancer is built with a balancer builder registered with balancerName. The
+// given ClientConn is passed to the newly built balancer instead of the
+// one passed to balancergroup.New().
+//
+// TODO: Get rid of the existing Add() API and replace it with this.
+func (bg *BalancerGroup) AddWithClientConn(id, balancerName string, cc balancer.ClientConn) error {
+	bg.logger.Infof("Adding child policy of type %q for child %q", balancerName, id)
+	builder := balancer.Get(balancerName)
+	if builder == nil {
+		return fmt.Errorf("balancergroup: unregistered balancer name %q", balancerName)
+	}
+
+	// Store data in static map, and then check to see if bg is started.
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return fmt.Errorf("balancergroup: already closed")
+	}
+	var sbc *subBalancerWrapper
+	// Skip searching the cache if disabled.
+	if bg.deletedBalancerCache != nil {
+		if old, ok := bg.deletedBalancerCache.Remove(id); ok {
+			if bg.logger.V(2) {
+				bg.logger.Infof("Removing and reusing child policy of type %q for child %q from the balancer cache", balancerName, id)
+				bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+			}
+
+			sbc, _ = old.(*subBalancerWrapper)
+			if sbc != nil && sbc.builder != builder {
+				// If the sub-balancer in cache was built with a different
+				// balancer builder, don't use it, cleanup this old-balancer,
+				// and behave as sub-balancer is not found in cache.
+				//
+				// NOTE that this will also drop the cached addresses for this
+				// sub-balancer, which seems to be reasonable.
+				sbc.stopBalancer()
+				// cleanupSubConns must be done before the new balancer starts,
+				// otherwise new SubConns created by the new balancer might be
+				// removed by mistake.
+				bg.cleanupSubConns(sbc)
+				sbc = nil
+			}
+		}
+	}
+	if sbc == nil {
+		sbc = &subBalancerWrapper{
+			ClientConn: cc,
+			id:         id,
+			group:      bg,
+			builder:    builder,
+			buildOpts:  bg.buildOpts,
+		}
+		sbc.startBalancer()
+	} else {
+		// When brining back a sub-balancer from cache, re-send the cached
+		// picker and state.
+		sbc.updateBalancerStateWithCachedPicker()
+	}
+	bg.idToBalancerConfig[id] = sbc
+	return nil
+}
+
+// Add adds a balancer built by builder to the group, with given id.
+func (bg *BalancerGroup) Add(id string, builder balancer.Builder) {
+	bg.AddWithClientConn(id, builder.Name(), bg.cc)
+}
+
+// Remove removes the balancer with id from the group.
+//
+// But doesn't close the balancer. The balancer is kept in a cache, and will be
+// closed after timeout. Cleanup work (closing sub-balancer and removing
+// subconns) will be done after timeout.
+func (bg *BalancerGroup) Remove(id string) {
+	bg.removeInternal(id, true)
+}
+
+// RemoveImmediately removes and closes the balancer with id from the group
+// immediately.
+func (bg *BalancerGroup) RemoveImmediately(id string) {
+	bg.removeInternal(id, false)
+}
+
+func (bg *BalancerGroup) removeInternal(id string, withCaching bool) {
+	bg.logger.Infof("Removing child policy for child %q", id)
+
+	bg.outgoingMu.Lock()
+	if bg.outgoingClosed {
+		bg.outgoingMu.Unlock()
+		return
+	}
+
+	sbToRemove, ok := bg.idToBalancerConfig[id]
+	if !ok {
+		bg.logger.Errorf("Child policy for child %q does not exist in the balancer group", id)
+		bg.outgoingMu.Unlock()
+		return
+	}
+
+	// Unconditionally remove the sub-balancer config from the map.
+	delete(bg.idToBalancerConfig, id)
+
+	if withCaching {
+		if bg.deletedBalancerCache != nil {
+			if bg.logger.V(2) {
+				bg.logger.Infof("Adding child policy for child %q to the balancer cache", id)
+			}
+			bg.deletedBalancerCache.Add(id, sbToRemove, func() {
+				if bg.logger.V(2) {
+					bg.logger.Infof("Removing child policy for child %q from the balancer cache after timeout", id)
+					bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+				}
+
+				// A sub-balancer evicted from the timeout cache needs to closed
+				// and its subConns need to removed, unconditionally. There is a
+				// possibility that a sub-balancer might be removed (thereby
+				// moving it to the cache) around the same time that the
+				// balancergroup is closed, and by the time we get here the
+				// balancergroup might be closed.  Check for `outgoingStarted ==
+				// true` at that point can lead to a leaked sub-balancer.
+				bg.outgoingMu.Lock()
+				sbToRemove.stopBalancer()
+				bg.outgoingMu.Unlock()
+				bg.cleanupSubConns(sbToRemove)
+			})
+			if bg.logger.V(2) {
+				bg.logger.Infof("Number of items remaining in the balancer cache: %d", bg.deletedBalancerCache.Len())
+			}
+			bg.outgoingMu.Unlock()
+			return
+		}
+
+		// Fall through to remove the sub-balancer with immediate effect if we are not caching.
+		if bg.logger.V(2) {
+			bg.logger.Infof("Child policy for child %q was requested to be cached before eventual removal. No such cache exists. Removing right away.", id)
+		}
+	}
+
+	// Remove the sub-balancer with immediate effect if we are not caching.
+	sbToRemove.stopBalancer()
+	bg.outgoingMu.Unlock()
+	bg.cleanupSubConns(sbToRemove)
+}
+
+// bg.remove(id) doesn't do cleanup for the sub-balancer. This function does
+// cleanup after the timeout.
+func (bg *BalancerGroup) cleanupSubConns(config *subBalancerWrapper) {
+	bg.incomingMu.Lock()
+	defer bg.incomingMu.Unlock()
+	// Remove SubConns. This is only done after the balancer is
+	// actually closed.
+	//
+	// NOTE: if NewSubConn is called by this (closed) balancer later, the
+	// SubConn will be leaked. This shouldn't happen if the balancer
+	// implementation is correct. To make sure this never happens, we need to
+	// add another layer (balancer manager) between balancer group and the
+	// sub-balancers.
+	for sc, b := range bg.scToSubBalancer {
+		if b == config {
+			delete(bg.scToSubBalancer, sc)
+		}
+	}
+}
+
+// Following are actions from the parent grpc.ClientConn, forward to sub-balancers.
+
+// updateSubConnState forwards the update to cb and updates scToSubBalancer if
+// needed.
+func (bg *BalancerGroup) updateSubConnState(sc balancer.SubConn, state balancer.SubConnState, cb func(balancer.SubConnState)) {
+	bg.incomingMu.Lock()
+	if bg.incomingClosed {
+		bg.incomingMu.Unlock()
+		return
+	}
+	if _, ok := bg.scToSubBalancer[sc]; !ok {
+		bg.incomingMu.Unlock()
+		return
+	}
+	if state.ConnectivityState == connectivity.Shutdown {
+		// Only delete sc from the map when state changed to Shutdown.
+		delete(bg.scToSubBalancer, sc)
+	}
+	bg.incomingMu.Unlock()
+
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	if cb != nil {
+		cb(state)
+	}
+}
+
+// UpdateSubConnState handles the state for the subconn. It finds the
+// corresponding balancer and forwards the update.
+func (bg *BalancerGroup) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
+	bg.logger.Errorf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, state)
+}
+
+// UpdateClientConnState handles ClientState (including balancer config and
+// addresses) from resolver. It finds the balancer and forwards the update.
+func (bg *BalancerGroup) UpdateClientConnState(id string, s balancer.ClientConnState) error {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return nil
+	}
+	if config, ok := bg.idToBalancerConfig[id]; ok {
+		return config.updateClientConnState(s)
+	}
+	return nil
+}
+
+// ResolverError forwards resolver errors to all sub-balancers.
+func (bg *BalancerGroup) ResolverError(err error) {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	for _, config := range bg.idToBalancerConfig {
+		config.resolverError(err)
+	}
+}
+
+// Following are actions from sub-balancers, forward to ClientConn.
+
+// newSubConn: forward to ClientConn, and also create a map from sc to balancer,
+// so state update will find the right balancer.
+//
+// One note about removing SubConn: only forward to ClientConn, but not delete
+// from map. Delete sc from the map only when state changes to Shutdown. Since
+// it's just forwarding the action, there's no need for a removeSubConn()
+// wrapper function.
+func (bg *BalancerGroup) newSubConn(sbw *subBalancerWrapper, addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	// NOTE: if balancer with id was already removed, this should also return
+	// error. But since we call balancer.stopBalancer when removing the balancer, this
+	// shouldn't happen.
+	bg.incomingMu.Lock()
+	if bg.incomingClosed {
+		bg.incomingMu.Unlock()
+		return nil, fmt.Errorf("balancergroup: NewSubConn is called after balancer group is closed")
+	}
+	var sc balancer.SubConn
+	oldListener := opts.StateListener
+	opts.StateListener = func(state balancer.SubConnState) { bg.updateSubConnState(sc, state, oldListener) }
+	sc, err := sbw.ClientConn.NewSubConn(addrs, opts)
+	if err != nil {
+		bg.incomingMu.Unlock()
+		return nil, err
+	}
+	bg.scToSubBalancer[sc] = sbw
+	bg.incomingMu.Unlock()
+	return sc, nil
+}
+
+// updateBalancerState: forward the new state to balancer state aggregator. The
+// aggregator will create an aggregated picker and an aggregated connectivity
+// state, then forward to ClientConn.
+func (bg *BalancerGroup) updateBalancerState(id string, state balancer.State) {
+	bg.logger.Infof("Balancer state update from child %v, new state: %+v", id, state)
+
+	// Send new state to the aggregator, without holding the incomingMu.
+	// incomingMu is to protect all calls to the parent ClientConn, this update
+	// doesn't necessary trigger a call to ClientConn, and should already be
+	// protected by aggregator's mutex if necessary.
+	if bg.stateAggregator != nil {
+		bg.stateAggregator.UpdateState(id, state)
+	}
+}
+
+// Close closes the balancer. It stops sub-balancers, and removes the subconns.
+// When a BalancerGroup is closed, it can not receive further address updates.
+func (bg *BalancerGroup) Close() {
+	bg.incomingMu.Lock()
+	bg.incomingClosed = true
+	// Also remove all SubConns.
+	for sc := range bg.scToSubBalancer {
+		sc.Shutdown()
+		delete(bg.scToSubBalancer, sc)
+	}
+	bg.incomingMu.Unlock()
+
+	bg.outgoingMu.Lock()
+	// Setting `outgoingClosed` ensures that no entries are added to
+	// `deletedBalancerCache` after this point.
+	bg.outgoingClosed = true
+	bg.outgoingMu.Unlock()
+
+	// Clear(true) runs clear function to close sub-balancers in cache. It
+	// must be called out of outgoing mutex.
+	if bg.deletedBalancerCache != nil {
+		bg.deletedBalancerCache.Clear(true)
+	}
+
+	bg.outgoingMu.Lock()
+	for id, config := range bg.idToBalancerConfig {
+		config.stopBalancer()
+		delete(bg.idToBalancerConfig, id)
+	}
+	bg.outgoingMu.Unlock()
+}
+
+// ExitIdle should be invoked when the parent LB policy's ExitIdle is invoked.
+// It will trigger this on all sub-balancers, or reconnect their subconns if
+// not supported.
+func (bg *BalancerGroup) ExitIdle() {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	for _, config := range bg.idToBalancerConfig {
+		config.exitIdle()
+	}
+}
+
+// ExitIdleOne instructs the sub-balancer `id` to exit IDLE state, if
+// appropriate and possible.
+func (bg *BalancerGroup) ExitIdleOne(id string) {
+	bg.outgoingMu.Lock()
+	defer bg.outgoingMu.Unlock()
+	if bg.outgoingClosed {
+		return
+	}
+	if config := bg.idToBalancerConfig[id]; config != nil {
+		config.exitIdle()
+	}
+}
+
+// ParseConfig parses a child config list and returns a LB config for the
+// gracefulswitch Balancer.
+//
+// cfg is expected to be a json.RawMessage containing a JSON array of LB policy
+// names + configs as the format of the "loadBalancingConfig" field in
+// ServiceConfig.  It returns a type that should be passed to
+// UpdateClientConnState in the BalancerConfig field.
+func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+	return gracefulswitch.ParseConfig(cfg)
+}

internal/balancergroup/balancergroup_test.go

@@ -0,0 +1,930 @@
+/*
+ * Copyright 2019 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package balancergroup_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/balancer/weightedtarget/weightedaggregator"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+var (
+	rrBuilder            = balancer.Get(roundrobin.Name)
+	testBalancerIDs      = []string{"b1", "b2", "b3"}
+	testBackendAddrs     []resolver.Address
+	testBackendEndpoints []resolver.Endpoint
+)
+
+const testBackendAddrsCount = 12
+
+func init() {
+	for i := 0; i < testBackendAddrsCount; i++ {
+		addr := resolver.Address{Addr: fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i)}
+		testBackendAddrs = append(testBackendAddrs, addr)
+		testBackendEndpoints = append(testBackendEndpoints, resolver.Endpoint{Addresses: []resolver.Address{addr}})
+	}
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// Create a new balancer group, add balancer and backends.
+// - b1, weight 2, backends [0,1]
+// - b2, weight 1, backends [2,3]
+// Start the balancer group and check behavior.
+//
+// Close the balancer group.
+func (s) TestBalancerGroup_start_close(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]balancer.SubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Stop()
+	bg.Close()
+	for i := 0; i < 4; i++ {
+		(<-cc.ShutdownSubConnCh).UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Shutdown})
+	}
+}
+
+// Test that balancer group start() doesn't deadlock if the balancer calls back
+// into balancer group inline when it gets an update.
+//
+// The potential deadlock can happen if we
+//   - hold a lock and send updates to balancer (e.g. update resolved addresses)
+//   - the balancer calls back (NewSubConn or update picker) in line
+//
+// The callback will try to hold the same lock again, which will cause a
+// deadlock.
+//
+// This test starts the balancer group with a test balancer, will update picker
+// whenever it gets an address update. It's expected that start() doesn't block
+// because of deadlock.
+func (s) TestBalancerGroup_start_close_deadlock(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_start_close_deadlock"
+	stub.Register(balancerName, stub.BalancerFuncs{})
+	builder := balancer.Get(balancerName)
+
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], builder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], builder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[2:4]}})
+}
+
+// initBalancerGroupForCachingTest creates a balancer group, and initialize it
+// to be ready for caching tests.
+//
+// Two rr balancers are added to bg, each with 2 ready subConns. A sub-balancer
+// is removed later, so the balancer group returned has one sub-balancer in its
+// own map, and one sub-balancer in cache.
+func initBalancerGroupForCachingTest(t *testing.T, idleCacheTimeout time.Duration) (*weightedaggregator.Aggregator, *balancergroup.BalancerGroup, *testutils.BalancerClientConn, map[string]*testutils.TestSubConn) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: idleCacheTimeout,
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]*testutils.TestSubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Remove(testBalancerIDs[1])
+	bg.Remove(testBalancerIDs[1])
+	// Don't wait for SubConns to be removed after close, because they are only
+	// removed after close timeout.
+	for i := 0; i < 10; i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			t.Fatalf("Got request to shut down subconn %v, want no shut down subconn (because subconns were still in cache)", sc)
+		default:
+		}
+		time.Sleep(time.Millisecond)
+	}
+	// Test roundrobin on the with only sub-balancer0.
+	p2 := <-cc.NewPickerCh
+	want = []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p2)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	return gator, bg, cc, m1
+}
+
+// Test that if a sub-balancer is removed, and re-added within close timeout,
+// the subConns won't be re-created.
+func (s) TestBalancerGroup_locality_caching(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Turn down subconn for addr2, shouldn't get picker update because
+	// sub-balancer1 was removed.
+	addrToSC[testBackendAddrs[2].Addr].UpdateState(balancer.SubConnState{
+		ConnectivityState: connectivity.TransientFailure,
+		ConnectionError:   errors.New("test error"),
+	})
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewPickerCh:
+			t.Fatalf("Got new picker, want no new picker (because the sub-balancer was removed)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+
+	// Re-add sub-balancer-1, because subconns were in cache, no new subconns
+	// should be created. But a new picker will still be generated, with subconn
+	// states update to date.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		// addr2 is down, b2 only has addr3 in READY state.
+		addrToSC[testBackendAddrs[3].Addr], addrToSC[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewSubConnAddrsCh:
+			t.Fatalf("Got new subconn, want no new subconn (because subconns were still in cache)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+}
+
+// Sub-balancers are put in cache when they are shut down. If balancer group is
+// closed within close timeout, all subconns should still be removed
+// immediately.
+func (s) TestBalancerGroup_locality_caching_close_group(t *testing.T) {
+	_, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	bg.Close()
+	// The balancer group is closed. The subconns should be shutdown immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[0].Addr]: 1,
+		addrToSC[testBackendAddrs[1].Addr]: 1,
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Sub-balancers in cache will be closed if not re-added within timeout, and
+// subConns will be shut down.
+func (s) TestBalancerGroup_locality_caching_not_read_within_timeout(t *testing.T) {
+	_, _, cc, addrToSC := initBalancerGroupForCachingTest(t, time.Second)
+
+	// The sub-balancer is not re-added within timeout. The subconns should be
+	// shut down.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-ctx.Done():
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Wrap the rr builder, so it behaves the same, but has a different name.
+type noopBalancerBuilderWrapper struct {
+	balancer.Builder
+}
+
+func init() {
+	balancer.Register(&noopBalancerBuilderWrapper{Builder: rrBuilder})
+}
+
+func (*noopBalancerBuilderWrapper) Name() string {
+	return "noopBalancerBuilderWrapper"
+}
+
+// After removing a sub-balancer, re-add with same ID, but different balancer
+// builder. Old subconns should be shut down, and new subconns should be created.
+func (s) TestBalancerGroup_locality_caching_read_with_different_builder(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Re-add sub-balancer-1, but with a different balancer builder. The
+	// sub-balancer was still in cache, but can't be reused. This should cause
+	// old sub-balancer's subconns to be shut down immediately, and new
+	// subconns to be created.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], &noopBalancerBuilderWrapper{rrBuilder})
+
+	// The cached sub-balancer should be closed, and the subconns should be
+	// shut down immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[4:6]}})
+
+	newSCTimeout := time.After(time.Millisecond * 500)
+	scToAdd := map[string]int{
+		testBackendAddrs[4].Addr: 1,
+		testBackendAddrs[5].Addr: 1,
+	}
+	for i := 0; i < len(scToAdd); i++ {
+		select {
+		case addr := <-cc.NewSubConnAddrsCh:
+			c := scToAdd[addr[0].Addr]
+			if c == 0 {
+				t.Fatalf("Got newSubConn for %v when there's %d new expected", addr, c)
+			}
+			scToAdd[addr[0].Addr] = c - 1
+			sc := <-cc.NewSubConnCh
+			addrToSC[addr[0].Addr] = sc
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		case <-newSCTimeout:
+			t.Fatalf("timeout waiting for subConns (from new sub-balancer) to be newed")
+		}
+	}
+
+	// Test roundrobin on the new picker.
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		addrToSC[testBackendAddrs[4].Addr], addrToSC[testBackendAddrs[5].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+}
+
+// After removing a sub-balancer, it will be kept in cache. Make sure that this
+// sub-balancer's Close is called when the balancer group is closed.
+func (s) TestBalancerGroup_CloseStopsBalancerInCache(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_check_close"
+	closed := make(chan struct{})
+	stub.Register(balancerName, stub.BalancerFuncs{Close: func(_ *stub.BalancerData) {
+		close(closed)
+	}})
+	builder := balancer.Get(balancerName)
+
+	gator, bg, _, _ := initBalancerGroupForCachingTest(t, time.Second)
+
+	// Add balancer, and remove
+	gator.Add(testBalancerIDs[2], 1)
+	bg.Add(testBalancerIDs[2], builder)
+	gator.Remove(testBalancerIDs[2])
+	bg.Remove(testBalancerIDs[2])
+
+	// Immediately close balancergroup, before the cache timeout.
+	bg.Close()
+
+	// Make sure the removed child balancer is closed eventually.
+	select {
+	case <-closed:
+	case <-time.After(time.Second * 2):
+		t.Fatalf("timeout waiting for the child balancer in cache to be closed")
+	}
+}
+
+// TestBalancerGroupBuildOptions verifies that the balancer.BuildOptions passed
+// to the balancergroup at creation time is passed to child policies.
+func (s) TestBalancerGroupBuildOptions(t *testing.T) {
+	const (
+		balancerName = "stubBalancer-TestBalancerGroupBuildOptions"
+		userAgent    = "ua"
+	)
+
+	// Setup the stub balancer such that we can read the build options passed to
+	// it in the UpdateClientConnState method.
+	bOpts := balancer.BuildOptions{
+		DialCreds:       insecure.NewCredentials(),
+		ChannelzParent:  channelz.RegisterChannel(nil, "test channel"),
+		CustomUserAgent: userAgent,
+	}
+	stub.Register(balancerName, stub.BalancerFuncs{
+		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
+			if bd.BuildOptions.DialCreds != bOpts.DialCreds || bd.BuildOptions.ChannelzParent != bOpts.ChannelzParent || bd.BuildOptions.CustomUserAgent != bOpts.CustomUserAgent {
+				return fmt.Errorf("buildOptions in child balancer: %v, want %v", bd, bOpts)
+			}
+			return nil
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       bOpts,
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	// Add the stub balancer build above as a child policy.
+	balancerBuilder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], balancerBuilder)
+
+	// Send an empty clientConn state change. This should trigger the
+	// verification of the buildOptions being passed to the child policy.
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func (s) TestBalancerGroup_UpdateClientConnState_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	clientConnStateCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		UpdateClientConnState: func(_ *stub.BalancerData, _ balancer.ClientConnState) error {
+			clientConnStateCh <- struct{}{}
+			return nil
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
+		t.Fatalf("Expected nil error, got %v", err)
+	}
+
+	select {
+	case <-clientConnStateCh:
+		t.Fatalf("UpdateClientConnState was called after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestBalancerGroup_ResolverError_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	resolveErrorCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ResolverError: func(_ *stub.BalancerData, _ error) {
+			resolveErrorCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+
+	bg.ResolverError(errors.New("test error"))
+
+	select {
+	case <-resolveErrorCh:
+		t.Fatalf("ResolverError was called on sub-balancer after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestBalancerExitIdleOne(t *testing.T) {
+	const balancerName = "stub-balancer-test-balancergroup-exit-idle-one"
+	exitIdleCh := make(chan struct{}, 1)
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(*stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	// Add the stub balancer build above as a child policy.
+	builder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], builder)
+
+	// Call ExitIdleOne on the child policy.
+	bg.ExitIdleOne(testBalancerIDs[0])
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("Timeout when waiting for ExitIdle to be invoked on child policy")
+	case <-exitIdleCh:
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{})
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			close(exitIdleCh)
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdleOne(testBalancerIDs[0])
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on sub-balancer after BalancerGroup was closed")
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_NonExistentID(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.ExitIdleOne("non-existent-id")
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on wrong sub-balancer ID")
+	}
+}
+
+// TestBalancerGracefulSwitch tests the graceful switch functionality for a
+// child of the balancer group. At first, the child is configured as a round
+// robin load balancer, and thus should behave accordingly. The test then
+// gracefully switches this child to a custom type which only creates a SubConn
+// for the second passed in address and also only picks that created SubConn.
+// The new aggregated picker should reflect this change for the child.
+func (s) TestBalancerGracefulSwitch(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: gator,
+		Logger:          nil,
+	})
+	gator.Add(testBalancerIDs[0], 1)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+
+	defer bg.Close()
+
+	m1 := make(map[string]balancer.SubConn)
+	scs := make(map[balancer.SubConn]bool)
+	for i := 0; i < 2; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		scs[sc] = true
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatal(err)
+	}
+
+	// The balancer type for testBalancersIDs[0] is currently Round Robin. Now,
+	// change it to a balancer that has separate behavior logically (creating
+	// SubConn for second address in address list and always picking that
+	// SubConn), and see if the downstream behavior reflects that change.
+	childPolicyName := t.Name()
+	stub.Register(childPolicyName, stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			ccs.ResolverState.Endpoints = ccs.ResolverState.Endpoints[1:]
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+	})
+	cfgJSON := json.RawMessage(fmt.Sprintf(`[{%q: {}}]`, t.Name()))
+	lbCfg, err := balancergroup.ParseConfig(cfgJSON)
+	if err != nil {
+		t.Fatalf("ParseConfig(%s) failed: %v", string(cfgJSON), err)
+	}
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: testBackendEndpoints[2:4]},
+		BalancerConfig: lbCfg,
+	}); err != nil {
+		t.Fatalf("error updating ClientConn state: %v", err)
+	}
+
+	addrs := <-cc.NewSubConnAddrsCh
+	if addrs[0].Addr != testBackendAddrs[3].Addr {
+		// Verifies forwarded to new created balancer, as the wrapped pick first
+		// balancer will delete first address.
+		t.Fatalf("newSubConn called with wrong address, want: %v, got : %v", testBackendAddrs[3].Addr, addrs[0].Addr)
+	}
+	sc := <-cc.NewSubConnCh
+
+	// Update the pick first balancers SubConn as CONNECTING. This will cause
+	// the pick first balancer to UpdateState() with CONNECTING, which shouldn't send
+	// a Picker update back, as the Graceful Switch process is not complete.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer cancel()
+	select {
+	case <-cc.NewPickerCh:
+		t.Fatalf("No new picker should have been sent due to the Graceful Switch process not completing")
+	case <-ctx.Done():
+	}
+
+	// Update the pick first balancers SubConn as READY. This will cause
+	// the pick first balancer to UpdateState() with READY, which should send a
+	// Picker update back, as the Graceful Switch process is complete. This
+	// Picker should always pick the pick first's created SubConn which
+	// corresponds to address 3.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	p2 := <-cc.NewPickerCh
+	pr, err := p2.Pick(balancer.PickInfo{})
+	if err != nil {
+		t.Fatalf("error picking: %v", err)
+	}
+	if pr.SubConn != sc {
+		t.Fatalf("picker.Pick(), want %v, got %v", sc, pr.SubConn)
+	}
+
+	// The Graceful Switch process completing for the child should cause the
+	// SubConns for the balancer being gracefully switched from to get deleted.
+	ctx, cancel = context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for i := 0; i < 2; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("error waiting for Shutdown()")
+		case sc := <-cc.ShutdownSubConnCh:
+			// The SubConn shut down should have been one of the two created
+			// SubConns, and both should be deleted.
+			if ok := scs[sc]; ok {
+				delete(scs, sc)
+				continue
+			} else {
+				t.Fatalf("Shutdown called for wrong SubConn %v, want in %v", sc, scs)
+			}
+		}
+	}
+}
+
+func (s) TestBalancerExitIdle_All(t *testing.T) {
+	balancer1 := t.Name() + "-1"
+	balancer2 := t.Name() + "-2"
+
+	testID1, testID2 := testBalancerIDs[0], testBalancerIDs[1]
+
+	exitIdleCh1, exitIdleCh2 := make(chan struct{}, 1), make(chan struct{}, 1)
+
+	stub.Register(balancer1, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh1 <- struct{}{}
+		},
+	})
+
+	stub.Register(balancer2, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh2 <- struct{}{}
+		},
+	})
+
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	bg.Add(testID1, balancer.Get(balancer1))
+	bg.Add(testID2, balancer.Get(balancer2))
+
+	bg.ExitIdle()
+
+	errCh := make(chan error, 2)
+
+	go func() {
+		select {
+		case <-exitIdleCh1:
+			errCh <- nil
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer1")
+		}
+	}()
+
+	go func() {
+		select {
+		case <-exitIdleCh2:
+			errCh <- nil
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer2")
+		}
+	}()
+
+	for i := 0; i < 2; i++ {
+		if err := <-errCh; err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdle_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdle()
+
+	select {
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdle was called on sub-balancer even after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+// Tests verifies that the RemoveImmediately method of balancergroup closes the
+// child balancer immediately, while the Remove method does not close the child
+// balancer immediately, when the cache is enabled.
+func (s) TestBalancerGroup_RemoveImmediately(t *testing.T) {
+	// Channels to track child balancer creation and closure.
+	childLBCreated := make(chan string, 1)
+	childLBClosed := make(chan string, 1)
+
+	childLBName1 := strings.ToLower(t.Name()) + "-child-1"
+	t.Logf("Registering a child balancer with name %q", childLBName1)
+	stub.Register(childLBName1, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName1
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName1
+		},
+	})
+
+	childLBName2 := strings.ToLower(t.Name()) + "-child-2"
+	t.Logf("Registering a child balancer with name %q", childLBName2)
+	stub.Register(childLBName2, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName2
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName2
+		},
+	})
+
+	t.Logf("Creating a balancergroup with cache enabled")
+	tcc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      tcc,
+		BuildOpts:               balancer.BuildOptions{},
+		SubBalancerCloseTimeout: defaultTestTimeout,
+	})
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-1")
+	bg.AddWithClientConn("child-1", childLBName1, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-2")
+	bg.AddWithClientConn("child-2", childLBName2, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with immediate effect", "child-1")
+	bg.RemoveImmediately("child-1")
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with caching", "child-2")
+	bg.Remove("child-2")
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-childLBClosed:
+		t.Fatalf("Child LB closed when expected to be cached")
+	case <-sCtx.Done():
+	}
+
+	t.Logf("Closing the balancergroup, which should close the second child balancer")
+	bg.Close()
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+}

internal/balancergroup/balancerstateaggregator.go

@@ -0,0 +1,37 @@
+/*
+ *
+ * Copyright 2020 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package balancergroup
+
+import (
+	"google.golang.org/grpc/balancer"
+)
+
+// BalancerStateAggregator aggregates sub-picker and connectivity states into a
+// state.
+//
+// It takes care of merging sub-picker into one picker. The picking config is
+// passed directly from the parent to the aggregator implementation (instead
+// via balancer group).
+type BalancerStateAggregator interface {
+	// UpdateState updates the state of the id.
+	//
+	// It's up to the implementation whether this will trigger an update to the
+	// parent ClientConn.
+	UpdateState(id string, state balancer.State)
+}