Merge pull request #62 from riandyrn/feat/trace-sampled

Feat: Add Trace Sampled on Response Headers

config.go

@@ -8,17 +8,22 @@ import (
 	oteltrace "go.opentelemetry.io/otel/trace"
 )
 
-const defaultTraceResponseHeaderKey = "X-Trace-Id"
+// These defaults are used in `TraceHeaderConfig`.
+const (
+	DefaultTraceIDResponseHeaderKey      = "X-Trace-Id"
+	DefaultTraceSampledResponseHeaderKey = "X-Trace-Sampled"
+)
 
 // config is used to configure the mux middleware.
 type config struct {
-	TracerProvider          oteltrace.TracerProvider
-	Propagators             propagation.TextMapPropagator
-	ChiRoutes               chi.Routes
-	RequestMethodInSpanName bool
-	Filters                 []Filter
-	TraceResponseHeaderKey  string
-	PublicEndpointFn        func(r *http.Request) bool
+	TracerProvider                oteltrace.TracerProvider
+	Propagators                   propagation.TextMapPropagator
+	ChiRoutes                     chi.Routes
+	RequestMethodInSpanName       bool
+	Filters                       []Filter
+	TraceIDResponseHeaderKey      string
+	TraceSampledResponseHeaderKey string
+	PublicEndpointFn              func(r *http.Request) bool
 }
 
 // Option specifies instrumentation configuration options.
@@ -32,7 +37,7 @@ func (o optionFunc) apply(c *config) {
 	o(c)
 }
 
-// Filter is a predicate used to determine whether a given http.request should
+// Filter is a predicate used to determine whether a given http.Request should
 // be traced. A Filter must return true if the request should be traced.
 type Filter func(*http.Request) bool
 
@@ -95,12 +100,39 @@ func WithFilter(filter Filter) Option {
 // WithTraceIDResponseHeader enables adding trace id into response header.
 // It accepts a function that generates the header key name. If this parameter
 // function set to `nil` the default header key which is `X-Trace-Id` will be used.
+//
+// Deprecated: use `WithTraceResponseHeaders` instead.
 func WithTraceIDResponseHeader(headerKeyFunc func() string) Option {
-	return optionFunc(func(cfg *config) {
-		if headerKeyFunc == nil {
-			cfg.TraceResponseHeaderKey = defaultTraceResponseHeaderKey // use default trace header
-		} else {
-			cfg.TraceResponseHeaderKey = headerKeyFunc()
+	cfg := TraceHeaderConfig{
+		TraceIDHeader:      "",
+		TraceSampledHeader: "",
+	}
+	if headerKeyFunc != nil {
+		cfg.TraceIDHeader = headerKeyFunc()
+	}
+	return WithTraceResponseHeaders(cfg)
+}
+
+// TraceHeaderConfig is configuration for trace headers in the response.
+type TraceHeaderConfig struct {
+	TraceIDHeader      string // if non-empty overrides the default of X-Trace-ID
+	TraceSampledHeader string // if non-empty overrides the default of X-Trace-Sampled
+}
+
+// WithTraceResponseHeaders configures the response headers for trace information.
+// It accepts a TraceHeaderConfig struct that contains the keys for the Trace ID
+// and Trace Sampled headers. If the provided keys are empty, default values will
+// be used for the respective headers.
+func WithTraceResponseHeaders(cfg TraceHeaderConfig) Option {
+	return optionFunc(func(c *config) {
+		c.TraceIDResponseHeaderKey = cfg.TraceIDHeader
+		if c.TraceIDResponseHeaderKey == "" {
+			c.TraceIDResponseHeaderKey = DefaultTraceIDResponseHeaderKey
+		}
+
+		c.TraceSampledResponseHeaderKey = cfg.TraceSampledHeader
+		if c.TraceSampledResponseHeaderKey == "" {
+			c.TraceSampledResponseHeaderKey = DefaultTraceSampledResponseHeaderKey
 		}
 	})
 }
@@ -138,7 +170,7 @@ func WithPublicEndpoint() Option {
 // incoming span context. Otherwise, the generated span will be set as the
 // child span of the incoming span context.
 //
-// Essentially it has the same functionality as WithPublicEndpoint but with
+// Essentially it has the same functionality as `WithPublicEndpoint` but with
 // more flexibility.
 func WithPublicEndpointFn(fn func(r *http.Request) bool) Option {
 	return optionFunc(func(cfg *config) {

middleware.go

@@ -2,6 +2,7 @@ package otelchi
 
 import (
 	"net/http"
+	"strconv"
 	"sync"
 
 	"github.com/felixge/httpsnoop"
@@ -40,29 +41,31 @@ func Middleware(serverName string, opts ...Option) func(next http.Handler) http.
 
 	return func(handler http.Handler) http.Handler {
 		return traceware{
-			serverName:             serverName,
-			tracer:                 tracer,
-			propagators:            cfg.Propagators,
-			handler:                handler,
-			chiRoutes:              cfg.ChiRoutes,
-			reqMethodInSpanName:    cfg.RequestMethodInSpanName,
-			filters:                cfg.Filters,
-			traceResponseHeaderKey: cfg.TraceResponseHeaderKey,
-			publicEndpointFn:       cfg.PublicEndpointFn,
+			serverName:                    serverName,
+			tracer:                        tracer,
+			propagators:                   cfg.Propagators,
+			handler:                       handler,
+			chiRoutes:                     cfg.ChiRoutes,
+			reqMethodInSpanName:           cfg.RequestMethodInSpanName,
+			filters:                       cfg.Filters,
+			traceIDResponseHeaderKey:      cfg.TraceIDResponseHeaderKey,
+			traceSampledResponseHeaderKey: cfg.TraceSampledResponseHeaderKey,
+			publicEndpointFn:              cfg.PublicEndpointFn,
 		}
 	}
 }
 
 type traceware struct {
-	serverName             string
-	tracer                 oteltrace.Tracer
-	propagators            propagation.TextMapPropagator
-	handler                http.Handler
-	chiRoutes              chi.Routes
-	reqMethodInSpanName    bool
-	filters                []Filter
-	traceResponseHeaderKey string
-	publicEndpointFn       func(r *http.Request) bool
+	serverName                    string
+	tracer                        oteltrace.Tracer
+	propagators                   propagation.TextMapPropagator
+	handler                       http.Handler
+	chiRoutes                     chi.Routes
+	reqMethodInSpanName           bool
+	filters                       []Filter
+	traceIDResponseHeaderKey      string
+	traceSampledResponseHeaderKey string
+	publicEndpointFn              func(r *http.Request) bool
 }
 
 type recordingResponseWriter struct {
@@ -175,9 +178,10 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx, span := tw.tracer.Start(ctx, spanName, spanOpts...)
 	defer span.End()
 
-	// put trace_id to response header only when WithTraceResponseHeaderKey is used
-	if len(tw.traceResponseHeaderKey) > 0 && span.SpanContext().HasTraceID() {
-		w.Header().Add(tw.traceResponseHeaderKey, span.SpanContext().TraceID().String())
+	// put trace_id to response header only when `WithTraceIDResponseHeader` is used
+	if len(tw.traceIDResponseHeaderKey) > 0 && span.SpanContext().HasTraceID() {
+		w.Header().Add(tw.traceIDResponseHeaderKey, span.SpanContext().TraceID().String())
+		w.Header().Add(tw.traceSampledResponseHeaderKey, strconv.FormatBool(span.SpanContext().IsSampled()))
 	}
 
 	// get recording response writer