147 lines
4.6 KiB
Go
147 lines
4.6 KiB
Go
/*
|
|
*
|
|
* Copyright 2014, Google Inc.
|
|
* All rights reserved.
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are
|
|
* met:
|
|
*
|
|
* * Redistributions of source code must retain the above copyright
|
|
* notice, this list of conditions and the following disclaimer.
|
|
* * Redistributions in binary form must reproduce the above
|
|
* copyright notice, this list of conditions and the following disclaimer
|
|
* in the documentation and/or other materials provided with the
|
|
* distribution.
|
|
* * Neither the name of Google Inc. nor the names of its
|
|
* contributors may be used to endorse or promote products derived from
|
|
* this software without specific prior written permission.
|
|
*
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
|
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
|
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
|
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
|
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
|
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
|
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
|
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
|
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
|
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
*
|
|
*/
|
|
|
|
/*
|
|
Package benchmark implements the building blocks to setup end-to-end gRPC benchmarks.
|
|
*/
|
|
package benchmark
|
|
|
|
import (
|
|
"io"
|
|
"math"
|
|
"net"
|
|
|
|
"golang.org/x/net/context"
|
|
"google.golang.org/grpc"
|
|
testpb "google.golang.org/grpc/benchmark/grpc_testing"
|
|
"google.golang.org/grpc/grpclog"
|
|
)
|
|
|
|
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,
|
|
}
|
|
}
|
|
|
|
type testServer struct {
|
|
}
|
|
|
|
func (s *testServer) UnaryCall(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
|
|
return &testpb.SimpleResponse{
|
|
Payload: newPayload(in.ResponseType, int(in.ResponseSize)),
|
|
}, nil
|
|
}
|
|
|
|
func (s *testServer) StreamingCall(stream testpb.TestService_StreamingCallServer) error {
|
|
for {
|
|
in, err := stream.Recv()
|
|
if err == io.EOF {
|
|
// read done.
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := stream.Send(&testpb.SimpleResponse{
|
|
Payload: newPayload(in.ResponseType, int(in.ResponseSize)),
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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()
|
|
}
|
|
}
|
|
|
|
// DoUnaryCall performs an unary RPC with given stub and request and response sizes.
|
|
func DoUnaryCall(tc testpb.TestServiceClient, reqSize, respSize int) {
|
|
pl := newPayload(testpb.PayloadType_COMPRESSABLE, reqSize)
|
|
req := &testpb.SimpleRequest{
|
|
ResponseType: pl.Type,
|
|
ResponseSize: int32(respSize),
|
|
Payload: pl,
|
|
}
|
|
if _, err := tc.UnaryCall(context.Background(), req); err != nil {
|
|
grpclog.Fatal("/TestService/UnaryCall RPC failed: ", err)
|
|
}
|
|
}
|
|
|
|
// 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)
|
|
if err != nil {
|
|
grpclog.Fatalf("NewClientConn(%q) failed to create a ClientConn %v", addr, err)
|
|
}
|
|
return conn
|
|
}
|