fix: use correct response writer for tracing middle (#3693)

coderd/tracing/httpmw.go

@@ -8,6 +8,7 @@ import (
 	"github.com/go-chi/chi/v5"
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
+	"go.opentelemetry.io/otel/trace"
 )
 
 // HTTPMW adds tracing to http routes.
@@ -27,28 +28,37 @@ func HTTPMW(tracerProvider *sdktrace.TracerProvider, name string) func(http.Hand
 			wrw := middleware.NewWrapResponseWriter(rw, r.ProtoMajor)
 
 			// pass the span through the request context and serve the request to the next middleware
-			next.ServeHTTP(rw, r)
+			next.ServeHTTP(wrw, r)
 
-			// set the resource name as we get it only once the handler is executed
+			// capture response data
-			route := chi.RouteContext(r.Context()).RoutePattern()
+			EndHTTPSpan(r, wrw.Status())
-			if route != "" {
-				span.SetName(fmt.Sprintf("%s %s", r.Method, route))
-			}
-			span.SetName(fmt.Sprintf("%s %s", r.Method, route))
-			span.SetAttributes(semconv.NetAttributesFromHTTPRequest("tcp", r)...)
-			span.SetAttributes(semconv.EndUserAttributesFromHTTPRequest(r)...)
-			span.SetAttributes(semconv.HTTPServerAttributesFromHTTPRequest("", route, r)...)
-			span.SetAttributes(semconv.HTTPRouteKey.String(route))
-
-			// set the status code
-			status := wrw.Status()
-			// 0 status means one has not yet been sent in which case net/http library will write StatusOK
-			if status == 0 {
-				status = http.StatusOK
-			}
-			span.SetAttributes(semconv.HTTPStatusCodeKey.Int(status))
-			spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(status)
-			span.SetStatus(spanStatus, spanMessage)
 		})
 	}
 }
+
+// EndHTTPSpan captures request and response data after the handler is done.
+func EndHTTPSpan(r *http.Request, status int) {
+	span := trace.SpanFromContext(r.Context())
+
+	// set the resource name as we get it only once the handler is executed
+	route := chi.RouteContext(r.Context()).RoutePattern()
+	if route != "" {
+		span.SetName(fmt.Sprintf("%s %s", r.Method, route))
+	}
+	span.SetName(fmt.Sprintf("%s %s", r.Method, route))
+	span.SetAttributes(semconv.NetAttributesFromHTTPRequest("tcp", r)...)
+	span.SetAttributes(semconv.EndUserAttributesFromHTTPRequest(r)...)
+	span.SetAttributes(semconv.HTTPServerAttributesFromHTTPRequest("", route, r)...)
+	span.SetAttributes(semconv.HTTPRouteKey.String(route))
+
+	// 0 status means one has not yet been sent in which case net/http library will write StatusOK
+	if status == 0 {
+		status = http.StatusOK
+	}
+	span.SetAttributes(semconv.HTTPStatusCodeKey.Int(status))
+	spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(status)
+	span.SetStatus(spanStatus, spanMessage)
+
+	// finally end span
+	span.End()
+}

coderd/workspaceagents.go

@@ -14,7 +14,6 @@ import (
 	"github.com/google/uuid"
 	"github.com/hashicorp/yamux"
 	"github.com/tabbed/pqtype"
-	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/xerrors"
 	"inet.af/netaddr"
 	"nhooyr.io/websocket"
@@ -27,6 +26,7 @@ import (
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/rbac"
+	"github.com/coder/coder/coderd/tracing"
 	"github.com/coder/coder/coderd/turnconn"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/peer"
@@ -111,7 +111,7 @@ func (api *API) workspaceAgentDial(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	// end span so we don't get long lived trace data
-	trace.SpanFromContext(ctx).End()
+	tracing.EndHTTPSpan(r, 200)
 
 	err = peerbroker.ProxyListen(ctx, session, peerbroker.ProxyOptions{
 		ChannelID: workspaceAgent.ID.String(),
@@ -276,7 +276,7 @@ func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	// end span so we don't get long lived trace data
-	trace.SpanFromContext(ctx).End()
+	tracing.EndHTTPSpan(r, 200)
 
 	api.Logger.Info(ctx, "accepting agent", slog.F("resource", resource), slog.F("agent", workspaceAgent))
 
@@ -365,8 +365,8 @@ func (api *API) workspaceAgentTurn(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	ctx, wsNetConn := websocketNetConn(r.Context(), wsConn, websocket.MessageBinary)
-	defer wsNetConn.Close()          // Also closes conn.
+	defer wsNetConn.Close()     // Also closes conn.
-	trace.SpanFromContext(ctx).End() // end span so we don't get long lived trace data
+	tracing.EndHTTPSpan(r, 200) // end span so we don't get long lived trace data
 
 	api.Logger.Debug(ctx, "accepting turn connection", slog.F("remote-address", r.RemoteAddr), slog.F("local-address", localAddress))
 	select {
@@ -581,7 +581,7 @@ func (api *API) workspaceAgentWireguardListener(rw http.ResponseWriter, r *http.
 	defer subCancel()
 
 	// end span so we don't get long lived trace data
-	trace.SpanFromContext(ctx).End()
+	tracing.EndHTTPSpan(r, 200)
 
 	// Wait for the connection to close or the client to send a message.
 	//nolint:dogsled

coderd/workspaceapps.go

@@ -16,6 +16,7 @@ import (
 	"github.com/coder/coder/coderd/httpapi"
 	"github.com/coder/coder/coderd/httpmw"
 	"github.com/coder/coder/coderd/rbac"
+	"github.com/coder/coder/coderd/tracing"
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/site"
 )
@@ -177,5 +178,9 @@ func (api *API) workspaceAppsProxyPath(rw http.ResponseWriter, r *http.Request)
 		r.Header.Add("Cookie", httpapi.StripCoderCookies(cookieHeader))
 	}
 	proxy.Transport = conn.HTTPTransport()
+
+	// end span so we don't get long lived trace data
+	tracing.EndHTTPSpan(r, 200)
+
 	proxy.ServeHTTP(rw, r)
 }