Eagerly provisions span so it is visible to httpclient interceptors (#…

…427)

Thanks @TilinC for reporting

instrumentation/http-tests/src/main/java/brave/http/ITHttp.java

@@ -18,7 +18,7 @@ public abstract class ITHttp {
   @Rule public ExpectedException thrown = ExpectedException.none();
   @Rule public MockWebServer server = new MockWebServer();
 
-  ConcurrentLinkedDeque<Span> spans = new ConcurrentLinkedDeque<>();
+  protected ConcurrentLinkedDeque<Span> spans = new ConcurrentLinkedDeque<>();
 
   protected CurrentTraceContext currentTraceContext = new StrictCurrentTraceContext();
   protected HttpTracing httpTracing;

instrumentation/http/src/main/java/brave/http/HttpClientHandler.java

@@ -29,6 +29,7 @@
  *
  * @param <Req> the native http request type of the client.
  * @param <Resp> the native http response type of the client.
+ * @since 4.3
  */
 public final class HttpClientHandler<Req, Resp> {
 
@@ -73,7 +74,26 @@ public Span handleSend(TraceContext.Injector<Req> injector, Req request) {
    * @see HttpClientParser#request(HttpAdapter, Object, SpanCustomizer)
    */
   public <C> Span handleSend(TraceContext.Injector<C> injector, C carrier, Req request) {
-    Span span = nextSpan(request);
+    return handleSend(injector, carrier, request, nextSpan(request));
+  }
+
+  /**
+   * Like {@link #handleSend(TraceContext.Injector, Object)}, except explicitly controls the span
+   * representing the request.
+   *
+   * @since 4.4
+   */
+  public Span handleSend(TraceContext.Injector<Req> injector, Req request, Span span) {
+    return handleSend(injector, request, request, span);
+  }
+
+  /**
+   * Like {@link #handleSend(TraceContext.Injector, Object, Object)}, except explicitly controls the
+   * span representing the request.
+   *
+   * @since 4.4
+   */
+  public <C> Span handleSend(TraceContext.Injector<C> injector, C carrier, Req request, Span span) {
     injector.inject(span.context(), carrier);
     if (span.isNoop()) return span;
 
@@ -87,8 +107,13 @@ public <C> Span handleSend(TraceContext.Injector<C> injector, C carrier, Req req
     return span.start();
   }
 
-  /** Creates a potentially noop span representing this request */
-  Span nextSpan(Req request) {
+  /**
+   * Creates a potentially noop span representing this request. This is used when you need to
+   * provision a span in a different scope than where the request is executed.
+   *
+   * @since 4.4
+   */
+  public Span nextSpan(Req request) {
     TraceContext parent = currentTraceContext.get();
     if (parent != null) return tracer.newChild(parent);
 

instrumentation/httpasyncclient/src/test/java/brave/httpasyncclient/ITTracingHttpAsyncClientBuilder.java

@@ -3,11 +3,17 @@
 import brave.http.ITHttpClient;
 import java.io.IOException;
 import java.net.URI;
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.RecordedRequest;
+import org.apache.http.HttpRequestInterceptor;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
 import org.apache.http.nio.entity.NStringEntity;
 import org.apache.http.util.EntityUtils;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
 
 public class ITTracingHttpAsyncClientBuilder extends ITHttpClient<CloseableHttpAsyncClient> {
 
@@ -38,4 +44,21 @@ protected void post(CloseableHttpAsyncClient client, String pathIncludingQuery,
   @Override protected void getAsync(CloseableHttpAsyncClient client, String pathIncludingQuery) {
     client.execute(new HttpGet(URI.create(url(pathIncludingQuery))), null);
   }
+
+  @Test public void currentSpanVisibleToUserFilters() throws Exception {
+    server.enqueue(new MockResponse());
+    closeClient(client);
+
+    client = TracingHttpAsyncClientBuilder.create(httpTracing)
+        .addInterceptorFirst((HttpRequestInterceptor) (request, context) ->
+            request.setHeader("my-id", currentTraceContext.get().traceIdString())
+        ).build();
+    client.start();
+
+    get(client, "/foo");
+
+    RecordedRequest request = server.takeRequest();
+    assertThat(request.getHeader("x-b3-traceId"))
+        .isEqualTo(request.getHeader("my-id"));
+  }
 }

instrumentation/httpclient/src/main/java/brave/httpclient/TracingHttpClientBuilder.java

@@ -16,6 +16,9 @@
 import org.apache.http.client.methods.HttpRequestWrapper;
 import org.apache.http.impl.client.HttpClientBuilder;
 import org.apache.http.impl.execchain.ClientExecChain;
+import org.apache.http.impl.execchain.MainClientExec;
+import org.apache.http.protocol.HttpProcessor;
+import org.apache.http.protocol.HttpRequestExecutor;
 import zipkin.Endpoint;
 
 public final class TracingHttpClientBuilder extends HttpClientBuilder {
@@ -40,16 +43,16 @@ public static HttpClientBuilder create(HttpTracing httpTracing) {
   }
 
   /**
-   * main exec is the first in the execution chain, so last to execute. This creates a concrete http
-   * request, so this is where the span is created.
+   * protocol exec is the last in the execution chain, so first to execute. We eagerly create a span
+   * here so that user interceptors can see it.
    */
-  @Override protected ClientExecChain decorateMainExec(ClientExecChain exec) {
-    return (route, request, context, execAware) -> {
-      Span span = handler.handleSend(injector, request);
+  @Override protected ClientExecChain decorateProtocolExec(ClientExecChain protocolExec) {
+    return (route, request, clientContext, execAware) -> {
+      Span span = handler.nextSpan(request);
       CloseableHttpResponse response = null;
       Throwable error = null;
       try (SpanInScope ws = tracer.withSpanInScope(span)) {
-        return response = exec.execute(route, request, context, execAware);
+        return response = protocolExec.execute(route, request, clientContext, execAware);
       } catch (IOException | HttpException | RuntimeException | Error e) {
         error = e;
         throw e;
@@ -59,6 +62,17 @@ public static HttpClientBuilder create(HttpTracing httpTracing) {
     };
   }
 
+  /**
+   * Main exec is the first in the execution chain, so last to execute. This creates a concrete http
+   * request, so this is where the span is started.
+   */
+  @Override protected ClientExecChain decorateMainExec(ClientExecChain exec) {
+    return (route, request, context, execAware) -> {
+      handler.handleSend(injector, request, tracer.currentSpan());
+      return exec.execute(route, request, context, execAware);
+    };
+  }
+
   static final class HttpAdapter
       extends brave.http.HttpClientAdapter<HttpRequestWrapper, HttpResponse> {
 

instrumentation/httpclient/src/test/java/brave/httpclient/ITTracingHttpClientBuilder.java

@@ -3,11 +3,17 @@
 import brave.http.ITHttpClient;
 import java.io.IOException;
 import java.net.URI;
+import okhttp3.mockwebserver.MockResponse;
+import okhttp3.mockwebserver.RecordedRequest;
+import org.apache.http.HttpRequestInterceptor;
 import org.apache.http.client.methods.HttpGet;
 import org.apache.http.client.methods.HttpPost;
 import org.apache.http.entity.StringEntity;
 import org.apache.http.impl.client.CloseableHttpClient;
 import org.junit.AssumptionViolatedException;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
 
 public class ITTracingHttpClientBuilder extends ITHttpClient<CloseableHttpClient> {
 
@@ -34,4 +40,20 @@ public class ITTracingHttpClientBuilder extends ITHttpClient<CloseableHttpClient
   @Override protected void getAsync(CloseableHttpClient client, String pathIncludingQuery) {
     throw new AssumptionViolatedException("This is not an async library");
   }
+
+  @Test public void currentSpanVisibleToUserFilters() throws Exception {
+    server.enqueue(new MockResponse());
+    closeClient(client);
+
+    client = TracingHttpClientBuilder.create(httpTracing).disableAutomaticRetries()
+        .addInterceptorFirst((HttpRequestInterceptor) (request, context) ->
+            request.setHeader("my-id", currentTraceContext.get().traceIdString())
+        ).build();
+
+    get(client, "/foo");
+
+    RecordedRequest request = server.takeRequest();
+    assertThat(request.getHeader("x-b3-traceId"))
+        .isEqualTo(request.getHeader("my-id"));
+  }
 }