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

Code before:
	defer es.childMu.Unlock()

	es.inhibitChildUpdates.Store(true)
	defer func() {
		es.inhibitChildUpdates.Store(false)
		es.updateState()
	}()
	var ret error

	children := es.children.Load()
	newChildren := resolver.NewEndpointMap()

	// Update/Create new children.
	for _, endpoint := range state.ResolverState.Endpoints {
		if _, ok := newChildren.Get(endpoint); ok {
			// Endpoint child was already created, continue to avoid duplicate
			// update.
			continue
		}
		var bal *balancerWrapper
		if child, ok := children.Get(endpoint); ok {
			bal = child.(*balancerWrapper)
			// Endpoint attributes may have changes, update the stored endpoint.
			es.mu.Lock()
			bal.childState.Endpoint = endpoint
			es.mu.Unlock()
		} else {
			bal = &balancerWrapper{
				childState: ChildState{Endpoint: endpoint},
				ClientConn: es.cc,
				es:         es,
			}
			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.

Review comment (suggestion):
Note that these updates are racy with respect to calls back into `UpdateState` and what is observed from `ChildStatesFromPicker`.  I'm not sure this is significant, since we will get an eventually-consistent `UpdateState` call after this, and also asynchronous calls out of the child should be pretty rare, since notifications into the child are always synchronous.  The `inhibit` atomic would also cause us to skip these updates most of the time, but if the call to `UpdateState` checks the flag before `UpdateCCS` sets it, then it is possible.

I'm also a little more worried about lock ordering.  But `es.childMu`->`es.mu` is already happening by virtue of calls into the child holding `childMu` and those resulting in synchronous calls into `cc` that takes `es.mu`.  So maybe the mutexes could use some documenting?

- `childMu` must be held on any call into any child.  (to prevent simultaneous calls into any single child)
- `mu` must not be held on any call into a child.  (since synchronous calls back from the child may require taking `mu`)
- `childMu` may not be taken if `mu` is held.  (the ordering implied by the above)


Code after:
// for endpoints that are no longer present. It also updates all the children,
// and sends a single synchronous update of the childrens' aggregated state at
// the end of the UpdateClientConnState operation. If any endpoint has no
// addresses it will ignore that endpoint. Otherwise, returns first error found
// from a child, but fully processes the new update.
func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState) error {
	es.childMu.Lock()
	defer es.childMu.Unlock()

	es.inhibitChildUpdates.Store(true)
	defer func() {
		es.inhibitChildUpdates.Store(false)
		es.updateState()
	}()
	var ret error

	children := es.children.Load()
	newChildren := resolver.NewEndpointMap()

	// Update/Create new children.
	for _, endpoint := range state.ResolverState.Endpoints {
		if _, ok := newChildren.Get(endpoint); ok {
			// Endpoint child was already created, continue to avoid duplicate
			// update.
			continue
		}
		var childBalancer *balancerWrapper
		if val, ok := children.Get(endpoint); ok {
			childBalancer = val.(*balancerWrapper)
			// Endpoint attributes may have changed, update the stored endpoint.
			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