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

Code before:
	// connectivityState stores the most recent connectivity state delivered
	// to the LB policy. It is a pointer to enable distinguishing when a SubConn
	// transitions to a state and then returns to the same state. It is atomic
	// because it is read when a health listener is registered by the LB policy.
	connectivityState atomic.Pointer[connectivity.State]
}

// updateState is invoked by grpc to push a subConn state update to the
// underlying balancer.
func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolver.Address, err error) {
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Invalidate all producers on any state change.
		acbw.closeProducers()

		// Even though it is optional for balancers, gracefulswitch ensures
		// opts.StateListener is set, so this cannot ever be nil.
		// TODO: delete this comment when UpdateSubConnState is removed.
		scs := balancer.SubConnState{ConnectivityState: s, ConnectionError: err}
		if s == connectivity.Ready {
			setConnectedAddress(&scs, curAddr)
		}
		// Invalidate the health listener by changing the state.
		acbw.connectivityState.Store(&scs.ConnectivityState)
		// A race may occur if a health listener is registered soon after the
		// connectivity state is set but before the stateListener is called.
		// Two cases may arise:
		// 1. The new state is not READY: RegisterHealthListener has checks to
		//    ensure no updates are sent when the connectivity state is not
		//    READY.
		// 2. The new state is READY: This means that the old state wasn't Ready.
		//    The RegisterHealthListener API mentions that a health listener
		//    must not be registered when a SubConn is not ready to avoid such
		//    races. When this happens, the LB policy would get health updates
		//    on the old listener. When the LB policy registers a new listener
		//    on receiving the connectivity update, the health updates will be
		//    sent to the new health listener.
		acbw.stateListener(scs)
	})
}

func (acbw *acBalancerWrapper) String() string {
	return fmt.Sprintf("SubConn(id:%d)", acbw.ac.channelz.ID)
}

func (acbw *acBalancerWrapper) UpdateAddresses(addrs []resolver.Address) {
	acbw.ac.updateAddrs(addrs)
}


Review comment (suggestion):
I got confused for a long time as to why the registered listener was not being stored in the balancer wrapper and how we would send any updates to it. And then I realized that:
- we support a health listener to be registered *only* when the subchannel is Ready
- when a health listener is registered and the subchannel is Ready, we schedule a callback that checks current state, ensures that the subchannel is still Ready, and sends that same state to the listener.
- when the subchannel state changes, we don't send any updates on the health listener, because at this point, the LB policy gets a subchannel state change on the connectivity state listener and it needs to wait till the connectivity state becomes Ready to register another health listener
- Updates on the health listener happen only when some petiole policy registers a health producer (either client side health checking or OD).

I feel this behavior needs to be clarified in the balancer wrapper's RegisterHealthListener. 

Code after:

	// Access to healthData is protected by healthMu.
	healthMu sync.Mutex
	// healthData is stored as a pointer to detect when the health listener is
	// dropped or updated. This is required as closures can't be compared for
	// equality.
	healthData *healthData
}

// healthData holds data related to health state reporting.
type healthData struct {
	// connectivityState stores the most recent connectivity state delivered
	// to the LB policy. This is stored to avoid sending updates when the
	// SubConn has already exited connectivity state READY.
	connectivityState connectivity.State
}

func newHealthData(s connectivity.State) *healthData {
	return &healthData{connectivityState: s}
}

// updateState is invoked by grpc to push a subConn state update to the
// underlying balancer.
func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolver.Address, err error) {
	acbw.ccb.serializer.TrySchedule(func(ctx context.Context) {
		if ctx.Err() != nil || acbw.ccb.balancer == nil {
			return
		}
		// Invalidate all producers on any state change.
		acbw.closeProducers()

		// Even though it is optional for balancers, gracefulswitch ensures
		// opts.StateListener is set, so this cannot ever be nil.
		// TODO: delete this comment when UpdateSubConnState is removed.
		scs := balancer.SubConnState{ConnectivityState: s, ConnectionError: err}
		if s == connectivity.Ready {
			setConnectedAddress(&scs, curAddr)
		}
		// Invalidate the health listener by updating the healthData.
		acbw.healthMu.Lock()
		// A race may occur if a health listener is registered soon after the
		// connectivity state is set but before the stateListener is called.
		// Two cases may arise:
		// 1. The new state is not READY: RegisterHealthListener has checks to
		//    ensure no updates are sent when the connectivity state is not
		//    READY.
		// 2. The new state is READY: This means that the old state wasn't Ready.
		//    The RegisterHealthListener API mentions that a health listener
		//    must not be registered when a SubConn is not ready to avoid such
		//    races. When this happens, the LB policy would get health updates
		//    on the old listener. When the LB policy registers a new listener