balancer_test.go

Documentation: google.golang.org/grpc/xds/internal/balancer/wrrlocality

     1  /*
     2   *
     3   * Copyright 2023 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package wrrlocality
    20  
    21  import (
    22  	"context"
    23  	"encoding/json"
    24  	"errors"
    25  	"strings"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/google/go-cmp/cmp"
    30  	"google.golang.org/grpc/balancer"
    31  	"google.golang.org/grpc/balancer/roundrobin"
    32  	"google.golang.org/grpc/balancer/weightedtarget"
    33  	"google.golang.org/grpc/internal/balancer/stub"
    34  	"google.golang.org/grpc/internal/grpctest"
    35  	internalserviceconfig "google.golang.org/grpc/internal/serviceconfig"
    36  	"google.golang.org/grpc/internal/testutils"
    37  	"google.golang.org/grpc/resolver"
    38  	"google.golang.org/grpc/serviceconfig"
    39  	"google.golang.org/grpc/xds/internal"
    40  )
    41  
    42  const (
    43  	defaultTestTimeout = 5 * time.Second
    44  )
    45  
    46  type s struct {
    47  	grpctest.Tester
    48  }
    49  
    50  func Test(t *testing.T) {
    51  	grpctest.RunSubTests(t, s{})
    52  }
    53  
    54  func (s) TestParseConfig(t *testing.T) {
    55  	const errParseConfigName = "errParseConfigBalancer"
    56  	stub.Register(errParseConfigName, stub.BalancerFuncs{
    57  		ParseConfig: func(json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
    58  			return nil, errors.New("some error")
    59  		},
    60  	})
    61  
    62  	parser := bb{}
    63  	tests := []struct {
    64  		name    string
    65  		input   string
    66  		wantCfg serviceconfig.LoadBalancingConfig
    67  		wantErr string
    68  	}{
    69  		{
    70  			name:  "happy-case-round robin-child",
    71  			input: `{"childPolicy": [{"round_robin": {}}]}`,
    72  			wantCfg: &LBConfig{
    73  				ChildPolicy: &internalserviceconfig.BalancerConfig{
    74  					Name: roundrobin.Name,
    75  				},
    76  			},
    77  		},
    78  		{
    79  			name:    "invalid-json",
    80  			input:   "{{invalidjson{{",
    81  			wantErr: "invalid character",
    82  		},
    83  
    84  		{
    85  			name:    "child-policy-field-isn't-set",
    86  			input:   `{}`,
    87  			wantErr: "child policy field must be set",
    88  		},
    89  		{
    90  			name:    "child-policy-type-is-empty",
    91  			input:   `{"childPolicy": []}`,
    92  			wantErr: "invalid loadBalancingConfig: no supported policies found in []",
    93  		},
    94  		{
    95  			name:    "child-policy-empty-config",
    96  			input:   `{"childPolicy": [{"": {}}]}`,
    97  			wantErr: "invalid loadBalancingConfig: no supported policies found in []",
    98  		},
    99  		{
   100  			name:    "child-policy-type-isn't-registered",
   101  			input:   `{"childPolicy": [{"doesNotExistBalancer": {"cluster": "test_cluster"}}]}`,
   102  			wantErr: "invalid loadBalancingConfig: no supported policies found in [doesNotExistBalancer]",
   103  		},
   104  		{
   105  			name:    "child-policy-config-is-invalid",
   106  			input:   `{"childPolicy": [{"errParseConfigBalancer": {"cluster": "test_cluster"}}]}`,
   107  			wantErr: "error parsing loadBalancingConfig for policy \"errParseConfigBalancer\"",
   108  		},
   109  	}
   110  	for _, test := range tests {
   111  		t.Run(test.name, func(t *testing.T) {
   112  			gotCfg, gotErr := parser.ParseConfig(json.RawMessage(test.input))
   113  			// Substring match makes this very tightly coupled to the
   114  			// internalserviceconfig.BalancerConfig error strings. However, it
   115  			// is important to distinguish the different types of error messages
   116  			// possible as the parser has a few defined buckets of ways it can
   117  			// error out.
   118  			if (gotErr != nil) != (test.wantErr != "") {
   119  				t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
   120  			}
   121  			if gotErr != nil && !strings.Contains(gotErr.Error(), test.wantErr) {
   122  				t.Fatalf("ParseConfig(%v) = %v, wantErr %v", test.input, gotErr, test.wantErr)
   123  			}
   124  			if test.wantErr != "" {
   125  				return
   126  			}
   127  			if diff := cmp.Diff(gotCfg, test.wantCfg); diff != "" {
   128  				t.Fatalf("ParseConfig(%v) got unexpected output, diff (-got +want): %v", test.input, diff)
   129  			}
   130  		})
   131  	}
   132  }
   133  
   134  // TestUpdateClientConnState tests the UpdateClientConnState method of the
   135  // wrr_locality_experimental balancer. This UpdateClientConn operation should
   136  // take the localities and their weights in the addresses passed in, alongside
   137  // the endpoint picking policy defined in the Balancer Config and construct a
   138  // weighted target configuration corresponding to these inputs.
   139  func (s) TestUpdateClientConnState(t *testing.T) {
   140  	// Configure the stub balancer defined below as the child policy of
   141  	// wrrLocalityBalancer.
   142  	cfgCh := testutils.NewChannel()
   143  	oldWeightedTargetName := weightedTargetName
   144  	defer func() {
   145  		weightedTargetName = oldWeightedTargetName
   146  	}()
   147  	weightedTargetName = "fake_weighted_target"
   148  	stub.Register("fake_weighted_target", stub.BalancerFuncs{
   149  		ParseConfig: func(c json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
   150  			var cfg weightedtarget.LBConfig
   151  			if err := json.Unmarshal(c, &cfg); err != nil {
   152  				return nil, err
   153  			}
   154  			return &cfg, nil
   155  		},
   156  		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
   157  			wtCfg, ok := ccs.BalancerConfig.(*weightedtarget.LBConfig)
   158  			if !ok {
   159  				return errors.New("child received config that was not a weighted target config")
   160  			}
   161  			defer cfgCh.Send(wtCfg)
   162  			return nil
   163  		},
   164  	})
   165  
   166  	builder := balancer.Get(Name)
   167  	if builder == nil {
   168  		t.Fatalf("balancer.Get(%q) returned nil", Name)
   169  	}
   170  	tcc := testutils.NewBalancerClientConn(t)
   171  	bal := builder.Build(tcc, balancer.BuildOptions{})
   172  	defer bal.Close()
   173  	wrrL := bal.(*wrrLocalityBalancer)
   174  
   175  	// Create the addresses with two localities with certain locality weights.
   176  	// This represents what addresses the wrr_locality balancer will receive in
   177  	// UpdateClientConnState.
   178  	addr1 := resolver.Address{
   179  		Addr: "locality-1",
   180  	}
   181  	addr1 = internal.SetLocalityID(addr1, internal.LocalityID{
   182  		Region:  "region-1",
   183  		Zone:    "zone-1",
   184  		SubZone: "subzone-1",
   185  	})
   186  	addr1 = SetAddrInfo(addr1, AddrInfo{LocalityWeight: 2})
   187  
   188  	addr2 := resolver.Address{
   189  		Addr: "locality-2",
   190  	}
   191  	addr2 = internal.SetLocalityID(addr2, internal.LocalityID{
   192  		Region:  "region-2",
   193  		Zone:    "zone-2",
   194  		SubZone: "subzone-2",
   195  	})
   196  	addr2 = SetAddrInfo(addr2, AddrInfo{LocalityWeight: 1})
   197  	addrs := []resolver.Address{addr1, addr2}
   198  
   199  	err := wrrL.UpdateClientConnState(balancer.ClientConnState{
   200  		BalancerConfig: &LBConfig{
   201  			ChildPolicy: &internalserviceconfig.BalancerConfig{
   202  				Name: "round_robin",
   203  			},
   204  		},
   205  		ResolverState: resolver.State{
   206  			Addresses: addrs,
   207  		},
   208  	})
   209  	if err != nil {
   210  		t.Fatalf("Unexpected error from UpdateClientConnState: %v", err)
   211  	}
   212  
   213  	// Note that these inline strings declared as the key in Targets built from
   214  	// Locality ID are not exactly what is shown in the example in the gRFC.
   215  	// However, this is an implementation detail that does not affect
   216  	// correctness (confirmed with Java team). The important thing is to get
   217  	// those three pieces of information region, zone, and subzone down to the
   218  	// child layer.
   219  	wantWtCfg := &weightedtarget.LBConfig{
   220  		Targets: map[string]weightedtarget.Target{
   221  			"{\"region\":\"region-1\",\"zone\":\"zone-1\",\"subZone\":\"subzone-1\"}": {
   222  				Weight: 2,
   223  				ChildPolicy: &internalserviceconfig.BalancerConfig{
   224  					Name: "round_robin",
   225  				},
   226  			},
   227  			"{\"region\":\"region-2\",\"zone\":\"zone-2\",\"subZone\":\"subzone-2\"}": {
   228  				Weight: 1,
   229  				ChildPolicy: &internalserviceconfig.BalancerConfig{
   230  					Name: "round_robin",
   231  				},
   232  			},
   233  		},
   234  	}
   235  
   236  	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
   237  	defer cancel()
   238  	cfg, err := cfgCh.Receive(ctx)
   239  	if err != nil {
   240  		t.Fatalf("No signal received from UpdateClientConnState() on the child: %v", err)
   241  	}
   242  
   243  	gotWtCfg, ok := cfg.(*weightedtarget.LBConfig)
   244  	if !ok {
   245  		// Shouldn't happen - only sends a config on this channel.
   246  		t.Fatalf("Unexpected config type: %T", gotWtCfg)
   247  	}
   248  
   249  	if diff := cmp.Diff(gotWtCfg, wantWtCfg); diff != "" {
   250  		t.Fatalf("Child received unexpected config, diff (-got, +want): %v", diff)
   251  	}
   252  }
   253
View as plain text