Repository: go-kratos/kratos
Pull request #1059 (2021-06-15 02:52:52)

transport: add transport kind

iamharvey: 👍👍

Patch:
diff --git a/middleware/logging/logging.go b/middleware/logging/logging.go
index af4e6a8c..d985e9bb 100644
--- a/middleware/logging/logging.go
+++ b/middleware/logging/logging.go
@@ -23,7 +23,7 @@ func Server(logger log.Logger) middleware.Middleware {
 			)
 			startTime := time.Now()
 			if info, ok := transport.FromServerContext(ctx); ok {
-				kind = info.Kind()
+				kind = info.Kind().String()
 				operation = info.Operation()
 			}
 			reply, err = handler(ctx, req)
@@ -59,7 +59,7 @@ func Client(logger log.Logger) middleware.Middleware {
 			)
 			startTime := time.Now()
 			if info, ok := transport.FromClientContext(ctx); ok {
-				kind = info.Kind()
+				kind = info.Kind().String()
 				operation = info.Operation()
 			}
 			reply, err = handler(ctx, req)
diff --git a/middleware/metrics/metrics.go b/middleware/metrics/metrics.go
index 780b9110..64ecaed0 100644
--- a/middleware/metrics/metrics.go
+++ b/middleware/metrics/metrics.go
@@ -51,7 +51,7 @@ func Server(opts ...Option) middleware.Middleware {
 			)
 			startTime := time.Now()
 			if info, ok := transport.FromServerContext(ctx); ok {
-				kind = info.Kind()
+				kind = info.Kind().String()
 				operation = info.Operation()
 			}
 			reply, err := handler(ctx, req)
@@ -86,7 +86,7 @@ func Client(opts ...Option) middleware.Middleware {
 			)
 			startTime := time.Now()
 			if info, ok := transport.FromClientContext(ctx); ok {
-				kind = info.Kind()
+				kind = info.Kind().String()
 				operation = info.Operation()
 			}
 			reply, err := handler(ctx, req)
diff --git a/middleware/tracing/tracing.go b/middleware/tracing/tracing.go
index bc5e70db..d978dc18 100644
--- a/middleware/tracing/tracing.go
+++ b/middleware/tracing/tracing.go
@@ -38,7 +38,7 @@ func Server(opts ...Option) middleware.Middleware {
 		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
 			if tr, ok := transport.FromServerContext(ctx); ok {
 				var span trace.Span
-				ctx, span = tracer.Start(ctx, tr.Kind(), tr.Operation(), tr.Header())
+				ctx, span = tracer.Start(ctx, tr.Kind().String(), tr.Operation(), tr.Header())
 				defer func() { tracer.End(ctx, span, err) }()
 			}
 			return handler(ctx, req)
@@ -53,7 +53,7 @@ func Client(opts ...Option) middleware.Middleware {
 		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
 			if tr, ok := transport.FromClientContext(ctx); ok {
 				var span trace.Span
-				ctx, span = tracer.Start(ctx, tr.Kind(), tr.Operation(), tr.Header())
+				ctx, span = tracer.Start(ctx, tr.Kind().String(), tr.Operation(), tr.Header())
 				defer func() { tracer.End(ctx, span, err) }()
 			}
 			return handler(ctx, req)
diff --git a/transport/grpc/transport.go b/transport/grpc/transport.go
index 80814f58..7a8f47bd 100644
--- a/transport/grpc/transport.go
+++ b/transport/grpc/transport.go
@@ -17,8 +17,8 @@ type Transport struct {
 }
 
 // Kind returns the transport kind.
-func (tr *Transport) Kind() string {
-	return "grpc"
+func (tr *Transport) Kind() transport.Kind {
+	return transport.KindGRPC
 }
 
 // Endpoint returns the transport endpoint.
diff --git a/transport/http/transport.go b/transport/http/transport.go
index 07a42741..1789c737 100644
--- a/transport/http/transport.go
+++ b/transport/http/transport.go
@@ -21,8 +21,8 @@ type Transport struct {
 }
 
 // Kind returns the transport kind.
-func (tr *Transport) Kind() string {
-	return "http"
+func (tr *Transport) Kind() transport.Kind {
+	return transport.KindHTTP
 }
 
 // Endpoint returns the transport endpoint.
diff --git a/transport/transport.go b/transport/transport.go
index fa814d5f..7d378311 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -31,12 +31,25 @@ type Header interface {
 
 // Transporter is transport context value interface.
 type Transporter interface {
-	Kind() string
+	Kind() Kind
 	Endpoint() string
 	Operation() string
 	Header() Header
 }
 
+// Kind defines the type of Transport
+type Kind string
+
+func (k Kind) String() string {
+	return string(k)
+}
+
+// Defines a set of transport kind
+const (
+	KindGRPC Kind = "grpc"
+	KindHTTP Kind = "http"
+)
+
 type serverTransportKey struct{}
 type clientTransportKey struct{}
 


Test patch:
diff --git a/middleware/logging/logging_test.go b/middleware/logging/logging_test.go
index e08d9db5..254bf9f4 100644
--- a/middleware/logging/logging_test.go
+++ b/middleware/logging/logging_test.go
@@ -16,12 +16,12 @@ var (
 )
 
 type Transport struct {
-	kind      string
+	kind      transport.Kind
 	endpoint  string
 	operation string
 }
 
-func (tr *Transport) Kind() string {
+func (tr *Transport) Kind() transport.Kind {
 	return tr.kind
 }
 
@@ -52,21 +52,21 @@ func TestHTTP(t *testing.T) {
 			Server,
 			err,
 			func() context.Context {
-				return transport.NewServerContext(context.Background(), &Transport{kind: "http", endpoint: "endpoint", operation: "/package.service/method"})
+				return transport.NewServerContext(context.Background(), &Transport{kind: transport.KindHTTP, endpoint: "endpoint", operation: "/package.service/method"})
 			}(),
 		},
 		{"http-server@succ",
 			Server,
 			nil,
 			func() context.Context {
-				return transport.NewServerContext(context.Background(), &Transport{kind: "http", endpoint: "endpoint", operation: "/package.service/method"})
+				return transport.NewServerContext(context.Background(), &Transport{kind: transport.KindHTTP, endpoint: "endpoint", operation: "/package.service/method"})
 			}(),
 		},
 		{"http-client@succ",
 			Client,
 			nil,
 			func() context.Context {
-				return transport.NewClientContext(context.Background(), &Transport{kind: "http", endpoint: "endpoint", operation: "/package.service/method"})
+				return transport.NewClientContext(context.Background(), &Transport{kind: transport.KindHTTP, endpoint: "endpoint", operation: "/package.service/method"})
 
 			}(),
 		},
@@ -74,7 +74,7 @@ func TestHTTP(t *testing.T) {
 			Client,
 			err,
 			func() context.Context {
-				return transport.NewClientContext(context.Background(), &Transport{kind: "http", endpoint: "endpoint", operation: "/package.service/method"})
+				return transport.NewClientContext(context.Background(), &Transport{kind: transport.KindHTTP, endpoint: "endpoint", operation: "/package.service/method"})
 			}(),
 		},
 	}