diff --git a/rest/handler/tracinghandler.go b/rest/handler/tracinghandler.go index 43ad9d2c..d3151569 100644 --- a/rest/handler/tracinghandler.go +++ b/rest/handler/tracinghandler.go @@ -1,24 +1,18 @@ package handler import ( - "bufio" - "errors" - "net" "net/http" "sync" "github.com/zeromicro/go-zero/core/lang" "github.com/zeromicro/go-zero/core/trace" + "github.com/zeromicro/go-zero/rest/internal/response" "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. @@ -26,40 +20,6 @@ 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 { @@ -89,21 +49,12 @@ func TracingHandler(serviceName, path string) func(http.Handler) http.Handler { // convenient for tracking error messages propagator.Inject(spanCtx, propagation.HeaderCarrier(w.Header())) - trw := &traceResponseWriter{ - w: w, - code: http.StatusOK, - } + + trw := &response.WithCodeResponseWriter{Writer: 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)) - } + span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(trw.Code)...) + span.SetStatus(semconv.SpanStatusFromHTTPStatusCodeAndSpanKind(trw.Code, oteltrace.SpanKindServer)) }) } } diff --git a/rest/handler/tracinghandler_test.go b/rest/handler/tracinghandler_test.go index d4aebe2a..79a94c4d 100644 --- a/rest/handler/tracinghandler_test.go +++ b/rest/handler/tracinghandler_test.go @@ -148,46 +148,3 @@ func TestTraceResponseWriter(t *testing.T) { }) } } - -func TestTraceHandler_traceResponseWriter(t *testing.T) { - writer := &traceResponseWriter{ - w: httptest.NewRecorder(), - } - assert.NotPanics(t, func() { - writer.Hijack() - }) - - writer = &traceResponseWriter{ - w: mockedHijackable{httptest.NewRecorder()}, - } - assert.NotPanics(t, func() { - writer.Hijack() - }) - - writer = &traceResponseWriter{ - w: httptest.NewRecorder(), - } - writer.WriteHeader(http.StatusBadRequest) - assert.NotNil(t, writer.Header()) - - writer = &traceResponseWriter{ - w: httptest.NewRecorder(), - } - assert.NotPanics(t, func() { - writer.Flush() - }) - - writer = &traceResponseWriter{ - w: mockedFlusher{httptest.NewRecorder()}, - } - assert.NotPanics(t, func() { - writer.Flush() - }) -} - -type mockedFlusher struct { - *httptest.ResponseRecorder -} - -func (m mockedFlusher) Flush() { -} diff --git a/rest/httpc/requests.go b/rest/httpc/requests.go index 4ddb15a8..3685d5c6 100644 --- a/rest/httpc/requests.go +++ b/rest/httpc/requests.go @@ -190,7 +190,7 @@ func request(r *http.Request, cli client) (*http.Response, error) { } span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(resp.StatusCode)...) - span.SetStatus(semconv.SpanStatusFromHTTPStatusCode(resp.StatusCode)) + span.SetStatus(semconv.SpanStatusFromHTTPStatusCodeAndSpanKind(resp.StatusCode, oteltrace.SpanKindClient)) return resp, err }