From 0211c8a1fc1142c2e2c0ba3ddbfa1b6726e4641f Mon Sep 17 00:00:00 2001 From: Pulak Bhowmick Date: Thu, 26 Oct 2023 16:13:15 +0600 Subject: [PATCH 1/3] improve trace middleware writer --- pkg/middleware/trace.go | 21 ++++----------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/pkg/middleware/trace.go b/pkg/middleware/trace.go index 9b11768f..ebfd2a04 100644 --- a/pkg/middleware/trace.go +++ b/pkg/middleware/trace.go @@ -25,6 +25,7 @@ import ( "net/http" "sync" + "github.com/go-chi/chi/v5/middleware" "github.com/optimizely/agent/config" "github.com/rs/zerolog/log" "go.opentelemetry.io/otel" @@ -80,16 +81,6 @@ func (gen *traceIDGenerator) NewIDs(ctx context.Context) (trace.TraceID, trace.S return tid, sid } -type statusRecorder struct { - http.ResponseWriter - statusCode int -} - -func (r *statusRecorder) WriteHeader(code int) { - r.statusCode = code - r.ResponseWriter.WriteHeader(code) -} - func AddTracing(conf config.TracingConfig, tracerName, spanName string) func(http.Handler) http.Handler { return func(next http.Handler) http.Handler { fn := func(w http.ResponseWriter, r *http.Request) { @@ -107,15 +98,11 @@ func AddTracing(conf config.TracingConfig, tracerName, spanName string) func(htt attribute.String(OptlySDKHeader, r.Header.Get(OptlySDKHeader)), ) - rec := &statusRecorder{ - ResponseWriter: w, - statusCode: http.StatusOK, - } - - next.ServeHTTP(rec, r.WithContext(ctx)) + next.ServeHTTP(w, r.WithContext(ctx)) + statusCode := middleware.NewWrapResponseWriter(w, r.ProtoMajor).Status() span.SetAttributes( - semconv.HTTPStatusCodeKey.Int(rec.statusCode), + semconv.HTTPStatusCodeKey.Int(statusCode), ) } return http.HandlerFunc(fn) From 1cf0cdf1ff8ba5c3b327106f5b8e251e35c07f48 Mon Sep 17 00:00:00 2001 From: Pulak Bhowmick Date: Thu, 26 Oct 2023 16:23:57 +0600 Subject: [PATCH 2/3] fix failing acceptance test --- tests/acceptance/test_acceptance/test_odp_redis.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/acceptance/test_acceptance/test_odp_redis.py b/tests/acceptance/test_acceptance/test_odp_redis.py index a58e2db2..7e4b5fdc 100644 --- a/tests/acceptance/test_acceptance/test_odp_redis.py +++ b/tests/acceptance/test_acceptance/test_odp_redis.py @@ -52,6 +52,7 @@ def test_redis_save(session_override_sdk_key_odp): """ expected_segments = ["atsbugbashsegmenthaspurchased", "atsbugbashsegmentdob"] + expected_segments_rev = ["atsbugbashsegmentdob", "atsbugbashsegmenthaspurchased"] uId = "fs_user_id-$-matjaz-user-1" r = redis.Redis(host='localhost', port=6379, db=0) # clean redis before testing since several tests use same user_id @@ -72,7 +73,7 @@ def test_redis_save(session_override_sdk_key_odp): params=params) # Check saved segments - assert json.loads(json.dumps(expected_segments)) == json.loads(r.get(uId)) + assert json.loads(json.dumps(expected_segments)) == json.loads(r.get(uId)) or json.loads(json.dumps(expected_segments_rev)) == json.loads(r.get(uId)) assert json.loads(json.dumps(expected_redis_save)) == resp.json() assert resp.status_code == 200, resp.text From a78adfb40e558f70f845a15b52b3483b0a08cb1e Mon Sep 17 00:00:00 2001 From: Pulak Bhowmick Date: Thu, 26 Oct 2023 16:46:56 +0600 Subject: [PATCH 3/3] fix bug --- pkg/middleware/trace.go | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pkg/middleware/trace.go b/pkg/middleware/trace.go index ebfd2a04..6caf33d3 100644 --- a/pkg/middleware/trace.go +++ b/pkg/middleware/trace.go @@ -98,11 +98,12 @@ func AddTracing(conf config.TracingConfig, tracerName, spanName string) func(htt attribute.String(OptlySDKHeader, r.Header.Get(OptlySDKHeader)), ) - next.ServeHTTP(w, r.WithContext(ctx)) + respWriter := middleware.NewWrapResponseWriter(w, r.ProtoMajor) + + next.ServeHTTP(respWriter, r.WithContext(ctx)) - statusCode := middleware.NewWrapResponseWriter(w, r.ProtoMajor).Status() span.SetAttributes( - semconv.HTTPStatusCodeKey.Int(statusCode), + semconv.HTTPStatusCodeKey.Int(respWriter.Status()), ) } return http.HandlerFunc(fn)