fix issue zeromicro#469 (zeromicro#471)

ScreamWolf · Feb 17, 2021 · 9602494 · 9602494
1 parent 38abfb8
commit 9602494
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 60 deletions.
diff --git a/zrpc/internal/rpcpubserver.go b/zrpc/internal/rpcpubserver.go
@@ -1,10 +1,22 @@
 package internal
 
-import "github.com/tal-tech/go-zero/core/discov"
+import (
+	"os"
+	"strings"
+
+	"github.com/tal-tech/go-zero/core/discov"
+	"github.com/tal-tech/go-zero/core/netx"
+)
+
+const (
+	allEths  = "0.0.0.0"
+	envPodIp = "POD_IP"
+)
 
 func NewRpcPubServer(etcdEndpoints []string, etcdKey, listenOn string, opts ...ServerOption) (Server, error) {
 	registerEtcd := func() error {
-		pubClient := discov.NewPublisher(etcdEndpoints, etcdKey, listenOn)
+		pubListenOn := figureOutListenOn(listenOn)
+		pubClient := discov.NewPublisher(etcdEndpoints, etcdKey, pubListenOn)
 		return pubClient.KeepAlive()
 	}
 	server := keepAliveServer{
@@ -27,3 +39,25 @@ func (ags keepAliveServer) Start(fn RegisterFn) error {
 
 	return ags.Server.Start(fn)
 }
+
+func figureOutListenOn(listenOn string) string {
+	fields := strings.Split(listenOn, ":")
+	if len(fields) == 0 {
+		return listenOn
+	}
+
+	host := fields[0]
+	if len(host) > 0 && host != allEths {
+		return listenOn
+	}
+
+	ip := os.Getenv(envPodIp)
+	if len(ip) == 0 {
+		ip = netx.InternalIp()
+	}
+	if len(ip) == 0 {
+		return listenOn
+	}
+
+	return strings.Join(append([]string{ip}, fields[1:]...), ":")
+}
diff --git a/zrpc/internal/rpcpubserver_test.go b/zrpc/internal/rpcpubserver_test.go
@@ -0,0 +1,33 @@
+package internal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tal-tech/go-zero/core/netx"
+)
+
+func TestFigureOutListenOn(t *testing.T) {
+	tests := []struct {
+		input  string
+		expect string
+	}{
+		{
+			input:  "192.168.0.5:1234",
+			expect: "192.168.0.5:1234",
+		},
+		{
+			input:  "0.0.0.0:8080",
+			expect: netx.InternalIp() + ":8080",
+		},
+		{
+			input:  ":8080",
+			expect: netx.InternalIp() + ":8080",
+		},
+	}
+
+	for _, test := range tests {
+		val := figureOutListenOn(test.input)
+		assert.Equal(t, test.expect, val)
+	}
+}
diff --git a/zrpc/server.go b/zrpc/server.go
@@ -2,25 +2,17 @@ package zrpc
 
 import (
 	"log"
-	"os"
-	"strings"
 	"time"
 
 	"github.com/tal-tech/go-zero/core/load"
 	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/netx"
 	"github.com/tal-tech/go-zero/core/stat"
 	"github.com/tal-tech/go-zero/zrpc/internal"
 	"github.com/tal-tech/go-zero/zrpc/internal/auth"
 	"github.com/tal-tech/go-zero/zrpc/internal/serverinterceptors"
 	"google.golang.org/grpc"
 )
 
-const (
-	allEths  = "0.0.0.0"
-	envPodIp = "POD_IP"
-)
-
 type RpcServer struct {
 	server   internal.Server
 	register internal.RegisterFn
@@ -44,8 +36,7 @@ func NewServer(c RpcServerConf, register internal.RegisterFn) (*RpcServer, error
 	var server internal.Server
 	metrics := stat.NewMetrics(c.ListenOn)
 	if c.HasEtcd() {
-		listenOn := figureOutListenOn(c.ListenOn)
-		server, err = internal.NewRpcPubServer(c.Etcd.Hosts, c.Etcd.Key, listenOn, internal.WithMetrics(metrics))
+		server, err = internal.NewRpcPubServer(c.Etcd.Hosts, c.Etcd.Key, c.ListenOn, internal.WithMetrics(metrics))
 		if err != nil {
 			return nil, err
 		}
@@ -92,28 +83,6 @@ func (rs *RpcServer) Stop() {
 	logx.Close()
 }
 
-func figureOutListenOn(listenOn string) string {
-	fields := strings.Split(listenOn, ":")
-	if len(fields) == 0 {
-		return listenOn
-	}
-
-	host := fields[0]
-	if len(host) > 0 && host != allEths {
-		return listenOn
-	}
-
-	ip := os.Getenv(envPodIp)
-	if len(ip) == 0 {
-		ip = netx.InternalIp()
-	}
-	if len(ip) == 0 {
-		return listenOn
-	}
-
-	return strings.Join(append([]string{ip}, fields[1:]...), ":")
-}
-
 func setupInterceptors(server internal.Server, c RpcServerConf, metrics *stat.Metrics) error {
 	if c.CpuThreshold > 0 {
 		shedder := load.NewAdaptiveShedder(load.WithCpuThreshold(c.CpuThreshold))

diff --git a/zrpc/server_test.go b/zrpc/server_test.go
@@ -7,7 +7,6 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/discov"
 	"github.com/tal-tech/go-zero/core/logx"
-	"github.com/tal-tech/go-zero/core/netx"
 	"github.com/tal-tech/go-zero/core/service"
 	"github.com/tal-tech/go-zero/core/stat"
 	"github.com/tal-tech/go-zero/core/stores/redis"
@@ -16,31 +15,6 @@ import (
 	"google.golang.org/grpc"
 )
 
-func TestFigureOutListenOn(t *testing.T) {
-	tests := []struct {
-		input  string
-		expect string
-	}{
-		{
-			input:  "192.168.0.5:1234",
-			expect: "192.168.0.5:1234",
-		},
-		{
-			input:  "0.0.0.0:8080",
-			expect: netx.InternalIp() + ":8080",
-		},
-		{
-			input:  ":8080",
-			expect: netx.InternalIp() + ":8080",
-		},
-	}
-
-	for _, test := range tests {
-		val := figureOutListenOn(test.input)
-		assert.Equal(t, test.expect, val)
-	}
-}
-
 func TestServer_setupInterceptors(t *testing.T) {
 	server := new(mockedServer)
 	err := setupInterceptors(server, RpcServerConf{