func closeLoop() { s := stats.NewStats(256) conn := benchmark.NewClientConn(*server) tc := testpb.NewTestServiceClient(conn) // Warm up connection. for i := 0; i < 100; i++ { caller(tc) } ch := make(chan int, *maxConcurrentRPCs*4) var ( mu sync.Mutex wg sync.WaitGroup ) wg.Add(*maxConcurrentRPCs) // Distribute RPCs over maxConcurrentCalls workers. for i := 0; i < *maxConcurrentRPCs; i++ { go func() { for _ = range ch { start := time.Now() caller(tc) elapse := time.Since(start) mu.Lock() s.Add(elapse) mu.Unlock() } wg.Done() }() } // Stop the client when time is up. done := make(chan struct{}) go func() { <-time.After(time.Duration(*duration) * time.Second) close(done) }() ok := true for ok { select { case ch <- 0: case <-done: ok = false } } close(ch) wg.Wait() conn.Close() grpclog.Println(s.String()) }
func run(b *testing.B, maxConcurrentCalls int, caller func(testpb.TestServiceClient)) { s := stats.AddStats(b, 38) b.StopTimer() target, stopper := StartServer() defer stopper() conn := NewClientConn(target) tc := testpb.NewTestServiceClient(conn) // Warm up connection. for i := 0; i < 10; i++ { caller(tc) } ch := make(chan int, maxConcurrentCalls*4) var ( mu sync.Mutex wg sync.WaitGroup ) wg.Add(maxConcurrentCalls) // Distribute the b.N calls over maxConcurrentCalls workers. for i := 0; i < maxConcurrentCalls; i++ { go func() { for _ = range ch { start := time.Now() caller(tc) elapse := time.Since(start) mu.Lock() s.Add(elapse) mu.Unlock() } wg.Done() }() } b.StartTimer() for i := 0; i < b.N; i++ { ch <- i } b.StopTimer() close(ch) wg.Wait() conn.Close() }
func main() { flag.Parse() serverAddr := net.JoinHostPort(*serverHost, strconv.Itoa(*serverPort)) var opts []grpc.DialOption if *useTLS { var sn string if *tlsServerName != "" { sn = *tlsServerName } var creds credentials.TransportAuthenticator if *caFile != "" { var err error creds, err = credentials.NewClientTLSFromFile(*caFile, sn) if err != nil { grpclog.Fatalf("Failed to create TLS credentials %v", err) } } else { creds = credentials.NewClientTLSFromCert(nil, sn) } opts = append(opts, grpc.WithTransportCredentials(creds)) if *testCase == "compute_engine_creds" { opts = append(opts, grpc.WithPerRPCCredentials(credentials.NewComputeEngine())) } else if *testCase == "service_account_creds" { jwtCreds, err := credentials.NewServiceAccountFromFile(*serviceAccountKeyFile, *oauthScope) if err != nil { grpclog.Fatalf("Failed to create JWT credentials: %v", err) } opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds)) } } conn, err := grpc.Dial(serverAddr, opts...) if err != nil { grpclog.Fatalf("Fail to dial: %v", err) } defer conn.Close() tc := testpb.NewTestServiceClient(conn) switch *testCase { case "empty_unary": doEmptyUnaryCall(tc) case "large_unary": doLargeUnaryCall(tc) case "client_streaming": doClientStreaming(tc) case "server_streaming": doServerStreaming(tc) case "ping_pong": doPingPong(tc) case "compute_engine_creds": if !*useTLS { grpclog.Fatalf("TLS is not enabled. TLS is required to execute compute_engine_creds test case.") } doComputeEngineCreds(tc) case "service_account_creds": if !*useTLS { grpclog.Fatalf("TLS is not enabled. TLS is required to execute service_account_creds test case.") } doServiceAccountCreds(tc) case "cancel_after_begin": doCancelAfterBegin(tc) case "cancel_after_first_response": doCancelAfterFirstResponse(tc) default: grpclog.Fatal("Unsupported test case: ", *testCase) } }