feat: trace http.status_code (#2708)

* feat: trace http.status_code

* feat: implements http.Flusher & http.Hijacker for traceResponseWriter

* test: delete notTracingSpans after test

* feat: trace http.status_code

* feat: implements http.Flusher & http.Hijacker for traceResponseWriter

* test: delete notTracingSpans after test

* refactor: update trace handler span message

* fix: code conflict

rest/handler/tracinghandler.go

@@ -1,17 +1,26 @@
 package handler
 
 import (
+	"bufio"
+	"errors"
+	"net"
 	"net/http"
 	"sync"
 
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/trace"
 	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
+const (
+	traceKeyStatusCode = "http.status_code"
+)
+
 var notTracingSpans sync.Map
 
 // DontTraceSpan disable tracing for the specified span name.
@@ -19,6 +28,40 @@ func DontTraceSpan(spanName string) {
 	notTracingSpans.Store(spanName, lang.Placeholder)
 }
 
+type traceResponseWriter struct {
+	w    http.ResponseWriter
+	code int
+}
+
+// Flush implements the http.Flusher interface.
+func (w *traceResponseWriter) Flush() {
+	if flusher, ok := w.w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+// Hijack implements the http.Hijacker interface.
+func (w *traceResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	if hijacked, ok := w.w.(http.Hijacker); ok {
+		return hijacked.Hijack()
+	}
+
+	return nil, nil, errors.New("server doesn't support hijacking")
+}
+
+func (w *traceResponseWriter) Header() http.Header {
+	return w.w.Header()
+}
+
+func (w *traceResponseWriter) Write(data []byte) (int, error) {
+	return w.w.Write(data)
+}
+
+func (w *traceResponseWriter) WriteHeader(statusCode int) {
+	w.w.WriteHeader(statusCode)
+	w.code = statusCode
+}
+
 // TracingHandler return a middleware that process the opentelemetry.
 func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
@@ -48,7 +91,21 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler {
 
 			// convenient for tracking error messages
 			propagator.Inject(spanCtx, propagation.HeaderCarrier(w.Header()))
-			next.ServeHTTP(w, r.WithContext(spanCtx))
+			trw := &traceResponseWriter{
+				w:    w,
+				code: http.StatusOK,
+			}
+			next.ServeHTTP(trw, r.WithContext(spanCtx))
+
+			span.SetAttributes(attribute.KeyValue{
+				Key:   traceKeyStatusCode,
+				Value: attribute.IntValue(trw.code),
+			})
+			if trw.code >= http.StatusBadRequest {
+				span.SetStatus(codes.Error, http.StatusText(trw.code))
+			} else {
+				span.SetStatus(codes.Ok, http.StatusText(trw.code))
+			}
 		})
 	}
 }