86 lines
2.5 KiB
Go
86 lines
2.5 KiB
Go
package observability
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"go.opentelemetry.io/otel"
|
|
"go.opentelemetry.io/otel/attribute"
|
|
"go.opentelemetry.io/otel/codes"
|
|
"go.opentelemetry.io/otel/propagation"
|
|
sdktrace "go.opentelemetry.io/otel/sdk/trace"
|
|
"go.opentelemetry.io/otel/trace"
|
|
)
|
|
|
|
// TracingMiddleware creates a middleware that adds OpenTelemetry tracing to HTTP requests.
|
|
func TracingMiddleware(next http.Handler, tp *sdktrace.TracerProvider) http.Handler {
|
|
if tp == nil {
|
|
// If tracing is not enabled, pass through without modification
|
|
return next
|
|
}
|
|
|
|
// Set up W3C Trace Context propagation
|
|
otel.SetTextMapPropagator(propagation.TraceContext{})
|
|
|
|
tracer := tp.Tracer("llm-gateway")
|
|
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Extract trace context from incoming request headers
|
|
ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
|
|
|
|
// Start a new span
|
|
ctx, span := tracer.Start(ctx, "HTTP "+r.Method+" "+r.URL.Path,
|
|
trace.WithSpanKind(trace.SpanKindServer),
|
|
trace.WithAttributes(
|
|
attribute.String("http.method", r.Method),
|
|
attribute.String("http.route", r.URL.Path),
|
|
attribute.String("http.scheme", r.URL.Scheme),
|
|
attribute.String("http.host", r.Host),
|
|
attribute.String("http.user_agent", r.Header.Get("User-Agent")),
|
|
),
|
|
)
|
|
defer span.End()
|
|
|
|
// Add request ID to span if present
|
|
if requestID := r.Header.Get("X-Request-ID"); requestID != "" {
|
|
span.SetAttributes(attribute.String("http.request_id", requestID))
|
|
}
|
|
|
|
// Create a response writer wrapper to capture status code
|
|
wrapped := &statusResponseWriter{
|
|
ResponseWriter: w,
|
|
statusCode: http.StatusOK,
|
|
}
|
|
|
|
// Inject trace context into request for downstream services
|
|
r = r.WithContext(ctx)
|
|
|
|
// Call the next handler
|
|
next.ServeHTTP(wrapped, r)
|
|
|
|
// Record the status code in the span
|
|
span.SetAttributes(attribute.Int("http.status_code", wrapped.statusCode))
|
|
|
|
// Set span status based on HTTP status code
|
|
if wrapped.statusCode >= 400 {
|
|
span.SetStatus(codes.Error, http.StatusText(wrapped.statusCode))
|
|
} else {
|
|
span.SetStatus(codes.Ok, "")
|
|
}
|
|
})
|
|
}
|
|
|
|
// statusResponseWriter wraps http.ResponseWriter to capture the status code.
|
|
type statusResponseWriter struct {
|
|
http.ResponseWriter
|
|
statusCode int
|
|
}
|
|
|
|
func (w *statusResponseWriter) WriteHeader(statusCode int) {
|
|
w.statusCode = statusCode
|
|
w.ResponseWriter.WriteHeader(statusCode)
|
|
}
|
|
|
|
func (w *statusResponseWriter) Write(b []byte) (int, error) {
|
|
return w.ResponseWriter.Write(b)
|
|
}
|