Merge branch 'master' of https://github.com/pvsravani/grpc-go into stubserver-test

Author: pvsravani

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,26 +48,30 @@ 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, true)
+// Options are the options to configure the behaviour of the
+// endpointsharding balancer.
+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
 }
 
-// NewBalancerWithoutAutoReconnect 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 NewBalancerWithoutAutoReconnect(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
-	return newBlanacer(cc, opts, false)
-}
+// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
+// type as the balancer.Builder.Build method.
+type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
 
-func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, autoReconnect bool) balancer.Balancer {
+// 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 ChildBuilderFunc, esOpts Options) balancer.Balancer {
 	es := &endpointSharding{
-		cc:                  cc,
-		bOpts:               opts,
-		enableAutoReconnect: autoReconnect,
+		cc:           cc,
+		bOpts:        opts,
+		esOpts:       esOpts,
+		childBuilder: childBuilder,
 	}
 	es.children.Store(resolver.NewEndpointMap())
 	return es
@@ -97,9 +81,10 @@ func newBlanacer(cc balancer.ClientConn, opts balancer.BuildOptions, autoReconne
 // 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
-	enableAutoReconnect bool
+	cc           balancer.ClientConn
+	bOpts        balancer.BuildOptions
+	esOpts       Options
+	childBuilder ChildBuilderFunc
 
 	// 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
@@ -160,7 +145,7 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 				es:         es,
 			}
 			childBalancer.childState.Balancer = childBalancer
-			childBalancer.child = gracefulswitch.NewBalancer(childBalancer, es.bOpts)
+			childBalancer.child = es.childBuilder(childBalancer, es.bOpts)
 		}
 		newChildren.Set(endpoint, childBalancer)
 		if err := childBalancer.updateClientConnStateLocked(balancer.ClientConnState{
@@ -334,7 +319,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.enableAutoReconnect {
+	if state.ConnectivityState == connectivity.Idle && !bw.es.esOpts.DisableAutoReconnect {
 		bw.ExitIdle()
 	}
 	bw.es.updateState()
@@ -370,12 +355,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)
-}

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).Build, 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,13 +180,11 @@ func (s) TestEndpointShardingReconnectDisabled(t *testing.T) {
 	name := strings.ReplaceAll(strings.ToLower(t.Name()), "/", "")
 	bf := stub.BalancerFuncs{
 		Init: func(bd *stub.BalancerData) {
-			bd.Data = endpointsharding.NewBalancerWithoutAutoReconnect(bd.ClientConn, bd.BuildOptions)
+			epOpts := endpointsharding.Options{DisableAutoReconnect: true}
+			bd.Data = endpointsharding.NewBalancer(bd.ClientConn, bd.BuildOptions, balancer.Get(pickfirstleaf.Name).Build, 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,44 +26,35 @@
 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() {
-	balancer.Register(builder{})
-}
-
 var (
-	// PickfirstConfig is the LB policy config json for a pick_first load
-	// balancer that is lazily initialized.
-	PickfirstConfig = fmt.Sprintf("{\"childPolicy\": [{%q: {}}]}", pickfirstleaf.Name)
-	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] "
 )
 
-type builder struct{}
+// ChildBuilderFunc creates a new balancer with the ClientConn. It has the same
+// type as the balancer.Builder.Build method.
+type ChildBuilderFunc func(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer
 
-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 ChildBuilderFunc) balancer.Balancer {
 	b := &lazyBalancer{
 		cc:           cc,
 		buildOptions: bOpts,
+		childBuilder: childBuilder,
 	}
 	b.logger = internalgrpclog.NewPrefixLogger(logger, fmt.Sprintf(logPrefix, b))
 	cc.UpdateState(balancer.State{
@@ -77,16 +68,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 ChildBuilderFunc
 
 	// The following fields are accessed while handling calls to the idlePicker
 	// and when handling ClientConn state updates. They are guarded by a mutex.
@@ -119,12 +107,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)
 	}
@@ -148,7 +130,7 @@ func (lb *lazyBalancer) ExitIdle() {
 		}
 		return
 	}
-	lb.delegate = gracefulswitch.NewBalancer(lb.cc, lb.buildOptions)
+	lb.delegate = lb.childBuilder(lb.cc, lb.buildOptions)
 	if lb.latestClientConnState != nil {
 		if err := lb.delegate.UpdateClientConnState(*lb.latestClientConnState); err != nil {
 			if err == balancer.ErrBadResolverState {
@@ -165,25 +147,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 {