Repository: grpc/grpc-go
File: xds/internal/balancer/clusterimpl/tests/balancer_test.go (Go)

Code before:
		t.Fatalf("Failure while waiting for load to be reported: %v", err)
	}

	loadStats := req.(*fakeserver.Request).Req.(*v3lrspb.LoadStatsRequest)

	// As pickfirst is configured, only one server should receive all the requests.
	// Take down the server which received all the requests initially.
	var initialRegion string
	for _, load := range loadStats.ClusterStats {
		for _, locality := range load.UpstreamLocalityStats {
			if locality.TotalSuccessfulRequests != 11 {
				continue
			}
			initialRegion = locality.Locality.Region
		}
	}

	if initialRegion == "" {
		t.Fatalf("None of the servers received all the requests based on load stats: %v", loadStats)
	}

	if initialRegion == "region-1" {
		server1.Stop()
	} else {
		server2.Stop()
	}

	// Send 10 more RPCs, now pickfirst will send them to a different server.
	for i := 0; i < 10; i++ {
		client.EmptyCall(ctx, &testpb.Empty{})
	}

	// Trigger load reporting.
	tickChan <- time.Now()

	req, err = mgmtServer.LRSServer.LRSRequestChan.Receive(ctx)
	if err != nil {
		t.Error(err)
	}

	loadStats = req.(*fakeserver.Request).Req.(*v3lrspb.LoadStatsRequest)
	var finalRegion string
	for _, load := range loadStats.ClusterStats {
		for _, locality := range load.UpstreamLocalityStats {
			if locality.TotalSuccessfulRequests != 10 {
				continue
			}
			finalRegion = locality.Locality.Region
		}
	}


Review comment (suggestion):
Given that we are creating an endpoints resource with `locality1` containing `port1` and `locality2` containing `port2`, we should expect the address list passed to `pickfirst` to be `[localhost:port1, localhost:port2]`. Is that not the case? If so, we need to figure out why?

If that is the case, we expect traffic to be sent to `localhost:port1`, but given our current implementation loads to be reported on `locality2`. We should verify both of these with a comment saying, once pickfirst is changed to create one subchannel per endpoint, we should no longer need this.

You can verify where traffic is being sent using the `peer.Peer` call option.

And once you verify that traffic is being sent to `localhost:port1`, you can stop `server1`.

Code after:
	server1.Stop()

	// Wait for the balancer to pick up the server state change.
	testutils.AwaitState(ctx, t, cc, connectivity.Idle)
	if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
		t.Fatalf("rpc EmptyCall() failed: %v", err)
	}

	// Verify that the request was sent to server 2.
	if got, want := peer.Addr.String(), server2.Address; got != want {
		t.Errorf("peer.Addr = %q, want = %q", got, want)
	}

	// Wait for load to be reported for locality of server 2.
	if err := waitForSuccessfulLoadReport(ctx, mgmtServer.LRSServer, "region-2"); err != nil {
		t.Fatalf("Server 2 did not receive load due to error: %v", err)
	}
}

// waitForSuccessfulLoadReport waits for a successful request to be reported for
// the specified locality region.
func waitForSuccessfulLoadReport(ctx context.Context, lrsServer *fakeserver.Server, region string) error {
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case req := <-lrsServer.LRSRequestChan.C:
			loadStats := req.(*fakeserver.Request).Req.(*v3lrspb.LoadStatsRequest)
			for _, load := range loadStats.ClusterStats {
				for _, locality := range load.UpstreamLocalityStats {
					if locality.TotalSuccessfulRequests > 0 && locality.Locality.Region == region {
						return nil
					}
				}
			}
		}
	}
}