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

Code before:
		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)
	}
}

// RegisterHealthListener stores 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 a
// the currently registered health listener, acbw updates the connectivity state
// stored by the object. 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()
	// 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 := &healthData{connectivityState: connectivity.Ready}
	acbw.healthData = hd
	if listener == nil {
		return
	}

	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Don't send updated if a new listener is registered.
		acbw.healthMu.Lock()
		curHD := acbw.healthData

Review comment (suggestion):
I feel that this comment needs some updating.

> // RegisterHealthListener stores a health listener from the LB policy. 

The health listener provided by the LB policy is not really stored anywhere.

> // 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. 

There is only one health update ever sent to the listener.

> // To invalidate a the currently registered health listener, acbw updates 
> // the connectivity state stored by the object. 

This is only true if the previously registered health listener was not yet invoked.

> // If a nil listener is registered, the active health listener is dropped.

Same here.

I think we need to make it clear in this comment that the implementation only over invokes the registered health listener once, and that happens only if the following conditions are met:
- the connectivity state is `Ready` when the health listener is registered
- the connectivity state did not change from the time the health listener was registered to the time the acbw got around to invoke it
- no other health listener was registered after this one was registered, but before acbw got around to invoke it.

And also that this health listener is only ever invoked with state `Ready`.

Code after:
		// 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)
	}
}

// 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()
	// 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
	}

	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()