record per-rpc rtt in benchmark_test

benchmark/benchmark_test.go

@@ -24,25 +24,30 @@ func run(b *testing.B, maxConcurrentCalls int, caller func(testpb.TestServiceCli
 	}
 
 	ch := make(chan int, maxConcurrentCalls*4)
-	var wg sync.WaitGroup
+	var (
+		mu sync.Mutex
+		wg sync.WaitGroup
+	)
 	wg.Add(maxConcurrentCalls)
 
 	// Distribute the b.N calls over maxConcurrentCalls workers.
 	for i := 0; i < maxConcurrentCalls; i++ {
 		go func() {
 			for _ = range ch {
+				start := time.Now()
 				caller(tc)
+				elapse := time.Since(start)
+				mu.Lock()
+				s.Add(elapse)
+				mu.Unlock()
 			}
 			wg.Done()
 		}()
 	}
 	for i := 0; i < b.N; i++ {
 		b.StartTimer()
-		start := time.Now()
 		ch <- i
-		elapsed := time.Since(start)
 		b.StopTimer()
-		s.Add(elapsed)
 	}
 	close(ch)
 	wg.Wait()