mirror of
https://github.com/zeromicro/go-zero.git
synced 2025-01-23 09:00:20 +08:00
refactor: use opentelemetry's standard api to track http status code (#2760)
This commit is contained in:
parent
e71c505e94
commit
fc9b3ffdc1
@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
@ -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() {
|
||||
}
|
||||
|
@ -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
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user