Move health producer activation to petiole policy

balancer/balancer.go

@@ -487,5 +487,4 @@ type HealthCheckOptions struct {
 	DisableHealthCheckDialOpt bool
 	ServiceName               func() string
 	HealthCheckFunc           internal.HealthChecker
-	EnableHealthCheck         bool
 }

balancer/pickfirstleaf/pickfirstleaf.go

@@ -141,12 +141,6 @@ func (b *pickfirstBalancer) newSCData(addr resolver.Address) (*scData, error) {
 	if err != nil {
 		return nil, err
 	}
-	// Start the health check service if its configured.
-	if internal.EnableHealthCheckViaProducer != nil {
-		sd.closeHealthProducers = internal.EnableHealthCheckViaProducer.(func(balancer.HealthCheckOptions, balancer.SubConn) func())(b.healthCheckOpts, sc)
-	} else if b.healthCheckOpts.EnableHealthCheck {
-		b.logger.Errorf("Health check is requested but health check function is not set.")
-	}
 	sd.subConn = sc
 	return sd, nil
 }

health/producer.go

@@ -138,7 +138,7 @@ func (p *healthServiceProducer) startHealthCheckLocked() {
 	marker := &struct{}{}
 	p.currentAttemptMarker = marker
 	serviceName := p.opts.ServiceName()
-	if p.opts.DisableHealthCheckDialOpt || !p.opts.EnableHealthCheck || serviceName == "" {
+	if p.opts.DisableHealthCheckDialOpt || serviceName == "" {
 		p.healthState = balancer.SubConnState{ConnectivityState: connectivity.Ready}
 		return
 	}

internal/internal.go

@@ -33,7 +33,7 @@ var (
 	WithHealthCheckFunc any // func (HealthChecker) DialOption
 	// HealthCheckFunc is used to provide client-side LB channel health checking
 	HealthCheckFunc              HealthChecker
-	EnableHealthCheckViaProducer any // func(balancer.HealthCheckOptions, balancer.SubConn) func()
+	EnableHealthCheckViaProducer any // func(balancer.HealthCheckOptions, balancer.SubConn) balancer.SubConn
 	// BalancerUnregister is exported by package balancer to unregister a balancer.
 	BalancerUnregister func(name string)
 	// KeepaliveMinPingTime is the minimum ping interval.  This must be 10s by

test/healthcheck_test.go

@@ -69,13 +69,37 @@ func init() {
 	}
 }
 
+// wraps the clientconn to enable health checking in the created subconns.
+type wrappedCC struct {
+	b *healthCheckingPetiolePolicy
+	balancer.ClientConn
+}
+
+func (w *wrappedCC) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	sc, err := w.ClientConn.NewSubConn(addrs, opts)
+	if err != nil {
+		return sc, err
+	}
+	if internal.EnableHealthCheckViaProducer == nil {
+		return sc, nil
+	}
+	cleanup := internal.EnableHealthCheckViaProducer.(func(balancer.HealthCheckOptions, balancer.SubConn) func())(w.b.opts, sc)
+	w.b.cleanups = append(w.b.cleanups, cleanup)
+	return sc, nil
+}
+
 type healthCheckingPetiolePolicyBuilder struct{}
 
-func (b *healthCheckingPetiolePolicyBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	opts.HealthCheckOptions.EnableHealthCheck = true
-	return &healthCheckingPetiolePolicy{
-		balancer.Get(pickfirst.Name).Build(cc, opts),
+func (bb *healthCheckingPetiolePolicyBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	wcc := &wrappedCC{
+		ClientConn: cc,
+	}
+	b := &healthCheckingPetiolePolicy{
+		Balancer: balancer.Get(pickfirst.Name).Build(wcc, opts),
+		opts:     opts.HealthCheckOptions,
 	}
+	wcc.b = b
+	return b
 }
 
 func (b *healthCheckingPetiolePolicyBuilder) Name() string {
@@ -84,6 +108,15 @@ func (b *healthCheckingPetiolePolicyBuilder) Name() string {
 
 type healthCheckingPetiolePolicy struct {
 	balancer.Balancer
+	cleanups []func()
+	opts     balancer.HealthCheckOptions
+}
+
+func (p *healthCheckingPetiolePolicy) Close() {
+	for _, cl := range p.cleanups {
+		cl()
+	}
+	p.Balancer.Close()
 }
 
 func newTestHealthServer() *testHealthServer {