From 82d44e0953dd9e27c201b590d0600464f1ce6d19 Mon Sep 17 00:00:00 2001
From: Matheus Nogueira <matheus.nogueira2008@gmail.com>
Date: Wed, 13 Dec 2023 16:23:39 -0300
Subject: [PATCH] telemetry(server): add more info about response in HTTP spans
 (#3440)

* telemetry(server): add more info from response in http endpoints

* fix: only add body attribute if it's a 5xx error

* remove set status in case of 5xx error
---
 server/http/custom_routes.go               | 21 ++++++++++++----
 server/http/middleware/metrics.go          | 28 ++++++++++++++++++----
 server/resourcemanager/resource_manager.go | 20 +++++++++++++---
 3 files changed, 57 insertions(+), 12 deletions(-)

diff --git a/server/http/custom_routes.go b/server/http/custom_routes.go
index 59894986b2..027a155cc4 100644
--- a/server/http/custom_routes.go
+++ b/server/http/custom_routes.go
@@ -9,6 +9,7 @@ import (
 	"strconv"
 
 	"github.com/gorilla/mux"
+	"github.com/kubeshop/tracetest/server/http/middleware"
 	"github.com/kubeshop/tracetest/server/openapi"
 	"github.com/kubeshop/tracetest/server/resourcemanager"
 	"go.opentelemetry.io/otel"
@@ -201,18 +202,30 @@ func (c *customController) instrumentRoute(name string, route string, f http.Han
 		}
 		headersJson, _ := json.Marshal(headers)
 
-		span.SetAttributes(
+		newRequest := r.WithContext(ctx)
+		responseWriter := middleware.NewStatusCodeCapturerWriter(w)
+
+		f(responseWriter, newRequest)
+
+		responseBody := responseWriter.Body()
+
+		attributes := []attribute.KeyValue{
 			attribute.String(string(semconv.HTTPMethodKey), r.Method),
 			attribute.String(string(semconv.HTTPRouteKey), route),
 			attribute.String(string(semconv.HTTPTargetKey), r.URL.String()),
 			attribute.String("http.request.params", string(paramsJson)),
 			attribute.String("http.request.query", string(queryStringJson)),
 			attribute.String("http.request.headers", string(headersJson)),
-		)
+			attribute.Int("http.response.status_code", responseWriter.StatusCode()),
+		}
 
-		newRequest := r.WithContext(ctx)
+		if responseWriter.StatusCode() >= 500 {
+			span.RecordError(fmt.Errorf("faulty server response"))
+
+			attributes = append(attributes, attribute.String("http.response.body", string(responseBody)))
+		}
 
-		f(w, newRequest)
+		span.SetAttributes(attributes...)
 	}
 }
 
diff --git a/server/http/middleware/metrics.go b/server/http/middleware/metrics.go
index 78a5dacebd..479f2d5eb0 100644
--- a/server/http/middleware/metrics.go
+++ b/server/http/middleware/metrics.go
@@ -46,16 +46,34 @@ var _ http.Handler = &httpMetricMiddleware{}
 
 type responseWriter struct {
 	http.ResponseWriter
-	statusCode int
+	responseBody []byte
+	statusCode   int
 }
 
 func NewStatusCodeCapturerWriter(w http.ResponseWriter) *responseWriter {
-	return &responseWriter{w, http.StatusOK}
+	return &responseWriter{
+		w,
+		[]byte{},
+		http.StatusOK,
+	}
+}
+
+func (w *responseWriter) WriteHeader(code int) {
+	w.statusCode = code
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *responseWriter) Write(body []byte) (int, error) {
+	w.responseBody = body
+	return w.ResponseWriter.Write(body)
+}
+
+func (w *responseWriter) StatusCode() int {
+	return w.statusCode
 }
 
-func (lrw *responseWriter) WriteHeader(code int) {
-	lrw.statusCode = code
-	lrw.ResponseWriter.WriteHeader(code)
+func (w *responseWriter) Body() []byte {
+	return w.responseBody
 }
 
 func NewMetricMiddleware(meter metric.Meter) mux.MiddlewareFunc {
diff --git a/server/resourcemanager/resource_manager.go b/server/resourcemanager/resource_manager.go
index 9d421aa646..1e35063d42 100644
--- a/server/resourcemanager/resource_manager.go
+++ b/server/resourcemanager/resource_manager.go
@@ -19,6 +19,7 @@ import (
 	"golang.org/x/exp/slices"
 
 	"github.com/gorilla/mux"
+	"github.com/kubeshop/tracetest/server/http/middleware"
 	"github.com/kubeshop/tracetest/server/pkg/id"
 	"github.com/kubeshop/tracetest/server/pkg/validation"
 )
@@ -220,16 +221,29 @@ func (m *manager[T]) instrumentRoute(route *mux.Route) {
 		}
 		headersJson, _ := json.Marshal(headers)
 
-		span.SetAttributes(
+		responseWriter := middleware.NewStatusCodeCapturerWriter(w)
+
+		originalHandler.ServeHTTP(responseWriter, r.WithContext(ctx))
+
+		responseBody := responseWriter.Body()
+
+		attributes := []attribute.KeyValue{
 			attribute.String(string(semconv.HTTPMethodKey), r.Method),
 			attribute.String(string(semconv.HTTPRouteKey), pathTemplate),
 			attribute.String(string(semconv.HTTPTargetKey), r.URL.String()),
 			attribute.String("http.request.params", string(paramsJson)),
 			attribute.String("http.request.query", string(queryStringJson)),
 			attribute.String("http.request.headers", string(headersJson)),
-		)
+			attribute.Int("http.response.status_code", responseWriter.StatusCode()),
+		}
+
+		if responseWriter.StatusCode() >= 500 {
+			span.RecordError(fmt.Errorf("faulty server response"))
+
+			attributes = append(attributes, attribute.String("http.response.body", string(responseBody)))
+		}
 
-		originalHandler.ServeHTTP(w, r.WithContext(ctx))
+		span.SetAttributes(attributes...)
 	})
 
 	route.Handler(newHandler)