feat: make trace provider in loadtest, add tracing to sdk (#4939)

coderd/activitybump_test.go

@@ -88,7 +88,7 @@ func TestWorkspaceActivityBump(t *testing.T) {
 		require.NoError(t, err)
 		defer conn.Close()
 
-		sshConn, err := conn.SSHClient()
+		sshConn, err := conn.SSHClient(ctx)
 		require.NoError(t, err)
 		_ = sshConn.Close()
 

coderd/templates_test.go

@@ -633,7 +633,7 @@ func TestTemplateMetrics(t *testing.T) {
 		_ = conn.Close()
 	}()
 
-	sshConn, err := conn.SSHClient()
+	sshConn, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	_ = sshConn.Close()
 

coderd/tracing/exporter.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/go-logr/logr"
+	"github.com/hashicorp/go-multierror"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
@@ -82,11 +83,23 @@ func TracerProvider(ctx context.Context, service string, opts TracerOpts) (*sdkt
 	otel.SetLogger(logr.Discard())
 
 	return tracerProvider, func(ctx context.Context) error {
-		for _, close := range closers {
-			_ = close(ctx)
+		var merr error
+		err := tracerProvider.ForceFlush(ctx)
+		if err != nil {
+			merr = multierror.Append(merr, xerrors.Errorf("tracerProvider.ForceFlush(): %w", err))
 		}
-		_ = tracerProvider.Shutdown(ctx)
-		return nil
+		for i, closer := range closers {
+			err = closer(ctx)
+			if err != nil {
+				merr = multierror.Append(merr, xerrors.Errorf("closer() %d: %w", i, err))
+			}
+		}
+		err = tracerProvider.Shutdown(ctx)
+		if err != nil {
+			merr = multierror.Append(merr, xerrors.Errorf("tracerProvider.Shutdown(): %w", err))
+		}
+
+		return merr
 	}, nil
 }
 

coderd/tracing/httpmw.go

@@ -6,6 +6,8 @@ import (
 	"net/http"
 
 	"github.com/go-chi/chi/v5"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.11.0"
 	"go.opentelemetry.io/otel/trace"
 )
@@ -23,11 +25,27 @@ func Middleware(tracerProvider trace.TracerProvider) func(http.Handler) http.Han
 				return
 			}
 
+			// Extract the trace context from the request headers.
+			tmp := otel.GetTextMapPropagator()
+			hc := propagation.HeaderCarrier(r.Header)
+			ctx := tmp.Extract(r.Context(), hc)
+
 			// start span with default span name. Span name will be updated to "method route" format once request finishes.
-			ctx, span := tracer.Start(r.Context(), fmt.Sprintf("%s %s", r.Method, r.RequestURI))
+			ctx, span := tracer.Start(ctx, fmt.Sprintf("%s %s", r.Method, r.RequestURI))
 			defer span.End()
 			r = r.WithContext(ctx)
 
+			if span.SpanContext().HasTraceID() && span.SpanContext().HasSpanID() {
+				// Technically these values are included in the Traceparent
+				// header, but they are easier to read for humans this way.
+				rw.Header().Set("X-Trace-ID", span.SpanContext().TraceID().String())
+				rw.Header().Set("X-Span-ID", span.SpanContext().SpanID().String())
+
+				// Inject the trace context into the response headers.
+				hc := propagation.HeaderCarrier(rw.Header())
+				tmp.Inject(ctx, hc)
+			}
+
 			sw, ok := rw.(*StatusWriter)
 			if !ok {
 				panic(fmt.Sprintf("ResponseWriter not a *tracing.StatusWriter; got %T", rw))
@@ -62,6 +80,37 @@ func EndHTTPSpan(r *http.Request, status int, span trace.Span) {
 	span.End()
 }
 
-func StartSpan(ctx context.Context, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
-	return trace.SpanFromContext(ctx).TracerProvider().Tracer(TracerName).Start(ctx, FuncNameSkip(1), opts...)
+type tracerNameKey struct{}
+
+// SetTracerName sets the tracer name that will be used by all spans created
+// from the context.
+func SetTracerName(ctx context.Context, tracerName string) context.Context {
+	return context.WithValue(ctx, tracerNameKey{}, tracerName)
+}
+
+// GetTracerName returns the tracer name from the context, or TracerName if none
+// is set.
+func GetTracerName(ctx context.Context) string {
+	if tracerName, ok := ctx.Value(tracerNameKey{}).(string); ok {
+		return tracerName
+	}
+
+	return TracerName
+}
+
+// StartSpan calls StartSpanWithName with the name set to the caller's function
+// name.
+func StartSpan(ctx context.Context, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
+	return StartSpanWithName(ctx, FuncNameSkip(1), opts...)
+}
+
+// StartSpanWithName starts a new span with the given name from the context. If
+// a tracer name was set on the context (or one of its parents), it will be used
+// as the tracer name instead of the default TracerName.
+func StartSpanWithName(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
+	tracerName := GetTracerName(ctx)
+	return trace.SpanFromContext(ctx).
+		TracerProvider().
+		Tracer(tracerName).
+		Start(ctx, name, opts...)
 }

coderd/workspaceagents.go

@@ -247,7 +247,7 @@ func (api *API) workspaceAgentPTY(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 	defer release()
-	ptNetConn, err := agentConn.ReconnectingPTY(reconnect.String(), uint16(height), uint16(width), r.URL.Query().Get("command"))
+	ptNetConn, err := agentConn.ReconnectingPTY(ctx, reconnect.String(), uint16(height), uint16(width), r.URL.Query().Get("command"))
 	if err != nil {
 		_ = conn.Close(websocket.StatusInternalError, httpapi.WebsocketCloseSprintf("dial: %s", err))
 		return

coderd/workspaceagents_test.go

@@ -260,7 +260,7 @@ func TestWorkspaceAgentTailnet(t *testing.T) {
 	})
 	require.NoError(t, err)
 	defer conn.Close()
-	sshClient, err := conn.SSHClient()
+	sshClient, err := conn.SSHClient(ctx)
 	require.NoError(t, err)
 	session, err := sshClient.NewSession()
 	require.NoError(t, err)