From 8e4b6f3bb2c5a064aad15914dbd560a8770d2f10 Mon Sep 17 00:00:00 2001
From: "renzheng.wang" <renzheng.wang@ucloud.cn>
Date: Wed, 16 Oct 2019 02:07:44 +0800
Subject: [PATCH] support get the actual listened address after grpc server is
 started

---
 go.mod                            |  1 +
 pkg/net/rpc/warden/server.go      | 19 +++++++++++++++++++
 pkg/net/rpc/warden/server_test.go | 14 +++++++++++++-
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index ce078a8e7..475ff1c43 100644
--- a/go.mod
+++ b/go.mod
@@ -26,6 +26,7 @@ require (
 	github.com/konsorten/go-windows-terminal-sequences v1.0.2 // indirect
 	github.com/leodido/go-urn v1.1.0 // indirect
 	github.com/mattn/go-colorable v0.1.2 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/montanaflynn/stats v0.5.0
 	github.com/openzipkin/zipkin-go v0.2.1
 	github.com/otokaze/mock v0.0.0-20190125081256-8282b7a7c7c3
diff --git a/pkg/net/rpc/warden/server.go b/pkg/net/rpc/warden/server.go
index 26b1f184d..9a8518674 100644
--- a/pkg/net/rpc/warden/server.go
+++ b/pkg/net/rpc/warden/server.go
@@ -316,6 +316,25 @@ func (s *Server) Start() (*Server, error) {
 	return s, nil
 }
 
+// StartWithAddr create a new goroutine run server with configured listen addr
+// will panic if any error happend
+// return server itself and the actually listened address (if configured listen
+// port is zero, the os will allocate an unused port)
+func (s *Server) StartWithAddr() (*Server, net.Addr, error) {
+	lis, err := net.Listen(s.conf.Network, s.conf.Addr)
+	if err != nil {
+		return nil, nil, err
+	}
+	log.Info("warden: start grpc listen addr: %v", lis.Addr())
+	reflection.Register(s.server)
+	go func() {
+		if err := s.Serve(lis); err != nil {
+			panic(err)
+		}
+	}()
+	return s, lis.Addr(), nil
+}
+
 // Serve accepts incoming connections on the listener lis, creating a new
 // ServerTransport and service goroutine for each.
 // Serve will return a non-nil error unless Stop or GracefulStop is called.
diff --git a/pkg/net/rpc/warden/server_test.go b/pkg/net/rpc/warden/server_test.go
index b6786290c..ee4cd664e 100644
--- a/pkg/net/rpc/warden/server_test.go
+++ b/pkg/net/rpc/warden/server_test.go
@@ -193,7 +193,7 @@ func runClient(ctx context.Context, cc *ClientConfig, t *testing.T, name string,
 	return
 }
 
-func TestMain(t *testing.T) {
+func TestMain(m *testing.M) {
 	log.Init(nil)
 }
 
@@ -594,3 +594,15 @@ func TestMetadata(t *testing.T) {
 	_, err := cli.SayHello(ctx, &pb.HelloRequest{Name: "test"})
 	assert.Nil(t, err)
 }
+
+func TestStartWithAddr(t *testing.T) {
+	configuredAddr := "127.0.0.1:0"
+	server = NewServer(&ServerConfig{Addr: configuredAddr, Timeout: xtime.Duration(time.Second)})
+	if _, realAddr, err := server.StartWithAddr(); err == nil && realAddr != nil {
+		assert.NotEqual(t, realAddr.String(), configuredAddr)
+	} else {
+		assert.NotNil(t, realAddr)
+		assert.Nil(t, err)
+	}
+}
+