Fix tests

balancer/endpointsharding/endpointsharding.go

@@ -86,7 +86,6 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 	if len(state.ResolverState.Endpoints) == 0 {
 		return errors.New("endpoints list is empty")
 	}
-
 	// Check/return early if any endpoints have no addresses.
 	// TODO: make this configurable if needed.
 	for i, endpoint := range state.ResolverState.Endpoints {
@@ -147,7 +146,6 @@ func (es *endpointSharding) UpdateClientConnState(state balancer.ClientConnState
 		}
 	}
 	es.children.Store(newChildren)
-
 	return ret
 }
 
@@ -233,7 +231,6 @@ func (es *endpointSharding) updateState() {
 		aggState = connectivity.TransientFailure
 		pickers = []balancer.Picker{base.NewErrPicker(errors.New("no children to pick from"))}
 	} // No children (resolver error before valid update).
-
 	p := &pickerWithChildStates{
 		pickers:     pickers,
 		childStates: childStates,

examples/examples_test.sh

@@ -110,7 +110,7 @@ declare -A EXPECTED_SERVER_OUTPUT=(
     ["features/authz"]="unary echoing message \"hello world\""
     ["features/cancellation"]="server: error receiving from stream: rpc error: code = Canceled desc = context canceled"
     ["features/compression"]="UnaryEcho called with message \"compress\""
-    ["features/customloadbalancer"]="serving on localhost:20000"
+    ["features/customloadbalancer"]="serving on localhost:50051"
     ["features/deadline"]=""
     ["features/encryption/TLS"]=""
     ["features/error_details"]=""

examples/features/customloadbalancer/README.md

@@ -40,13 +40,13 @@ useful functionality, such as Sticky Transient Failure, Happy Eyeballs, and
 Health Checking.
 
 ```
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20001)
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20001)
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20000)
-this is examples/customloadbalancing (from localhost:20001)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50051)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50051)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50050)
+this is examples/customloadbalancing (from localhost:50051)
 ```

examples/features/customloadbalancer/client/customroundrobin/customroundrobin.go

@@ -23,20 +23,16 @@ import (
 	"fmt"
 	"sync/atomic"
 
+	_ "google.golang.org/grpc" // to register pick_first
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/balancer/endpointsharding"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/serviceconfig"
 )
 
-var gracefulSwitchPickFirst serviceconfig.LoadBalancingConfig
-
 func init() {
 	balancer.Register(customRoundRobinBuilder{})
-	// Hardcode a pick first with no shuffling, since this is a petiole, and
-	// that is what petiole policies will interact with.
-	gracefulSwitchPickFirst, _ = endpointsharding.ParseConfig(json.RawMessage(endpointsharding.PickFirstConfig))
 }
 
 const customRRName = "custom_round_robin"
@@ -103,6 +99,7 @@ func (crr *customRoundRobin) UpdateClientConnState(state balancer.ClientConnStat
 	// A call to UpdateClientConnState should always produce a new Picker.  That
 	// is guaranteed to happen since the aggregator will always call
 	// UpdateChildState in its UpdateClientConnState.
+	gracefulSwitchPickFirst, _ := endpointsharding.ParseConfig(json.RawMessage(endpointsharding.PickFirstConfig))
 	return crr.Balancer.UpdateClientConnState(balancer.ClientConnState{
 		BalancerConfig: gracefulSwitchPickFirst,
 		ResolverState:  state.ResolverState,

examples/features/customloadbalancer/client/main.go

@@ -28,7 +28,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	_ "google.golang.org/grpc/examples/features/customloadbalancer/client/customroundrobin" // To register custom_round_robin.
-	"google.golang.org/grpc/examples/features/proto/echo"
+	pb "google.golang.org/grpc/examples/features/proto/echo"
 	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/resolver"
@@ -37,8 +37,8 @@ import (
 )
 
 var (
-	addr1 = "localhost:20000"
-	addr2 = "localhost:20001"
+	addr1 = "localhost:50050"
+	addr2 = "localhost:50051"
 )
 
 func main() {
@@ -64,7 +64,7 @@ func main() {
 	}
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
 	defer cancel()
-	ec := echo.NewEchoClient(cc)
+	ec := pb.NewEchoClient(cc)
 	if err := waitForDistribution(ctx, ec); err != nil {
 		log.Fatalf(err.Error())
 	}
@@ -74,7 +74,7 @@ func main() {
 // waitForDistribution makes RPC's on the echo client until 3 RPC's follow the
 // same 1:2 address ratio for the peer. Returns an error if fails to do so
 // before context timeout.
-func waitForDistribution(ctx context.Context, ec echo.EchoClient) error {
+func waitForDistribution(ctx context.Context, ec pb.EchoClient) error {
 	for {
 		results := make(map[string]uint32)
 	InnerLoop:
@@ -87,13 +87,13 @@ func waitForDistribution(ctx context.Context, ec echo.EchoClient) error {
 				res := make(map[string]uint32)
 				for j := 0; j < 3; j++ {
 					var peer peer.Peer
-					r, err := ec.UnaryEcho(ctx, &echo.EchoRequest{Message: "this is examples/customloadbalancing"}, grpc.Peer(&peer))
+					r, err := ec.UnaryEcho(ctx, &pb.EchoRequest{Message: "this is examples/customloadbalancing"}, grpc.Peer(&peer))
 					if err != nil {
 						return fmt.Errorf("UnaryEcho failed: %v", err)
 					}
 					fmt.Println(r)
 					peerAddr := peer.Addr.String()
-					if strings.HasSuffix(peerAddr, addr1) || strings.HasSuffix(peerAddr, addr2) {
+					if !strings.HasSuffix(peerAddr, "50050") && !strings.HasSuffix(peerAddr, "50051") {
 						return fmt.Errorf("peer address was not one of %v or %v, got: %v", addr1, addr2, peerAddr)
 					}
 					res[peerAddr]++

examples/features/customloadbalancer/server/main.go

@@ -26,20 +26,20 @@ import (
 	"sync"
 
 	"google.golang.org/grpc"
-	"google.golang.org/grpc/examples/features/proto/echo"
+	pb "google.golang.org/grpc/examples/features/proto/echo"
 )
 
 var (
-	addrs = []string{"localhost:20000", "localhost:20001"}
+	addrs = []string{"localhost:50050", "localhost:50051"}
 )
 
 type echoServer struct {
-	echo.UnimplementedEchoServer
+	pb.UnimplementedEchoServer
 	addr string
 }
 
-func (s *echoServer) UnaryEcho(ctx context.Context, req *echo.EchoRequest) (*echo.EchoResponse, error) {
-	return &echo.EchoResponse{Message: fmt.Sprintf("%s (from %s)", req.Message, s.addr)}, nil
+func (s *echoServer) UnaryEcho(ctx context.Context, req *pb.EchoRequest) (*pb.EchoResponse, error) {
+	return &pb.EchoResponse{Message: fmt.Sprintf("%s (from %s)", req.Message, s.addr)}, nil
 }
 
 func main() {
@@ -50,7 +50,7 @@ func main() {
 			log.Fatalf("failed to listen: %v", err)
 		}
 		s := grpc.NewServer()
-		echo.RegisterEchoServer(s, &echoServer{
+		pb.RegisterEchoServer(s, &echoServer{
 			addr: addr,
 		})
 		log.Printf("serving on %s\n", addr)

internal/balancer/gracefulswitch/config.go

@@ -75,7 +75,6 @@ func ParseConfig(cfg json.RawMessage) (serviceconfig.LoadBalancingConfig, error)
 		if err != nil {
 			return nil, fmt.Errorf("error parsing config for policy %q: %v", name, err)
 		}
-
 		return &lbConfig{childBuilder: builder, childConfig: cfg}, nil
 	}
 

internal/balancer/gracefulswitch/gracefulswitch.go

@@ -169,7 +169,6 @@ func (gsb *Balancer) latestBalancer() *balancerWrapper {
 func (gsb *Balancer) UpdateClientConnState(state balancer.ClientConnState) error {
 	// The resolver data is only relevant to the most recent LB Policy.
 	balToUpdate := gsb.latestBalancer()
-
 	gsbCfg, ok := state.BalancerConfig.(*lbConfig)
 	if ok {
 		// Switch to the child in the config unless it is already active.