Setup PublicEndpointFn option (#2342)

* setup PublicEndpointFn option

* note that WithPublicEndpoint takes precedence

instrumentation/net/http/otelhttp/config.go

@@ -38,6 +38,7 @@ type config struct {
 	Propagators       propagation.TextMapPropagator
 	SpanStartOptions  []trace.SpanStartOption
 	PublicEndpoint    bool
+	PublicEndpointFn  func(*http.Request) bool
 	ReadEvent         bool
 	WriteEvent        bool
 	Filters           []Filter
@@ -108,7 +109,17 @@ func WithMeterProvider(provider metric.MeterProvider) Option {
 func WithPublicEndpoint() Option {
 	return optionFunc(func(c *config) {
 		c.PublicEndpoint = true
-		c.SpanStartOptions = append(c.SpanStartOptions, trace.WithNewRoot())
+	})
+}
+
+// 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
 	})
 }
 

instrumentation/net/http/otelhttp/handler.go

@@ -52,6 +52,7 @@ type Handler struct {
 	counters          map[string]syncint64.Counter
 	valueRecorders    map[string]syncfloat64.Histogram
 	publicEndpoint    bool
+	publicEndpointFn  func(*http.Request) bool
 }
 
 func defaultHandlerFormatter(operation string, _ *http.Request) string {
@@ -88,6 +89,7 @@ func (h *Handler) configure(c *config) {
 	h.filters = c.Filters
 	h.spanNameFormatter = c.SpanNameFormatter
 	h.publicEndpoint = c.PublicEndpoint
+	h.publicEndpointFn = c.PublicEndpointFn
 }
 
 func handleErr(err error) {
@@ -125,11 +127,21 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	opts := append([]trace.SpanStartOption{
+	ctx := h.propagators.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+	opts := h.spanStartOptions
+	if h.publicEndpoint || (h.publicEndpointFn != nil && h.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}))
+		}
+	}
+
+	opts = append([]trace.SpanStartOption{
 		trace.WithAttributes(semconv.NetAttributesFromHTTPRequest("tcp", r)...),
 		trace.WithAttributes(semconv.EndUserAttributesFromHTTPRequest(r)...),
 		trace.WithAttributes(semconv.HTTPServerAttributesFromHTTPRequest(h.operation, "", r)...),
-	}, h.spanStartOptions...) // start with the configured options
+	}, opts...) // start with the configured options
 
 	tracer := h.tracer
 
@@ -141,13 +153,6 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	ctx := h.propagators.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
-	if h.publicEndpoint {
-		// 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}))
-		}
-	}
 	ctx, span := tracer.Start(ctx, h.spanNameFormatter(h.operation, r), opts...)
 	defer span.End()
 

instrumentation/net/http/otelhttp/test/handler_test.go

@@ -190,6 +190,90 @@ func TestWithPublicEndpoint(t *testing.T) {
 	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{}
+
+	for _, tt := range []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")
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			spanRecorder := tracetest.NewSpanRecorder()
+			provider := sdktrace.NewTracerProvider(
+				sdktrace.WithSpanProcessor(spanRecorder),
+			)
+
+			h := otelhttp.NewHandler(
+				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					s := trace.SpanFromContext(r.Context())
+					tt.handlerAssert(t, s.SpanContext())
+				}), "test_handler",
+				otelhttp.WithPublicEndpointFn(tt.fn),
+				otelhttp.WithPropagators(prop),
+				otelhttp.WithTracerProvider(provider),
+			)
+
+			r, err := http.NewRequest(http.MethodGet, "http://localhost/", nil)
+			require.NoError(t, err)
+
+			sc := trace.NewSpanContext(remoteSpan)
+			ctx := trace.ContextWithSpanContext(context.Background(), sc)
+			prop.Inject(ctx, propagation.HeaderCarrier(r.Header))
+
+			rr := httptest.NewRecorder()
+			h.ServeHTTP(rr, r)
+			assert.Equal(t, 200, rr.Result().StatusCode)
+
+			// Recorded span should be linked with an incoming span context.
+			assert.NoError(t, spanRecorder.ForceFlush(ctx))
+			spans := spanRecorder.Ended()
+			tt.spansAssert(t, sc, spans)
+		})
+	}
+}
+
 func TestSpanStatus(t *testing.T) {
 	testCases := []struct {
 		httpStatusCode int