balancer/resolver: add loadBalancingConfig and pre-parsing support (#…

…2732)

balancer/balancer.go

@@ -22,6 +22,7 @@ package balancer
 
 import (
  "context"
+ "encoding/json"
  "errors"
  "net"
  "strings"
@@ -31,6 +32,7 @@ import (
  "google.golang.org/grpc/internal"
  "google.golang.org/grpc/metadata"
  "google.golang.org/grpc/resolver"
+ "google.golang.org/grpc/serviceconfig"
 )
 
 var (
@@ -39,7 +41,10 @@ var (
 )
 
 // Register registers the balancer builder to the balancer map. b.Name
-// (lowercased) will be used as the name registered with this builder.
+// (lowercased) will be used as the name registered with this builder. If the
+// Builder implements ConfigParser, ParseConfig will be called when new service
+// configs are received by the resolver, and the result will be provided to the
+// Balancer in UpdateClientConnState.
 //
 // NOTE: this function must only be called during initialization time (i.e. in
 // an init() function), and is not thread-safe. If multiple Balancers are
@@ -172,6 +177,14 @@ type Builder interface {
  Name() string
 }
 
+// ConfigParser parses load balancer configs.
+type ConfigParser interface {
+ // ParseConfig parses the JSON load balancer config provided into an
+ // internal form or returns an error if the config is invalid. For future
+ // compatibility reasons, unknown fields in the config should be ignored.
+ ParseConfig(LoadBalancingConfigJSON json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
+}
+
 // PickOptions contains addition information for the Pick operation.
 type PickOptions struct {
  // FullMethodName is the method name that NewClientStream() is called
@@ -270,7 +283,7 @@ type Balancer interface {
  // non-nil error to gRPC.
  //
  // Deprecated: if V2Balancer is implemented by the Balancer,
- // UpdateResolverState will be called instead.
+ // UpdateClientConnState will be called instead.
  HandleResolvedAddrs([]resolver.Address, error)
  // Close closes the balancer. The balancer is not required to call
  // ClientConn.RemoveSubConn for its existing SubConns.
@@ -283,14 +296,23 @@ type SubConnState struct {
  // TODO: add last connection error
 }
 
+// ClientConnState describes the state of a ClientConn relevant to the
+// balancer.
+type ClientConnState struct {
+ ResolverState resolver.State
+ // The parsed load balancing configuration returned by the builder's
+ // ParseConfig method, if implemented.
+ BalancerConfig serviceconfig.LoadBalancingConfig
+}
+
 // V2Balancer is defined for documentation purposes. If a Balancer also
-// implements V2Balancer, its UpdateResolverState method will be called instead
-// of HandleResolvedAddrs and its UpdateSubConnState will be called instead of
-// HandleSubConnStateChange.
+// implements V2Balancer, its UpdateClientConnState method will be called
+// instead of HandleResolvedAddrs and its UpdateSubConnState will be called
+// instead of HandleSubConnStateChange.
 type V2Balancer interface {
- // UpdateResolverState is called by gRPC when the state of the resolver
+ // UpdateClientConnState is called by gRPC when the state of the ClientConn
  // changes.
- UpdateResolverState(resolver.State)
+ UpdateClientConnState(ClientConnState)
  // UpdateSubConnState is called by gRPC when the state of a SubConn
  // changes.
  UpdateSubConnState(SubConn, SubConnState)

balancer/base/balancer.go

@@ -70,13 +70,13 @@ func (b *baseBalancer) HandleResolvedAddrs(addrs []resolver.Address, err error)
  panic("not implemented")
 }
 
-func (b *baseBalancer) UpdateResolverState(s resolver.State) {
- // TODO: handle s.Err (log if not nil) once implemented.
- // TODO: handle s.ServiceConfig?
- grpclog.Infoln("base.baseBalancer: got new resolver state: ", s)
+func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) {
+ // TODO: handle s.ResolverState.Err (log if not nil) once implemented.
+ // TODO: handle s.ResolverState.ServiceConfig?
+ grpclog.Infoln("base.baseBalancer: got new ClientConn state: ", s)
  // addrsSet is the set converted from addrs, it's used for quick lookup of an address.
  addrsSet := make(map[resolver.Address]struct{})
- for _, a := range s.Addresses {
+ for _, a := range s.ResolverState.Addresses {
  addrsSet[a] = struct{}{}
  if _, ok := b.subConns[a]; !ok {
  // a is a new address (not existing in b.subConns).

balancer/grpclb/grpclb.go

@@ -408,11 +408,11 @@ func (lb *lbBalancer) HandleResolvedAddrs(addrs []resolver.Address, err error) {
  panic("not used")
 }
 
-func (lb *lbBalancer) handleServiceConfig(sc string) {
+func (lb *lbBalancer) handleServiceConfig(gc *grpclbServiceConfig) {
  lb.mu.Lock()
  defer lb.mu.Unlock()
 
- newUsePickFirst := childIsPickFirst(sc)
+ newUsePickFirst := childIsPickFirst(gc)
  if lb.usePickFirst == newUsePickFirst {
  return
  }
@@ -422,13 +422,14 @@ func (lb *lbBalancer) handleServiceConfig(sc string) {
  lb.refreshSubConns(lb.backendAddrs, lb.inFallback, newUsePickFirst)
 }
 
-func (lb *lbBalancer) UpdateResolverState(rs resolver.State) {
+func (lb *lbBalancer) UpdateClientConnState(ccs balancer.ClientConnState) {
  if grpclog.V(2) {
- grpclog.Infof("lbBalancer: UpdateResolverState: %+v", rs)
+ grpclog.Infof("lbBalancer: UpdateClientConnState: %+v", ccs)
  }
- lb.handleServiceConfig(rs.ServiceConfig)
+ gc, _ := ccs.BalancerConfig.(*grpclbServiceConfig)
+ lb.handleServiceConfig(gc)
 
- addrs := rs.Addresses
+ addrs := ccs.ResolverState.Addresses
  if len(addrs) <= 0 {
  return
  }

balancer/grpclb/grpclb_config.go

@@ -23,53 +23,32 @@ import (
 
  "google.golang.org/grpc"
  "google.golang.org/grpc/balancer/roundrobin"
+ "google.golang.org/grpc/serviceconfig"
 )
 
-type serviceConfig struct {
- LoadBalancingConfig *[]map[string]*grpclbServiceConfig
-}
+const (
+ roundRobinName = roundrobin.Name
+ pickFirstName = grpc.PickFirstBalancerName
+)
 
 type grpclbServiceConfig struct {
+ serviceconfig.LoadBalancingConfig
  ChildPolicy *[]map[string]json.RawMessage
 }
 
-func parseFullServiceConfig(s string) *serviceConfig {
- var ret serviceConfig
- err := json.Unmarshal([]byte(s), &ret)
- if err != nil {
- return nil
- }
- return &ret
-}
-
-func parseServiceConfig(s string) *grpclbServiceConfig {
- parsedSC := parseFullServiceConfig(s)
- if parsedSC == nil {
- return nil
+func (b *lbBuilder) ParseConfig(lbConfig json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
+ ret := &grpclbServiceConfig{}
+ if err := json.Unmarshal(lbConfig, ret); err != nil {
+ return nil, err
  }
- lbConfigs := parsedSC.LoadBalancingConfig
- if lbConfigs == nil {
- return nil
- }
- for _, lbC := range *lbConfigs {
- if v, ok := lbC[grpclbName]; ok {
- return v
- }
- }
- return nil
+ return ret, nil
 }
 
-const (
- roundRobinName = roundrobin.Name
- pickFirstName = grpc.PickFirstBalancerName
-)
-
-func childIsPickFirst(s string) bool {
- parsedSC := parseServiceConfig(s)
- if parsedSC == nil {
+func childIsPickFirst(sc *grpclbServiceConfig) bool {
+ if sc == nil {
  return false
  }
- childConfigs := parsedSC.ChildPolicy
+ childConfigs := sc.ChildPolicy
  if childConfigs == nil {
  return false
  }

balancer/grpclb/grpclb_config_test.go

@@ -20,72 +20,31 @@ package grpclb
 
 import (
  "encoding/json"
+ "errors"
+ "fmt"
  "reflect"
+ "strings"
  "testing"
-)
 
-func Test_parseFullServiceConfig(t *testing.T) {
- tests := []struct {
- name string
- s string
- want *serviceConfig
- }{
- {
- name: "empty",
- s: "",
- want: nil,
- },
- {
- name: "success1",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`,
- want: &serviceConfig{
- LoadBalancingConfig: &[]map[string]*grpclbServiceConfig{
- {"grpclb": &grpclbServiceConfig{
- ChildPolicy: &[]map[string]json.RawMessage{
- {"pick_first": json.RawMessage("{}")},
- },
- }},
- },
- },
- },
- {
- name: "success2",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}}]}`,
- want: &serviceConfig{
- LoadBalancingConfig: &[]map[string]*grpclbServiceConfig{
- {"grpclb": &grpclbServiceConfig{
- ChildPolicy: &[]map[string]json.RawMessage{
- {"round_robin": json.RawMessage("{}")},
- {"pick_first": json.RawMessage("{}")},
- },
- }},
- },
- },
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- if got := parseFullServiceConfig(tt.s); !reflect.DeepEqual(got, tt.want) {
- t.Errorf("parseFullServiceConfig() = %+v, want %+v", got, tt.want)
- }
- })
- }
-}
+ "google.golang.org/grpc/serviceconfig"
+)
 
-func Test_parseServiceConfig(t *testing.T) {
+func Test_Parse(t *testing.T) {
  tests := []struct {
- name string
- s string
- want *grpclbServiceConfig
+ name string
+ s string
+ want serviceconfig.LoadBalancingConfig
+ wantErr error
  }{
  {
- name: "empty",
- s: "",
- want: nil,
+ name: "empty",
+ s: "",
+ want: nil,
+ wantErr: errors.New("unexpected end of JSON input"),
  },
  {
  name: "success1",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`,
+ s: `{"childPolicy":[{"pick_first":{}}]}`,
  want: &grpclbServiceConfig{
  ChildPolicy: &[]map[string]json.RawMessage{
  {"pick_first": json.RawMessage("{}")},
@@ -94,24 +53,19 @@ func Test_parseServiceConfig(t *testing.T) {
  },
  {
  name: "success2",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}}]}`,
+ s: `{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}`,
  want: &grpclbServiceConfig{
  ChildPolicy: &[]map[string]json.RawMessage{
  {"round_robin": json.RawMessage("{}")},
  {"pick_first": json.RawMessage("{}")},
  },
  },
  },
- {
- name: "no_grpclb",
- s: `{"loadBalancingConfig":[{"notgrpclb":{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}}]}`,
- want: nil,
- },
  }
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
- if got := parseServiceConfig(tt.s); !reflect.DeepEqual(got, tt.want) {
- t.Errorf("parseFullServiceConfig() = %+v, want %+v", got, tt.want)
+ if got, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.s)); !reflect.DeepEqual(got, tt.want) || !strings.Contains(fmt.Sprint(err), fmt.Sprint(tt.wantErr)) {
+ t.Errorf("parseFullServiceConfig() = %+v, %+v, want %+v, <contains %q>", got, err, tt.want, tt.wantErr)
  }
  })
  }
@@ -123,30 +77,29 @@ func Test_childIsPickFirst(t *testing.T) {
  s string
  want bool
  }{
- {
- name: "invalid",
- s: "",
- want: false,
- },
  {
  name: "pickfirst_only",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}}]}}]}`,
+ s: `{"childPolicy":[{"pick_first":{}}]}`,
  want: true,
  },
  {
  name: "pickfirst_before_rr",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"pick_first":{}},{"round_robin":{}}]}}]}`,
+ s: `{"childPolicy":[{"pick_first":{}},{"round_robin":{}}]}`,
  want: true,
  },
  {
  name: "rr_before_pickfirst",
- s: `{"loadBalancingConfig":[{"grpclb":{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}}]}`,
+ s: `{"childPolicy":[{"round_robin":{}},{"pick_first":{}}]}`,
  want: false,
  },
  }
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
- if got := childIsPickFirst(tt.s); got != tt.want {
+ gc, err := (&lbBuilder{}).ParseConfig(json.RawMessage(tt.s))
+ if err != nil {
+ t.Fatalf("Parse(%v) = _, %v; want _, nil", tt.s, err)
+ }
+ if got := childIsPickFirst(gc.(*grpclbServiceConfig)); got != tt.want {
  t.Errorf("childIsPickFirst() = %v, want %v", got, tt.want)
  }
  })