roundrobin: Delegate subchannel creation to pickfirst

balancer/endpointsharding/endpointsharding_test.go

@@ -16,7 +16,7 @@
  *
  */
 
-package endpointsharding
+package endpointsharding_test
 
 import (
 	"context"
@@ -28,6 +28,7 @@ import (
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/endpointsharding"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/internal"
@@ -55,7 +56,7 @@ var logger = grpclog.Component("endpoint-sharding-test")
 
 func init() {
 	var err error
-	gracefulSwitchPickFirst, err = ParseConfig(json.RawMessage(PickFirstConfig))
+	gracefulSwitchPickFirst, err = endpointsharding.ParseConfig(json.RawMessage(endpointsharding.PickFirstConfig))
 	if err != nil {
 		logger.Fatal(err)
 	}
@@ -75,7 +76,7 @@ func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptio
 		ClientConn: cc,
 		bOpts:      opts,
 	}
-	fp.Balancer = NewBalancer(fp, opts)
+	fp.Balancer = endpointsharding.NewBalancer(fp, opts)
 	return fp
 }
 
@@ -105,7 +106,7 @@ func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) err
 }
 
 func (fp *fakePetiole) UpdateState(state balancer.State) {
-	childStates := ChildStatesFromPicker(state.Picker)
+	childStates := endpointsharding.ChildStatesFromPicker(state.Picker)
 	// Both child states should be present in the child picker. States and
 	// picker change over the lifecycle of test, but there should always be two.
 	if len(childStates) != 2 {

balancer/roundrobin/roundrobin.go

@@ -22,60 +22,61 @@
 package roundrobin
 
 import (
-	rand "math/rand/v2"
-	"sync/atomic"
+	"encoding/json"
+	"fmt"
 
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/grpclog"
+	internalgrpclog "google.golang.org/grpc/internal/grpclog"
+	"google.golang.org/grpc/serviceconfig"
 )
 
 // Name is the name of round_robin balancer.
 const Name = "round_robin"
 
-var logger = grpclog.Component("roundrobin")
-
-// newBuilder creates a new roundrobin balancer builder.
-func newBuilder() balancer.Builder {
-	return base.NewBalancerBuilder(Name, &rrPickerBuilder{}, base.Config{HealthCheck: true})
-}
+var (
+	logger = grpclog.Component("roundrobin")
+	// endpointSharding which specifies pick first children.
+	endpointShardingLBConfig serviceconfig.LoadBalancingConfig
+)
 
 func init() {
-	balancer.Register(newBuilder())
+	var err error
+	endpointShardingLBConfig, err = endpointsharding.ParseConfig(json.RawMessage(endpointsharding.PickFirstConfig))
+	if err != nil {
+		logger.Fatal(err)
+	}
+	balancer.Register(builder{})
 }
 
-type rrPickerBuilder struct{}
+type builder struct{}
 
-func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
-	logger.Infof("roundrobinPicker: Build called with info: %v", info)
-	if len(info.ReadySCs) == 0 {
-		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
-	}
-	scs := make([]balancer.SubConn, 0, len(info.ReadySCs))
-	for sc := range info.ReadySCs {
-		scs = append(scs, sc)
-	}
-	return &rrPicker{
-		subConns: scs,
-		// Start at a random index, as the same RR balancer rebuilds a new
-		// picker when SubConn states change, and we don't want to apply excess
-		// load to the first server in the list.
-		next: uint32(rand.IntN(len(scs))),
-	}
+func (bb builder) Name() string {
+	return Name
 }
 
-type rrPicker struct {
-	// subConns is the snapshot of the roundrobin balancer when this picker was
-	// created. The slice is immutable. Each Get() will do a round robin
-	// selection from it and return the selected SubConn.
-	subConns []balancer.SubConn
-	next     uint32
+func (bb builder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
+	bal := &rrBalancer{
+		cc:       cc,
+		Balancer: endpointsharding.NewBalancer(cc, opts),
+	}
+	bal.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf("[%p] ", bal))
+	bal.logger.Infof("Created")
+	return bal
 }
 
-func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
-	subConnsLen := uint32(len(p.subConns))
-	nextIndex := atomic.AddUint32(&p.next, 1)
+type rrBalancer struct {
+	balancer.Balancer
+	cc     balancer.ClientConn
+	logger *internalgrpclog.PrefixLogger
+}
 
-	sc := p.subConns[nextIndex%subConnsLen]
-	return balancer.PickResult{SubConn: sc}, nil
+func (b *rrBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
+	// Enable the health listener in pickfirst children for client side health
+	// checks and outlier detection, if configured.
+	ccs.ResolverState = pickfirstleaf.EnableHealthListener(ccs.ResolverState)
+	ccs.BalancerConfig = endpointShardingLBConfig
+	return b.Balancer.UpdateClientConnState(ccs)
 }

balancer/weightedroundrobin/balancer.go

@@ -405,6 +405,7 @@ func (b *wrrBalancer) Close() {
 			ew.stopORCAListener()
 		}
 	}
+	b.child.Close()
 }
 
 func (b *wrrBalancer) ExitIdle() {