diff --git a/changelog/6.5.0-rc1/pr-790.v2.yml b/changelog/6.5.0-rc1/pr-790.v2.yml new file mode 100644 index 000000000..a52088d47 --- /dev/null +++ b/changelog/6.5.0-rc1/pr-790.v2.yml @@ -0,0 +1,6 @@ +type: feature +feature: + description: User agent specified inside `ForUserAgent` http header will be propagated + to requests made in the same trace. + links: + - https://github.com/palantir/tracing-java/pull/790 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..96e0fb4df 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 @@ -30,6 +30,8 @@ public interface TraceHttpHeaders { String SPAN_ID = "X-B3-SpanId"; String IS_SAMPLED = "X-B3-Sampled"; // Boolean (either “1” or “0”, can be absent) + String FOR_USER_AGENT = "For-User-Agent"; + /** * Field is no longer used by the tracing library. * diff --git a/tracing-api/src/main/java/com/palantir/tracing/api/TraceTags.java b/tracing-api/src/main/java/com/palantir/tracing/api/TraceTags.java index 74370f7dc..1f3cbf494 100644 --- a/tracing-api/src/main/java/com/palantir/tracing/api/TraceTags.java +++ b/tracing-api/src/main/java/com/palantir/tracing/api/TraceTags.java @@ -62,6 +62,8 @@ public final class TraceTags { public static final String HTTP_REQUEST_ID = "http.request_id"; /** The User-Agent as it is sent (raw format). */ public static final String HTTP_USER_AGENT = "http.useragent"; + /** The User-Agent propagated across service boundaries as it is sent (raw format). */ + public static final String HTTP_FOR_USER_AGENT = "http.for_useragent"; /** The version of HTTP used for the request. */ public static final String HTTP_VERSION = "http.version"; 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..018c5c695 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 @@ -16,6 +16,7 @@ package com.palantir.tracing.jersey; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; import com.palantir.tracing.Observability; import com.palantir.tracing.TagTranslator; @@ -35,6 +36,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; @@ -55,6 +57,9 @@ public final class TraceEnrichingFilter implements ContainerRequestFilter, Conta public static final String SAMPLED_PROPERTY_NAME = "com.palantir.tracing.sampled"; + @VisibleForTesting + static final String FETCH_USER_AGENT_HEADER = "Fetch-User-Agent"; + @Context @SuppressWarnings("NullAway") // instantiated using by Jersey using reflection private ExtendedUriInfo uriInfo; @@ -68,6 +73,7 @@ public void filter(ContainerRequestContext requestContext) throws IOException { // The following strings are all nullable String traceId = requestContext.getHeaderString(TraceHttpHeaders.TRACE_ID); String spanId = requestContext.getHeaderString(TraceHttpHeaders.SPAN_ID); + Optional forUserAgent = getForUserAgent(requestContext); // Set up thread-local span that inherits state from HTTP headers if (Strings.isNullOrEmpty(traceId)) { @@ -75,15 +81,25 @@ public void filter(ContainerRequestContext requestContext) throws IOException { Tracer.initTraceWithSpan( getObservabilityFromHeader(requestContext), Tracers.randomId(), + forUserAgent, operation, SpanType.SERVER_INCOMING); } else if (spanId == null) { Tracer.initTraceWithSpan( - getObservabilityFromHeader(requestContext), traceId, operation, SpanType.SERVER_INCOMING); + getObservabilityFromHeader(requestContext), + traceId, + forUserAgent, + operation, + SpanType.SERVER_INCOMING); } else { // caller's span is this span's parent. Tracer.initTraceWithSpan( - getObservabilityFromHeader(requestContext), traceId, operation, spanId, SpanType.SERVER_INCOMING); + getObservabilityFromHeader(requestContext), + traceId, + forUserAgent, + operation, + spanId, + SpanType.SERVER_INCOMING); } // Give asynchronous downstream handlers access to the trace id @@ -131,6 +147,18 @@ private static Observability getObservabilityFromHeader(ContainerRequestContext } } + private static Optional getForUserAgent(ContainerRequestContext context) { + String forUserAgent = context.getHeaderString(TraceHttpHeaders.FOR_USER_AGENT); + if (forUserAgent != null) { + return Optional.of(forUserAgent); + } + String fetchUserAgent = context.getHeaderString(FETCH_USER_AGENT_HEADER); + if (fetchUserAgent != null) { + return Optional.of(fetchUserAgent); + } + return Optional.ofNullable(context.getHeaderString(HttpHeaders.USER_AGENT)); + } + private String getPathTemplate() { return Optional.ofNullable(uriInfo) .map(ExtendedUriInfo::getMatchedModelResource) diff --git a/tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java b/tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java index 65e3ad3cd..33fde0729 100644 --- a/tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java +++ b/tracing-jersey/src/test/java/com/palantir/tracing/jersey/TraceEnrichingFilterTest.java @@ -25,6 +25,7 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.palantir.tracing.InternalTracers; import com.palantir.tracing.TraceSampler; import com.palantir.tracing.Tracer; import com.palantir.tracing.Tracers; @@ -45,6 +46,7 @@ import javax.ws.rs.client.Entity; import javax.ws.rs.client.WebTarget; import javax.ws.rs.container.ContainerRequestContext; +import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; import javax.ws.rs.core.StreamingOutput; @@ -238,6 +240,34 @@ public void testFilter_setsMdcIfTraceIdHeaderIsPresent() throws Exception { verify(request).setProperty(eq(TraceEnrichingFilter.REQUEST_ID_PROPERTY_NAME), anyString()); } + @Test + public void testFilter_setsUserAgentAsForUserAgent() throws Exception { + when(request.getHeaderString(TraceHttpHeaders.TRACE_ID)).thenReturn("traceId"); + when(request.getHeaderString(HttpHeaders.USER_AGENT)).thenReturn("userAgent"); + TraceEnrichingFilter.INSTANCE.filter(request); + + assertThat(InternalTracers.getForUserAgent()).contains("userAgent"); + } + + @Test + public void testFilter_setsFetchUserAgentAsForUserAgent() throws Exception { + when(request.getHeaderString(TraceHttpHeaders.TRACE_ID)).thenReturn("traceId"); + when(request.getHeaderString(TraceEnrichingFilter.FETCH_USER_AGENT_HEADER)) + .thenReturn("fetchUserAgent"); + TraceEnrichingFilter.INSTANCE.filter(request); + + assertThat(InternalTracers.getForUserAgent()).contains("fetchUserAgent"); + } + + @Test + public void testFilter_propagatesProvidedForUserAgent() throws Exception { + when(request.getHeaderString(TraceHttpHeaders.TRACE_ID)).thenReturn("traceId"); + when(request.getHeaderString(TraceHttpHeaders.FOR_USER_AGENT)).thenReturn("forUserAgent"); + TraceEnrichingFilter.INSTANCE.filter(request); + + assertThat(InternalTracers.getForUserAgent()).contains("forUserAgent"); + } + @Test public void testFilter_createsReceiveAndSendEvents() throws Exception { target.path("/trace").request().header(TraceHttpHeaders.TRACE_ID, "").get(); 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..1caed043e 100644 --- a/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java +++ b/tracing-okhttp3/src/main/java/com/palantir/tracing/OkhttpTraceInterceptor2.java @@ -17,7 +17,6 @@ package com.palantir.tracing; import com.palantir.logsafe.exceptions.SafeRuntimeException; -import com.palantir.tracing.api.TraceHttpHeaders; import java.io.Closeable; import java.io.IOException; import java.util.function.Function; @@ -44,14 +43,24 @@ public Response intercept(Chain chain) throws IOException { Request request = chain.request(); try (Closeable span = createNetworkCallSpan.apply(request)) { - TraceMetadata metadata = Tracer.maybeGetTraceMetadata() - .orElseThrow(() -> new SafeRuntimeException("Trace with no spans in progress")); - - return chain.proceed(request.newBuilder() - .header(TraceHttpHeaders.TRACE_ID, Tracer.getTraceId()) - .header(TraceHttpHeaders.SPAN_ID, metadata.getSpanId()) - .header(TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0") - .build()); + if (!Tracer.hasTraceId()) { + throw new SafeRuntimeException("Trace with no spans in progress"); + } + + Request.Builder requestBuilder = request.newBuilder(); + + Tracers.addTracingHeaders(requestBuilder, EnrichingFunction.INSTANCE); + + return chain.proceed(requestBuilder.build()); + } + } + + private enum EnrichingFunction implements TracingHeadersEnrichingFunction { + INSTANCE; + + @Override + public void addHeader(String headerName, String headerValue, Request.Builder state) { + state.header(headerName, headerValue); } } } 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..2dc3d3d56 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 @@ -17,9 +17,9 @@ package com.palantir.tracing.okhttp3; import com.palantir.tracing.Tracer; -import com.palantir.tracing.api.OpenSpan; +import com.palantir.tracing.Tracers; +import com.palantir.tracing.TracingHeadersEnrichingFunction; import com.palantir.tracing.api.SpanType; -import com.palantir.tracing.api.TraceHttpHeaders; import java.io.IOException; import okhttp3.Interceptor; import okhttp3.Request; @@ -48,11 +48,9 @@ public Response intercept(Chain chain) throws IOException { request = request.newBuilder().removeHeader(PATH_TEMPLATE_HEADER).build(); } - OpenSpan span = Tracer.startSpan(spanName, SpanType.CLIENT_OUTGOING); - Request.Builder tracedRequest = request.newBuilder() - .header(TraceHttpHeaders.TRACE_ID, Tracer.getTraceId()) - .header(TraceHttpHeaders.SPAN_ID, span.getSpanId()) - .header(TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0"); + Tracer.fastStartSpan(spanName, SpanType.CLIENT_OUTGOING); + Request.Builder tracedRequest = request.newBuilder(); + Tracers.addTracingHeaders(tracedRequest, EnrichingFunction.INSTANCE); Response response; try { @@ -63,4 +61,13 @@ public Response intercept(Chain chain) throws IOException { return response; } + + private enum EnrichingFunction implements TracingHeadersEnrichingFunction { + INSTANCE; + + @Override + public void addHeader(String headerName, String headerValue, Request.Builder state) { + state.header(headerName, headerValue); + } + } } 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..889825a39 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", Optional.of(forUserAgent), "operation", "parent", SpanType.SERVER_INCOMING); + 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..4ed9586ec 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 @@ -16,6 +16,7 @@ package com.palantir.tracing.undertow; +import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Strings; import com.palantir.logsafe.SafeArg; import com.palantir.logsafe.exceptions.SafeIllegalStateException; @@ -32,6 +33,7 @@ import io.undertow.server.HttpServerExchange; import io.undertow.util.AttachmentKey; import io.undertow.util.HeaderMap; +import io.undertow.util.Headers; import io.undertow.util.HttpString; import java.util.Optional; @@ -48,6 +50,11 @@ 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); + // Tracing headers for obtaining for constructing forUserAgent. + private static final HttpString FOR_USER_AGENT = HttpString.tryFromString(TraceHttpHeaders.FOR_USER_AGENT); + + @VisibleForTesting + static final HttpString FETCH_USER_AGENT = HttpString.tryFromString("Fetch-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. @@ -77,7 +84,8 @@ private static DetachedSpan initializeRequestTrace( String maybeTraceId = requestHeaders.getFirst(TRACE_ID); boolean newTraceId = maybeTraceId == null; String traceId = newTraceId ? Tracers.randomId() : maybeTraceId; - DetachedSpan detachedSpan = detachedSpan(operationName, newTraceId, traceId, requestHeaders); + Optional forUserAgent = getForUserAgent(requestHeaders); + DetachedSpan detachedSpan = detachedSpan(operationName, newTraceId, traceId, forUserAgent, requestHeaders); setExchangeState(exchange, detachedSpan, traceId, translator); return detachedSpan; } @@ -92,7 +100,7 @@ private static void setExchangeState( boolean isSampled = InternalTracers.isSampled(detachedSpan); exchange.putAttachment(TracingAttachments.IS_SAMPLED, isSampled); Optional requestId = InternalTracers.getRequestId(detachedSpan); - if (!requestId.isPresent()) { + if (requestId.isEmpty()) { throw new SafeIllegalStateException("No requestId is set", SafeArg.of("span", detachedSpan)); } exchange.putAttachment(TracingAttachments.REQUEST_ID, requestId.get()); @@ -102,10 +110,15 @@ private static void setExchangeState( } private static DetachedSpan detachedSpan( - String operationName, boolean newTrace, String traceId, HeaderMap requestHeaders) { + String operationName, + boolean newTrace, + String traceId, + Optional forUserAgent, + HeaderMap requestHeaders) { return DetachedSpan.start( getObservabilityFromHeader(requestHeaders), traceId, + forUserAgent, newTrace ? Optional.empty() : Optional.ofNullable(requestHeaders.getFirst(SPAN_ID)), operationName, SpanType.SERVER_INCOMING); @@ -148,5 +161,17 @@ private static Observability getObservabilityFromHeader(HeaderMap headers) { } } + private static Optional getForUserAgent(HeaderMap requestHeaders) { + String forUserAgent = requestHeaders.getFirst(FOR_USER_AGENT); + if (forUserAgent != null) { + return Optional.of(forUserAgent); + } + String fetchUserAgent = requestHeaders.getFirst(FETCH_USER_AGENT); + if (fetchUserAgent != null) { + return Optional.of(fetchUserAgent); + } + return Optional.ofNullable(requestHeaders.getFirst(Headers.USER_AGENT)); + } + private UndertowTracing() {} } diff --git a/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java b/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java index fcd3500ba..330bf0c3a 100644 --- a/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java +++ b/tracing-undertow/src/test/java/com/palantir/tracing/undertow/TracedOperationHandlerTest.java @@ -24,6 +24,8 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.palantir.tracing.DetachedSpan; +import com.palantir.tracing.InternalTracers; import com.palantir.tracing.TraceSampler; import com.palantir.tracing.Tracer; import com.palantir.tracing.Tracers; @@ -33,6 +35,7 @@ import io.undertow.server.HttpHandler; import io.undertow.server.HttpServerExchange; import io.undertow.util.HeaderMap; +import io.undertow.util.Headers; import io.undertow.util.HttpString; import java.util.List; import java.util.concurrent.atomic.AtomicReference; @@ -125,6 +128,39 @@ public void whenParentSpanIsGiven_usesParentSpan() throws Exception { assertThat(span.getSpanId()).isNotEqualTo(parentSpanId); } + @Test + public void whenOnlyUserAgentIsProvided_setsItAsForUserAgent() throws Exception { + setRequestTraceId(traceId); + setUserAgent("userAgent"); + handler.handleRequest(exchange); + + DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN); + assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("userAgent"); + } + + @Test + public void whenFetchUserAgentIsProvided_setsItAsForUserAgent() throws Exception { + setRequestTraceId(traceId); + setUserAgent("userAgent"); + setFetchUserAgent("fetchUserAgent"); + handler.handleRequest(exchange); + + DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN); + assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("fetchUserAgent"); + } + + @Test + public void whenForUserAgentIsProvided_propagateItFurther() throws Exception { + setRequestTraceId(traceId); + setUserAgent("userAgent"); + setFetchUserAgent("fetchUserAgent"); + setForUserAgent("forUserAgent"); + handler.handleRequest(exchange); + + DetachedSpan detachedSpan = exchange.getAttachment(UndertowTracing.REQUEST_SPAN); + assertThat(InternalTracers.getForUserAgent(detachedSpan)).contains("forUserAgent"); + } + @Test public void whenTraceIsAlreadySampled_doesNotCallSampler() throws Exception { exchange.getRequestHeaders().put(HttpString.tryFromString(TraceHttpHeaders.IS_SAMPLED), "1"); @@ -184,10 +220,30 @@ public void populatesSlf4jMdc() throws Exception { } private void setRequestTraceId(String theTraceId) { - exchange.getRequestHeaders().put(HttpString.tryFromString(TraceHttpHeaders.TRACE_ID), theTraceId); + setHeader(TraceHttpHeaders.TRACE_ID, theTraceId); } private void setRequestSpanId(String spanId) { - exchange.getRequestHeaders().put(HttpString.tryFromString(TraceHttpHeaders.SPAN_ID), spanId); + setHeader(TraceHttpHeaders.SPAN_ID, spanId); + } + + private void setUserAgent(String userAgent) { + setHeader(Headers.USER_AGENT, userAgent); + } + + private void setFetchUserAgent(String fetchUserAgent) { + setHeader(UndertowTracing.FETCH_USER_AGENT, fetchUserAgent); + } + + private void setForUserAgent(String forUserAgent) { + setHeader(TraceHttpHeaders.FOR_USER_AGENT, forUserAgent); + } + + private void setHeader(String headerName, String headerValue) { + setHeader(HttpString.tryFromString(headerName), headerValue); + } + + private void setHeader(HttpString headerName, String headerValue) { + exchange.getRequestHeaders().put(headerName, headerValue); } } diff --git a/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java b/tracing/src/main/java/com/palantir/tracing/DetachedSpan.java index 066c02678..dc53a05f5 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 start(observability, traceId, Optional.empty(), parentSpanId, operation, type); + } + + /** + * 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 forUserAgent, + Optional parentSpanId, + @Safe String operation, + SpanType type) { + return Tracer.detachInternal(observability, traceId, forUserAgent, parentSpanId, operation, type); } /** diff --git a/tracing/src/main/java/com/palantir/tracing/InternalTracers.java b/tracing/src/main/java/com/palantir/tracing/InternalTracers.java index ec4a1dced..7342b773e 100644 --- a/tracing/src/main/java/com/palantir/tracing/InternalTracers.java +++ b/tracing/src/main/java/com/palantir/tracing/InternalTracers.java @@ -31,10 +31,22 @@ public static boolean isSampled(DetachedSpan detachedSpan) { return Tracer.isSampled(detachedSpan); } - /** Returns true if the provided detachedSpan is sampled. */ + /** Returns requestId of the provided detachedSpan. */ public static Optional getRequestId(DetachedSpan detachedSpan) { return Optional.ofNullable(Tracer.getRequestId(detachedSpan)); } + /** + * Returns the forUserAgent propagated inside the trace. + */ + public static Optional getForUserAgent() { + return Tracer.getForUserAgent(); + } + + /** Returns the forUserAgent in the provided detachedSpan. */ + public static Optional getForUserAgent(DetachedSpan detachedSpan) { + return Optional.ofNullable(Tracer.getForUserAgent(detachedSpan)); + } + private InternalTracers() {} } diff --git a/tracing/src/main/java/com/palantir/tracing/Trace.java b/tracing/src/main/java/com/palantir/tracing/Trace.java index 9fbb88622..07d458eff 100644 --- a/tracing/src/main/java/com/palantir/tracing/Trace.java +++ b/tracing/src/main/java/com/palantir/tracing/Trace.java @@ -129,12 +129,19 @@ final Optional getRequestId() { return Optional.ofNullable(traceState.requestId()); } + /** + * The user agent propagated across this trace. + */ + final Optional getForUserAgent() { + return Optional.ofNullable(traceState.forUserAgent()); + } + /** Returns a copy of this Trace which can be independently mutated. */ abstract Trace deepCopy(); @Deprecated static Trace of(boolean isObservable, String traceId, Optional requestId) { - return of(isObservable, TraceState.of(traceId, requestId)); + return of(isObservable, TraceState.of(traceId, requestId, Optional.empty())); } static Trace of(boolean isObservable, TraceState traceState) { diff --git a/tracing/src/main/java/com/palantir/tracing/TraceState.java b/tracing/src/main/java/com/palantir/tracing/TraceState.java index 9f39885a4..c359fd8fd 100644 --- a/tracing/src/main/java/com/palantir/tracing/TraceState.java +++ b/tracing/src/main/java/com/palantir/tracing/TraceState.java @@ -36,15 +36,20 @@ final class TraceState implements Serializable { @Nullable private final String requestId; - static TraceState of(String traceId, Optional requestId) { + @Nullable + private final String forUserAgent; + + static TraceState of(String traceId, Optional requestId, Optional forUserAgent) { checkArgument(!Strings.isNullOrEmpty(traceId), "traceId must be non-empty"); checkNotNull(requestId, "requestId should be not-null"); - return new TraceState(traceId, requestId.orElse(null)); + checkNotNull(forUserAgent, "forUserAgent should be not-null"); + return new TraceState(traceId, requestId.orElse(null), forUserAgent.orElse(null)); } - private TraceState(String traceId, @Nullable String requestId) { + private TraceState(String traceId, @Nullable String requestId, @Nullable String forUserAgent) { this.traceId = traceId; this.requestId = requestId; + this.forUserAgent = forUserAgent; } /** @@ -66,8 +71,20 @@ String requestId() { return requestId; } + /** + * The user agent propagated throughout the duration of this trace. + */ + @Nullable + String forUserAgent() { + return forUserAgent; + } + @Override public String toString() { - return "TraceState{traceId='" + traceId + "', requestId='" + requestId + "'}"; + return "TraceState{" + + "traceId='" + traceId + "', " + + "requestId='" + requestId + "', " + + "forUserAgent='" + forUserAgent + + "'}"; } } diff --git a/tracing/src/main/java/com/palantir/tracing/Tracer.java b/tracing/src/main/java/com/palantir/tracing/Tracer.java index 428908424..d27f9e9d6 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracer.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracer.java @@ -72,9 +72,17 @@ private Tracer() {} * Creates a new trace, but does not set it as the current trace. */ private static Trace createTrace(Observability observability, String traceId, Optional requestId) { + return createTrace(observability, traceId, requestId, Optional.empty()); + } + + /** + * 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, TraceState.of(traceId, requestId)); + return Trace.of(observable, TraceState.of(traceId, requestId, forUserAgent)); } private static boolean shouldObserve(Observability observability) { @@ -159,7 +167,10 @@ public static void initTrace(Observability observability, String traceId) { /** * 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()}. + * + * @deprecated Use {@link #initTraceWithSpan(Observability, String, Optional, String, String, SpanType)} */ + @Deprecated public static void initTraceWithSpan( Observability observability, String traceId, @Safe String operation, String parentSpanId, SpanType type) { setTrace(createTrace( @@ -182,6 +193,43 @@ 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, + Optional forUserAgent, + @Safe String operation, + String parentSpanId, + SpanType type) { + 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()}. + */ + public static void initTraceWithSpan( + Observability observability, + String traceId, + Optional forUserAgent, + @Safe String operation, + SpanType type) { + 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, @@ -252,12 +300,13 @@ static DetachedSpan detachInternal(@Safe String operation, SpanType type) { static DetachedSpan detachInternal( Observability observability, String traceId, + Optional forUserAgent, Optional parentSpanId, @Safe String operation, SpanType type) { 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, forUserAgent, parentSpanId, operation, type); } /** @@ -268,12 +317,13 @@ static DetachedSpan detachInternal( Observability observability, String traceId, Optional requestId, + Optional forUserAgent, Optional parentSpanId, @Safe String operation, SpanType type) { // The current trace has no impact on this function, a new trace is spawned and existing thread state // is not modified. - TraceState traceState = TraceState.of(traceId, requestId); + TraceState traceState = TraceState.of(traceId, requestId, forUserAgent); return shouldObserve(observability) ? new SampledDetachedSpan(operation, type, traceState, parentSpanId) : new UnsampledDetachedSpan(traceState, parentSpanId); @@ -314,7 +364,7 @@ private static TraceState getTraceState(@Nullable Trace maybeCurrentTrace, SpanT if (maybeCurrentTrace != null) { return maybeCurrentTrace.getTraceState(); } - return TraceState.of(Tracers.randomId(), getRequestIdForSpan(newSpanType)); + return TraceState.of(Tracers.randomId(), getRequestIdForSpan(newSpanType), Optional.empty()); } private static Optional getRequestIdForSpan(SpanType newSpanType) { @@ -750,6 +800,28 @@ public static String getTraceId() { return checkNotNull(currentTrace.get(), "There is no trace").getTraceId(); } + /** + * Returns the forUserAgent propagated inside the trace. + */ + static Optional getForUserAgent() { + Trace trace = currentTrace.get(); + return trace == null ? Optional.empty() : trace.getForUserAgent(); + } + + /** + * Returns the forUserAgent propagated inside the trace. + */ + @Nullable + static String getForUserAgent(DetachedSpan detachedSpan) { + if (detachedSpan instanceof SampledDetachedSpan) { + return ((SampledDetachedSpan) detachedSpan).traceState.forUserAgent(); + } + if (detachedSpan instanceof UnsampledDetachedSpan) { + return ((UnsampledDetachedSpan) detachedSpan).traceState.forUserAgent(); + } + throw new SafeIllegalStateException("Unknown span type", SafeArg.of("detachedSpan", detachedSpan)); + } + /** * Clears the current trace id and returns it if present. */ diff --git a/tracing/src/main/java/com/palantir/tracing/Tracers.java b/tracing/src/main/java/com/palantir/tracing/Tracers.java index 97d10115d..32b9ea89e 100644 --- a/tracing/src/main/java/com/palantir/tracing/Tracers.java +++ b/tracing/src/main/java/com/palantir/tracing/Tracers.java @@ -21,6 +21,7 @@ import com.google.common.util.concurrent.MoreExecutors; import com.palantir.logsafe.Preconditions; import com.palantir.tracing.api.SpanType; +import com.palantir.tracing.api.TraceHttpHeaders; import java.util.Map; import java.util.Optional; import java.util.concurrent.Callable; @@ -262,6 +263,23 @@ public static > U wrapListenableFuture( return result; } + public static void addTracingHeaders( + T state, TracingHeadersEnrichingFunction tracingHeadersEnrichingFunction) { + Optional maybeTraceMetadata = Tracer.maybeGetTraceMetadata(); + if (maybeTraceMetadata.isEmpty()) { + return; + } + TraceMetadata traceMetadata = maybeTraceMetadata.get(); + tracingHeadersEnrichingFunction.addHeader(TraceHttpHeaders.TRACE_ID, traceMetadata.getTraceId(), state); + tracingHeadersEnrichingFunction.addHeader(TraceHttpHeaders.SPAN_ID, traceMetadata.getSpanId(), state); + tracingHeadersEnrichingFunction.addHeader( + TraceHttpHeaders.IS_SAMPLED, Tracer.isTraceObservable() ? "1" : "0", state); + Optional forUserAgent = Tracer.getForUserAgent(); + if (forUserAgent.isPresent()) { + tracingHeadersEnrichingFunction.addHeader(TraceHttpHeaders.FOR_USER_AGENT, forUserAgent.get(), state); + } + } + private static final class ListenableFutureSpanListener implements Runnable { private final DetachedSpan span; diff --git a/tracing/src/main/java/com/palantir/tracing/TracingHeadersEnrichingFunction.java b/tracing/src/main/java/com/palantir/tracing/TracingHeadersEnrichingFunction.java new file mode 100644 index 000000000..2347062ec --- /dev/null +++ b/tracing/src/main/java/com/palantir/tracing/TracingHeadersEnrichingFunction.java @@ -0,0 +1,22 @@ +/* + * (c) Copyright 2021 Palantir Technologies Inc. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.palantir.tracing; + +/** Function used to adapt the tracing library to any arbitrary request type. */ +public interface TracingHeadersEnrichingFunction { + void addHeader(String headerName, String headerValue, T state); +} diff --git a/tracing/src/test/java/com/palantir/tracing/TraceTest.java b/tracing/src/test/java/com/palantir/tracing/TraceTest.java index 0f3059082..6ae26bd0f 100644 --- a/tracing/src/test/java/com/palantir/tracing/TraceTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TraceTest.java @@ -28,19 +28,19 @@ public final class TraceTest { @Test public void constructTrace_emptyTraceId() { - assertThatThrownBy(() -> Trace.of(false, TraceState.of("", Optional.empty()))) + assertThatThrownBy(() -> Trace.of(false, TraceState.of("", Optional.empty(), Optional.empty()))) .isInstanceOf(IllegalArgumentException.class); } @Test public void testToString() { - Trace trace = Trace.of(true, TraceState.of("traceId", Optional.empty())); + Trace trace = Trace.of(true, TraceState.of("traceId", Optional.empty(), Optional.empty())); OpenSpan span = trace.startSpan("operation", SpanType.LOCAL); assertThat(trace.toString()) .isEqualTo("Trace{" + "stack=[" + span + "], " + "isObservable=true, " - + "state=TraceState{traceId='traceId', requestId='null'}}") + + "state=TraceState{traceId='traceId', requestId='null', forUserAgent='null'}}") .contains(span.getOperation()) .contains(span.getSpanId()); } diff --git a/tracing/src/test/java/com/palantir/tracing/TracerTest.java b/tracing/src/test/java/com/palantir/tracing/TracerTest.java index 80acb3a1b..feeb2a64c 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracerTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracerTest.java @@ -146,14 +146,14 @@ public void testDoesNotNotifyObserversWhenCompletingNonexistingSpan() throws Exc public void testObserversAreInvokedOnObservableTracesOnly() throws Exception { Tracer.subscribe("1", observer1); - Tracer.setTrace(Trace.of(true, TraceState.of(Tracers.randomId(), Optional.empty()))); + Tracer.setTrace(Trace.of(true, TraceState.of(Tracers.randomId(), Optional.empty(), Optional.empty()))); Span span = startAndCompleteSpan(); verify(observer1).consume(span); span = startAndCompleteSpan(); verify(observer1).consume(span); verifyNoMoreInteractions(observer1); - Tracer.setTrace(Trace.of(false, TraceState.of(Tracers.randomId(), Optional.empty()))); + Tracer.setTrace(Trace.of(false, TraceState.of(Tracers.randomId(), Optional.empty(), Optional.empty()))); startAndFastCompleteSpan(); // not sampled, see above verifyNoMoreInteractions(observer1); } @@ -164,7 +164,7 @@ public void testCountsSpansWhenTraceIsNotObservable() throws Exception { assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isNull(); assertThat(Tracer.hasTraceId()).isFalse(); assertThat(Tracer.hasUnobservableTrace()).isFalse(); - Tracer.setTrace(Trace.of(false, TraceState.of(traceId, Optional.empty()))); + Tracer.setTrace(Trace.of(false, TraceState.of(traceId, Optional.empty(), Optional.empty()))); // Unsampled trace should still apply thread state assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo(traceId); assertThat(Tracer.hasTraceId()).isTrue(); @@ -218,7 +218,7 @@ public void testTraceCopyIsIndependent() throws Exception { @Test public void testSetTraceSetsCurrentTraceAndMdcTraceIdKey() throws Exception { Tracer.fastStartSpan("operation"); - Tracer.setTrace(Trace.of(true, TraceState.of("newTraceId", Optional.empty()))); + Tracer.setTrace(Trace.of(true, TraceState.of("newTraceId", Optional.empty(), Optional.empty()))); assertThat(Tracer.getTraceId()).isEqualTo("newTraceId"); assertThat(MDC.get(Tracers.TRACE_ID_KEY)).isEqualTo("newTraceId"); assertThat(Tracer.completeSpan()).isEmpty(); @@ -227,7 +227,7 @@ public void testSetTraceSetsCurrentTraceAndMdcTraceIdKey() throws Exception { @Test public void testSetTraceSetsMdcTraceSampledKeyWhenObserved() { - Tracer.setTrace(Trace.of(true, TraceState.of("observedTraceId", Optional.empty()))); + Tracer.setTrace(Trace.of(true, TraceState.of("observedTraceId", Optional.empty(), Optional.empty()))); assertThat(MDC.get(Tracers.TRACE_SAMPLED_KEY)).isEqualTo("1"); assertThat(Tracer.completeSpan()).isEmpty(); assertThat(MDC.get(Tracers.TRACE_SAMPLED_KEY)).isNull(); @@ -235,7 +235,7 @@ public void testSetTraceSetsMdcTraceSampledKeyWhenObserved() { @Test public void testSetTraceMissingMdcTraceSampledKeyWhenNotObserved() { - Tracer.setTrace(Trace.of(false, TraceState.of("notObservedTraceId", Optional.empty()))); + Tracer.setTrace(Trace.of(false, TraceState.of("notObservedTraceId", Optional.empty(), Optional.empty()))); assertThat(MDC.get(Tracers.TRACE_SAMPLED_KEY)).isNull(); assertThat(Tracer.completeSpan()).isEmpty(); assertThat(MDC.get(Tracers.TRACE_SAMPLED_KEY)).isNull(); @@ -336,7 +336,7 @@ public void testObserversThrow() { @Test public void testGetAndClearTraceIfPresent() { - Trace trace = Trace.of(true, TraceState.of("newTraceId", Optional.empty())); + Trace trace = Trace.of(true, TraceState.of("newTraceId", Optional.empty(), Optional.empty())); Tracer.setTrace(trace); Optional nonEmptyTrace = Tracer.getAndClearTraceIfPresent(); diff --git a/tracing/src/test/java/com/palantir/tracing/TracersTest.java b/tracing/src/test/java/com/palantir/tracing/TracersTest.java index 7122b5e83..fcb01adb0 100644 --- a/tracing/src/test/java/com/palantir/tracing/TracersTest.java +++ b/tracing/src/test/java/com/palantir/tracing/TracersTest.java @@ -31,8 +31,10 @@ import com.palantir.tracing.api.SpanType; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.UUID; @@ -61,7 +63,8 @@ public void before() { Tracer.setSampler(AlwaysSampler.INSTANCE); // Initialize a new trace for each test - Tracer.initTraceWithSpan(Observability.SAMPLE, "defaultTraceId", "rootOperation", SpanType.LOCAL); + Tracer.initTraceWithSpan( + Observability.SAMPLE, "defaultTraceId", Optional.of("forUserAgent"), "rootOperation", SpanType.LOCAL); } @After @@ -890,6 +893,22 @@ public void testExecutorStackDepth_runnable() { .hasSize(rawStackTrace.get().length + 1); } + @Test + public void testAddTracingHeaders_populates() { + Map headers = new HashMap<>(); + TracingHeadersEnrichingFunction> enrichingFunction = + (headerName, headerValue, state) -> state.put(headerName, headerValue); + + Tracers.addTracingHeaders(headers, enrichingFunction); + + assertThat(headers) + .containsAllEntriesOf(Map.of( + "For-User-Agent", "forUserAgent", + "X-B3-TraceId", "defaultTraceId", + "X-B3-Sampled", "1")); + assertThat(headers).containsKey("X-B3-SpanId"); + } + private static Callable newTraceExpectingCallable(String expectedOperation) { final Set seenTraceIds = new HashSet<>(); seenTraceIds.add(Tracer.getTraceId());