Draft For-User-Agent implementation

.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<java.lang.String>)\
         \ @ 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<java.lang.String> com.palantir.tracing.TraceMetadata::getForUserAgent()"
+      justification: "Adding a method will not break anyone"

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

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

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

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 {

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

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 {

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,13 +107,15 @@ public DeferredTracer(@Safe String operation, @Safe Map<String, String> 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;
             this.isObservable = false;
             this.parentSpanId = null;
             this.operation = null;
             this.metadata = null;
+            this.forUserAgent = null;
         }
     }
 
@@ -134,7 +139,8 @@ CloseableTrace withTrace() {
 
         Optional<Trace> 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 {

tracing/src/main/java/com/palantir/tracing/DetachedSpan.java

@@ -67,7 +67,23 @@ static DetachedSpan start(
             Optional<String> 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<String> parentSpanId,
+            @Safe String operation,
+            SpanType type,
+            Optional<String> forUserAgent) {
+        return Tracer.detachInternal(observability, traceId, parentSpanId, operation, type, forUserAgent);
     }
 
     /**

tracing/src/main/java/com/palantir/tracing/Trace.java

@@ -48,10 +48,13 @@ public abstract class Trace {
 
     private final Optional<String> requestId;
 
-    private Trace(String traceId, Optional<String> requestId) {
+    private final Optional<String> forUserAgent;
+
+    private Trace(String traceId, Optional<String> requestId, Optional<String> 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<String> getRequestId() {
         return requestId;
     }
 
+    final Optional<String> 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<String> 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<String> requestId, Optional<String> forUserAgent) {
+        return isObservable
+                ? new Sampled(traceId, requestId, forUserAgent)
+                : new Unsampled(traceId, requestId, forUserAgent);
     }
 
     private static final class Sampled extends Trace {
 
         private final Deque<OpenSpan> stack;
 
-        private Sampled(ArrayDeque<OpenSpan> stack, String traceId, Optional<String> requestId) {
-            super(traceId, requestId);
+        private Sampled(
+                ArrayDeque<OpenSpan> stack, String traceId, Optional<String> requestId, Optional<String> forUserAgent) {
+            super(traceId, requestId, forUserAgent);
             this.stack = stack;
         }
 
-        private Sampled(String traceId, Optional<String> requestId) {
-            this(new ArrayDeque<>(), traceId, requestId);
+        private Sampled(String traceId, Optional<String> requestId, Optional<String> 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<String> requestId) {
-            super(traceId, requestId);
+        private Unsampled(
+                int numberOfSpans, String traceId, Optional<String> requestId, Optional<String> forUserAgent) {
+            super(traceId, requestId, forUserAgent);
             this.numberOfSpans = numberOfSpans;
             validateNumberOfSpans();
         }
 
-        private Unsampled(String traceId, Optional<String> requestId) {
-            this(0, traceId, requestId);
+        private Unsampled(String traceId, Optional<String> requestId, Optional<String> 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. */

tracing/src/main/java/com/palantir/tracing/TraceMetadata.java

@@ -50,6 +50,9 @@ default Optional<String> getOriginatingSpanId() {
         return Optional.empty();
     }
 
+    /** Corresponds to {@link com.palantir.tracing.api.TraceHttpHeaders#FOR_USER_AGENT}. */
+    Optional<String> getForUserAgent();
+
     static Builder builder() {
         return new Builder();
     }