Repository: grpc/grpc-go
File: xds/internal/balancer/clusterimpl/clusterimpl.go (Go)

Code before:
	// it.
	bb := balancer.Get(newConfig.ChildPolicy.Name)
	if bb == nil {
		return fmt.Errorf("child policy %q not registered", newConfig.ChildPolicy.Name)
	}

	if b.xdsClient == nil {
		c := xdsclient.FromResolverState(s.ResolverState)
		if c == nil {
			return balancer.ErrBadResolverState
		}
		b.xdsClient = c
	}

	// Update load reporting config. This needs to be done before updating the
	// child policy because we need the loadStore from the updated client to be
	// passed to the ccWrapper, so that the next picker from the child policy
	// will pick up the new loadStore.
	if err := b.updateLoadStore(newConfig); err != nil {
		return err
	}

	// Build config for the gracefulswitch balancer. It is safe to ignore JSON
	// marshaling errors here, since the config was already validated as part of
	// ParseConfig().
	cfg := []map[string]any{{newConfig.ChildPolicy.Name: newConfig.ChildPolicy.Config}}
	cfgJSON, _ := json.Marshal(cfg)
	parsedCfg, err := gracefulswitch.ParseConfig(cfgJSON)
	if err != nil {
		return err
	}

	b.config = newConfig
	// Addresses and sub-balancer config are sent to sub-balancer.
	err = b.child.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  s.ResolverState,
		BalancerConfig: parsedCfg,
	})

	b.childStateMu.Lock()
	b.inhibitPickerUpdates = false
	childState := b.childState
	b.childStateMu.Unlock()
	b.telemetryLabels = newConfig.TelemetryLabels
	b.configMu.Lock()
	dc := b.handleDropAndRequestCount(newConfig)
	b.configMu.Unlock()
	if dc != nil && childState.Picker != nil {
		b.ClientConn.UpdateState(balancer.State{
			ConnectivityState: childState.ConnectivityState,
			Picker:            b.newPicker(dc),

Review comment (suggestion):
I would do something like this here:

```golang
func (b *clusterImplBalancer) updateClientConnState(s balancer.ClientConnState) error {
	defer clientConnUpdateHook()

	if b.logger.V(2) {
		b.logger.Infof("Received configuration: %s", pretty.ToJSON(s.BalancerConfig))
	}
	newConfig, ok := s.BalancerConfig.(*LBConfig)
	if !ok {
		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
	}

	// Need to check for potential errors at the beginning of this function, so
	// that on errors, we reject the whole config, instead of applying part of
	// it.
	bb := balancer.Get(newConfig.ChildPolicy.Name)
	if bb == nil {
		return fmt.Errorf("child policy %q not registered", newConfig.ChildPolicy.Name)
	}

	if b.xdsClient == nil {
		c := xdsclient.FromResolverState(s.ResolverState)
		if c == nil {
			return balancer.ErrBadResolverState
		}
		b.xdsClient = c
	}

	// Update load reporting config. This needs to be done before updating the
	// child policy because we need the loadStore from the updated client to be
	// passed to the ccWrapper, so that the next picker from the child policy
	// will pick up the new loadStore.
	if err := b.updateLoadStore(newConfig); err != nil {
		return err
	}

	// Build config for the gracefulswitch balancer. It is safe to ignore JSON
	// marshaling errors here, since the config was already validated as part of
	// ParseConfig().
	cfg := []map[string]any{{newConfig.ChildPolicy.Name: newConfig.ChildPolicy.Config}}
	cfgJSON, _ := json.Marshal(cfg)
	parsedCfg, err := gracefulswitch.ParseConfig(cfgJSON)
	if err != nil {
		return err
	}

	// Addresses and sub-balancer config are sent to sub-balancer.
	err = b.child.UpdateClientConnState(balancer.ClientConnState{
		ResolverState:  s.ResolverState,
		BalancerConfig: parsedCfg,
	})

	b.mu.Lock()
	b.telemetryLabels = newConfig.TelemetryLabels
	if b.handleDropAndRequestCount(newConfig) && b.childState.Picker != nil {
		b.ClientConn.UpdateState(balancer.State{
			ConnectivityState: b.childState.ConnectivityState,
			Picker:            b.newPickerLocked(),
		})
	}
	b.inhibitPickerUpdates = false
	b.mu.Unlock()

	newPickerUpdated()
	return err
}

func (b *clusterImplBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
	b.mu.Lock()
	b.inhibitPickerUpdates = true
	b.mu.Unlock()
	return b.updateClientConnState(s)
}
```

Code after:
	}

	if stopOldLoadReport {
		if b.cancelLoadReport != nil {
			b.cancelLoadReport()
			b.cancelLoadReport = nil
			if !startNewLoadReport {
				// If a new LRS stream will be started later, no need to update
				// it to nil here.
				b.loadWrapper.UpdateLoadStore(nil)
			}
		}
	}
	if startNewLoadReport {
		var loadStore *load.Store
		if b.xdsClient != nil {
			loadStore, b.cancelLoadReport = b.xdsClient.ReportLoad(b.lrsServer)
		}
		b.loadWrapper.UpdateLoadStore(loadStore)
	}

	return nil
}

func (b *clusterImplBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
	defer clientConnUpdateHook()

	b.mu.Lock()
	b.inhibitPickerUpdates = true
	b.mu.Unlock()
	if b.logger.V(2) {
		b.logger.Infof("Received configuration: %s", pretty.ToJSON(s.BalancerConfig))
	}
	newConfig, ok := s.BalancerConfig.(*LBConfig)
	if !ok {
		return fmt.Errorf("unexpected balancer config with type: %T", s.BalancerConfig)
	}

	// Need to check for potential errors at the beginning of this function, so
	// that on errors, we reject the whole config, instead of applying part of
	// it.
	bb := balancer.Get(newConfig.ChildPolicy.Name)
	if bb == nil {
		return fmt.Errorf("child policy %q not registered", newConfig.ChildPolicy.Name)
	}

	if b.xdsClient == nil {
		c := xdsclient.FromResolverState(s.ResolverState)
		if c == nil {
			return balancer.ErrBadResolverState
		}