github.com/stackdocker/rkt@v0.10.1-0.20151109095037-1aa827478248/Godeps/_workspace/src/google.golang.org/grpc/benchmark/client/main.go (about) 1 package main 2 3 import ( 4 "flag" 5 "math" 6 "net" 7 "net/http" 8 _ "net/http/pprof" 9 "sync" 10 "time" 11 12 "github.com/coreos/rkt/Godeps/_workspace/src/golang.org/x/net/context" 13 "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc" 14 "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/benchmark" 15 testpb "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/benchmark/grpc_testing" 16 "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/benchmark/stats" 17 "github.com/coreos/rkt/Godeps/_workspace/src/google.golang.org/grpc/grpclog" 18 ) 19 20 var ( 21 server = flag.String("server", "", "The server address") 22 maxConcurrentRPCs = flag.Int("max_concurrent_rpcs", 1, "The max number of concurrent RPCs") 23 duration = flag.Int("duration", math.MaxInt32, "The duration in seconds to run the benchmark client") 24 trace = flag.Bool("trace", true, "Whether tracing is on") 25 rpcType = flag.Int("rpc_type", 0, 26 `Configure different client rpc type. Valid options are: 27 0 : unary call; 28 1 : streaming call.`) 29 ) 30 31 func unaryCaller(client testpb.TestServiceClient) { 32 benchmark.DoUnaryCall(client, 1, 1) 33 } 34 35 func streamCaller(client testpb.TestServiceClient, stream testpb.TestService_StreamingCallClient) { 36 benchmark.DoStreamingRoundTrip(client, stream, 1, 1) 37 } 38 39 func buildConnection() (s *stats.Stats, conn *grpc.ClientConn, tc testpb.TestServiceClient) { 40 s = stats.NewStats(256) 41 conn = benchmark.NewClientConn(*server) 42 tc = testpb.NewTestServiceClient(conn) 43 return s, conn, tc 44 } 45 46 func closeLoopUnary() { 47 s, conn, tc := buildConnection() 48 49 for i := 0; i < 100; i++ { 50 unaryCaller(tc) 51 } 52 ch := make(chan int, *maxConcurrentRPCs*4) 53 var ( 54 mu sync.Mutex 55 wg sync.WaitGroup 56 ) 57 wg.Add(*maxConcurrentRPCs) 58 59 for i := 0; i < *maxConcurrentRPCs; i++ { 60 go func() { 61 for _ = range ch { 62 start := time.Now() 63 unaryCaller(tc) 64 elapse := time.Since(start) 65 mu.Lock() 66 s.Add(elapse) 67 mu.Unlock() 68 } 69 wg.Done() 70 }() 71 } 72 // Stop the client when time is up. 73 done := make(chan struct{}) 74 go func() { 75 <-time.After(time.Duration(*duration) * time.Second) 76 close(done) 77 }() 78 ok := true 79 for ok { 80 select { 81 case ch <- 0: 82 case <-done: 83 ok = false 84 } 85 } 86 close(ch) 87 wg.Wait() 88 conn.Close() 89 grpclog.Println(s.String()) 90 91 } 92 93 func closeLoopStream() { 94 s, conn, tc := buildConnection() 95 stream, err := tc.StreamingCall(context.Background()) 96 if err != nil { 97 grpclog.Fatalf("%v.StreamingCall(_) = _, %v", tc, err) 98 } 99 for i := 0; i < 100; i++ { 100 streamCaller(tc, stream) 101 } 102 ch := make(chan int, *maxConcurrentRPCs*4) 103 var ( 104 mu sync.Mutex 105 wg sync.WaitGroup 106 ) 107 wg.Add(*maxConcurrentRPCs) 108 // Distribute RPCs over maxConcurrentCalls workers. 109 for i := 0; i < *maxConcurrentRPCs; i++ { 110 go func() { 111 for _ = range ch { 112 start := time.Now() 113 streamCaller(tc, stream) 114 elapse := time.Since(start) 115 mu.Lock() 116 s.Add(elapse) 117 mu.Unlock() 118 } 119 wg.Done() 120 }() 121 } 122 // Stop the client when time is up. 123 done := make(chan struct{}) 124 go func() { 125 <-time.After(time.Duration(*duration) * time.Second) 126 close(done) 127 }() 128 ok := true 129 for ok { 130 select { 131 case ch <- 0: 132 case <-done: 133 ok = false 134 } 135 } 136 close(ch) 137 wg.Wait() 138 conn.Close() 139 grpclog.Println(s.String()) 140 } 141 142 func main() { 143 flag.Parse() 144 grpc.EnableTracing = *trace 145 go func() { 146 lis, err := net.Listen("tcp", ":0") 147 if err != nil { 148 grpclog.Fatalf("Failed to listen: %v", err) 149 } 150 grpclog.Println("Client profiling address: ", lis.Addr().String()) 151 if err := http.Serve(lis, nil); err != nil { 152 grpclog.Fatalf("Failed to serve: %v", err) 153 } 154 }() 155 switch *rpcType { 156 case 0: 157 closeLoopUnary() 158 case 1: 159 closeLoopStream() 160 } 161 }