refactor: minor changes

Author: wdsjk

benchmark/worker/benchmark_client.go

@@ -264,9 +264,6 @@ func (bc *benchmarkClient) unaryLoop(ctx context.Context, conns []*grpc.ClientCo
 			// Create histogram for each goroutine.
 			idx := ic*rpcCountPerConn + j
 			bc.lockingHistograms[idx].histogram = stats.NewHistogram(bc.histogramOptions)
-			if ctx.Err() != nil {
-				return
-			}
 			// Start goroutine on the created mutex and histogram.
 			go func(idx int) {
 				// TODO: do warm up if necessary.
@@ -277,7 +274,10 @@ func (bc *benchmarkClient) unaryLoop(ctx context.Context, conns []*grpc.ClientCo
 					for {
 						start := time.Now()
 						if err := benchmark.DoUnaryCall(ctx, client, reqSize, respSize); err != nil {
-							return
+							if status.Code(err) == codes.Canceled {
+								return
+							}
+							continue
 						}
 						elapse := time.Since(start)
 						bc.lockingHistograms[idx].add(int64(elapse))
@@ -311,9 +311,6 @@ func (bc *benchmarkClient) streamingLoop(ctx context.Context, conns []*grpc.Clie
 			idx := ic*rpcCountPerConn + j
 			bc.lockingHistograms[idx].histogram = stats.NewHistogram(bc.histogramOptions)
 			if poissonLambda == nil { // Closed loop.
-				if stream.Context().Err() != nil {
-					return
-				}
 				// Start goroutine on the created mutex and histogram.
 				go func(idx int) {
 					// TODO: do warm up if necessary.
@@ -342,7 +339,6 @@ func (bc *benchmarkClient) streamingLoop(ctx context.Context, conns []*grpc.Clie
 func (bc *benchmarkClient) poissonUnary(ctx context.Context, client testgrpc.BenchmarkServiceClient, idx int, reqSize int, respSize int, lambda float64) {
 	go func() {
 		start := time.Now()
-
 		if err := benchmark.DoUnaryCall(ctx, client, reqSize, respSize); err != nil {
 			return
 		}

benchmark/worker/main.go

@@ -139,7 +139,9 @@ func (s *workerServer) RunServer(stream testgrpc.WorkerService_RunServerServer)
 
 func (s *workerServer) RunClient(stream testgrpc.WorkerService_RunClientServer) error {
 	var bc *benchmarkClient
+	ctx, cancel := context.WithCancel(stream.Context())
 	defer func() {
+		cancel()
 		// Shut down benchmark client when stream ends.
 		logger.Infof("shutting down benchmark client")
 		if bc != nil {
@@ -163,9 +165,6 @@ func (s *workerServer) RunClient(stream testgrpc.WorkerService_RunClientServer)
 				logger.Infof("client setup received when client already exists, shutting down the existing client")
 				bc.shutdown()
 			}
-
-			ctx, cancel := context.WithCancel(stream.Context())
-			defer cancel()
 			bc, err = startBenchmarkClient(ctx, t.Setup)
 			if err != nil {
 				return err