feat(http/transport): allow passing a function to customize trace propagation
This commit is contained in:
parent
c6c33c7024
commit
d373c1bc7f
1 changed files with 20 additions and 4 deletions
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue