func doCancelAfterFirstResponse(tc testpb.TestServiceClient) {
ctx, cancel := context.WithCancel(context.Background())
stream, err := tc.FullDuplexCall(ctx)
if err != nil {
grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
}
respParam := []*testpb.ResponseParameters{
{
Size: proto.Int32(31415),
},
}
pl := newPayload(testpb.PayloadType_COMPRESSABLE, 27182)
req := &testpb.StreamingOutputCallRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseParameters: respParam,
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("%v.Send(%v) = %v", stream, req, err)
}
if _, err := stream.Recv(); err != nil {
grpclog.Fatalf("%v.Recv() = %v", stream, err)
}
cancel()
if _, err := stream.Recv(); grpc.Code(err) != codes.Canceled {
grpclog.Fatalf("%v compleled with error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
}
grpclog.Println("CancelAfterFirstResponse done")
}
func doPerRPCCreds(tc testpb.TestServiceClient) {
jsonKey := getServiceAccountJSONKey()
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
FillOauthScope: proto.Bool(true),
}
token := getToken()
kv := map[string]string{"authorization": token.TokenType + " " + token.AccessToken}
ctx := metadata.NewContext(context.Background(), metadata.MD{"authorization": []string{kv["authorization"]}})
reply, err := tc.UnaryCall(ctx, req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
user := reply.GetUsername()
scope := reply.GetOauthScope()
if !strings.Contains(string(jsonKey), user) {
grpclog.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
}
if !strings.Contains(*oauthScope, scope) {
grpclog.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, *oauthScope)
}
grpclog.Println("PerRPCCreds done")
}
func doClientStreaming(tc testpb.TestServiceClient) {
stream, err := tc.StreamingInputCall(context.Background())
if err != nil {
grpclog.Fatalf("%v.StreamingInputCall(_) = _, %v", tc, err)
}
var sum int
for _, s := range reqSizes {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, s)
req := &testpb.StreamingInputCallRequest{
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("%v.Send(%v) = %v", stream, req, err)
}
sum += s
grpclog.Printf("Sent a request of size %d, aggregated size %d", s, sum)
}
reply, err := stream.CloseAndRecv()
if err != nil {
grpclog.Fatalf("%v.CloseAndRecv() got error %v, want %v", stream, err, nil)
}
if reply.GetAggregatedPayloadSize() != int32(sum) {
grpclog.Fatalf("%v.CloseAndRecv().GetAggregatePayloadSize() = %v; want %v", stream, reply.GetAggregatedPayloadSize(), sum)
}
grpclog.Println("ClientStreaming done")
}
// loadFeatures loads features from a JSON file.
func (s *routeGuideServer) loadFeatures(filePath string) {
file, err := ioutil.ReadFile(filePath)
if err != nil {
grpclog.Fatalf("Failed to load default features: %v", err)
}
if err := json.Unmarshal(file, &s.savedFeatures); err != nil {
grpclog.Fatalf("Failed to load default features: %v", err)
}
}
func getToken() *oauth2.Token {
jsonKey := getServiceAccountJSONKey()
config, err := google.JWTConfigFromJSON(jsonKey, *oauthScope)
if err != nil {
grpclog.Fatalf("Failed to get the config: %v", err)
}
token, err := config.TokenSource(context.Background()).Token()
if err != nil {
grpclog.Fatalf("Failed to get the token: %v", err)
}
return token
}
func doCancelAfterBegin(tc testpb.TestServiceClient) {
ctx, cancel := context.WithCancel(metadata.NewContext(context.Background(), testMetadata))
stream, err := tc.StreamingInputCall(ctx)
if err != nil {
grpclog.Fatalf("%v.StreamingInputCall(_) = _, %v", tc, err)
}
cancel()
_, err = stream.CloseAndRecv()
if grpc.Code(err) != codes.Canceled {
grpclog.Fatalf("%v.CloseAndRecv() got error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
}
grpclog.Println("CancelAfterBegin done")
}
func doEmptyStream(tc testpb.TestServiceClient) {
stream, err := tc.FullDuplexCall(context.Background())
if err != nil {
grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
}
if err := stream.CloseSend(); err != nil {
grpclog.Fatalf("%v.CloseSend() got %v, want %v", stream, err, nil)
}
if _, err := stream.Recv(); err != io.EOF {
grpclog.Fatalf("%v failed to complete the empty stream test: %v", stream, err)
}
grpclog.Println("Emptystream done")
}
// DoStreamingRoundTrip performs a round trip for a single streaming rpc.
func DoStreamingRoundTrip(tc testpb.TestServiceClient, stream testpb.TestService_StreamingCallClient, reqSize, respSize int) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, reqSize)
req := &testpb.SimpleRequest{
ResponseType: pl.Type,
ResponseSize: int32(respSize),
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("StreamingCall(_).Send: %v", err)
}
if _, err := stream.Recv(); err != nil {
grpclog.Fatalf("StreamingCall(_).Recv: %v", err)
}
}
// NewClientConn creates a gRPC client connection to addr.
func NewClientConn(addr string) *grpc.ClientConn {
conn, err := grpc.Dial(addr, grpc.WithInsecure())
if err != nil {
grpclog.Fatalf("NewClientConn(%q) failed to create a ClientConn %v", addr, err)
}
return conn
}
func getServiceAccountJSONKey() []byte {
jsonKey, err := ioutil.ReadFile(*serviceAccountKeyFile)
if err != nil {
grpclog.Fatalf("Failed to read the service account key file: %v", err)
}
return jsonKey
}
func newPayload(t testpb.PayloadType, size int) *testpb.Payload {
if size < 0 {
grpclog.Fatalf("Requested a response with invalid length %d", size)
}
body := make([]byte, size)
switch t {
case testpb.PayloadType_COMPRESSABLE:
case testpb.PayloadType_UNCOMPRESSABLE:
grpclog.Fatalf("PayloadType UNCOMPRESSABLE is not supported")
default:
grpclog.Fatalf("Unsupported payload type: %d", t)
}
return &testpb.Payload{
Type: t,
Body: body,
}
}
// RegisterService register a service and its implementation to the gRPC
// server. Called from the IDL generated code. This must be called before
// invoking Serve.
func (s *Server) RegisterService(sd *ServiceDesc, ss interface{}) {
ht := reflect.TypeOf(sd.HandlerType).Elem()
st := reflect.TypeOf(ss)
if !st.Implements(ht) {
grpclog.Fatalf("grpc: Server.RegisterService found the handler of type %v that does not satisfy %v", st, ht)
}
s.register(sd, ss)
}
func main() {
flag.Parse()
go func() {
lis, err := net.Listen("tcp", ":0")
if err != nil {
grpclog.Fatalf("Failed to listen: %v", err)
}
grpclog.Println("Server profiling address: ", lis.Addr().String())
if err := http.Serve(lis, nil); err != nil {
grpclog.Fatalf("Failed to serve: %v", err)
}
}()
addr, stopper := benchmark.StartServer(":0") // listen on all interfaces
grpclog.Println("Server Address: ", addr)
<-time.After(time.Duration(*duration) * time.Second)
stopper()
}
func main() {
flag.Parse()
lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
if err != nil {
grpclog.Fatalf("failed to listen: %v", err)
}
var opts []grpc.ServerOption
if *tls {
creds, err := credentials.NewServerTLSFromFile(*certFile, *keyFile)
if err != nil {
grpclog.Fatalf("Failed to generate credentials %v", err)
}
opts = []grpc.ServerOption{grpc.Creds(creds)}
}
grpcServer := grpc.NewServer(opts...)
pb.RegisterRouteGuideServer(grpcServer, newServer())
grpcServer.Serve(lis)
}
func doServerStreaming(tc testpb.TestServiceClient) {
respParam := make([]*testpb.ResponseParameters, len(respSizes))
for i, s := range respSizes {
respParam[i] = &testpb.ResponseParameters{
Size: proto.Int32(int32(s)),
}
}
req := &testpb.StreamingOutputCallRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseParameters: respParam,
}
stream, err := tc.StreamingOutputCall(context.Background(), req)
if err != nil {
grpclog.Fatalf("%v.StreamingOutputCall(_) = _, %v", tc, err)
}
var rpcStatus error
var respCnt int
var index int
for {
reply, err := stream.Recv()
if err != nil {
rpcStatus = err
break
}
t := reply.GetPayload().GetType()
if t != testpb.PayloadType_COMPRESSABLE {
grpclog.Fatalf("Got the reply of type %d, want %d", t, testpb.PayloadType_COMPRESSABLE)
}
size := len(reply.GetPayload().GetBody())
if size != int(respSizes[index]) {
grpclog.Fatalf("Got reply body of length %d, want %d", size, respSizes[index])
}
index++
respCnt++
}
if rpcStatus != io.EOF {
grpclog.Fatalf("Failed to finish the server streaming rpc: %v", err)
}
if respCnt != len(respSizes) {
grpclog.Fatalf("Got %d reply, want %d", len(respSizes), respCnt)
}
grpclog.Println("ServerStreaming done")
}
func doEmptyUnaryCall(tc testpb.TestServiceClient) {
reply, err := tc.EmptyCall(context.Background(), &testpb.Empty{})
if err != nil {
grpclog.Fatal("/TestService/EmptyCall RPC failed: ", err)
}
if !proto.Equal(&testpb.Empty{}, reply) {
grpclog.Fatalf("/TestService/EmptyCall receives %v, want %v", reply, testpb.Empty{})
}
grpclog.Println("EmptyUnaryCall done")
}
func main() {
flag.Parse()
p := strconv.Itoa(*port)
lis, err := net.Listen("tcp", ":"+p)
if err != nil {
grpclog.Fatalf("failed to listen: %v", err)
}
var opts []grpc.ServerOption
if *useTLS {
creds, err := credentials.NewServerTLSFromFile(*certFile, *keyFile)
if err != nil {
grpclog.Fatalf("Failed to generate credentials %v", err)
}
opts = []grpc.ServerOption{grpc.Creds(creds)}
}
server := grpc.NewServer(opts...)
testpb.RegisterTestServiceServer(server, &testServer{})
server.Serve(lis)
}
func main() {
flag.Parse()
grpc.EnableTracing = *trace
go func() {
lis, err := net.Listen("tcp", ":0")
if err != nil {
grpclog.Fatalf("Failed to listen: %v", err)
}
grpclog.Println("Client profiling address: ", lis.Addr().String())
if err := http.Serve(lis, nil); err != nil {
grpclog.Fatalf("Failed to serve: %v", err)
}
}()
switch *rpcType {
case 0:
closeLoopUnary()
case 1:
closeLoopStream()
}
}
// StartServer starts a gRPC server serving a benchmark service on the given
// address, which may be something like "localhost:0". It returns its listen
// address and a function to stop the server.
func StartServer(addr string) (string, func()) {
lis, err := net.Listen("tcp", addr)
if err != nil {
grpclog.Fatalf("Failed to listen: %v", err)
}
s := grpc.NewServer(grpc.MaxConcurrentStreams(math.MaxUint32))
testpb.RegisterTestServiceServer(s, &testServer{})
go s.Serve(lis)
return lis.Addr().String(), func() {
s.Stop()
}
}
func doTimeoutOnSleepingServer(tc testpb.TestServiceClient) {
ctx, _ := context.WithTimeout(context.Background(), 1*time.Millisecond)
stream, err := tc.FullDuplexCall(ctx)
if err != nil {
if grpc.Code(err) == codes.DeadlineExceeded {
grpclog.Println("TimeoutOnSleepingServer done")
return
}
grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
}
pl := newPayload(testpb.PayloadType_COMPRESSABLE, 27182)
req := &testpb.StreamingOutputCallRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("%v.Send(%v) = %v", stream, req, err)
}
if _, err := stream.Recv(); grpc.Code(err) != codes.DeadlineExceeded {
grpclog.Fatalf("%v.Recv() = _, %v, want error code %d", stream, err, codes.DeadlineExceeded)
}
grpclog.Println("TimeoutOnSleepingServer done")
}
// SendHeader sends header metadata. It may be called at most once from a unary
// RPC handler. The ctx is the RPC handler's Context or one derived from it.
func SendHeader(ctx context.Context, md metadata.MD) error {
if md.Len() == 0 {
return nil
}
stream, ok := transport.StreamFromContext(ctx)
if !ok {
return fmt.Errorf("grpc: failed to fetch the stream from the context %v", ctx)
}
t := stream.ServerTransport()
if t == nil {
grpclog.Fatalf("grpc: SendHeader: %v has no ServerTransport to send header metadata.", stream)
}
return t.WriteHeader(stream, md)
}
func (s *Server) sendResponse(t transport.ServerTransport, stream *transport.Stream, msg interface{}, pf payloadFormat, opts *transport.Options) error {
p, err := encode(s.opts.codec, msg, pf)
if err != nil {
// This typically indicates a fatal issue (e.g., memory
// corruption or hardware faults) the application program
// cannot handle.
//
// TODO(zhaoq): There exist other options also such as only closing the
// faulty stream locally and remotely (Other streams can keep going). Find
// the optimal option.
grpclog.Fatalf("grpc: Server failed to encode response %v", err)
}
return t.Write(stream, p, opts)
}
func doComputeEngineCreds(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
FillOauthScope: proto.Bool(true),
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
user := reply.GetUsername()
scope := reply.GetOauthScope()
if user != *defaultServiceAccount {
grpclog.Fatalf("Got user name %q, want %q.", user, *defaultServiceAccount)
}
if !strings.Contains(*oauthScope, scope) {
grpclog.Fatalf("Got OAuth scope %q which is NOT a substring of %q.", scope, *oauthScope)
}
grpclog.Println("ComputeEngineCreds done")
}
func closeLoopStream() {
s, conn, tc := buildConnection()
stream, err := tc.StreamingCall(context.Background())
if err != nil {
grpclog.Fatalf("%v.StreamingCall(_) = _, %v", tc, err)
}
for i := 0; i < 100; i++ {
streamCaller(tc, stream)
}
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()
streamCaller(tc, stream)
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 doLargeUnaryCall(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
t := reply.GetPayload().GetType()
s := len(reply.GetPayload().GetBody())
if t != testpb.PayloadType_COMPRESSABLE || s != largeRespSize {
grpclog.Fatalf("Got the reply with type %d len %d; want %d, %d", t, s, testpb.PayloadType_COMPRESSABLE, largeRespSize)
}
grpclog.Println("LargeUnaryCall done")
}
func doJWTTokenCreds(tc testpb.TestServiceClient) {
pl := newPayload(testpb.PayloadType_COMPRESSABLE, largeReqSize)
req := &testpb.SimpleRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseSize: proto.Int32(int32(largeRespSize)),
Payload: pl,
FillUsername: proto.Bool(true),
}
reply, err := tc.UnaryCall(context.Background(), req)
if err != nil {
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
}
jsonKey := getServiceAccountJSONKey()
user := reply.GetUsername()
if !strings.Contains(string(jsonKey), user) {
grpclog.Fatalf("Got user name %q which is NOT a substring of %q.", user, jsonKey)
}
grpclog.Println("JWTtokenCreds done")
}
func doPingPong(tc testpb.TestServiceClient) {
stream, err := tc.FullDuplexCall(context.Background())
if err != nil {
grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
}
var index int
for index < len(reqSizes) {
respParam := []*testpb.ResponseParameters{
{
Size: proto.Int32(int32(respSizes[index])),
},
}
pl := newPayload(testpb.PayloadType_COMPRESSABLE, reqSizes[index])
req := &testpb.StreamingOutputCallRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseParameters: respParam,
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("%v.Send(%v) = %v", stream, req, err)
}
reply, err := stream.Recv()
if err != nil {
grpclog.Fatalf("%v.Recv() = %v", stream, err)
}
t := reply.GetPayload().GetType()
if t != testpb.PayloadType_COMPRESSABLE {
grpclog.Fatalf("Got the reply of type %d, want %d", t, testpb.PayloadType_COMPRESSABLE)
}
size := len(reply.GetPayload().GetBody())
if size != int(respSizes[index]) {
grpclog.Fatalf("Got reply body of length %d, want %d", size, respSizes[index])
}
index++
}
if err := stream.CloseSend(); err != nil {
grpclog.Fatalf("%v.CloseSend() got %v, want %v", stream, err, nil)
}
if _, err := stream.Recv(); err != io.EOF {
grpclog.Fatalf("%v failed to complele the ping pong test: %v", stream, err)
}
grpclog.Println("Pingpong done")
}
func (s *Server) register(sd *ServiceDesc, ss interface{}) {
s.mu.Lock()
defer s.mu.Unlock()
s.printf("RegisterService(%q)", sd.ServiceName)
if _, ok := s.m[sd.ServiceName]; ok {
grpclog.Fatalf("grpc: Server.RegisterService found duplicate service registration for %q", sd.ServiceName)
}
srv := &service{
server: ss,
md: make(map[string]*MethodDesc),
sd: make(map[string]*StreamDesc),
}
for i := range sd.Methods {
d := &sd.Methods[i]
srv.md[d.MethodName] = d
}
for i := range sd.Streams {
d := &sd.Streams[i]
srv.sd[d.StreamName] = d
}
s.m[sd.ServiceName] = srv
}
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 *testCA {
var err error
creds, err = credentials.NewClientTLSFromFile(testCAFile, 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(oauth.NewComputeEngine()))
} else if *testCase == "service_account_creds" {
jwtCreds, err := oauth.NewServiceAccountFromFile(*serviceAccountKeyFile, *oauthScope)
if err != nil {
grpclog.Fatalf("Failed to create JWT credentials: %v", err)
}
opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds))
} else if *testCase == "jwt_token_creds" {
jwtCreds, err := oauth.NewJWTAccessFromFile(*serviceAccountKeyFile)
if err != nil {
grpclog.Fatalf("Failed to create JWT credentials: %v", err)
}
opts = append(opts, grpc.WithPerRPCCredentials(jwtCreds))
} else if *testCase == "oauth2_auth_token" {
opts = append(opts, grpc.WithPerRPCCredentials(oauth.NewOauthAccess(getToken())))
}
} else {
opts = append(opts, grpc.WithInsecure())
}
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 "empty_stream":
doEmptyStream(tc)
case "timeout_on_sleeping_server":
doTimeoutOnSleepingServer(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 "jwt_token_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute jwt_token_creds test case.")
}
doJWTTokenCreds(tc)
case "per_rpc_creds":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute per_rpc_creds test case.")
}
doPerRPCCreds(tc)
case "oauth2_auth_token":
if !*useTLS {
grpclog.Fatalf("TLS is not enabled. TLS is required to execute oauth2_auth_token test case.")
}
doOauth2TokenCreds(tc)
case "cancel_after_begin":
doCancelAfterBegin(tc)
case "cancel_after_first_response":
doCancelAfterFirstResponse(tc)
default:
grpclog.Fatal("Unsupported test case: ", *testCase)
}
}