From 92ba6a3209147eaf8c7a5be55a4da87c1099eb2e Mon Sep 17 00:00:00 2001
From: songzhibin97 <49082129+songzhibin97@users.noreply.github.com>
Date: Sat, 21 May 2022 16:25:13 +0800
Subject: [PATCH] Style:trace (#2017)

* Update metadata.go

* Update span.go

* Update statsHandler.go
---
 middleware/tracing/metadata.go     | 18 +++++++++---------
 middleware/tracing/span.go         | 16 ++++++++++------
 middleware/tracing/statsHandler.go | 17 ++++++++++-------
 3 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/middleware/tracing/metadata.go b/middleware/tracing/metadata.go
index 7e8cbe951..01e511329 100644
--- a/middleware/tracing/metadata.go
+++ b/middleware/tracing/metadata.go
@@ -26,16 +26,16 @@ func (b Metadata) Inject(ctx context.Context, carrier propagation.TextMapCarrier
 // Extract returns a copy of parent with the metadata from the carrier added.
 func (b Metadata) Extract(parent context.Context, carrier propagation.TextMapCarrier) context.Context {
 	name := carrier.Get(serviceHeader)
-	if name != "" {
-		if md, ok := metadata.FromServerContext(parent); ok {
-			md.Set(serviceHeader, name)
-		} else {
-			md := metadata.New()
-			md.Set(serviceHeader, name)
-			parent = metadata.NewServerContext(parent, md)
-		}
+	if name == "" {
+		return parent
 	}
-
+	if md, ok := metadata.FromServerContext(parent); ok {
+		md.Set(serviceHeader, name)
+		return parent
+	}
+	md := metadata.New()
+	md.Set(serviceHeader, name)
+	parent = metadata.NewServerContext(parent, md)
 	return parent
 }
 
diff --git a/middleware/tracing/span.go b/middleware/tracing/span.go
index b60433b6b..38bb35f42 100644
--- a/middleware/tracing/span.go
+++ b/middleware/tracing/span.go
@@ -22,10 +22,12 @@ func setClientSpan(ctx context.Context, span trace.Span, m interface{}) {
 	var remote string
 	var operation string
 	var rpcKind string
-	if tr, ok := transport.FromClientContext(ctx); ok {
+	tr, ok := transport.FromClientContext(ctx)
+	if ok {
 		operation = tr.Operation()
 		rpcKind = tr.Kind().String()
-		if tr.Kind() == transport.KindHTTP {
+		switch tr.Kind() {
+		case transport.KindHTTP:
 			if ht, ok := tr.(http.Transporter); ok {
 				method := ht.Request().Method
 				route := ht.PathTemplate()
@@ -35,7 +37,7 @@ func setClientSpan(ctx context.Context, span trace.Span, m interface{}) {
 				attrs = append(attrs, semconv.HTTPTargetKey.String(path))
 				remote = ht.Request().Host
 			}
-		} else if tr.Kind() == transport.KindGRPC {
+		case transport.KindGRPC:
 			remote, _ = parseTarget(tr.Endpoint())
 		}
 	}
@@ -57,10 +59,12 @@ func setServerSpan(ctx context.Context, span trace.Span, m interface{}) {
 	var remote string
 	var operation string
 	var rpcKind string
-	if tr, ok := transport.FromServerContext(ctx); ok {
+	tr, ok := transport.FromServerContext(ctx)
+	if ok {
 		operation = tr.Operation()
 		rpcKind = tr.Kind().String()
-		if tr.Kind() == transport.KindHTTP {
+		switch tr.Kind() {
+		case transport.KindHTTP:
 			if ht, ok := tr.(http.Transporter); ok {
 				method := ht.Request().Method
 				route := ht.PathTemplate()
@@ -70,7 +74,7 @@ func setServerSpan(ctx context.Context, span trace.Span, m interface{}) {
 				attrs = append(attrs, semconv.HTTPTargetKey.String(path))
 				remote = ht.Request().RemoteAddr
 			}
-		} else if tr.Kind() == transport.KindGRPC {
+		case transport.KindGRPC:
 			if p, ok := peer.FromContext(ctx); ok {
 				remote = p.Addr.String()
 			}
diff --git a/middleware/tracing/statsHandler.go b/middleware/tracing/statsHandler.go
index b5cb1233b..11b12679d 100644
--- a/middleware/tracing/statsHandler.go
+++ b/middleware/tracing/statsHandler.go
@@ -22,13 +22,16 @@ func (c *ClientHandler) TagConn(ctx context.Context, cti *stats.ConnTagInfo) con
 
 // HandleRPC implements per-RPC tracing and stats instrumentation.
 func (c *ClientHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
-	if _, ok := rs.(*stats.OutHeader); ok {
-		if p, ok := peer.FromContext(ctx); ok {
-			remoteAddr := p.Addr.String()
-			if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
-				span.SetAttributes(peerAttr(remoteAddr)...)
-			}
-		}
+	if _, ok := rs.(*stats.OutHeader); !ok {
+		return
+	}
+	p, ok := peer.FromContext(ctx)
+	if !ok {
+		return
+	}
+	remoteAddr := p.Addr.String()
+	if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
+		span.SetAttributes(peerAttr(remoteAddr)...)
 	}
 }