Repository: grpc/grpc-go
File: balancer/endpointsharding/endpointsharding.go (Go)

Code before:
			}
			bal.childState.Balancer = bal
			bal.Balancer = gracefulswitch.NewBalancer(bal, es.bOpts)
		}
		newChildren.Set(endpoint, bal)
		if err := bal.UpdateClientConnState(balancer.ClientConnState{
			BalancerConfig: state.BalancerConfig,
			ResolverState: resolver.State{
				Endpoints:  []resolver.Endpoint{endpoint},
				Attributes: state.ResolverState.Attributes,
			},
		}); err != nil && ret == nil {
			// Return first error found, and always commit full processing of
			// updating children. If desired to process more specific errors
			// across all endpoints, caller should make these specific
			// validations, this is a current limitation for simplicity sake.
			ret = err
		}
	}
	// Delete old children that are no longer present.
	for _, e := range children.Keys() {
		child, _ := children.Get(e)
		bal := child.(*balancerWrapper)
		if _, ok := newChildren.Get(e); !ok {
			bal.Close()
			bal.isClosed = true
		}
	}
	es.children.Store(newChildren)
	if newChildren.Len() == 0 {
		return balancer.ErrBadResolverState
	}
	return ret
}

// ResolverError forwards the resolver error to all of the endpointSharding's
// children and sends a single synchronous update of the childStates at the end
// of the ResolverError operation.
func (es *endpointSharding) ResolverError(err error) {
	es.childMu.Lock()
	defer es.childMu.Unlock()
	es.inhibitChildUpdates.Store(true)
	defer func() {
		es.inhibitChildUpdates.Store(false)
		es.updateState()
	}()
	children := es.children.Load()
	for _, child := range children.Values() {
		bal := child.(balancer.Balancer)
		bal.ResolverError(err)
	}

Review comment (refactor):
Can we move this into `balancerWrapper`'s `Close` method?

Also since calls into a `balancer.Balancer` must not happen concurrently, I don't think it requires a lock.

Code after:
			es.mu.Lock()
			childBalancer.childState.Endpoint = endpoint
			es.mu.Unlock()
		} else {
			childBalancer = &balancerWrapper{
				childState: ChildState{Endpoint: endpoint},
				ClientConn: es.cc,
				es:         es,
			}
			childBalancer.childState.Balancer = childBalancer
			childBalancer.child = gracefulswitch.NewBalancer(childBalancer, es.bOpts)
		}
		newChildren.Set(endpoint, childBalancer)
		if err := childBalancer.updateClientConnStateLocked(balancer.ClientConnState{
			BalancerConfig: state.BalancerConfig,
			ResolverState: resolver.State{
				Endpoints:  []resolver.Endpoint{endpoint},
				Attributes: state.ResolverState.Attributes,
			},
		}); err != nil && ret == nil {
			// Return first error found, and always commit full processing of
			// updating children. If desired to process more specific errors
			// across all endpoints, caller should make these specific
			// validations, this is a current limitation for simplicity sake.
			ret = err
		}
	}
	// Delete old children that are no longer present.
	for _, e := range children.Keys() {
		child, _ := children.Get(e)
		if _, ok := newChildren.Get(e); !ok {
			child.(*balancerWrapper).closeLocked()
		}
	}
	es.children.Store(newChildren)
	if newChildren.Len() == 0 {
		return balancer.ErrBadResolverState
	}
	return ret
}

// ResolverError forwards the resolver error to all of the endpointSharding's
// children and sends a single synchronous update of the childStates at the end
// of the ResolverError operation.
func (es *endpointSharding) ResolverError(err error) {
	es.childMu.Lock()
	defer es.childMu.Unlock()
	es.inhibitChildUpdates.Store(true)
	defer func() {
		es.inhibitChildUpdates.Store(false)
		es.updateState()