Introduce ForUserAgent

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

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.
      *

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";
 

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,22 +73,33 @@ 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<String> forUserAgent = getForUserAgent(requestContext);
 
         // Set up thread-local span that inherits state from HTTP headers
         if (Strings.isNullOrEmpty(traceId)) {
             // HTTP request did not indicate a trace; initialize trace state and create a span.
             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<String> 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)

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();

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<Request.Builder> {
+        INSTANCE;
+
+        @Override
+        public void addHeader(String headerName, String headerValue, Request.Builder state) {
+            state.header(headerName, headerValue);
         }
     }
 }

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<Request.Builder> {
+        INSTANCE;
+
+        @Override
+        public void addHeader(String headerName, String headerValue, Request.Builder state) {
+            state.header(headerName, headerValue);
+        }
+    }
 }

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);

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<String> 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<String> 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<String> 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<String> 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() {}
 }