Repository: go-kratos/kratos
Pull request #972 (2021-05-31 06:55:12)

add feat endpoint

tonybase: LGTM

Patch:
diff --git a/app.go b/app.go
index a0dd9d02..f832d6b1 100644
--- a/app.go
+++ b/app.go
@@ -116,22 +116,27 @@ func (a *App) Stop() error {
 }
 
 func (a *App) buildInstance() (*registry.ServiceInstance, error) {
-	if len(a.opts.endpoints) == 0 {
+	var endpoints []string
+	for _, e := range a.opts.endpoints {
+		endpoints = append(endpoints, e.String())
+	}
+	if len(endpoints) == 0 {
 		for _, srv := range a.opts.servers {
 			if r, ok := srv.(transport.Endpointer); ok {
 				e, err := r.Endpoint()
 				if err != nil {
 					return nil, err
 				}
-				a.opts.endpoints = append(a.opts.endpoints, e)
+				endpoints = append(endpoints, e.String())
 			}
 		}
 	}
+
 	return &registry.ServiceInstance{
 		ID:        a.opts.id,
 		Name:      a.opts.name,
 		Version:   a.opts.version,
 		Metadata:  a.opts.metadata,
-		Endpoints: a.opts.endpoints,
+		Endpoints: endpoints,
 	}, nil
 }
diff --git a/options.go b/options.go
index 2bc8ba66..cf2275d9 100644
--- a/options.go
+++ b/options.go
@@ -2,6 +2,7 @@ package kratos
 
 import (
 	"context"
+	"net/url"
 	"os"
 
 	"github.com/go-kratos/kratos/v2/log"
@@ -18,7 +19,7 @@ type options struct {
 	name      string
 	version   string
 	metadata  map[string]string
-	endpoints []string
+	endpoints []*url.URL
 
 	ctx  context.Context
 	sigs []os.Signal
@@ -50,7 +51,7 @@ func Metadata(md map[string]string) Option {
 }
 
 // Endpoint with service endpoint.
-func Endpoint(endpoints ...string) Option {
+func Endpoint(endpoints ...*url.URL) Option {
 	return func(o *options) { o.endpoints = endpoints }
 }
 
diff --git a/transport/grpc/context.go b/transport/grpc/context.go
index 3ff5610c..0347a78c 100644
--- a/transport/grpc/context.go
+++ b/transport/grpc/context.go
@@ -1,6 +1,9 @@
 package grpc
 
-import "context"
+import (
+	"context"
+	"net/url"
+)
 
 // ServerInfo represent gRPC server information.
 type ServerInfo struct {
@@ -8,6 +11,8 @@ type ServerInfo struct {
 	Server interface{}
 	// FullMethod is the full RPC method string, i.e., /package.service/method.
 	FullMethod string
+	// Endpoint is a real address to registry endpoint.
+	Endpoint *url.URL
 }
 
 type serverKey struct{}
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index 33124dbb..40f7fc48 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -2,8 +2,8 @@ package grpc
 
 import (
 	"context"
-	"fmt"
 	"net"
+	"net/url"
 	"strings"
 	"time"
 
@@ -75,6 +75,7 @@ type Server struct {
 	lis        net.Listener
 	network    string
 	address    string
+	endpoint   *url.URL
 	timeout    time.Duration
 	log        *log.Helper
 	middleware middleware.Middleware
@@ -118,19 +119,24 @@ func NewServer(opts ...ServerOption) *Server {
 // Endpoint return a real address to registry endpoint.
 // examples:
 //   grpc://127.0.0.1:9000?isSecure=false
-func (s *Server) Endpoint() (string, error) {
+func (s *Server) Endpoint() (*url.URL, error) {
 	if s.lis == nil && strings.HasSuffix(s.address, ":0") {
 		lis, err := net.Listen(s.network, s.address)
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 		s.lis = lis
 	}
 	addr, err := host.Extract(s.address, s.lis)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
-	return fmt.Sprintf("grpc://%s", addr), nil
+	u := &url.URL{
+		Scheme: "grpc",
+		Host:   addr,
+	}
+	s.endpoint = u
+	return u, nil
 }
 
 // Start start the gRPC server.
@@ -161,7 +167,7 @@ func (s *Server) unaryServerInterceptor() grpc.UnaryServerInterceptor {
 		ctx, cancel := ic.Merge(ctx, s.ctx)
 		defer cancel()
 		ctx = transport.NewContext(ctx, transport.Transport{Kind: transport.KindGRPC})
-		ctx = NewServerContext(ctx, ServerInfo{Server: info.Server, FullMethod: info.FullMethod})
+		ctx = NewServerContext(ctx, ServerInfo{Server: info.Server, FullMethod: info.FullMethod, Endpoint: s.endpoint})
 		if s.timeout > 0 {
 			var cancel context.CancelFunc
 			ctx, cancel = context.WithTimeout(ctx, s.timeout)
diff --git a/transport/http/context.go b/transport/http/context.go
index eecb7e2e..68c57d9a 100644
--- a/transport/http/context.go
+++ b/transport/http/context.go
@@ -3,12 +3,14 @@ package http
 import (
 	"context"
 	"net/http"
+	"net/url"
 )
 
 // ServerInfo represent HTTP server information.
 type ServerInfo struct {
 	Request  *http.Request
 	Response http.ResponseWriter
+	Endpoint *url.URL
 }
 
 type serverKey struct{}
diff --git a/transport/http/server.go b/transport/http/server.go
index e5f2e7a2..10afbb57 100644
--- a/transport/http/server.go
+++ b/transport/http/server.go
@@ -3,9 +3,9 @@ package http
 import (
 	"context"
 	"errors"
-	"fmt"
 	"net"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -54,13 +54,14 @@ func Logger(logger log.Logger) ServerOption {
 // Server is an HTTP server wrapper.
 type Server struct {
 	*http.Server
-	ctx     context.Context
-	lis     net.Listener
-	network string
-	address string
-	timeout time.Duration
-	router  *mux.Router
-	log     *log.Helper
+	ctx      context.Context
+	lis      net.Listener
+	network  string
+	address  string
+	endpoint *url.URL
+	timeout  time.Duration
+	router   *mux.Router
+	log      *log.Helper
 }
 
 // NewServer creates an HTTP server by options.
@@ -99,7 +100,7 @@ func (s *Server) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 	ctx, cancel := ic.Merge(req.Context(), s.ctx)
 	defer cancel()
 	ctx = transport.NewContext(ctx, transport.Transport{Kind: transport.KindHTTP})
-	ctx = NewServerContext(ctx, ServerInfo{Request: req, Response: res})
+	ctx = NewServerContext(ctx, ServerInfo{Request: req, Response: res, Endpoint: s.endpoint})
 	if s.timeout > 0 {
 		ctx, cancel = context.WithTimeout(ctx, s.timeout)
 		defer cancel()
@@ -110,19 +111,24 @@ func (s *Server) ServeHTTP(res http.ResponseWriter, req *http.Request) {
 // Endpoint return a real address to registry endpoint.
 // examples:
 //   http://127.0.0.1:8000?isSecure=false
-func (s *Server) Endpoint() (string, error) {
+func (s *Server) Endpoint() (*url.URL, error) {
 	if s.lis == nil && strings.HasSuffix(s.address, ":0") {
 		lis, err := net.Listen(s.network, s.address)
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 		s.lis = lis
 	}
 	addr, err := host.Extract(s.address, s.lis)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
-	return fmt.Sprintf("http://%s", addr), nil
+	u := &url.URL{
+		Scheme: "http",
+		Host:   addr,
+	}
+	s.endpoint = u
+	return u, nil
 }
 
 // Start start the HTTP server.
diff --git a/transport/transport.go b/transport/transport.go
index aee456d1..aec86d1b 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -2,6 +2,7 @@ package transport
 
 import (
 	"context"
+	"net/url"
 
 	// init encoding
 	_ "github.com/go-kratos/kratos/v2/encoding/json"
@@ -18,7 +19,7 @@ type Server interface {
 
 // Endpointer is registry endpoint.
 type Endpointer interface {
-	Endpoint() (string, error)
+	Endpoint() (*url.URL, error)
 }
 
 // Transport is transport context value.


Test patch:
diff --git a/transport/grpc/server_test.go b/transport/grpc/server_test.go
index a101648a..c2901dc0 100644
--- a/transport/grpc/server_test.go
+++ b/transport/grpc/server_test.go
@@ -15,7 +15,7 @@ func TestServer(t *testing.T) {
 	ctx := context.Background()
 	ctx = context.WithValue(ctx, testKey{}, "test")
 	srv := NewServer()
-	if e, err := srv.Endpoint(); err != nil || e == "" {
+	if e, err := srv.Endpoint(); err != nil || e == nil {
 		t.Fatal(e, err)
 	}
 
diff --git a/transport/http/server_test.go b/transport/http/server_test.go
index d5efb141..90da0969 100644
--- a/transport/http/server_test.go
+++ b/transport/http/server_test.go
@@ -32,7 +32,7 @@ func TestServer(t *testing.T) {
 	srv := NewServer()
 	srv.HandleFunc("/index", fn)
 
-	if e, err := srv.Endpoint(); err != nil || e == "" {
+	if e, err := srv.Endpoint(); err != nil || e == nil {
 		t.Fatal(e, err)
 	}