Repository: grpc/grpc-go
File: benchmark/worker/benchmark_client.go (Go)

Code before:

		lastResetTime:   time.Now(),
		closeConns:      closeConns,
		rusageLastReset: syscall.GetRusage(),
	}

	if err = performRPCs(ctx, config, conns, bc); err != nil {
		// Close all connections if performRPCs failed.
		closeConns()
		return nil, err
	}

	return bc, nil
}

func (bc *benchmarkClient) unaryLoop(ctx context.Context, conns []*grpc.ClientConn, rpcCountPerConn int, reqSize int, respSize int, poissonLambda *float64) {
	for ic, conn := range conns {
		client := testgrpc.NewBenchmarkServiceClient(conn)
		// For each connection, create rpcCountPerConn goroutines to do rpc.
		for j := range rpcCountPerConn {
			// 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.
				// Now relying on worker client to reserve time to do warm up.
				// The worker client needs to wait for some time after client is created,
				// before starting benchmark.
				if poissonLambda == nil { // Closed loop.
					for {
						start := time.Now()
						if err := benchmark.DoUnaryCall(ctx, client, reqSize, respSize); err != nil {
							return
						}
						elapse := time.Since(start)
						bc.lockingHistograms[idx].add(int64(elapse))
					}
				} else { // Open loop.
					timeBetweenRPCs := time.Duration((rand.ExpFloat64() / *poissonLambda) * float64(time.Second))
					time.AfterFunc(timeBetweenRPCs, func() {
						bc.poissonUnary(ctx, client, idx, reqSize, respSize, *poissonLambda)
					})
				}
			}(idx)
		}
	}
}

Review comment (bug):
I think we can avoid the context cancellation check here also. This check happens at the beginning of the benchmark run and is unlikely that the context is already cancelled. In case of early cancellation, the client may do some throwaway work, but the test will eventually return when the RPCs fail.

Code after:

		lastResetTime:   time.Now(),
		closeConns:      closeConns,
		rusageLastReset: syscall.GetRusage(),
	}

	if err = performRPCs(ctx, config, conns, bc); err != nil {
		// Close all connections if performRPCs failed.
		closeConns()
		return nil, err
	}

	return bc, nil
}

func (bc *benchmarkClient) unaryLoop(ctx context.Context, conns []*grpc.ClientConn, rpcCountPerConn int, reqSize int, respSize int, poissonLambda *float64) {
	for ic, conn := range conns {
		client := testgrpc.NewBenchmarkServiceClient(conn)
		// For each connection, create rpcCountPerConn goroutines to do rpc.
		for j := range rpcCountPerConn {
			// Create histogram for each goroutine.
			idx := ic*rpcCountPerConn + j
			bc.lockingHistograms[idx].histogram = stats.NewHistogram(bc.histogramOptions)
			// Start goroutine on the created mutex and histogram.
			go func(idx int) {
				// TODO: do warm up if necessary.
				// Now relying on worker client to reserve time to do warm up.
				// The worker client needs to wait for some time after client is created,
				// before starting benchmark.
				if poissonLambda == nil { // Closed loop.
					for {
						if ctx.Err() != nil {
							break
						}
						start := time.Now()
						if err := benchmark.DoUnaryCall(client, reqSize, respSize); err != nil {
							continue
						}
						elapse := time.Since(start)
						bc.lockingHistograms[idx].add(int64(elapse))
					}
				} else { // Open loop.
					timeBetweenRPCs := time.Duration((rand.ExpFloat64() / *poissonLambda) * float64(time.Second))
					time.AfterFunc(timeBetweenRPCs, func() {
						bc.poissonUnary(client, idx, reqSize, respSize, *poissonLambda)
					})
				}
			}(idx)
		}
	}
}