Address comments by @trustin

site/src/sphinx/advanced-zipkin.rst

@@ -3,4 +3,4 @@
 Zipkin integration
 ==================
 
-TBW - See :api:`HttpTracingService` and :api:`HttpTracingClient`
+TBW - See :api:`TracingService` and :api:`TracingClient`

zipkin/src/main/java/com/linecorp/armeria/client/tracing/HttpTracingClient.java

@@ -54,55 +54,45 @@
  *
  * <p>This decorator puts trace data into HTTP headers. The specifications of header names and its values
  * correspond to <a href="http://zipkin.io/">Zipkin</a>.
+ *
+ * @deprecated Use {@link TracingClient}.
  */
+@Deprecated
 public final class HttpTracingClient extends SimpleDecoratingClient<HttpRequest, HttpResponse> {
 
     private static final Logger logger = LoggerFactory.getLogger(HttpTracingClient.class);
 
     /**
      * Returns a new builder.
+     *
+     * @deprecated Use {@link TracingClient#builder()}.
      */
-    public static HttpTracingClientBuilder builder() {
-        return new HttpTracingClientBuilder();
-    }
-
-    /**
-     * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance.
-     */
-    public static Function<Client<HttpRequest, HttpResponse>, HttpTracingClient> decorator(Tracing tracing) {
-        return decorator(tracing, null);
-    }
-
-    /**
-     * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance
-     * and the remote service name.
-     */
-    public static Function<Client<HttpRequest, HttpResponse>, HttpTracingClient> decorator(
-            Tracing tracing, @Nullable String remoteServiceName) {
-        checkTracing(tracing);
-        return delegate -> new HttpTracingClient(delegate, tracing, remoteServiceName);
+    @Deprecated
+    public static TracingClientBuilder builder() {
+        return new TracingClientBuilder();
     }
 
     /**
      * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance.
      *
-     * @deprecated Use {@link #decorator(Tracing)}.
+     * @deprecated Use {@link TracingClient#newDecorator(Tracing)}.
      */
     @Deprecated
     public static Function<Client<HttpRequest, HttpResponse>, HttpTracingClient> newDecorator(Tracing tracing) {
-        return decorator(tracing, null);
+        return newDecorator(tracing, null);
     }
 
     /**
      * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance
      * and the remote service name.
      *
-     * @deprecated Use {@link #decorator(Tracing, String)}.
+     * @deprecated Use {@link TracingClient#newDecorator(Tracing, String)}.
      */
     @Deprecated
     public static Function<Client<HttpRequest, HttpResponse>, HttpTracingClient> newDecorator(
             Tracing tracing, @Nullable String remoteServiceName) {
-        return decorator(tracing, remoteServiceName);
+        checkTracing(tracing);
+        return delegate -> new HttpTracingClient(delegate, tracing, remoteServiceName);
     }
 
     private final Tracer tracer;

zipkin/src/main/java/com/linecorp/armeria/client/tracing/TracingClient.java

@@ -0,0 +1,176 @@
+/*
+ * Copyright 2019 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you 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:
+ *
+ *   https://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.linecorp.armeria.client.tracing;
+
+import static com.linecorp.armeria.common.tracing.RequestContextCurrentTraceContext.ensureScopeUsesRequestContext;
+
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+import java.util.function.Function;
+
+import javax.annotation.Nullable;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.linecorp.armeria.client.Client;
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.SimpleDecoratingClient;
+import com.linecorp.armeria.common.HttpRequest;
+import com.linecorp.armeria.common.HttpResponse;
+import com.linecorp.armeria.common.RequestHeadersBuilder;
+import com.linecorp.armeria.common.logging.RequestLog;
+import com.linecorp.armeria.common.logging.RequestLogAvailability;
+import com.linecorp.armeria.common.tracing.RequestContextCurrentTraceContext;
+import com.linecorp.armeria.internal.tracing.AsciiStringKeyFactory;
+import com.linecorp.armeria.internal.tracing.SpanContextUtil;
+import com.linecorp.armeria.internal.tracing.SpanTags;
+
+import brave.Span;
+import brave.Span.Kind;
+import brave.Tracer;
+import brave.Tracer.SpanInScope;
+import brave.Tracing;
+import brave.propagation.TraceContext;
+
+/**
+ * Decorates a {@link Client} to trace outbound {@link HttpRequest}s using
+ * <a href="http://zipkin.io/">Zipkin</a>.
+ *
+ * <p>This decorator puts trace data into HTTP headers. The specifications of header names and its values
+ * correspond to <a href="http://zipkin.io/">Zipkin</a>.
+ */
+public final class TracingClient extends SimpleDecoratingClient<HttpRequest, HttpResponse> {
+
+    private static final Logger logger = LoggerFactory.getLogger(TracingClient.class);
+
+    /**
+     * Returns a new builder.
+     */
+    public static TracingClientBuilder builder() {
+        return new TracingClientBuilder();
+    }
+
+    /**
+     * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance.
+     */
+    public static Function<Client<HttpRequest, HttpResponse>, TracingClient> newDecorator(Tracing tracing) {
+        return newDecorator(tracing, null);
+    }
+
+    /**
+     * Creates a new tracing {@link Client} decorator using the specified {@link Tracing} instance
+     * and the remote service name.
+     */
+    public static Function<Client<HttpRequest, HttpResponse>, TracingClient> newDecorator(
+            Tracing tracing, @Nullable String remoteServiceName) {
+        checkTracing(tracing);
+        return delegate -> new TracingClient(delegate, tracing, remoteServiceName);
+    }
+
+    private final Tracer tracer;
+    private final TraceContext.Injector<RequestHeadersBuilder> injector;
+    @Nullable
+    private final String remoteServiceName;
+
+    /**
+     * Creates a new instance.
+     */
+    TracingClient(Client<HttpRequest, HttpResponse> delegate, Tracing tracing,
+                  @Nullable String remoteServiceName) {
+        super(delegate);
+        tracer = tracing.tracer();
+        injector = tracing.propagationFactory().create(AsciiStringKeyFactory.INSTANCE)
+                          .injector(RequestHeadersBuilder::set);
+        this.remoteServiceName = remoteServiceName;
+    }
+
+    @Override
+    public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Exception {
+        final Span span = tracer.nextSpan();
+
+        // Inject the headers.
+        final RequestHeadersBuilder newHeaders = req.headers().toBuilder();
+        injector.inject(span.context(), newHeaders);
+        req = HttpRequest.of(req, newHeaders.build());
+        ctx.updateRequest(req);
+
+        // For no-op spans, we only need to inject into headers and don't set any other attributes.
+        if (span.isNoop()) {
+            return delegate().execute(ctx, req);
+        }
+
+        final String method = ctx.method().name();
+        span.kind(Kind.CLIENT).name(method);
+        ctx.log().addListener(log -> SpanContextUtil.startSpan(span, log),
+                              RequestLogAvailability.REQUEST_START);
+
+        // Ensure the trace context propagates to children
+        ctx.onChild(RequestContextCurrentTraceContext::copy);
+
+        ctx.log().addListener(log -> {
+            SpanTags.logWireSend(span, log.requestFirstBytesTransferredTimeNanos(), log);
+
+            // If the client timed-out the request, we will have never received any response data at all.
+            if (log.isAvailable(RequestLogAvailability.RESPONSE_FIRST_BYTES_TRANSFERRED)) {
+                SpanTags.logWireReceive(span, log.responseFirstBytesTransferredTimeNanos(), log);
+            }
+
+            finishSpan(span, log);
+        }, RequestLogAvailability.COMPLETE);
+
+        try (SpanInScope ignored = tracer.withSpanInScope(span)) {
+            return delegate().execute(ctx, req);
+        }
+    }
+
+    private void finishSpan(Span span, RequestLog log) {
+        setRemoteEndpoint(span, log);
+        SpanContextUtil.closeSpan(span, log);
+    }
+
+    private void setRemoteEndpoint(Span span, RequestLog log) {
+        final SocketAddress remoteAddress = log.context().remoteAddress();
+        final InetAddress address;
+        final int port;
+        if (remoteAddress instanceof InetSocketAddress) {
+            final InetSocketAddress socketAddress = (InetSocketAddress) remoteAddress;
+            address = socketAddress.getAddress();
+            port = socketAddress.getPort();
+        } else {
+            address = null;
+            port = 0;
+        }
+        if (remoteServiceName != null) {
+            span.remoteServiceName(remoteServiceName);
+        }
+        if (address != null) {
+            span.remoteIpAndPort(address.getHostAddress(), port);
+        }
+    }
+
+    static void checkTracing(Tracing tracing) {
+        try {
+            ensureScopeUsesRequestContext(tracing);
+        } catch (IllegalStateException e) {
+            logger.warn("{} - it is appropriate to ignore this warning if this client is not being used " +
+                        "inside an Armeria server (e.g., this is a normal spring-mvc tomcat server).",
+                        e.getMessage());
+        }
+    }
+}

zipkin/src/main/java/com/linecorp/armeria/client/tracing/HttpTracingClientBuilder.java → zipkin/src/main/java/com/linecorp/armeria/client/tracing/TracingClientBuilder.java

@@ -16,7 +16,7 @@
 
 package com.linecorp.armeria.client.tracing;
 
-import static com.linecorp.armeria.client.tracing.HttpTracingClient.checkTracing;
+import static com.linecorp.armeria.client.tracing.TracingClient.checkTracing;
 import static java.util.Objects.requireNonNull;
 
 import java.util.function.Function;
@@ -31,40 +31,40 @@
 import brave.Tracing;
 
 /**
- * Builds a new {@link HttpTracingClient} or its decorator function.
+ * Builds a new {@link TracingClient} or its decorator function.
  */
-public final class HttpTracingClientBuilder extends AbstractTracingBuilder<HttpTracingClientBuilder> {
+public final class TracingClientBuilder extends AbstractTracingBuilder<TracingClientBuilder> {
 
     @Nullable
     private String remoteServiceName;
 
     /**
      * Creates a new instance.
      */
-    HttpTracingClientBuilder() {}
+    TracingClientBuilder() {}
 
     /**
      * Sets the remote service name.
      */
-    public HttpTracingClientBuilder remoteServiceName(String remoteServiceName) {
+    public TracingClientBuilder remoteServiceName(String remoteServiceName) {
         this.remoteServiceName = requireNonNull(remoteServiceName, "remoteServiceName");
         return this;
     }
 
     /**
-     * Returns a newly-created {@link HttpTracingClient} based on the properties of this builder.
+     * Returns a newly-created {@link TracingClient} based on the properties of this builder.
      */
-    public HttpTracingClient build(Client<HttpRequest, HttpResponse> delegate) {
+    public TracingClient build(Client<HttpRequest, HttpResponse> delegate) {
         final Tracing tracing = buildTracing();
         checkTracing(tracing);
-        return new HttpTracingClient(delegate, tracing, remoteServiceName);
+        return new TracingClient(delegate, tracing, remoteServiceName);
     }
 
     /**
-     * Returns a newly-created decorator that decorates a {@link Client} with a new {@link HttpTracingClient}
+     * Returns a newly-created decorator that decorates a {@link Client} with a new {@link TracingClient}
      * based on the properties of this builder.
      */
-    public Function<Client<HttpRequest, HttpResponse>, HttpTracingClient> decorator() {
+    public Function<Client<HttpRequest, HttpResponse>, TracingClient> decorator() {
         return this::build;
     }
 }

zipkin/src/main/java/com/linecorp/armeria/common/tracing/AbstractTracingBuilder.java

@@ -19,6 +19,9 @@
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import javax.annotation.Nullable;
 
 import brave.Clock;
@@ -27,6 +30,7 @@
 import brave.Tracing.Builder;
 import brave.handler.FinishedSpanHandler;
 import brave.propagation.CurrentTraceContext;
+import brave.propagation.CurrentTraceContext.ScopeDecorator;
 import brave.propagation.Propagation;
 import brave.propagation.Propagation.Factory;
 import brave.sampler.Sampler;
@@ -46,7 +50,7 @@ public abstract class AbstractTracingBuilder<B extends AbstractTracingBuilder<B>
     @Nullable
     private Tracing tracing;
 
-    private boolean isCurrentTraceContextSet;
+    private final List<ScopeDecorator> scopeDecorators = new ArrayList<>();
 
     @SuppressWarnings("unchecked")
     private B self() {
@@ -56,10 +60,14 @@ private B self() {
     /**
      * Sets the {@link Tracing}.
      *
-     * @throws IllegalStateException if other setters for building a new {@link Tracing} are already invoked
+     * @throws IllegalStateException if other setters(e.g., localServiceName, spanReporter, etc.) for
+     *                               building a new {@link Tracing} are already invoked
      */
     public B tracing(Tracing tracing) {
-        checkState(tracingBuilder == null, "tracingBuilder is already set.");
+        if (tracingBuilder != null) {
+            throw new IllegalStateException("cannot set tracing after other setters(e.g., localServiceName, " +
+                                            " spanReporter, etc.) for building a new Tracing are invoked");
+        }
         this.tracing = requireNonNull(tracing, "tracing");
         return self();
     }
@@ -146,19 +154,14 @@ public B sampler(Sampler sampler) {
     }
 
     /**
-     * Sets the {@link CurrentTraceContext}. If this is not set and {@link #tracing(Tracing)} is not invoked,
-     * {@link RequestContextCurrentTraceContext} will be set by default.
+     * Adds the {@link ScopeDecorator}. Use this to add features such as thread checks or log correlation
+     * fields when a scope is created or closed.
      *
      * @throws IllegalStateException if {@link #tracing(Tracing)} is already set
-     *
-     * @see Builder#currentTraceContext(CurrentTraceContext)
      */
-    public B currentTraceContext(CurrentTraceContext currentTraceContext) {
+    public B addScopeDecorator(ScopeDecorator scopeDecorator) {
         checkTracingAndBuilder();
-        if (!isCurrentTraceContextSet) {
-            isCurrentTraceContextSet = true;
-        }
-        tracingBuilder.currentTraceContext(currentTraceContext);
+        scopeDecorators.add(requireNonNull(scopeDecorator, "scopeDecorator"));
         return self();
     }
 
@@ -234,9 +237,14 @@ protected final Tracing buildTracing() {
         if (tracingBuilder == null) {
             tracingBuilder = Tracing.newBuilder();
         }
-        if (!isCurrentTraceContextSet) {
+        if (scopeDecorators.isEmpty()) {
             tracingBuilder.currentTraceContext(RequestContextCurrentTraceContext.DEFAULT);
+        } else {
+            final CurrentTraceContext.Builder builder = RequestContextCurrentTraceContext.builder();
+            scopeDecorators.forEach(builder::addScopeDecorator);
+            tracingBuilder.currentTraceContext(builder.build());
         }
+
         return tracingBuilder.build();
     }