Remove graceful switch from lazy and fix endpointsharding usages

Author: arjan-bal

balancer/endpointsharding/endpointsharding.go

@@ -26,37 +26,17 @@
 package endpointsharding
 
 import (
-	"encoding/json"
 	"errors"
-	"fmt"
 	rand "math/rand/v2"
 	"sync"
 	"sync/atomic"
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/base"
-	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
-	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/resolver"
-	"google.golang.org/grpc/serviceconfig"
 )
 
-var (
-	// PickFirstConfig is a pick first config without shuffling enabled.
-	PickFirstConfig serviceconfig.LoadBalancingConfig
-	logger          = grpclog.Component("endpoint-sharding")
-)
-
-func init() {
-	var err error
-	PickFirstConfig, err = ParseConfig(json.RawMessage(fmt.Sprintf("[{%q: {}}]", pickfirstleaf.Name)))
-	if err != nil {
-		logger.Fatal(err)
-	}
-}
-
 // ChildState is the balancer state of a child along with the endpoint which
 // identifies the child balancer.
 type ChildState struct {
@@ -68,36 +48,26 @@ type ChildState struct {
 	Balancer balancer.ExitIdler
 }
 
-// NewBalancer returns a load balancing policy that manages homogeneous child
-// policies each owning a single endpoint. The balancer will automatically call
-// ExitIdle on its children if they report IDLE connectivity state.
-func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	return newBlanacer(cc, opts, BuildOptions{})
-}
-
-// BuildOptions are the options to configure the behaviour of the
+// Options are the options to configure the behaviour of the
 // endpointsharding balancer.
-type BuildOptions struct {
+type Options struct {
+	// DisableAutoReconnect allows the balancer to keep child balancer in the
+	// IDLE state until they are explicitly triggered to exit using the
+	// ChildState obtained from the endpointsharding picker. When set to false,
+	// the endpointsharding balancer will automatically call ExitIdle on child
+	// connections that report IDLE.
 	DisableAutoReconnect bool
-	ChildBuilder         balancer.Builder
 }
 
-// NewBalancerWithOpts returns a load balancing policy that manages
-// homogeneous child policies each owning a single endpoint. The balancer will
-// allow children to remain in IDLE state until triggered to exit idle state
-// using the ChildState obtained using the endpointsharding picker.
-func NewBalancerWithOpts(cc balancer.ClientConn, opts balancer.BuildOptions, epOpts BuildOptions) balancer.Balancer {
-	return newBlanacer(cc, opts, epOpts)
-}
-
-func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, epOpts BuildOptions) balancer.Balancer {
-	if epOpts.ChildBuilder == nil {
-		epOpts.ChildBuilder = &gracefulSwitchBuilder{}
-	}
+// NewBalancer returns a load balancing policy that manages homogeneous child
+// policies each owning a single endpoint. The endpointsharding balancer
+// forwards the LoadBalancingConfig in ClientConn state updates to its children.
+func NewBalancer(cc balancer.ClientConn, opts balancer.BuildOptions, childBuilder balancer.Builder, esOpts Options) balancer.Balancer {
 	es := &endpointSharding{
-		cc:     cc,
-		bOpts:  opts,
-		epOpts: epOpts,
+		cc:           cc,
+		bOpts:        opts,
+		esOpts:       esOpts,
+		childBuilder: childBuilder,
 	}
 	es.children.Store(resolver.NewEndpointMap())
 	return es
@@ -107,9 +77,10 @@ func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, epOpts Buil
 // balancer with child config for every unique Endpoint received. It updates the
 // child states on any update from parent or child.
 type endpointSharding struct {
-	cc     balancer.ClientConn
-	bOpts  balancer.BuildOptions
-	epOpts BuildOptions
+	cc           balancer.ClientConn
+	bOpts        balancer.BuildOptions
+	esOpts       Options
+	childBuilder balancer.Builder
 
 	// childMu synchronizes calls to any single child. It must be held for all
 	// calls into a child. To avoid deadlocks, do not acquire childMu while
@@ -170,7 +141,7 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 				es:         es,
 			}
 			childBalancer.childState.Balancer = childBalancer
-			childBalancer.child = es.epOpts.ChildBuilder.Build(childBalancer, es.bOpts)
+			childBalancer.child = es.childBuilder.Build(childBalancer, es.bOpts)
 		}
 		newChildren.Set(endpoint, childBalancer)
 		if err := childBalancer.updateClientConnStateLocked(balancer.ClientConnState{
@@ -344,7 +315,7 @@ func (bw *balancerWrapper) UpdateState(state balancer.State) {
 	bw.es.mu.Lock()
 	bw.childState.State = state
 	bw.es.mu.Unlock()
-	if state.ConnectivityState == connectivity.Idle && !bw.es.epOpts.DisableAutoReconnect {
+	if state.ConnectivityState == connectivity.Idle && !bw.es.esOpts.DisableAutoReconnect {
 		bw.ExitIdle()
 	}
 	bw.es.updateState()
@@ -380,23 +351,3 @@ func (bw *balancerWrapper) closeLocked() {
 	bw.child.Close()
 	bw.isClosed = true
 }
-
-// ParseConfig parses a child config list and returns an LB config to use with
-// the endpointsharding balancer.
-//
-// cfg is expected to be a JSON array of LB policy names + configs as the
-// format of the loadBalancingConfig field in ServiceConfig.
-func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
-	return gracefulswitch.ParseConfig(cfg)
-}
-
-type gracefulSwitchBuilder struct {
-}
-
-func (b *gracefulSwitchBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	return gracefulswitch.NewBalancer(cc, opts)
-}
-
-func (b *gracefulSwitchBuilder) Name() string {
-	return "gracefulswitch"
-}

balancer/endpointsharding/endpointsharding_test.go

@@ -30,6 +30,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/endpointsharding"
+	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/grpclog"
@@ -80,7 +81,7 @@ func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptio
 		ClientConn: cc,
 		bOpts:      opts,
 	}
-	fp.Balancer = endpointsharding.NewBalancer(fp, opts)
+	fp.Balancer = endpointsharding.NewBalancer(fp, opts, balancer.Get(pickfirstleaf.Name), endpointsharding.Options{})
 	return fp
 }
 
@@ -103,10 +104,7 @@ func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) err
 		return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
 	}
 
