From 1db6590e400ee69e2c15b9312255795625862ab0 Mon Sep 17 00:00:00 2001 From: Zach Reyes <39203661+zasweq@users.noreply.github.com> Date: Wed, 22 May 2024 16:26:02 -0400 Subject: [PATCH] grpc: Move Pick First Balancer to separate package (#7255) --- balancer/grpclb/grpclb_config.go | 4 ++-- .../pickfirst/pickfirst.go | 18 +++++++++++++----- balancer/rls/balancer_test.go | 6 ++++-- .../rls/internal/test/e2e/rls_child_policy.go | 4 ++-- clientconn.go | 1 - internal/balancergroup/balancergroup_test.go | 4 ++-- service_config.go | 5 +++-- test/balancer_switching_test.go | 19 ++++++++++--------- test/balancer_test.go | 3 ++- test/resolver_update_test.go | 3 ++- .../clustermanager/clustermanager_test.go | 4 ++-- .../e2e_test/clustermanager_test.go | 3 ++- .../xdslbregistry/converter/converter.go | 4 ++-- 13 files changed, 46 insertions(+), 32 deletions(-) rename pickfirst.go => balancer/pickfirst/pickfirst.go (96%) diff --git a/balancer/grpclb/grpclb_config.go b/balancer/grpclb/grpclb_config.go index 8942c31310af..96a57c8c70c8 100644 --- a/balancer/grpclb/grpclb_config.go +++ b/balancer/grpclb/grpclb_config.go @@ -21,14 +21,14 @@ package grpclb import ( "encoding/json" - "google.golang.org/grpc" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/balancer/roundrobin" "google.golang.org/grpc/serviceconfig" ) const ( roundRobinName = roundrobin.Name - pickFirstName = grpc.PickFirstBalancerName + pickFirstName = pickfirst.Name ) type grpclbServiceConfig struct { diff --git a/pickfirst.go b/balancer/pickfirst/pickfirst.go similarity index 96% rename from pickfirst.go rename to balancer/pickfirst/pickfirst.go index 8853626614e8..853f3885c021 100644 --- a/pickfirst.go +++ b/balancer/pickfirst/pickfirst.go @@ -16,7 +16,8 @@ * */ -package grpc +// Package pickfirst contains the pick_first load balancing policy. +package pickfirst import ( "encoding/json" @@ -25,6 +26,7 @@ import ( "google.golang.org/grpc/balancer" "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/grpclog" internalgrpclog "google.golang.org/grpc/internal/grpclog" "google.golang.org/grpc/internal/grpcrand" "google.golang.org/grpc/internal/pretty" @@ -32,10 +34,16 @@ import ( "google.golang.org/grpc/serviceconfig" ) +func init() { + balancer.Register(pickfirstBuilder{}) +} + +var logger = grpclog.Component("pick-first-lb") + const ( - // PickFirstBalancerName is the name of the pick_first balancer. - PickFirstBalancerName = "pick_first" - logPrefix = "[pick-first-lb %p] " + // Name is the name of the pick_first balancer. + Name = "pick_first" + logPrefix = "[pick-first-lb %p] " ) type pickfirstBuilder struct{} @@ -47,7 +55,7 @@ func (pickfirstBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) } func (pickfirstBuilder) Name() string { - return PickFirstBalancerName + return Name } type pfConfig struct { diff --git a/balancer/rls/balancer_test.go b/balancer/rls/balancer_test.go index 444b8b99d4a3..6930e3da14f5 100644 --- a/balancer/rls/balancer_test.go +++ b/balancer/rls/balancer_test.go @@ -30,6 +30,7 @@ import ( "github.com/google/go-cmp/cmp" "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/balancer/rls/internal/test/e2e" "google.golang.org/grpc/codes" "google.golang.org/grpc/connectivity" @@ -37,7 +38,6 @@ import ( "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/internal" "google.golang.org/grpc/internal/balancer/stub" - rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1" internalserviceconfig "google.golang.org/grpc/internal/serviceconfig" "google.golang.org/grpc/internal/testutils" rlstest "google.golang.org/grpc/internal/testutils/rls" @@ -46,6 +46,8 @@ import ( "google.golang.org/grpc/resolver/manual" "google.golang.org/grpc/serviceconfig" "google.golang.org/grpc/testdata" + + rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1" "google.golang.org/protobuf/types/known/durationpb" ) @@ -919,7 +921,7 @@ func (s) TestUpdateStatePauses(t *testing.T) { } stub.Register(childPolicyName, stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { - bd.Data = balancer.Get(grpc.PickFirstBalancerName).Build(bd.ClientConn, bd.BuildOptions) + bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions) }, ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) { cfg := &childPolicyConfig{} diff --git a/balancer/rls/internal/test/e2e/rls_child_policy.go b/balancer/rls/internal/test/e2e/rls_child_policy.go index 5a6e3e69175a..8742fa135886 100644 --- a/balancer/rls/internal/test/e2e/rls_child_policy.go +++ b/balancer/rls/internal/test/e2e/rls_child_policy.go @@ -23,8 +23,8 @@ import ( "errors" "fmt" - "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/internal/grpcsync" "google.golang.org/grpc/resolver" "google.golang.org/grpc/serviceconfig" @@ -68,7 +68,7 @@ type bb struct { func (bb bb) Name() string { return bb.name } func (bb bb) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer { - pf := balancer.Get(grpc.PickFirstBalancerName) + pf := balancer.Get(pickfirst.Name) b := &bal{ Balancer: pf.Build(cc, opts), bf: bb.bf, diff --git a/clientconn.go b/clientconn.go index 77f2d2cceaaa..29fc11698d10 100644 --- a/clientconn.go +++ b/clientconn.go @@ -693,7 +693,6 @@ func (cc *ClientConn) waitForResolvedAddrs(ctx context.Context) error { var emptyServiceConfig *ServiceConfig func init() { - balancer.Register(pickfirstBuilder{}) cfg := parseServiceConfig("{}") if cfg.Err != nil { panic(fmt.Sprintf("impossible error parsing empty service config: %v", cfg.Err)) diff --git a/internal/balancergroup/balancergroup_test.go b/internal/balancergroup/balancergroup_test.go index b182ec8489e2..9de47f54504f 100644 --- a/internal/balancergroup/balancergroup_test.go +++ b/internal/balancergroup/balancergroup_test.go @@ -23,8 +23,8 @@ import ( "testing" "time" - "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/balancer/roundrobin" "google.golang.org/grpc/balancer/weightedtarget/weightedaggregator" "google.golang.org/grpc/connectivity" @@ -602,7 +602,7 @@ func (s) TestBalancerGracefulSwitch(t *testing.T) { childPolicyName := t.Name() stub.Register(childPolicyName, stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { - bd.Data = balancer.Get(grpc.PickFirstBalancerName).Build(bd.ClientConn, bd.BuildOptions) + bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions) }, UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error { ccs.ResolverState.Addresses = ccs.ResolverState.Addresses[1:] diff --git a/service_config.go b/service_config.go index 9da8fc8027d1..b7a4f5960490 100644 --- a/service_config.go +++ b/service_config.go @@ -26,6 +26,7 @@ import ( "time" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/codes" "google.golang.org/grpc/internal" "google.golang.org/grpc/internal/balancer/gracefulswitch" @@ -183,12 +184,12 @@ func parseServiceConfig(js string) *serviceconfig.ParseResult { } c := rsc.LoadBalancingConfig if c == nil { - name := PickFirstBalancerName + name := pickfirst.Name if rsc.LoadBalancingPolicy != nil { name = *rsc.LoadBalancingPolicy } if balancer.Get(name) == nil { - name = PickFirstBalancerName + name = pickfirst.Name } cfg := []map[string]any{{name: struct{}{}}} strCfg, err := json.Marshal(cfg) diff --git a/test/balancer_switching_test.go b/test/balancer_switching_test.go index b70afebd4059..34fd871b65c7 100644 --- a/test/balancer_switching_test.go +++ b/test/balancer_switching_test.go @@ -26,12 +26,13 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/balancer" grpclbstate "google.golang.org/grpc/balancer/grpclb/state" + pickfirst "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/internal" "google.golang.org/grpc/internal/balancer/stub" "google.golang.org/grpc/internal/stubserver" "google.golang.org/grpc/internal/testutils/fakegrpclb" - "google.golang.org/grpc/internal/testutils/pickfirst" + pfutil "google.golang.org/grpc/internal/testutils/pickfirst" rrutil "google.golang.org/grpc/internal/testutils/roundrobin" "google.golang.org/grpc/resolver" "google.golang.org/grpc/resolver/manual" @@ -127,7 +128,7 @@ func (s) TestBalancerSwitch_Basic(t *testing.T) { r.UpdateState(resolver.State{Addresses: addrs}) ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout) defer cancel() - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { t.Fatal(err) } @@ -146,7 +147,7 @@ func (s) TestBalancerSwitch_Basic(t *testing.T) { Addresses: addrs, ServiceConfig: parseServiceConfig(t, r, pickFirstServiceConfig), }) - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { t.Fatal(err) } } @@ -195,7 +196,7 @@ func (s) TestBalancerSwitch_grpclbToPickFirst(t *testing.T) { // newly configured backends, as part of the balancer switch. emptyConfig := parseServiceConfig(t, r, `{}`) r.UpdateState(resolver.State{Addresses: addrs[1:], ServiceConfig: emptyConfig}) - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { t.Fatal(err) } } @@ -220,7 +221,7 @@ func (s) TestBalancerSwitch_pickFirstToGRPCLB(t *testing.T) { r.UpdateState(resolver.State{Addresses: addrs[1:]}) ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout) defer cancel() - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { t.Fatal(err) } @@ -245,7 +246,7 @@ func (s) TestBalancerSwitch_pickFirstToGRPCLB(t *testing.T) { // Switch to "pick_first" again by sending no grpclb server addresses. emptyConfig := parseServiceConfig(t, r, `{}`) r.UpdateState(resolver.State{Addresses: addrs[1:], ServiceConfig: emptyConfig}) - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[1]); err != nil { t.Fatal(err) } } @@ -340,7 +341,7 @@ func (s) TestBalancerSwitch_grpclbNotRegistered(t *testing.T) { r.UpdateState(grpclbstate.Set(state, &grpclbstate.State{BalancerAddresses: grpclbAddr})) ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout) defer cancel() - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { t.Fatal(err) } @@ -468,7 +469,7 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) { waitToProceed := make(chan struct{}) stub.Register(t.Name(), stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { - pf := balancer.Get(grpc.PickFirstBalancerName) + pf := balancer.Get(pickfirst.Name) bd.Data = pf.Build(bd.ClientConn, bd.BuildOptions) }, UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error { @@ -503,7 +504,7 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) { // underlying "pick_first" balancer which will result in a healthy picker // being reported to the channel. RPCs should start using the new balancer. close(waitToProceed) - if err := pickfirst.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { + if err := pfutil.CheckRPCsToBackend(ctx, cc, addrs[0]); err != nil { t.Fatal(err) } } diff --git a/test/balancer_test.go b/test/balancer_test.go index 721f390a76a6..de7ab5557e80 100644 --- a/test/balancer_test.go +++ b/test/balancer_test.go @@ -32,6 +32,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/attributes" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/codes" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials" @@ -847,7 +848,7 @@ func (s) TestMetadataInPickResult(t *testing.T) { stub.Register(t.Name(), stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { cc := &testCCWrapper{ClientConn: bd.ClientConn} - bd.Data = balancer.Get(grpc.PickFirstBalancerName).Build(cc, bd.BuildOptions) + bd.Data = balancer.Get(pickfirst.Name).Build(cc, bd.BuildOptions) }, UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error { bal := bd.Data.(balancer.Balancer) diff --git a/test/resolver_update_test.go b/test/resolver_update_test.go index 072e0b575ef8..103c732f61af 100644 --- a/test/resolver_update_test.go +++ b/test/resolver_update_test.go @@ -29,6 +29,7 @@ import ( "github.com/google/go-cmp/cmp" "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/internal" @@ -158,7 +159,7 @@ func (s) TestResolverUpdate_InvalidServiceConfigAfterGoodUpdate(t *testing.T) { ccUpdateCh := testutils.NewChannel() stub.Register(t.Name(), stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { - pf := balancer.Get(grpc.PickFirstBalancerName) + pf := balancer.Get(pickfirst.Name) bd.Data = pf.Build(bd.ClientConn, bd.BuildOptions) }, ParseConfig: func(lbCfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error) { diff --git a/xds/internal/balancer/clustermanager/clustermanager_test.go b/xds/internal/balancer/clustermanager/clustermanager_test.go index b998c1b35f29..1b3fa954b86f 100644 --- a/xds/internal/balancer/clustermanager/clustermanager_test.go +++ b/xds/internal/balancer/clustermanager/clustermanager_test.go @@ -25,8 +25,8 @@ import ( "time" "github.com/google/go-cmp/cmp" - "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/codes" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials/insecure" @@ -642,7 +642,7 @@ func TestClusterGracefulSwitch(t *testing.T) { childPolicyName := t.Name() stub.Register(childPolicyName, stub.BalancerFuncs{ Init: func(bd *stub.BalancerData) { - bd.Data = balancer.Get(grpc.PickFirstBalancerName).Build(bd.ClientConn, bd.BuildOptions) + bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions) }, UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error { bal := bd.Data.(balancer.Balancer) diff --git a/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go b/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go index e9fd3c389fef..443e08b3bf5a 100644 --- a/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go +++ b/xds/internal/balancer/clustermanager/e2e_test/clustermanager_test.go @@ -27,6 +27,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/balancer" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/codes" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials/insecure" @@ -181,7 +182,7 @@ func (s) TestConfigUpdate_ChildPolicyChange(t *testing.T) { // Create a wrapped pickfirst LB policy. When the endpoint picking policy on // the cluster resource is changed to pickfirst, this will allow us to // verify that load balancing configuration is pushed to it. - pfBuilder := balancer.Get(grpc.PickFirstBalancerName) + pfBuilder := balancer.Get(pickfirst.Name) internal.BalancerUnregister(pfBuilder.Name()) lbCfgCh := make(chan serviceconfig.LoadBalancingConfig, 1) diff --git a/xds/internal/xdsclient/xdslbregistry/converter/converter.go b/xds/internal/xdsclient/xdslbregistry/converter/converter.go index 076ae8644f88..3c48f1bdea3d 100644 --- a/xds/internal/xdsclient/xdslbregistry/converter/converter.go +++ b/xds/internal/xdsclient/xdslbregistry/converter/converter.go @@ -27,9 +27,9 @@ import ( "fmt" "strings" - "google.golang.org/grpc" "google.golang.org/grpc/balancer" "google.golang.org/grpc/balancer/leastrequest" + "google.golang.org/grpc/balancer/pickfirst" "google.golang.org/grpc/balancer/roundrobin" "google.golang.org/grpc/balancer/weightedroundrobin" "google.golang.org/grpc/internal/envconfig" @@ -110,7 +110,7 @@ func convertPickFirstProtoToServiceConfig(rawProto []byte, _ int) (json.RawMessa if err != nil { return nil, fmt.Errorf("error marshaling JSON for type %T: %v", pfCfg, err) } - return makeBalancerConfigJSON(grpc.PickFirstBalancerName, js), nil + return makeBalancerConfigJSON(pickfirst.Name, js), nil } func convertRoundRobinProtoToServiceConfig([]byte, int) (json.RawMessage, error) {