diff --git a/.palantir/revapi.yml b/.palantir/revapi.yml index f1212ff3a..60f43ad45 100644 --- a/.palantir/revapi.yml +++ b/.palantir/revapi.yml @@ -57,3 +57,8 @@ acceptedBreaks: old: "method com.palantir.tracing.api.OpenSpan.Builder com.palantir.tracing.api.ImmutableOpenSpan.Builder::originatingSpanId(java.util.Optional)\ \ @ com.palantir.tracing.api.OpenSpan.Builder" justification: "Type is not meant for external creation" + "6.4.0": + com.palantir.tracing:tracing: + - code: "java.method.addedToInterface" + new: "method java.util.Optional com.palantir.tracing.TraceMetadata::getForUserAgent()" + justification: "Adding a method will not break anyone" diff --git a/tracing-api/src/main/java/com/palantir/tracing/api/TraceHttpHeaders.java b/tracing-api/src/main/java/com/palantir/tracing/api/TraceHttpHeaders.java index 80ed5ecf2..40c5eb01a 100644 --- a/tracing-api/src/main/java/com/palantir/tracing/api/TraceHttpHeaders.java +++ b/tracing-api/src/main/java/com/palantir/tracing/api/TraceHttpHeaders.java @@ -37,4 +37,6 @@ public interface TraceHttpHeaders { */ @Deprecated String ORIGINATING_SPAN_ID = "X-OrigSpanId"; + + String FOR_USER_AGENT = "X-For-User-Agent"; } diff --git a/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java b/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java index ab51a60cf..dc7d89cc7 100644 --- a/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java +++ b/tracing-jersey/src/main/java/com/palantir/tracing/jersey/TraceEnrichingFilter.java @@ -35,6 +35,7 @@ import javax.ws.rs.container.ContainerResponseContext; import javax.ws.rs.container.ContainerResponseFilter; import javax.ws.rs.core.Context; +import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.MultivaluedMap; import javax.ws.rs.ext.Provider; import org.glassfish.jersey.server.ExtendedUriInfo; @@ -53,6 +54,8 @@ public final class TraceEnrichingFilter implements ContainerRequestFilter, Conta public static final String REQUEST_ID_PROPERTY_NAME = "com.palantir.tracing.requestId"; + public static final String FOR_USER_AGENT_PROPERTY_NAME = "com.palantir.tracing.forUserAgent"; + public static final String SAMPLED_PROPERTY_NAME = "com.palantir.tracing.sampled"; @Context @@ -76,22 +79,36 @@ public void filter(ContainerRequestContext requestContext) throws IOException { getObservabilityFromHeader(requestContext), Tracers.randomId(), operation, - SpanType.SERVER_INCOMING); + SpanType.SERVER_INCOMING, + getForUserAgentFromHeader(requestContext)); } else if (spanId == null) { Tracer.initTraceWithSpan( - getObservabilityFromHeader(requestContext), traceId, operation, SpanType.SERVER_INCOMING); + getObservabilityFromHeader(requestContext), + traceId, + operation, + SpanType.SERVER_INCOMING, + getForUserAgentFromHeader(requestContext)); } else { // caller's span is this span's parent. Tracer.initTraceWithSpan( - getObservabilityFromHeader(requestContext), traceId, operation, spanId, SpanType.SERVER_INCOMING); + getObservabilityFromHeader(requestContext), + traceId, + operation, + spanId, + SpanType.SERVER_INCOMING, + getForUserAgentFromHeader(requestContext)); } // Give asynchronous downstream handlers access to the trace id requestContext.setProperty(TRACE_ID_PROPERTY_NAME, Tracer.getTraceId()); requestContext.setProperty(SAMPLED_PROPERTY_NAME, Tracer.isTraceObservable()); - Tracer.maybeGetTraceMetadata() + Optional traceMetadata = Tracer.maybeGetTraceMetadata(); + traceMetadata .flatMap(TraceMetadata::getRequestId) .ifPresent(requestId -> requestContext.setProperty(REQUEST_ID_PROPERTY_NAME, requestId)); + traceMetadata + .flatMap(TraceMetadata::getForUserAgent) + .ifPresent(forUserAgent -> requestContext.setProperty(FOR_USER_AGENT_PROPERTY_NAME, forUserAgent)); } // Handles outgoing response @@ -131,6 +148,14 @@ private static Observability getObservabilityFromHeader(ContainerRequestContext } } + private static Optional getForUserAgentFromHeader(ContainerRequestContext requestContext) { + String forUserAgent = requestContext.getHeaderString(TraceHttpHeaders.FOR_USER_AGENT); + if (forUserAgent == null) { + return Optional.ofNullable(requestContext.getHeaderString(HttpHeaders.USER_AGENT)); + } + return Optional.of(forUserAgent); + } + private String getPathTemplate() { return Optional.ofNullable(uriInfo) .map(ExtendedUriInfo::getMatchedModelResource) diff --git a/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java b/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java index d51460559..ba2cdfb89 100644 --- a/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java +++ b/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java @@ -47,11 +47,17 @@ public Response intercept(Chain chain) throws IOException { TraceMetadata metadata = Tracer.maybeGetTraceMetadata() .orElseThrow(() -> new SafeRuntimeException("Trace with no spans in progress")); - return chain.proceed(request.newBuilder() + Request.Builder tracedRequest = request.newBuilder() .header(TraceHttpHeaders.TRACE_ID, Tracer.getTraceId()) .header(TraceHttpHeaders.SPAN_ID, metadata.getSpanId()) - .header(TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0") - .build()); + .header(TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0"); + if (metadata.getForUserAgent().isPresent()) { + tracedRequest.header( + TraceHttpHeaders.FOR_USER_AGENT, + metadata.getForUserAgent().get()); + } + + return chain.proceed(tracedRequest.build()); } } } diff --git a/tracing-okhttp3/src/main/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptor.java b/tracing-okhttp3/src/main/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptor.java index 70fa091e9..53d7bb5e8 100644 --- a/tracing-okhttp3/src/main/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptor.java +++ b/tracing-okhttp3/src/main/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptor.java @@ -53,6 +53,10 @@ public Response intercept(Chain chain) throws IOException { .header(TraceHttpHeaders.TRACE_ID, Tracer.getTraceId()) .header(TraceHttpHeaders.SPAN_ID, span.getSpanId()) .header(TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0"); + if (Tracer.getForUserAgent().isPresent()) { + tracedRequest.header( + TraceHttpHeaders.FOR_USER_AGENT, Tracer.getForUserAgent().get()); + } Response response; try { diff --git a/tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java b/tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java index 31ee58584..5fd681435 100644 --- a/tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java +++ b/tracing-okhttp3/src/test/java/com/palantir/tracing/okhttp3/OkhttpTraceInterceptorTest.java @@ -31,6 +31,7 @@ import com.palantir.tracing.api.SpanType; import com.palantir.tracing.api.TraceHttpHeaders; import java.io.IOException; +import java.util.Optional; import okhttp3.Interceptor; import okhttp3.Request; import org.junit.After; @@ -104,6 +105,28 @@ public void testPopulatesNewTrace_whenParentTraceIsPresent() throws IOException assertThat(intercepted.headers(TraceHttpHeaders.TRACE_ID)).containsOnly(traceId); } + @Test + public void testPopulatesNewTrace_whenForUserAgentIsPresent() throws IOException { + String forUserAgent = "forUserAgent"; + Tracer.initTraceWithSpan( + Observability.SAMPLE, "id", "operation", "parent", SpanType.SERVER_INCOMING, Optional.of(forUserAgent)); + String traceId = Tracer.getTraceId(); + try { + OkhttpTraceInterceptor.INSTANCE.intercept(chain); + } finally { + Tracer.fastCompleteSpan(); + } + + verify(chain).request(); + verify(chain).proceed(requestCaptor.capture()); + verifyNoMoreInteractions(chain); + + Request intercepted = requestCaptor.getValue(); + assertThat(intercepted.headers(TraceHttpHeaders.SPAN_ID)).isNotNull(); + assertThat(intercepted.headers(TraceHttpHeaders.TRACE_ID)).containsOnly(traceId); + assertThat(intercepted.headers(TraceHttpHeaders.FOR_USER_AGENT)).containsOnly(forUserAgent); + } + @Test public void testAddsIsSampledHeader_whenTraceIsObservable() throws IOException { Tracer.initTraceWithSpan(Observability.SAMPLE, Tracers.randomId(), "op", SpanType.LOCAL); diff --git a/tracing-undertow/src/main/java/com/palantir/tracing/undertow/UndertowTracing.java b/tracing-undertow/src/main/java/com/palantir/tracing/undertow/UndertowTracing.java index 3de412c64..b23d5e8c3 100644 --- a/tracing-undertow/src/main/java/com/palantir/tracing/undertow/UndertowTracing.java +++ b/tracing-undertow/src/main/java/com/palantir/tracing/undertow/UndertowTracing.java @@ -48,6 +48,7 @@ final class UndertowTracing { private static final HttpString TRACE_ID = HttpString.tryFromString(TraceHttpHeaders.TRACE_ID); private static final HttpString SPAN_ID = HttpString.tryFromString(TraceHttpHeaders.SPAN_ID); private static final HttpString IS_SAMPLED = HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED); + private static final HttpString FOR_USER_AGENT = HttpString.tryFromString(TraceHttpHeaders.FOR_USER_AGENT); // Consider moving this to TracingAttachments and making it public. For now it's well encapsulated // here because we expect the two handler implementations to be sufficient. @@ -108,7 +109,8 @@ private static DetachedSpan detachedSpan( traceId, newTrace ? Optional.empty() : Optional.ofNullable(requestHeaders.getFirst(SPAN_ID)), operationName, - SpanType.SERVER_INCOMING); + SpanType.SERVER_INCOMING, + Optional.ofNullable(requestHeaders.getFirst(FOR_USER_AGENT))); } private enum DetachedTraceCompletionListener implements ExchangeCompletionListener { diff --git a/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java b/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java index 44fdd65fa..0eb0de32b 100644 --- a/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java +++ b/tracing/src/main/java/com/palantir/tracing/DeferredTracer.java @@ -74,6 +74,9 @@ public final class DeferredTracer implements Serializable { @Nullable private final transient String requestId; + @Nullable + private final String forUserAgent; + /** * Deprecated. * @@ -104,6 +107,7 @@ public DeferredTracer(@Safe String operation, @Safe Map metadata this.parentSpanId = trace.top().map(OpenSpan::getSpanId).orElse(null); this.operation = operation; this.metadata = metadata; + this.forUserAgent = trace.getForUserAgent().orElse(null); } else { this.traceId = null; this.requestId = null; @@ -111,6 +115,7 @@ public DeferredTracer(@Safe String operation, @Safe Map metadata this.parentSpanId = null; this.operation = null; this.metadata = null; + this.forUserAgent = null; } } @@ -134,7 +139,8 @@ CloseableTrace withTrace() { Optional originalTrace = Tracer.getAndClearTraceIfPresent(); - Tracer.setTrace(Trace.of(isObservable, traceId, Optional.ofNullable(requestId))); + Tracer.setTrace( + Trace.of(isObservable, traceId, Optional.ofNullable(requestId), Optional.ofNullable(forUserAgent))); if (parentSpanId != null) { Tracer.fastStartSpan(operation, parentSpanId, SpanType.LOCAL); } else { diff --git a/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java b/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java index 066c02678..5e11ed20f 100644 --- a/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java +++ b/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java @@ -67,7 +67,23 @@ static DetachedSpan start( Optional parentSpanId, @Safe String operation, SpanType type) { - return Tracer.detachInternal(observability, traceId, parentSpanId, operation, type); + return Tracer.detachInternal(observability, traceId, parentSpanId, operation, type, Optional.empty()); + } + + /** + * Marks the beginning of a span, which you can {@link #complete} on any other thread. + * + * @see DetachedSpan#start(String) + */ + @CheckReturnValue + static DetachedSpan start( + Observability observability, + String traceId, + Optional parentSpanId, + @Safe String operation, + SpanType type, + Optional forUserAgent) { + return Tracer.detachInternal(observability, traceId, parentSpanId, operation, type, forUserAgent); } /** diff --git a/tracing/src/main/java/com/palantir/tracing/Trace.java b/tracing/src/main/java/com/palantir/tracing/Trace.java index 7ffbb3fad..fc7d0a54b 100644 --- a/tracing/src/main/java/com/palantir/tracing/Trace.java +++ b/tracing/src/main/java/com/palantir/tracing/Trace.java @@ -48,10 +48,13 @@ public abstract class Trace { private final Optional requestId; - private Trace(String traceId, Optional requestId) { + private final Optional forUserAgent; + + private Trace(String traceId, Optional requestId, Optional forUserAgent) { checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty"); this.traceId = traceId; this.requestId = checkNotNull(requestId, "requestId"); + this.forUserAgent = forUserAgent; } /** @@ -127,24 +130,37 @@ final Optional getRequestId() { return requestId; } + final Optional getForUserAgent() { + return forUserAgent; + } + /** Returns a copy of this Trace which can be independently mutated. */ abstract Trace deepCopy(); static Trace of(boolean isObservable, String traceId, Optional requestId) { - return isObservable ? new Sampled(traceId, requestId) : new Unsampled(traceId, requestId); + return isObservable + ? new Sampled(traceId, requestId, Optional.empty()) + : new Unsampled(traceId, requestId, Optional.empty()); + } + + static Trace of(boolean isObservable, String traceId, Optional requestId, Optional forUserAgent) { + return isObservable + ? new Sampled(traceId, requestId, forUserAgent) + : new Unsampled(traceId, requestId, forUserAgent); } private static final class Sampled extends Trace { private final Deque stack; - private Sampled(ArrayDeque stack, String traceId, Optional requestId) { - super(traceId, requestId); + private Sampled( + ArrayDeque stack, String traceId, Optional requestId, Optional forUserAgent) { + super(traceId, requestId, forUserAgent); this.stack = stack; } - private Sampled(String traceId, Optional requestId) { - this(new ArrayDeque<>(), traceId, requestId); + private Sampled(String traceId, Optional requestId, Optional forUserAgent) { + this(new ArrayDeque<>(), traceId, requestId, forUserAgent); } @Override @@ -186,7 +202,7 @@ boolean isObservable() { @Override Trace deepCopy() { - return new Sampled(new ArrayDeque<>(stack), getTraceId(), getRequestId()); + return new Sampled(new ArrayDeque<>(stack), getTraceId(), getRequestId(), getForUserAgent()); } @Override @@ -202,14 +218,15 @@ private static final class Unsampled extends Trace { */ private int numberOfSpans; - private Unsampled(int numberOfSpans, String traceId, Optional requestId) { - super(traceId, requestId); + private Unsampled( + int numberOfSpans, String traceId, Optional requestId, Optional forUserAgent) { + super(traceId, requestId, forUserAgent); this.numberOfSpans = numberOfSpans; validateNumberOfSpans(); } - private Unsampled(String traceId, Optional requestId) { - this(0, traceId, requestId); + private Unsampled(String traceId, Optional requestId, Optional forUserAgent) { + this(0, traceId, requestId, forUserAgent); } @Override @@ -254,7 +271,7 @@ boolean isObservable() { @Override Trace deepCopy() { - return new Unsampled(numberOfSpans, getTraceId(), getRequestId()); + return new Unsampled(numberOfSpans, getTraceId(), getRequestId(), getForUserAgent()); } /** Internal validation, this should never fail because {@link #pop()} only decrements positive values. */ diff --git a/tracing/src/main/java/com/palantir/tracing/TraceMetadata.java b/tracing/src/main/java/com/palantir/tracing/TraceMetadata.java index 7abd455aa..f661612c1 100644 --- a/tracing/src/main/java/com/palantir/tracing/TraceMetadata.java +++ b/tracing/src/main/java/com/palantir/tracing/TraceMetadata.java @@ -50,6 +50,9 @@ default Optional getOriginatingSpanId() { return Optional.empty(); } + /** Corresponds to {@link com.palantir.tracing.api.TraceHttpHeaders#FOR_USER_AGENT}. */ + Optional getForUserAgent(); + static Builder builder() { return new Builder(); } diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 6b1722bfa..8aa022775 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -68,6 +68,16 @@ private Tracer() {} // Thread-safe since stateless private static volatile TraceSampler sampler = RandomSampler.create(0.0005f); + /** + * Creates a new trace, but does not set it as the current trace. + */ + private static Trace createTrace( + Observability observability, String traceId, Optional requestId, Optional forUserAgent) { + checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty"); + boolean observable = shouldObserve(observability); + return Trace.of(observable, traceId, requestId, forUserAgent); + } + /** * Creates a new trace, but does not set it as the current trace. */ @@ -117,6 +127,7 @@ public static Optional maybeGetTraceMetadata() { return trace.top().map(openSpan -> TraceMetadata.builder() .spanId(openSpan.getSpanId()) .parentSpanId(openSpan.getParentSpanId()) + .forUserAgent(trace.getForUserAgent()) .traceId(trace.getTraceId()) .requestId(trace.getRequestId()) .build()); @@ -124,6 +135,7 @@ public static Optional maybeGetTraceMetadata() { return Optional.of(TraceMetadata.builder() .spanId(Tracers.randomId()) .parentSpanId(Optional.empty()) + .forUserAgent(trace.getForUserAgent()) .traceId(trace.getTraceId()) .requestId(trace.getRequestId()) .build()); @@ -154,6 +166,25 @@ public static void initTrace(Observability observability, String traceId) { setTrace(createTrace(observability, traceId, Optional.empty())); } + /** + * Initializes the current thread's trace with a root span, erasing any previously accrued open spans. + * The root span must eventually be completed using {@link #fastCompleteSpan()} or {@link #completeSpan()}. + */ + public static void initTraceWithSpan( + Observability observability, + String traceId, + @Safe String operation, + String parentSpanId, + SpanType type, + Optional forUserAgent) { + setTrace(createTrace( + observability, + traceId, + type == SpanType.SERVER_INCOMING ? Optional.of(Tracers.randomId()) : Optional.empty(), + forUserAgent)); + fastStartSpan(operation, parentSpanId, type); + } + /** * Initializes the current thread's trace with a root span, erasing any previously accrued open spans. * The root span must eventually be completed using {@link #fastCompleteSpan()} or {@link #completeSpan()}. @@ -180,6 +211,24 @@ public static void initTraceWithSpan( fastStartSpan(operation, type); } + /** + * Initializes the current thread's trace with a root span, erasing any previously accrued open spans. + * The root span must eventually be completed using {@link #fastCompleteSpan()} or {@link #completeSpan()}. + */ + public static void initTraceWithSpan( + Observability observability, + String traceId, + @Safe String operation, + SpanType type, + Optional forUserAgent) { + setTrace(createTrace( + observability, + traceId, + type == SpanType.SERVER_INCOMING ? Optional.of(Tracers.randomId()) : Optional.empty(), + forUserAgent)); + fastStartSpan(operation, type); + } + /** * Opens a new span for this thread's call trace, labeled with the provided operation and parent span. Only allowed * when the current trace is empty. If the return value is not used, prefer {@link Tracer#fastStartSpan(String, @@ -240,8 +289,8 @@ static DetachedSpan detachInternal(@Safe String operation, SpanType type) { Optional parentSpan = getParentSpanId(maybeCurrentTrace); Optional requestId = getRequestId(maybeCurrentTrace, type); return sampled - ? new SampledDetachedSpan(operation, type, traceId, requestId, parentSpan) - : new UnsampledDetachedSpan(traceId, requestId, Optional.empty()); + ? new SampledDetachedSpan(operation, type, traceId, requestId, parentSpan, Optional.empty()) + : new UnsampledDetachedSpan(traceId, requestId, Optional.empty(), Optional.empty()); } /** @@ -253,10 +302,11 @@ static DetachedSpan detachInternal( String traceId, Optional parentSpanId, @Safe String operation, - SpanType type) { + SpanType type, + Optional forUserAgent) { Optional requestId = type == SpanType.SERVER_INCOMING ? Optional.of(Tracers.randomId()) : Optional.empty(); - return detachInternal(observability, traceId, requestId, parentSpanId, operation, type); + return detachInternal(observability, traceId, requestId, parentSpanId, operation, type, forUserAgent); } /** @@ -269,12 +319,13 @@ static DetachedSpan detachInternal( Optional requestId, Optional parentSpanId, @Safe String operation, - SpanType type) { + SpanType type, + Optional forUserAgent) { // The current trace has no impact on this function, a new trace is spawned and existing thread state // is not modified. return shouldObserve(observability) - ? new SampledDetachedSpan(operation, type, traceId, requestId, parentSpanId) - : new UnsampledDetachedSpan(traceId, requestId, parentSpanId); + ? new SampledDetachedSpan(operation, type, traceId, requestId, parentSpanId, forUserAgent) + : new UnsampledDetachedSpan(traceId, requestId, parentSpanId, forUserAgent); } /** @@ -294,9 +345,9 @@ static Detached detachInternal() { if (maybeOpenSpan == null) { return NopDetached.INSTANCE; } - return new SampledDetached(traceId, requestId, maybeOpenSpan); + return new SampledDetached(traceId, requestId, maybeOpenSpan, trace.getForUserAgent()); } else { - return new UnsampledDetachedSpan(traceId, requestId, Optional.empty()); + return new UnsampledDetachedSpan(traceId, requestId, Optional.empty(), trace.getForUserAgent()); } } @@ -374,6 +425,7 @@ private static final class SampledDetachedSpan implements DetachedSpan { private final String traceId; private final Optional requestId; + private final Optional forUserAgent; private final OpenSpan openSpan; private volatile int completed; @@ -385,9 +437,11 @@ private static final class SampledDetachedSpan implements DetachedSpan { SpanType type, String traceId, Optional requestId, - Optional parentSpanId) { + Optional parentSpanId, + Optional forUserAgent) { this.traceId = traceId; this.requestId = requestId; + this.forUserAgent = forUserAgent; this.openSpan = OpenSpan.of(operation, Tracers.randomId(), type, parentSpanId); } @@ -399,9 +453,10 @@ private static CloseableSpan childSpan( String operationName, TagTranslator translator, T data, - SpanType type) { + SpanType type, + Optional forUserAgent) { Trace maybeCurrentTrace = currentTrace.get(); - setTrace(Trace.of(true, traceId, requestId)); + setTrace(Trace.of(true, traceId, requestId, forUserAgent)); Tracer.fastStartSpan(operationName, openSpan.getSpanId(), type); return TraceRestoringCloseableSpanWithMetadata.of(maybeCurrentTrace, translator, data); } @@ -410,18 +465,20 @@ private static CloseableSpan childSpan( @MustBeClosed public CloseableSpan childSpan( String operationName, TagTranslator translator, T data, SpanType type) { - return childSpan(traceId, openSpan, requestId, operationName, translator, data, type); + return childSpan(traceId, openSpan, requestId, operationName, translator, data, type, forUserAgent); } @Override public DetachedSpan childDetachedSpan(String operation, SpanType type) { - return new SampledDetachedSpan(operation, type, traceId, requestId, Optional.of(openSpan.getSpanId())); + return new SampledDetachedSpan( + operation, type, traceId, requestId, Optional.of(openSpan.getSpanId()), forUserAgent); } @MustBeClosed - private static CloseableSpan attach(String traceId, OpenSpan openSpan, Optional requestId) { + private static CloseableSpan attach( + String traceId, OpenSpan openSpan, Optional requestId, Optional forUserAgent) { Trace maybeCurrentTrace = currentTrace.get(); - Trace newTrace = Trace.of(true, traceId, requestId); + Trace newTrace = Trace.of(true, traceId, requestId, forUserAgent); // Push the DetachedSpan OpenSpan to provide the correct parent information // to child spans created within the context of this attach. // It is VITAL that this span is never completed, it exists only for attribution. @@ -435,7 +492,7 @@ private static CloseableSpan attach(String traceId, OpenSpan openSpan, Optional< @Override @MustBeClosed public CloseableSpan attach() { - return attach(traceId, openSpan, requestId); + return attach(traceId, openSpan, requestId, forUserAgent); } @Override @@ -471,29 +528,33 @@ private static final class SampledDetached implements Detached { private final String traceId; private final Optional requestId; private final OpenSpan openSpan; + private final Optional forUserAgent; - SampledDetached(String traceId, Optional requestId, OpenSpan openSpan) { + SampledDetached(String traceId, Optional requestId, OpenSpan openSpan, Optional forUserAgent) { this.traceId = traceId; this.requestId = requestId; this.openSpan = openSpan; + this.forUserAgent = forUserAgent; } @Override @MustBeClosed public CloseableSpan childSpan( String operationName, TagTranslator translator, T data, SpanType type) { - return SampledDetachedSpan.childSpan(traceId, openSpan, requestId, operationName, translator, data, type); + return SampledDetachedSpan.childSpan( + traceId, openSpan, requestId, operationName, translator, data, type, forUserAgent); } @Override public DetachedSpan childDetachedSpan(String operation, SpanType type) { - return new SampledDetachedSpan(operation, type, traceId, requestId, Optional.of(openSpan.getSpanId())); + return new SampledDetachedSpan( + operation, type, traceId, requestId, Optional.of(openSpan.getSpanId()), forUserAgent); } @Override @MustBeClosed public CloseableSpan attach() { - return SampledDetachedSpan.attach(traceId, openSpan, requestId); + return SampledDetachedSpan.attach(traceId, openSpan, requestId, forUserAgent); } @Override @@ -515,18 +576,24 @@ private static final class UnsampledDetachedSpan implements DetachedSpan { private final String traceId; private final Optional requestId; private final Optional parentSpanId; + private final Optional forUserAgent; - UnsampledDetachedSpan(String traceId, Optional requestId, Optional parentSpanId) { + UnsampledDetachedSpan( + String traceId, + Optional requestId, + Optional parentSpanId, + Optional forUserAgent) { this.traceId = traceId; this.requestId = requestId; this.parentSpanId = parentSpanId; + this.forUserAgent = forUserAgent; } @Override public CloseableSpan childSpan( String operationName, TagTranslator _translator, T _data, SpanType type) { Trace maybeCurrentTrace = currentTrace.get(); - setTrace(Trace.of(false, traceId, requestId)); + setTrace(Trace.of(false, traceId, requestId, forUserAgent)); if (parentSpanId.isPresent()) { Tracer.fastStartSpan(operationName, parentSpanId.get(), type); } else { @@ -769,6 +836,10 @@ public static String getTraceId() { return checkNotNull(currentTrace.get(), "There is no trace").getTraceId(); } + public static Optional getForUserAgent() { + return currentTrace.get().getForUserAgent(); + } + /** * Clears the current trace id and returns it if present. */ @@ -827,10 +898,20 @@ static void setTrace(Trace trace) { MDC.put(Tracers.TRACE_ID_KEY, trace.getTraceId()); setTraceSampledMdcIfObservable(trace.isObservable()); setTraceRequestId(trace.getRequestId()); + setForUserAgent(trace.getForUserAgent()); logSettingTrace(); } + private static void setForUserAgent(Optional forUserAgent) { + if (forUserAgent.isPresent()) { + MDC.put(Tracers.FOR_USER_AGENT_KEY, forUserAgent.get()); + } else { + // Ensure MDC state is cleared when there is no origin user header + MDC.remove(Tracers.FOR_USER_AGENT_KEY); + } + } + private static void setTraceSampledMdcIfObservable(boolean observable) { if (observable) { // Set to 1 to be consistent with values associated with http header key TraceHttpHeaders.IS_SAMPLED @@ -870,6 +951,7 @@ static void clearCurrentTrace() { MDC.remove(Tracers.TRACE_ID_KEY); MDC.remove(Tracers.TRACE_SAMPLED_KEY); MDC.remove(Tracers.REQUEST_ID_KEY); + MDC.remove(Tracers.FOR_USER_AGENT_KEY); } private static void logClearingTrace() { diff --git a/tracing/src/main/java/com/palantir/tracing/Tracers.java b/tracing/src/main/java/com/palantir/tracing/Tracers.java index 97d10115d..690895bb9 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracers.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracers.java @@ -48,6 +48,11 @@ public final class Tracers { */ public static final String REQUEST_ID_KEY = "_requestId"; + /** + * The Key under which origin user agents are inserted into SLF4J {@link org.slf4j.MDC MDCs}. + */ + public static final String FOR_USER_AGENT_KEY = "__forUserAgent"; + private static final String DEFAULT_ROOT_SPAN_OPERATION = "root"; private static final char[] HEX_DIGITS = { '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f' diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index 7a38d25fa..866eeaf6d 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracerTest.java @@ -647,8 +647,8 @@ public void testDetachedTraceBuildsUponExistingTrace() { public void testNewDetachedTrace() { try (CloseableTracer ignored = CloseableTracer.startSpan("test")) { String currentTraceId = Tracer.getTraceId(); - DetachedSpan span = - DetachedSpan.start(Observability.SAMPLE, "12345", Optional.empty(), "op", SpanType.LOCAL); + DetachedSpan span = DetachedSpan.start( + Observability.SAMPLE, "12345", Optional.empty(), "op", SpanType.LOCAL, Optional.empty()); try (CloseableSpan ignored2 = span.completeAndStartChild("foo")) { assertThat(Tracer.getTraceId()).isEqualTo("12345"); } @@ -662,8 +662,8 @@ public void testNewDetachedTrace() { public void testNewDetachedTrace_doesNotModifyCurrentState() { try (CloseableTracer ignored = CloseableTracer.startSpan("test")) { String currentTraceId = Tracer.getTraceId(); - DetachedSpan span = - DetachedSpan.start(Observability.SAMPLE, "12345", Optional.empty(), "op", SpanType.LOCAL); + DetachedSpan span = DetachedSpan.start( + Observability.SAMPLE, "12345", Optional.empty(), "op", SpanType.LOCAL, Optional.empty()); assertThat(Tracer.getTraceId()) .as("Current thread state should not be modified") .isEqualTo(currentTraceId);