Method TestBalancerSwitch_Graceful

test/balancer_switching_test.go:455–530

TestBalancerSwitch_Graceful tests the graceful switching of LB policies. It starts off by configuring "round_robin" on the channel and ensures that RPCs are successful. Then, it switches to a stub balancer which does not report a picker until instructed by the test do to so. At this point, the test

(t *testing.T)

// asks the new balancer to report a healthy picker and the test verifies that
// the RPCs get routed using the picker reported by the new balancer.
func (s) TestBalancerSwitch_Graceful(t *testing.T) {
	backends, cleanup := startBackendsForBalancerSwitch(t)
	defer cleanup()
	addrs := stubBackendsToResolverAddrs(backends)

	r := manual.NewBuilderWithScheme("whatever")
	cc, err := grpc.NewClient(r.Scheme()+":///test.server", grpc.WithTransportCredentials(insecure.NewCredentials()), grpc.WithResolvers(r))
	if err != nil {
		t.Fatalf("grpc.NewClient() failed: %v", err)
	}
	defer cc.Close()
	cc.Connect()
	// Push a resolver update with the service config specifying "round_robin".
	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
	defer cancel()
	r.UpdateState(resolver.State{
		Addresses:     addrs[1:],
		ServiceConfig: parseServiceConfig(t, r, rrServiceConfig),
	})
	client := testgrpc.NewTestServiceClient(cc)
	if err := rrutil.CheckRoundRobinRPCs(ctx, client, addrs[1:]); err != nil {
		t.Fatal(err)
	}

	// Register a stub balancer which uses a "pick_first" balancer underneath and
	// signals on a channel when it receives ClientConn updates. But it does not
	// forward the ccUpdate to the underlying "pick_first" balancer until the test
	// asks it to do so. This allows us to test the graceful switch functionality.
	// Until the test asks the stub balancer to forward the ccUpdate, RPCs should
	// get routed to the old balancer. And once the test gives the go ahead, RPCs
	// should get routed to the new balancer.
	ccUpdateCh := make(chan struct{})
	waitToProceed := make(chan struct{})
	stub.Register(t.Name(), stub.BalancerFuncs{
		Init: func(bd *stub.BalancerData) {
			pf := balancer.Get(pickfirst.Name)
			bd.ChildBalancer = pf.Build(bd.ClientConn, bd.BuildOptions)
		},
		Close: func(bd *stub.BalancerData) {
			bd.ChildBalancer.Close()
		},
		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
			close(ccUpdateCh)
			go func() {
				<-waitToProceed
				bd.ChildBalancer.UpdateClientConnState(ccs)
			}()
			return nil
		},
	})

	// Push a resolver update with the service config specifying our stub
	// balancer. We should see a trace event for this balancer switch. But RPCs
	// should still be routed to the old balancer since our stub balancer does not
	// report a ready picker until we ask it to do so.
	r.UpdateState(resolver.State{
		Addresses:     addrs[:1],
		ServiceConfig: r.CC().ParseServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%v": {}}]}`, t.Name())),