github.com/blixtra/rkt@v0.8.1-0.20160204105720-ab0d1add1a43/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  	"golang.org/x/net/context"
    13  	"google.golang.org/grpc"
    14  	"google.golang.org/grpc/benchmark"
    15  	testpb "google.golang.org/grpc/benchmark/grpc_testing"
    16  	"google.golang.org/grpc/benchmark/stats"
    17  	"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  	ch := make(chan int, *maxConcurrentRPCs*4)
    96  	var (
    97  		mu sync.Mutex
    98  		wg sync.WaitGroup
    99  	)
   100  	wg.Add(*maxConcurrentRPCs)
   101  	// Distribute RPCs over maxConcurrentCalls workers.
   102  	for i := 0; i < *maxConcurrentRPCs; i++ {
   103  		go func() {
   104  			stream, err := tc.StreamingCall(context.Background())
   105  			if err != nil {
   106  				grpclog.Fatalf("%v.StreamingCall(_) = _, %v", tc, err)
   107  			}
   108  			// Do some warm up.
   109  			for i := 0; i < 100; i++ {
   110  				streamCaller(tc, stream)
   111  			}
   112  			for range ch {
   113  				start := time.Now()
   114  				streamCaller(tc, stream)
   115  				elapse := time.Since(start)
   116  				mu.Lock()
   117  				s.Add(elapse)
   118  				mu.Unlock()
   119  			}
   120  			wg.Done()
   121  		}()
   122  	}
   123  	// Stop the client when time is up.
   124  	done := make(chan struct{})
   125  	go func() {
   126  		<-time.After(time.Duration(*duration) * time.Second)
   127  		close(done)
   128  	}()
   129  	ok := true
   130  	for ok {
   131  		select {
   132  		case ch <- 0:
   133  		case <-done:
   134  			ok = false
   135  		}
   136  	}
   137  	close(ch)
   138  	wg.Wait()
   139  	conn.Close()
   140  	grpclog.Println(s.String())
   141  }
   142  
   143  func main() {
   144  	flag.Parse()
   145  	grpc.EnableTracing = *trace
   146  	go func() {
   147  		lis, err := net.Listen("tcp", ":0")
   148  		if err != nil {
   149  			grpclog.Fatalf("Failed to listen: %v", err)
   150  		}
   151  		grpclog.Println("Client profiling address: ", lis.Addr().String())
   152  		if err := http.Serve(lis, nil); err != nil {
   153  			grpclog.Fatalf("Failed to serve: %v", err)
   154  		}
   155  	}()
   156  	switch *rpcType {
   157  	case 0:
   158  		closeLoopUnary()
   159  	case 1:
   160  		closeLoopStream()
   161  	}
   162  }