feat(http/transport): allow passing a function to customize trace propagation

This commit is contained in:
DevMiner 2024-08-24 13:48:07 +02:00
parent f560b15a57
commit 109459cad0

View file

@ -24,12 +24,28 @@ func WithLogger(l logr.Logger) HTTPTransportOpt {
} }
} }
type TracePropagator = func(r *http.Request) bool
func WithTracePropagation(propagator TracePropagator) HTTPTransportOpt {
return func(t *HTTPTransport) {
t.tracePropagator = propagator
}
}
func PropagateAllTraces(req *http.Request) bool {
return true
}
func PropagateNoTraces(req *http.Request) bool {
return false
}
type HTTPTransport struct { type HTTPTransport struct {
logger logr.Logger logger logr.Logger
telemetry *unitel.Telemetry telemetry *unitel.Telemetry
transport http.RoundTripper transport http.RoundTripper
forwardTrace bool tracePropagator TracePropagator
tracedRequestHeaders []string tracedRequestHeaders []string
tracedResponseHeaders []string tracedResponseHeaders []string
@ -37,13 +53,13 @@ type HTTPTransport struct {
tracer trace.Tracer tracer trace.Tracer
} }
func NewTracedTransport(t *unitel.Telemetry, inner http.RoundTripper, forwardTrace bool, tracedRequestHeaders []string, tracedResponseHeaders []string, opts ...HTTPTransportOpt) *HTTPTransport { func NewTracedTransport(t *unitel.Telemetry, inner http.RoundTripper, tracedRequestHeaders []string, tracedResponseHeaders []string, opts ...HTTPTransportOpt) *HTTPTransport {
transport := &HTTPTransport{ transport := &HTTPTransport{
logger: logr.Discard(), logger: logr.Discard(),
telemetry: t, telemetry: t,
transport: inner, transport: inner,
forwardTrace: forwardTrace, tracePropagator: PropagateNoTraces,
tracedRequestHeaders: tracedRequestHeaders, tracedRequestHeaders: tracedRequestHeaders,
tracedResponseHeaders: tracedResponseHeaders, tracedResponseHeaders: tracedResponseHeaders,
@ -69,7 +85,7 @@ func (t *HTTPTransport) RoundTrip(req *http.Request) (*http.Response, error) {
ctx := s.Context() ctx := s.Context()
req = req.WithContext(ctx) req = req.WithContext(ctx)
if t.forwardTrace { if t.tracePropagator(req) {
t.telemetry.InjectIntoHeaders(ctx, req.Header) t.telemetry.InjectIntoHeaders(ctx, req.Header)
} }