1 /* 2 * 3 * Copyright 2022 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 pickfirst contains helper functions to check for pickfirst load 20 // balancing of RPCs in tests. 21 package pickfirst 22 23 import ( 24 "context" 25 "fmt" 26 "time" 27 28 "google.golang.org/grpc" 29 "google.golang.org/grpc/peer" 30 "google.golang.org/grpc/resolver" 31 32 testgrpc "google.golang.org/grpc/interop/grpc_testing" 33 testpb "google.golang.org/grpc/interop/grpc_testing" 34 ) 35 36 // CheckRPCsToBackend makes a bunch of RPCs on the given ClientConn and verifies 37 // if the RPCs are routed to a peer matching wantAddr. 38 // 39 // Returns a non-nil error if context deadline expires before all RPCs begin to 40 // be routed to the peer matching wantAddr, or if the backend returns RPC errors. 41 func CheckRPCsToBackend(ctx context.Context, cc *grpc.ClientConn, wantAddr resolver.Address) error { 42 client := testgrpc.NewTestServiceClient(cc) 43 peer := &peer.Peer{} 44 // Make sure that 20 RPCs in a row reach the expected backend. Some 45 // tests switch from round_robin back to pick_first and call this 46 // function. None of our tests spin up more than 10 backends. So, 47 // waiting for 20 RPCs to reach a single backend would a decent 48 // indicator of having switched to pick_first. 49 count := 0 50 for { 51 time.Sleep(time.Millisecond) 52 if ctx.Err() != nil { 53 return fmt.Errorf("timeout waiting for RPC to be routed to %s", wantAddr.Addr) 54 } 55 if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(peer)); err != nil { 56 // Some tests remove backends and check if pick_first is happening across 57 // the remaining backends. In such cases, RPCs can initially fail on the 58 // connection using the removed backend. Just keep retrying and eventually 59 // the connection using the removed backend will shutdown and will be 60 // removed. 61 continue 62 } 63 if peer.Addr.String() != wantAddr.Addr { 64 count = 0 65 continue 66 } 67 count++ 68 if count > 20 { 69 break 70 } 71 } 72 // Make sure subsequent RPCs are all routed to the same backend. 73 for i := 0; i < 10; i++ { 74 if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(peer)); err != nil { 75 return fmt.Errorf("EmptyCall() = %v, want <nil>", err) 76 } 77 if gotAddr := peer.Addr.String(); gotAddr != wantAddr.Addr { 78 return fmt.Errorf("rpc sent to peer %q, want peer %q", gotAddr, wantAddr) 79 } 80 } 81 return nil 82 } 83