endpointsharding_test.go

Documentation: google.golang.org/grpc/balancer/endpointsharding

     1  /*
     2   *
     3   * Copyright 2024 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 endpointsharding
    20  
    21  import (
    22  	"context"
    23  	"encoding/json"
    24  	"fmt"
    25  	"log"
    26  	"testing"
    27  	"time"
    28  
    29  	"google.golang.org/grpc"
    30  	"google.golang.org/grpc/balancer"
    31  	"google.golang.org/grpc/credentials/insecure"
    32  	"google.golang.org/grpc/grpclog"
    33  	"google.golang.org/grpc/internal"
    34  	"google.golang.org/grpc/internal/grpctest"
    35  	"google.golang.org/grpc/internal/stubserver"
    36  	"google.golang.org/grpc/internal/testutils/roundrobin"
    37  	"google.golang.org/grpc/resolver"
    38  	"google.golang.org/grpc/resolver/manual"
    39  	"google.golang.org/grpc/serviceconfig"
    40  
    41  	testgrpc "google.golang.org/grpc/interop/grpc_testing"
    42  )
    43  
    44  type s struct {
    45  	grpctest.Tester
    46  }
    47  
    48  func Test(t *testing.T) {
    49  	grpctest.RunSubTests(t, s{})
    50  }
    51  
    52  var gracefulSwitchPickFirst serviceconfig.LoadBalancingConfig
    53  
    54  var logger = grpclog.Component("endpoint-sharding-test")
    55  
    56  func init() {
    57  	var err error
    58  	gracefulSwitchPickFirst, err = ParseConfig(json.RawMessage(PickFirstConfig))
    59  	if err != nil {
    60  		logger.Fatal(err)
    61  	}
    62  	balancer.Register(fakePetioleBuilder{})
    63  }
    64  
    65  const fakePetioleName = "fake_petiole"
    66  
    67  type fakePetioleBuilder struct{}
    68  
    69  func (fakePetioleBuilder) Name() string {
    70  	return fakePetioleName
    71  }
    72  
    73  func (fakePetioleBuilder) Build(cc balancer.ClientConn, opts balancer.BuildOptions) balancer.Balancer {
    74  	fp := &fakePetiole{
    75  		ClientConn: cc,
    76  		bOpts:      opts,
    77  	}
    78  	fp.Balancer = NewBalancer(fp, opts)
    79  	return fp
    80  }
    81  
    82  func (fakePetioleBuilder) ParseConfig(s json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
    83  	return nil, nil
    84  }
    85  
    86  // fakePetiole is a load balancer that wraps the endpointShardingBalancer, and
    87  // forwards ClientConnUpdates with a child config of graceful switch that wraps
    88  // pick first. It also intercepts UpdateState to make sure it can access the
    89  // child state maintained by EndpointSharding.
    90  type fakePetiole struct {
    91  	balancer.Balancer
    92  	balancer.ClientConn
    93  	bOpts balancer.BuildOptions
    94  }
    95  
    96  func (fp *fakePetiole) UpdateClientConnState(state balancer.ClientConnState) error {
    97  	if el := state.ResolverState.Endpoints; len(el) != 2 {
    98  		return fmt.Errorf("UpdateClientConnState wants two endpoints, got: %v", el)
    99  	}
   100  
   101  	return fp.Balancer.UpdateClientConnState(balancer.ClientConnState{
   102  		BalancerConfig: gracefulSwitchPickFirst,
   103  		ResolverState:  state.ResolverState,
   104  	})
   105  }
   106  
   107  func (fp *fakePetiole) UpdateState(state balancer.State) {
   108  	childStates := ChildStatesFromPicker(state.Picker)
   109  	// Both child states should be present in the child picker. States and
   110  	// picker change over the lifecycle of test, but there should always be two.
   111  	if len(childStates) != 2 {
   112  		logger.Fatal(fmt.Errorf("length of child states received: %v, want 2", len(childStates)))
   113  	}
   114  
   115  	fp.ClientConn.UpdateState(state)
   116  }
   117  
   118  // TestEndpointShardingBasic tests the basic functionality of the endpoint
   119  // sharding balancer. It specifies a petiole policy that is essentially a
   120  // wrapper around the endpoint sharder. Two backends are started, with each
   121  // backend's address specified in an endpoint. The petiole does not have a
   122  // special picker, so it should fallback to the default behavior, which is to
   123  // round_robin amongst the endpoint children that are in the aggregated state.
   124  // It also verifies the petiole has access to the raw child state in case it
   125  // wants to implement a custom picker.
   126  func (s) TestEndpointShardingBasic(t *testing.T) {
   127  	backend1 := stubserver.StartTestService(t, nil)
   128  	defer backend1.Stop()
   129  	backend2 := stubserver.StartTestService(t, nil)
   130  	defer backend2.Stop()
   131  
   132  	mr := manual.NewBuilderWithScheme("e2e-test")
   133  	defer mr.Close()
   134  
   135  	json := `{"loadBalancingConfig": [{"fake_petiole":{}}]}`
   136  	sc := internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(json)
   137  	mr.InitialState(resolver.State{
   138  		Endpoints: []resolver.Endpoint{
   139  			{Addresses: []resolver.Address{{Addr: backend1.Address}}},
   140  			{Addresses: []resolver.Address{{Addr: backend2.Address}}},
   141  		},
   142  		ServiceConfig: sc,
   143  	})
   144  
   145  	cc, err := grpc.Dial(mr.Scheme()+":///", grpc.WithResolvers(mr), grpc.WithTransportCredentials(insecure.NewCredentials()))
   146  	if err != nil {
   147  		log.Fatalf("Failed to dial: %v", err)
   148  	}
   149  	defer cc.Close()
   150  	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
   151  	defer cancel()
   152  	client := testgrpc.NewTestServiceClient(cc)
   153  	// Assert a round robin distribution between the two spun up backends. This
   154  	// requires a poll and eventual consistency as both endpoint children do not
   155  	// start in state READY.
   156  	if err = roundrobin.CheckRoundRobinRPCs(ctx, client, []resolver.Address{{Addr: backend1.Address}, {Addr: backend2.Address}}); err != nil {
   157  		t.Fatalf("error in expected round robin: %v", err)
   158  	}
   159  }
   160
View as plain text