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

Code before:
		// producer will not be double-closed.
		pData.refs--
		if pData.refs == 0 {
			defer pData.close() // Run outside the acbw mutex
			delete(acbw.producers, pb)
		}
		acbw.producersMu.Unlock()
	}
	return pData.producer, grpcsync.OnceFunc(unref)
}

func (acbw *acBalancerWrapper) closeProducers() {
	acbw.producersMu.Lock()
	defer acbw.producersMu.Unlock()
	for pb, pData := range acbw.producers {
		pData.refs = 0
		pData.close()
		delete(acbw.producers, pb)
	}
}

// healthProducerRegisterFn is a type alias for the health producer's function
// for registering listeners.
type healthProducerRegisterFn = func(context.Context, balancer.SubConn, string, func(balancer.SubConnState)) func()

// healthServiceOpts returns the options for client side health checking.
// It returns a nil registerHealthListenerFn if client side health checks are
// disabled.
// 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.
func (acbw *acBalancerWrapper) healthServiceOpts() (string, healthProducerRegisterFn) {
	if acbw.ccb.cc.dopts.disableHealthCheck {
		return "", nil
	}
	regHealthLisFn := internal.RegisterClientHealthCheckListener
	if regHealthLisFn == nil {
		// The health package is not imported.
		return "", nil
	}
	cfg := acbw.ac.cc.healthCheckConfig()
	if cfg == nil {
		return "", nil
	}
	return cfg.ServiceName, regHealthLisFn.(healthProducerRegisterFn)
}

// RegisterHealthListener accepts a health listener from the LB policy. It sends
// updates to the health listener as long as the SubConn's connectivity state

Review comment (suggestion):
Nit: instead of saying that it returns options for client side health checking, you might as well say that it returns the service name to be used for health checking and the function to register a health producer.

Code after:

	// Return a cleanup function wrapped in a OnceFunc to remove this reference
	// and delete the refCountedProducer from the map if the total reference
	// count goes to zero.
	unref := func() {
		acbw.producersMu.Lock()
		// If closeProducers has already closed this producer instance, refs is
		// set to 0, so the check after decrementing will never pass, and the
		// producer will not be double-closed.
		pData.refs--
		if pData.refs == 0 {
			defer pData.close() // Run outside the acbw mutex
			delete(acbw.producers, pb)
		}
		acbw.producersMu.Unlock()
	}
	return pData.producer, grpcsync.OnceFunc(unref)
}

func (acbw *acBalancerWrapper) closeProducers() {
	acbw.producersMu.Lock()
	defer acbw.producersMu.Unlock()
	for pb, pData := range acbw.producers {
		pData.refs = 0
		pData.close()
		delete(acbw.producers, pb)
	}
}

// healthProducerRegisterFn is a type alias for the health producer's function
// for registering listeners.
type healthProducerRegisterFn = func(context.Context, balancer.SubConn, string, func(balancer.SubConnState)) func()

// healthListenerRegFn returns a function to register a listener for health
// updates. If client side health checks are disabled, the registered listener
// will get a single READY (raw connectivity state) update.
//
// 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.
func (acbw *acBalancerWrapper) healthListenerRegFn() func(context.Context, func(balancer.SubConnState)) func() {
	if acbw.ccb.cc.dopts.disableHealthCheck {
		return noOpRegisterHealthListenerFn
	}
	regHealthLisFn := internal.RegisterClientHealthCheckListener
	if regHealthLisFn == nil {
		// The health package is not imported.
		return noOpRegisterHealthListenerFn
	}