From b08997ec479f3d848908c8749c2bcbbb74ac6cb1 Mon Sep 17 00:00:00 2001
From: Anton Audzei <anton.audzei@outreach.io>
Date: Mon, 5 Aug 2024 11:40:06 +0200
Subject: [PATCH] feat(trace): propagate WithPublicEndpointFn to otel
 NewHandler

---
 pkg/trace/propagation.go  | 23 +++++++++++++++++++----
 pkg/trace/roundtripper.go |  9 +++++++--
 pkg/trace/tracer.go       |  2 +-
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/pkg/trace/propagation.go b/pkg/trace/propagation.go
index 0835d247..ee4b312b 100644
--- a/pkg/trace/propagation.go
+++ b/pkg/trace/propagation.go
@@ -40,8 +40,17 @@ func (t *otelTracer) newTransport(old http.RoundTripper) http.RoundTripper {
 }
 
 type Handler struct {
-	handler   http.Handler
-	operation string
+	handler          http.Handler
+	operation        string
+	publicEndpointFn func(*http.Request) bool
+}
+
+type HandlerOption func(*Handler)
+
+func WithPublicEndpointFn(fn func(*http.Request) bool) HandlerOption {
+	return func(h *Handler) {
+		h.publicEndpointFn = fn
+	}
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -52,18 +61,24 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	force := r.Header.Get(HeaderForceTracing)
 	if force != "" {
 		startOptions = oteltrace.WithAttributes(attribute.String(fieldForceTrace, force))
-		handler = otelhttp.NewHandler(h.handler, h.operation, otelhttp.WithSpanOptions(startOptions))
+		handler = otelhttp.NewHandler(h.handler, h.operation,
+			otelhttp.WithSpanOptions(startOptions),
+			otelhttp.WithPublicEndpointFn(h.publicEndpointFn), // passing nil function is equivalent to "not configured"
+		)
 		r = r.WithContext(forceTracing(r.Context()))
 	}
 
 	handler.ServeHTTP(w, r)
 }
 
-func (t *otelTracer) newHandler(handler http.Handler, operation string) http.Handler {
+func (t *otelTracer) newHandler(handler http.Handler, operation string, opts ...HandlerOption) http.Handler {
 	h := Handler{
 		handler:   handler,
 		operation: operation,
 	}
+	for _, opt := range opts {
+		opt(&h)
+	}
 
 	return &h
 }
diff --git a/pkg/trace/roundtripper.go b/pkg/trace/roundtripper.go
index 95740134..bee1060f 100644
--- a/pkg/trace/roundtripper.go
+++ b/pkg/trace/roundtripper.go
@@ -35,6 +35,11 @@ func NewTransport(old http.RoundTripper) http.RoundTripper {
 // NewHandler creates a new handlers which reads propagated headers
 // from the current trace context.
 //
+// Supported options are:
+//   - WithPublicEndpointFn conditionally configure the Handler to link the span with an incoming span context
+//     instead of child association. Useful for public facing endpoints that don't need to attach to externally
+//     defined traces
+//
 // Usage:
 //
 //		  trace.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) *roundtripperState {
@@ -42,10 +47,10 @@ func NewTransport(old http.RoundTripper) http.RoundTripper {
 //			defer trace.End(r.Context())
 //			... do actual request handling ...
 //	   }), "my endpoint")
-func NewHandler(handler http.Handler, operation string) http.Handler {
+func NewHandler(handler http.Handler, operation string, opts ...HandlerOption) http.Handler {
 	if defaultTracer == nil {
 		return handler
 	}
 
-	return defaultTracer.newHandler(handler, operation)
+	return defaultTracer.newHandler(handler, operation, opts...)
 }
diff --git a/pkg/trace/tracer.go b/pkg/trace/tracer.go
index bf261f4e..a7bf2c81 100644
--- a/pkg/trace/tracer.go
+++ b/pkg/trace/tracer.go
@@ -46,7 +46,7 @@ type tracer interface {
 
 	newTransport(http.RoundTripper) http.RoundTripper
 
-	newHandler(handler http.Handler, operation string) http.Handler
+	newHandler(handler http.Handler, operation string, opts ...HandlerOption) http.Handler
 
 	toHeaders(ctx context.Context) map[string][]string