otelmux: Treat incoming trace ID as links rather than parents (#3661)

* Treat incoming trace ID as links rather than parents

* CHANGELOG.md updated

* Fix  double import

* Update CHANGELOG.md

Co-authored-by: Anthony Mirabella <a9@aneurysm9.com>

* Removed publicEndpointFn nil check in Middleware. Guard in ServeHTTP already prevents nil deref.

---------

Co-authored-by: Anthony Mirabella <a9@aneurysm9.com>
Co-authored-by: Tyler Yahn <MrAlias@users.noreply.github.com>
Co-authored-by: Robert Pająk <pellared@hotmail.com>

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Added
 
 - Add the new `go.opentelemetry.io/contrib/instrgen` package to provide auto-generated source code instrumentation. (#3068, #3108)
+- The `WithPublicEndpoint` and `WithPublicEndpointFn` options in `go.opentelemetry.io/contrib/instrumentation/github.com/gorilla/mux/otelmux`.
 
 ## [1.16.0/0.41.0/0.10.0] - 2023-04-28
 

instrumentation/github.com/gorilla/mux/otelmux/config.go

@@ -26,6 +26,8 @@ type config struct {
 	TracerProvider    oteltrace.TracerProvider
 	Propagators       propagation.TextMapPropagator
 	spanNameFormatter func(string, *http.Request) string
+	PublicEndpoint    bool
+	PublicEndpointFn  func(*http.Request) bool
 }
 
 // Option specifies instrumentation configuration options.
@@ -39,6 +41,26 @@ func (o optionFunc) apply(c *config) {
 	o(c)
 }
 
+// WithPublicEndpoint configures the Handler to link the span with an incoming
+// span context. If this option is not provided, then the association is a child
+// association instead of a link.
+func WithPublicEndpoint() Option {
+	return optionFunc(func(c *config) {
+		c.PublicEndpoint = true
+	})
+}
+
+// WithPublicEndpointFn runs with every request, and allows conditionnally
+// configuring the Handler to link the span with an incoming span context. If
+// this option is not provided or returns false, then the association is a
+// child association instead of a link.
+// Note: WithPublicEndpoint takes precedence over WithPublicEndpointFn.
+func WithPublicEndpointFn(fn func(*http.Request) bool) Option {
+	return optionFunc(func(c *config) {
+		c.PublicEndpointFn = fn
+	})
+}
+
 // WithPropagators specifies propagators to use for extracting
 // information from the HTTP requests. If none are specified, global
 // ones will be used.

instrumentation/github.com/gorilla/mux/otelmux/mux.go

@@ -27,7 +27,7 @@ import (
 	"go.opentelemetry.io/otel/propagation"
 	semconv "go.opentelemetry.io/otel/semconv/v1.17.0"
 	"go.opentelemetry.io/otel/semconv/v1.17.0/httpconv"
-	oteltrace "go.opentelemetry.io/otel/trace"
+	"go.opentelemetry.io/otel/trace"
 )
 
 const (
@@ -47,31 +47,36 @@ func Middleware(service string, opts ...Option) mux.MiddlewareFunc {
 	}
 	tracer := cfg.TracerProvider.Tracer(
 		tracerName,
-		oteltrace.WithInstrumentationVersion(SemVersion()),
+		trace.WithInstrumentationVersion(SemVersion()),
 	)
 	if cfg.Propagators == nil {
 		cfg.Propagators = otel.GetTextMapPropagator()
 	}
 	if cfg.spanNameFormatter == nil {
 		cfg.spanNameFormatter = defaultSpanNameFunc
 	}
+
 	return func(handler http.Handler) http.Handler {
 		return traceware{
 			service:           service,
 			tracer:            tracer,
 			propagators:       cfg.Propagators,
 			handler:           handler,
 			spanNameFormatter: cfg.spanNameFormatter,
+			publicEndpoint:    cfg.PublicEndpoint,
+			publicEndpointFn:  cfg.PublicEndpointFn,
 		}
 	}
 }
 
 type traceware struct {
 	service           string
-	tracer            oteltrace.Tracer
+	tracer            trace.Tracer
 	propagators       propagation.TextMapPropagator
 	handler           http.Handler
 	spanNameFormatter func(string, *http.Request) string
+	publicEndpoint    bool
+	publicEndpointFn  func(*http.Request) bool
 }
 
 type recordingResponseWriter struct {
@@ -136,15 +141,25 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 	}
-	opts := []oteltrace.SpanStartOption{
-		oteltrace.WithAttributes(httpconv.ServerRequest(tw.service, r)...),
-		oteltrace.WithSpanKind(oteltrace.SpanKindServer),
+
+	opts := []trace.SpanStartOption{
+		trace.WithAttributes(httpconv.ServerRequest(tw.service, r)...),
+		trace.WithSpanKind(trace.SpanKindServer),
+	}
+
+	if tw.publicEndpoint || (tw.publicEndpointFn != nil && tw.publicEndpointFn(r.WithContext(ctx))) {
+		opts = append(opts, trace.WithNewRoot())
+		// Linking incoming span context if any for public endpoint.
+		if s := trace.SpanContextFromContext(ctx); s.IsValid() && s.IsRemote() {
+			opts = append(opts, trace.WithLinks(trace.Link{SpanContext: s}))
+		}
 	}
+
 	if routeStr == "" {
 		routeStr = fmt.Sprintf("HTTP %s route not found", r.Method)
 	} else {
 		rAttr := semconv.HTTPRoute(routeStr)
-		opts = append(opts, oteltrace.WithAttributes(rAttr))
+		opts = append(opts, trace.WithAttributes(rAttr))
 	}
 	spanName := tw.spanNameFormatter(routeStr, r)
 	ctx, span := tw.tracer.Start(ctx, spanName, opts...)

instrumentation/github.com/gorilla/mux/otelmux/test/mux_test.go

@@ -15,6 +15,7 @@
 package test
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -27,6 +28,7 @@ import (
 	"go.opentelemetry.io/contrib/instrumentation/github.com/gorilla/mux/otelmux"
 	"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/sdk/trace/tracetest"
 	"go.opentelemetry.io/otel/trace"
@@ -160,3 +162,134 @@ func assertSpan(t *testing.T, span sdktrace.ReadOnlySpan, name string, kind trac
 		assert.Equal(t, got[want.Key], want.Value)
 	}
 }
+
+func TestWithPublicEndpoint(t *testing.T) {
+	sr := tracetest.NewSpanRecorder()
+	provider := sdktrace.NewTracerProvider()
+	provider.RegisterSpanProcessor(sr)
+
+	remoteSpan := trace.SpanContextConfig{
+		TraceID: trace.TraceID{0x01},
+		SpanID:  trace.SpanID{0x01},
+		Remote:  true,
+	}
+	prop := propagation.TraceContext{}
+
+	router := mux.NewRouter()
+	router.Use(otelmux.Middleware("foobar",
+		otelmux.WithPublicEndpoint(),
+		otelmux.WithPropagators(prop),
+		otelmux.WithTracerProvider(provider),
+	))
+	router.HandleFunc("/with/public/endpoint", func(w http.ResponseWriter, r *http.Request) {
+		s := trace.SpanFromContext(r.Context())
+		sc := s.SpanContext()
+
+		// Should be with new root trace.
+		assert.True(t, sc.IsValid())
+		assert.False(t, sc.IsRemote())
+		assert.NotEqual(t, remoteSpan.TraceID, sc.TraceID())
+	})
+
+	r0 := httptest.NewRequest("GET", "/with/public/endpoint", nil)
+	w := httptest.NewRecorder()
+
+	sc := trace.NewSpanContext(remoteSpan)
+	ctx := trace.ContextWithSpanContext(context.Background(), sc)
+	prop.Inject(ctx, propagation.HeaderCarrier(r0.Header))
+
+	router.ServeHTTP(w, r0)
+	assert.Equal(t, http.StatusOK, w.Result().StatusCode)
+
+	// Recorded span should be linked with an incoming span context.
+	assert.NoError(t, sr.ForceFlush(ctx))
+	done := sr.Ended()
+	require.Len(t, done, 1)
+	require.Len(t, done[0].Links(), 1, "should contain link")
+	require.True(t, sc.Equal(done[0].Links()[0].SpanContext), "should link incoming span context")
+}
+
+func TestWithPublicEndpointFn(t *testing.T) {
+	remoteSpan := trace.SpanContextConfig{
+		TraceID:    trace.TraceID{0x01},
+		SpanID:     trace.SpanID{0x01},
+		TraceFlags: trace.FlagsSampled,
+		Remote:     true,
+	}
+	prop := propagation.TraceContext{}
+
+	testdata := []struct {
+		name          string
+		fn            func(*http.Request) bool
+		handlerAssert func(*testing.T, trace.SpanContext)
+		spansAssert   func(*testing.T, trace.SpanContext, []sdktrace.ReadOnlySpan)
+	}{
+		{
+			name: "with the method returning true",
+			fn: func(r *http.Request) bool {
+				return true
+			},
+			handlerAssert: func(t *testing.T, sc trace.SpanContext) {
+				// Should be with new root trace.
+				assert.True(t, sc.IsValid())
+				assert.False(t, sc.IsRemote())
+				assert.NotEqual(t, remoteSpan.TraceID, sc.TraceID())
+			},
+			spansAssert: func(t *testing.T, sc trace.SpanContext, spans []sdktrace.ReadOnlySpan) {
+				require.Len(t, spans, 1)
+				require.Len(t, spans[0].Links(), 1, "should contain link")
+				require.True(t, sc.Equal(spans[0].Links()[0].SpanContext), "should link incoming span context")
+			},
+		},
+		{
+			name: "with the method returning false",
+			fn: func(r *http.Request) bool {
+				return false
+			},
+			handlerAssert: func(t *testing.T, sc trace.SpanContext) {
+				// Should have remote span as parent
+				assert.True(t, sc.IsValid())
+				assert.False(t, sc.IsRemote())
+				assert.Equal(t, remoteSpan.TraceID, sc.TraceID())
+			},
+			spansAssert: func(t *testing.T, _ trace.SpanContext, spans []sdktrace.ReadOnlySpan) {
+				require.Len(t, spans, 1)
+				require.Len(t, spans[0].Links(), 0, "should not contain link")
+			},
+		},
+	}
+
+	for _, tt := range testdata {
+		t.Run(tt.name, func(t *testing.T) {
+			sr := tracetest.NewSpanRecorder()
+			provider := sdktrace.NewTracerProvider()
+			provider.RegisterSpanProcessor(sr)
+
+			router := mux.NewRouter()
+			router.Use(otelmux.Middleware("foobar",
+				otelmux.WithPublicEndpointFn(tt.fn),
+				otelmux.WithPropagators(prop),
+				otelmux.WithTracerProvider(provider),
+			))
+			router.HandleFunc("/with/public/endpointfn", func(w http.ResponseWriter, r *http.Request) {
+				s := trace.SpanFromContext(r.Context())
+				tt.handlerAssert(t, s.SpanContext())
+			})
+
+			r0 := httptest.NewRequest("GET", "/with/public/endpointfn", nil)
+			w := httptest.NewRecorder()
+
+			sc := trace.NewSpanContext(remoteSpan)
+			ctx := trace.ContextWithSpanContext(context.Background(), sc)
+			prop.Inject(ctx, propagation.HeaderCarrier(r0.Header))
+
+			router.ServeHTTP(w, r0)
+			assert.Equal(t, http.StatusOK, w.Result().StatusCode)
+
+			// Recorded span should be linked with an incoming span context.
+			assert.NoError(t, sr.ForceFlush(ctx))
+			spans := sr.Ended()
+			tt.spansAssert(t, sc, spans)
+		})
+	}
+}