diff --git a/zrpc/internal/rpcpubserver.go b/zrpc/internal/rpcpubserver.go index 93490aa4..334853b5 100644 --- 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 new file mode 100644 index 00000000..196ef156 --- /dev/null +++ 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 index a6f8b0e1..f1557929 100644 --- a/zrpc/server.go +++ b/zrpc/server.go @@ -2,13 +2,10 @@ 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" @@ -16,11 +13,6 @@ import ( "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 index 8a4a257f..7d7a1f9a 100644 --- 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{