Repository: grpc/grpc-go
File: balancer_wrapper.go (Go)

Code before:
	}
	// Replace the health data to stop sending updates to any previously
	// registered health listeners.
	hd := newHealthData(connectivity.Ready)
	acbw.healthData = hd
	if listener == nil {
		return
	}

	// Client side health checking is enabled when all the following
	// conditions are satisfied:
	// 1. Health checking is not disabled using the dial option.
	// 2. The health package is imported.
	// 3. The health check config is present in the service config.
	healthCheckEnabled := !acbw.ccb.cc.dopts.disableHealthCheck
	regHealthLisFn := internal.RegisterClientHealthCheckListener
	if regHealthLisFn == nil {
		// The health package is not imported.
		healthCheckEnabled = false
	}
	var cfg *healthCheckConfig
	if healthCheckEnabled {
		// Avoid acquiring cc.mu unless necessary.
		cfg = acbw.ac.cc.healthCheckConfig()
		healthCheckEnabled = cfg != nil
	}

	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		if acbw.healthData != hd {
			return
		}
		if !healthCheckEnabled {
			listener(balancer.SubConnState{ConnectivityState: connectivity.Ready})
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return
				}
				acbw.healthMu.Lock()
				defer acbw.healthMu.Unlock()
				if acbw.healthData != hd {

Review comment (refactor):
What do you feel about making a helper method on `acBalancerWrapper` that returns `(string, bool)` to indicate whether health checking is enabled or not (and the first return value will contain the service name when the second value is `true`).

The reason why a helper would be nicer is because we wont have to successively rewrite a local variable as is done currently, and instead can return `false` whenever one of the conditions is not met.

Code after:
	}
	cfg := acbw.ac.cc.healthCheckConfig()
	if cfg == nil {
		return noOpRegisterHealthListenerFn
	}
	return func(ctx context.Context, listener func(balancer.SubConnState)) func() {
		return regHealthLisFn.(healthProducerRegisterFn)(ctx, acbw, cfg.ServiceName, listener)
	}
}

// RegisterHealthListener accepts a health listener from the LB policy. It sends
// updates to the health listener as long as the SubConn's connectivity state
// doesn't change and a new health listener is not registered. To invalidate
// the currently registered health listener, acbw updates the healthData. If a
// nil listener is registered, the active health listener is dropped.
func (acbw *acBalancerWrapper) RegisterHealthListener(listener func(balancer.SubConnState)) {
	acbw.healthMu.Lock()
	defer acbw.healthMu.Unlock()
	acbw.healthData.closeHealthProducer()
	// listeners should not be registered when the connectivity state
	// isn't Ready. This may happen when the balancer registers a listener
	// after the connectivityState is updated, but before it is notified
	// of the update.
	if acbw.healthData.connectivityState != connectivity.Ready {
		return
	}
	// Replace the health data to stop sending updates to any previously
	// registered health listeners.
	hd := newHealthData(connectivity.Ready)
	acbw.healthData = hd
	if listener == nil {
		return
	}

	registerFn := acbw.healthListenerRegFn()
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updates if a new listener is registered.
		acbw.healthMu.Lock()
		defer acbw.healthMu.Unlock()
		if acbw.healthData != hd {
			return
		}
		// Serialize the health updates from the health producer with
		// other calls into the LB policy.
		listenerWrapper := func(scs balancer.SubConnState) {
			acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
				if ctx.Err() != nil || acbw.ccb.balancer == nil {
					return