-	return fp.Balancer.UpdateClientConnState(balancer.ClientConnState{
-		BalancerConfig: endpointsharding.PickFirstConfig,
-		ResolverState:  state.ResolverState,
-	})
+	return fp.Balancer.UpdateClientConnState(state)
 }
 
 func (fp *fakePetiole) UpdateState(state balancer.State) {
@@ -182,14 +180,11 @@ func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
 	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			epOpts := endpointsharding.BuildOptions{DisableAutoReconnect: true}
-			bd.Data = endpointsharding.NewBalancerWithOpts(bd.ClientConn, bd.BuildOptions, epOpts)
+			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
+			bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name), epOpts)
 		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
-			return bd.Data.(balancer.Balancer).UpdateClientConnState(balancer.ClientConnState{
-				BalancerConfig: endpointsharding.PickFirstConfig,
-				ResolverState:  ccs.ResolverState,
-			})
+			return bd.Data.(balancer.Balancer).UpdateClientConnState(ccs)
 		},
 		Close: func(bd *stub.BalancerData) {
 			bd.Data.(balancer.Balancer).Close()

balancer/lazy/lazy.go

@@ -26,51 +26,31 @@
 package lazy
 
 import (
-	"encoding/json"
 	"fmt"
 	"sync"
 
 	"google.golang.org/grpc/balancer"
-	"google.golang.org/grpc/balancer/pickfirst/pickfirstleaf"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
-	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/resolver"
-	"google.golang.org/grpc/serviceconfig"
 
 	internalgrpclog "google.golang.org/grpc/internal/grpclog"
 )
 
-func init() {
-	pfCfgJSON := fmt.Sprintf("{\"childPolicy\": [{%q: {}}]}", pickfirstleaf.Name)
-	var err error
-	PickfirstConfig, err = Builder{}.ParseConfig(json.RawMessage(pfCfgJSON))
-	if err != nil {
-		logger.Fatalf("Failed to parse pickfirst config: %v", err)
-	}
-}
-
 var (
-	// PickfirstConfig is the LB policy config json for a pick_first load
-	// balancer that is lazily initialized.
-	PickfirstConfig serviceconfig.LoadBalancingConfig
-	logger          = grpclog.Component("lazy-lb")
+	logger = grpclog.Component("lazy-lb")
 )
 
 const (
-	// Name is the name of the lazy balancer.
-	Name      = "lazy"
 	logPrefix = "[lazy-lb %p] "
 )
 
-// Builder builds the lazy balancer. It is not registered in the global balancer
-// registry by default.
-type Builder struct{}
-
-func (Builder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balancer.Balancer {
+// NewBalancer is the constructor for the lazy balancer.
+func NewBalancer(cc balancer.ClientConn, bOpts balancer.BuildOptions, childBuilder balancer.Builder) balancer.Balancer {
 	b := &lazyBalancer{
 		cc:           cc,
 		buildOptions: bOpts,
+		childBuilder: childBuilder,
 	}
 	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
 	cc.UpdateState(balancer.State{
@@ -84,16 +64,13 @@ func (Builder) Build(cc balancer.ClientConn, bOpts balancer.BuildOptions) balanc
 	return b
 }
 
-func (Builder) Name() string {
-	return Name
-}
-
 type lazyBalancer struct {
 	// The following fields are initialized at build time and read-only after
 	// that and therefore do not need to be guarded by a mutex.
 	cc           balancer.ClientConn
 	buildOptions balancer.BuildOptions
 	logger       *internalgrpclog.PrefixLogger
+	childBuilder balancer.Builder
 
 	// The following fields are accessed while handling calls to the idlePicker
 	// and when handling ClientConn state updates. They are guarded by a mutex.
@@ -126,12 +103,6 @@ func (lb *lazyBalancer) ResolverError(err error) {
 func (lb *lazyBalancer) UpdateClientConnState(ccs balancer.ClientConnState) error {
 	lb.mu.Lock()
 	defer lb.mu.Unlock()
-	childLBCfg, ok := ccs.BalancerConfig.(lbCfg)
-	if !ok {
-		lb.logger.Errorf("Got LB config of unexpected type: %v", ccs.BalancerConfig)
-		return balancer.ErrBadResolverState
-	}
-	ccs.BalancerConfig = childLBCfg.childLBCfg
 	if lb.delegate != nil {
 		return lb.delegate.UpdateClientConnState(ccs)
 	}
@@ -155,7 +126,7 @@ func (lb *lazyBalancer) ExitIdle() {
 		}
 		return
 	}
-	lb.delegate = gracefulswitch.NewBalancer(lb.cc, lb.buildOptions)
+	lb.delegate = lb.childBuilder.Build(lb.cc, lb.buildOptions)
 	if lb.latestClientConnState != nil {
 		if err := lb.delegate.UpdateClientConnState(*lb.latestClientConnState); err != nil {
 			if err == balancer.ErrBadResolverState {
@@ -172,25 +143,6 @@ func (lb *lazyBalancer) ExitIdle() {
 	}
 }
 
-type lbCfg struct {
-	serviceconfig.LoadBalancingConfig
-	childLBCfg serviceconfig.LoadBalancingConfig
-}
-
-func (b Builder) ParseConfig(lbConfig json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
-	jsonReprsentation := &struct {
-		ChildPolicy json.RawMessage
-	}{}
-	if err := json.Unmarshal(lbConfig, jsonReprsentation); err != nil {
-		return nil, err
-	}
-	childCfg, err := gracefulswitch.ParseConfig(jsonReprsentation.ChildPolicy)
-	if err != nil {
-		return nil, err
-	}
-	return lbCfg{childLBCfg: childCfg}, nil
-}
-
 // idlePicker is used when the SubConn is IDLE and kicks the SubConn into
 // CONNECTING when Pick is called.
 type idlePicker struct {