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  }