Change StartServer back to return address rather than port number

benchmark/benchmark.go

@@ -93,8 +93,8 @@ func (s *testServer) StreamingCall(stream testpb.BenchmarkService_StreamingCallS
 
 // StartServer starts a gRPC server serving a benchmark service on the given
 // address, which may be something like "localhost:0". It returns its listen
-// port number and a function to stop the server.
-func StartServer(addr string, opts ...grpc.ServerOption) (int, func()) {
+// address and a function to stop the server.
+func StartServer(addr string, opts ...grpc.ServerOption) (string, func()) {
 	lis, err := net.Listen("tcp", addr)
 	if err != nil {
 		grpclog.Fatalf("Failed to listen: %v", err)
@@ -102,7 +102,7 @@ func StartServer(addr string, opts ...grpc.ServerOption) (int, func()) {
 	s := grpc.NewServer(opts...)
 	testpb.RegisterBenchmarkServiceServer(s, &testServer{})
 	go s.Serve(lis)
-	return lis.Addr().(*net.TCPAddr).Port, func() {
+	return lis.Addr().String(), func() {
 		s.Stop()
 	}
 }
@@ -133,8 +133,8 @@ func (s *genericTestServer) StreamingCall(stream testpb.BenchmarkService_Streami
 }
 
 // StartGenericServer starts a benchmark service server that supports custom codec.
-// It returns its listen port number and a function to stop the server.
-func StartGenericServer(addr string, reqSize, respSize int32, opts ...grpc.ServerOption) (int, func()) {
+// It returns its listen address and a function to stop the server.
+func StartGenericServer(addr string, reqSize, respSize int32, opts ...grpc.ServerOption) (string, func()) {
 	lis, err := net.Listen("tcp", addr)
 	if err != nil {
 		grpclog.Fatalf("Failed to listen: %v", err)
@@ -142,7 +142,7 @@ func StartGenericServer(addr string, reqSize, respSize int32, opts ...grpc.Serve
 	s := grpc.NewServer(opts...)
 	testpb.RegisterBenchmarkServiceServer(s, &genericTestServer{reqSize: reqSize, respSize: respSize})
 	go s.Serve(lis)
-	return lis.Addr().(*net.TCPAddr).Port, func() {
+	return lis.Addr().String(), func() {
 		s.Stop()
 	}
 }

benchmark/benchmark_test.go

@@ -2,7 +2,6 @@ package benchmark
 
 import (
 	"os"
-	"strconv"
 	"sync"
 	"testing"
 	"time"
@@ -17,9 +16,9 @@ import (
 func runUnary(b *testing.B, maxConcurrentCalls int) {
 	s := stats.AddStats(b, 38)
 	b.StopTimer()
-	targetPort, stopper := StartServer("localhost:0")
+	target, stopper := StartServer("localhost:0")
 	defer stopper()
-	conn := NewClientConn(":" + strconv.Itoa(targetPort))
+	conn := NewClientConn(target)
 	tc := testpb.NewBenchmarkServiceClient(conn)
 
 	// Warm up connection.
@@ -60,9 +59,9 @@ func runUnary(b *testing.B, maxConcurrentCalls int) {
 func runStream(b *testing.B, maxConcurrentCalls int) {
 	s := stats.AddStats(b, 38)
 	b.StopTimer()
-	targetPort, stopper := StartServer("localhost:0")
+	target, stopper := StartServer("localhost:0")
 	defer stopper()
-	conn := NewClientConn(":" + strconv.Itoa(targetPort))
+	conn := NewClientConn(target)
 	tc := testpb.NewBenchmarkServiceClient(conn)
 
 	// Warm up connection.

benchmark/worker/benchmark_server.go

@@ -36,6 +36,7 @@ package main
 import (
 	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 
@@ -107,15 +108,15 @@ func startBenchmarkServerWithSetup(setup *testpb.ServerConfig, serverPort int) (
 	}
 
 	grpclog.Printf(" - payload config: %v", setup.PayloadConfig)
-	var p int
+	var addr string
 	var close func()
 	if setup.PayloadConfig != nil {
 		switch payload := setup.PayloadConfig.Payload.(type) {
 		case *testpb.PayloadConfig_BytebufParams:
 			opts = append(opts, grpc.CustomCodec(byteBufCodec{}))
-			p, close = benchmark.StartGenericServer(":"+strconv.Itoa(port), payload.BytebufParams.ReqSize, payload.BytebufParams.RespSize, opts...)
+			addr, close = benchmark.StartGenericServer(":"+strconv.Itoa(port), payload.BytebufParams.ReqSize, payload.BytebufParams.RespSize, opts...)
 		case *testpb.PayloadConfig_SimpleParams:
-			p, close = benchmark.StartServer(":"+strconv.Itoa(port), opts...)
+			addr, close = benchmark.StartServer(":"+strconv.Itoa(port), opts...)
 		case *testpb.PayloadConfig_ComplexParams:
 			return nil, grpc.Errorf(codes.Unimplemented, "unsupported payload config: %v", setup.PayloadConfig)
 		default:
@@ -123,10 +124,15 @@ func startBenchmarkServerWithSetup(setup *testpb.ServerConfig, serverPort int) (
 		}
 	} else {
 		// Start protobuf server is payload config is nil.
-		p, close = benchmark.StartServer(":"+strconv.Itoa(port), opts...)
+		addr, close = benchmark.StartServer(":"+strconv.Itoa(port), opts...)
 	}
 
-	grpclog.Printf("benchmark server listening at port %v", p)
+	grpclog.Printf("benchmark server listening at %v", addr)
+	addrSplitted := strings.Split(addr, ":")
+	p, err := strconv.Atoi(addrSplitted[len(addrSplitted)-1])
+	if err != nil {
+		grpclog.Fatalf("failed to get port number from server address: %v", err)
+	}
 
 	return &benchmarkServer{port: p, cores: numOfCores, close: close, lastResetTime: time.Now()}, nil
 }