Java 分布式跟踪系统 Zipkin（三）：Brave 源码分析 —— Tracing

public class TraceDemo {

    public static void main(String[] args) {
	Sender sender = OkHttpSender.create("http://localhost:9411/api/v2/spans");
	AsyncReporter asyncReporter = AsyncReporter.builder(sender)
		.closeTimeout(500, TimeUnit.MILLISECONDS)
		.build(SpanBytesEncoder.JSON_V2);

	Tracing tracing = Tracing.newBuilder()
		.localServiceName("tracer-demo")
		.spanReporter(asyncReporter)
		.propagationFactory(ExtraFieldPropagation.newFactory(B3Propagation.FACTORY, "user-name"))
		.currentTraceContext(ThreadContextCurrentTraceContext.create())
		.build();
	Tracer tracer = tracing.tracer();
	// ...
    }
}

public static final class Tracing.Builder {
    String localServiceName;
    Endpoint localEndpoint;
    Reporter<zipkin2.Span> reporter;
    Clock clock;
    Sampler sampler = Sampler.ALWAYS_SAMPLE;
    CurrentTraceContext currentTraceContext = CurrentTraceContext.Default.inheritable();
    boolean traceId128Bit = false;
    boolean supportsJoin = true;
    Propagation.Factory propagationFactory = Propagation.Factory.B3;

    public Tracing build() {
      if (clock == null) clock = Platform.get();
      if (localEndpoint == null) {
        localEndpoint = Platform.get().localEndpoint();
        if (localServiceName != null) {
          localEndpoint = localEndpoint.toBuilder().serviceName(localServiceName).build();
        }
      }
      if (reporter == null) reporter = Platform.get();
      return new Default(this);
    }

    Builder() {
    }
}

@Override public void report(zipkin2.Span span) {
  if (!logger.isLoggable(Level.INFO)) return;
  if (span == null) throw new NullPointerException("span == null");
  logger.info(span.toString());
}

public abstract class Sampler {

  public static final Sampler ALWAYS_SAMPLE = new Sampler() {
    @Override public boolean isSampled(long traceId) {
      return true;
    }

    @Override public String toString() {
      return "AlwaysSample";
    }
  };

  public static final Sampler NEVER_SAMPLE = new Sampler() {
    @Override public boolean isSampled(long traceId) {
      return false;
    }

    @Override public String toString() {
      return "NeverSample";
    }
  };

  /** Returns true if the trace ID should be measured. */
  public abstract boolean isSampled(long traceId);

  /**
   * Returns a sampler, given a rate expressed as a percentage.
   *
   * <p>The sampler returned is good for low volumes of traffic (<100K requests), as it is precise.
   * If you have high volumes of traffic, consider {@link BoundarySampler}.
   *
   * @param rate minimum sample rate is 0.01, or 1% of traces
   */
  public static Sampler create(float rate) {
    return CountingSampler.create(rate);
  }
}

TraceContext newRootContext(SamplingFlags samplingFlags, List<Object> extra) {
  long nextId = Platform.get().randomLong();
  Boolean sampled = samplingFlags.sampled();
  if (sampled == null) sampled = sampler.isSampled(nextId);
  return TraceContext.newBuilder()
      .sampled(sampled)
      .traceIdHigh(traceId128Bit ? Platform.get().nextTraceIdHigh() : 0L).traceId(nextId)
      .spanId(nextId)
      .debug(samplingFlags.debug())
      .extra(extra).build();
}

public interface Injector<C> {
  void inject(TraceContext traceContext, C carrier);
}

public interface Extractor<C> {
  TraceContextOrSamplingFlags extract(C carrier);
}

public static final class Default extends CurrentTraceContext {
    static final ThreadLocal<TraceContext> DEFAULT = new ThreadLocal<>();
    // Inheritable as Brave 3's ThreadLocalServerClientAndLocalSpanState was inheritable
    static final InheritableThreadLocal<TraceContext> INHERITABLE = new InheritableThreadLocal<>();

    final ThreadLocal<TraceContext> local;

    /** @deprecated prefer {@link #create()} as it isn't inheritable, so can't leak contexts. */
    @Deprecated
    public Default() {
      this(INHERITABLE);
    }

    /** Uses a non-inheritable static thread local */
    public static CurrentTraceContext create() {
      return new Default(DEFAULT);
    }

    /**
     * Uses an inheritable static thread local which allows arbitrary calls to {@link
     * Thread#start()} to automatically inherit this context. This feature is available as it is was
     * the default in Brave 3, because some users couldn't control threads in their applications.
     *
     * <p>This can be a problem in scenarios such as thread pool expansion, leading to data being
     * recorded in the wrong span, or spans with the wrong parent. If you are impacted by this,
     * switch to {@link #create()}.
     */
    public static CurrentTraceContext inheritable() {
      return new Default(INHERITABLE);
    }

    Default(ThreadLocal<TraceContext> local) {
      if (local == null) throw new NullPointerException("local == null");
      this.local = local;
    }

    @Override public TraceContext get() {
      return local.get();
    }
}

public abstract Scope newScope(@Nullable TraceContext currentSpan);

/** A span remains in the scope it was bound to until close is called. */
public interface Scope extends Closeable {
  /** No exceptions are thrown when unbinding a span scope. */
  @Override void close();
}

try (Scope scope = newScope(invocationContext)) {
  // do somthing
}

@Override public Scope newScope(@Nullable TraceContext currentSpan) {
      final TraceContext previous = local.get();
      local.set(currentSpan);
      class DefaultCurrentTraceContextScope implements Scope {
        @Override public void close() {
          local.set(previous);
        }
      }
      return new DefaultCurrentTraceContextScope();
    }

TraceContext traceContext1;
TraceContext traceContext2;
try (Scope scope = newScope(traceContext1)) {
  // 1.此处CurrentTraceContext.get()能获得traceContext1
  try (Scope scope = newScope(traceContext2)) {
  // 2.此处CurrentTraceContext.get()能获得traceContext2
  }
  // 3.此处CurrentTraceContext.get()能获得traceContext1
}

/** Wraps the input so that it executes with the same context as now. */
public <C> Callable<C> wrap(Callable<C> task) {
  final TraceContext invocationContext = get();
  class CurrentTraceContextCallable implements Callable<C> {
    @Override public C call() throws Exception {
      try (Scope scope = newScope(invocationContext)) {
        return task.call();
      }
    }
  }
  return new CurrentTraceContextCallable();
}

/** Wraps the input so that it executes with the same context as now. */
public Runnable wrap(Runnable task) {
  final TraceContext invocationContext = get();
  class CurrentTraceContextRunnable implements Runnable {
    @Override public void run() {
      try (Scope scope = newScope(invocationContext)) {
        task.run();
      }
    }
  }
  return new CurrentTraceContextRunnable();
}

/**
 * Decorates the input such that the {@link #get() current trace context} at the time a task is
 * scheduled is made current when the task is executed.
 */
public Executor executor(Executor delegate) {
  class CurrentTraceContextExecutor implements Executor {
    @Override public void execute(Runnable task) {
      delegate.execute(CurrentTraceContext.this.wrap(task));
    }
  }
  return new CurrentTraceContextExecutor();
}

/**
 * Decorates the input such that the {@link #get() current trace context} at the time a task is
 * scheduled is made current when the task is executed.
 */
public ExecutorService executorService(ExecutorService delegate) {
  class CurrentTraceContextExecutorService extends brave.internal.WrappingExecutorService {

    @Override protected ExecutorService delegate() {
      return delegate;
    }

    @Override protected <C> Callable<C> wrap(Callable<C> task) {
      return CurrentTraceContext.this.wrap(task);
    }

    @Override protected Runnable wrap(Runnable task) {
      return CurrentTraceContext.this.wrap(task);
    }
  }
  return new CurrentTraceContextExecutorService();
}

public final class ThreadContextCurrentTraceContext extends CurrentTraceContext {
  public static ThreadContextCurrentTraceContext create() {
    return create(CurrentTraceContext.Default.inheritable());
  }

  public static ThreadContextCurrentTraceContext create(CurrentTraceContext delegate) {
    return new ThreadContextCurrentTraceContext(delegate);
  }

  final CurrentTraceContext delegate;

  ThreadContextCurrentTraceContext(CurrentTraceContext delegate) {
    if (delegate == null) throw new NullPointerException("delegate == null");
    this.delegate = delegate;
  }

  @Override public TraceContext get() {
    return delegate.get();
  }

  @Override public Scope newScope(@Nullable TraceContext currentSpan) {
    final String previousTraceId = ThreadContext.get("traceId");
    final String previousSpanId = ThreadContext.get("spanId");

    if (currentSpan != null) {
      ThreadContext.put("traceId", currentSpan.traceIdString());
      ThreadContext.put("spanId", HexCodec.toLowerHex(currentSpan.spanId()));
    } else {
      ThreadContext.remove("traceId");
      ThreadContext.remove("spanId");
    }

    Scope scope = delegate.newScope(currentSpan);
    class ThreadContextCurrentTraceContextScope implements Scope {
      @Override public void close() {
        scope.close();
        ThreadContext.put("traceId", previousTraceId);
        ThreadContext.put("spanId", previousSpanId);
      }
    }
    return new ThreadContextCurrentTraceContextScope();
  }
}

interface Setter<C, K> {
  void put(C carrier, K key, String value);
}
interface Getter<C, K> {
  @Nullable String get(C carrier, K key);
}

<C> TraceContext.Injector<C> injector(Setter<C, K> setter);
<C> TraceContext.Extractor<C> extractor(Getter<C, K> getter);

abstract class Factory {
    public static final Factory B3 = B3Propagation.FACTORY;

    public boolean supportsJoin() {
      return false;
    }

    public boolean requires128BitTraceId() {
      return false;
    }

    public abstract <K> Propagation<K> create(KeyFactory<K> keyFactory);
}

interface KeyFactory<K> {
    KeyFactory<String> STRING = name -> name;

    K create(String name);
}

@Override public <C> TraceContext.Injector<C> injector(Setter<C, K> setter) {
  if (setter == null) throw new NullPointerException("setter == null");
  return new B3Injector<>(this, setter);
}

static final class B3Injector<C, K> implements TraceContext.Injector<C> {
  final B3Propagation<K> propagation;
  final Setter<C, K> setter;

  B3Injector(B3Propagation<K> propagation, Setter<C, K> setter) {
    this.propagation = propagation;
    this.setter = setter;
  }

  @Override public void inject(TraceContext traceContext, C carrier) {
    setter.put(carrier, propagation.traceIdKey, traceContext.traceIdString());
    setter.put(carrier, propagation.spanIdKey, toLowerHex(traceContext.spanId()));
    if (traceContext.parentId() != null) {
      setter.put(carrier, propagation.parentSpanIdKey, toLowerHex(traceContext.parentId()));
    }
    if (traceContext.debug()) {
      setter.put(carrier, propagation.debugKey, "1");
    } else if (traceContext.sampled() != null) {
      setter.put(carrier, propagation.sampledKey, traceContext.sampled() ? "1" : "0");
    }
  }
}

@Override public <C> TraceContext.Extractor<C> extractor(Getter<C, K> getter) {
  if (getter == null) throw new NullPointerException("getter == null");
  return new B3Extractor(this, getter);
}

static final class B3Extractor<C, K> implements TraceContext.Extractor<C> {
  final B3Propagation<K> propagation;
  final Getter<C, K> getter;

  B3Extractor(B3Propagation<K> propagation, Getter<C, K> getter) {
    this.propagation = propagation;
    this.getter = getter;
  }

  @Override public TraceContextOrSamplingFlags extract(C carrier) {
    if (carrier == null) throw new NullPointerException("carrier == null");

    String traceId = getter.get(carrier, propagation.traceIdKey);
    String sampled = getter.get(carrier, propagation.sampledKey);
    String debug = getter.get(carrier, propagation.debugKey);
    if (traceId == null && sampled == null && debug == null) {
      return TraceContextOrSamplingFlags.EMPTY;
    }

    // Official sampled value is 1, though some old instrumentation send true
    Boolean sampledV = sampled != null
        ? sampled.equals("1") || sampled.equalsIgnoreCase("true")
        : null;
    boolean debugV = "1".equals(debug);

    String spanId = getter.get(carrier, propagation.spanIdKey);
    if (spanId == null) { // return early if there's no span ID
      return TraceContextOrSamplingFlags.create(
          debugV ? SamplingFlags.DEBUG : SamplingFlags.Builder.build(sampledV)
      );
    }

    TraceContext.Builder result = TraceContext.newBuilder().sampled(sampledV).debug(debugV);
    result.traceIdHigh(
        traceId.length() == 32 ? lowerHexToUnsignedLong(traceId, 0) : 0
    );
    result.traceId(lowerHexToUnsignedLong(traceId));
    result.spanId(lowerHexToUnsignedLong(spanId));
    String parentSpanIdString = getter.get(carrier, propagation.parentSpanIdKey);
    if (parentSpanIdString != null) result.parentId(lowerHexToUnsignedLong(parentSpanIdString));
    return TraceContextOrSamplingFlags.create(result.build());
  }
}

static final class ExtraFieldInjector<C, K> implements Injector<C> {
  final Injector<C> delegate;
  final Propagation.Setter<C, K> setter;
  final Map<String, K> nameToKey;

  ExtraFieldInjector(Injector<C> delegate, Setter<C, K> setter, Map<String, K> nameToKey) {
    this.delegate = delegate;
    this.setter = setter;
    this.nameToKey = nameToKey;
  }

  @Override public void inject(TraceContext traceContext, C carrier) {
    for (Object extra : traceContext.extra()) {
      if (extra instanceof Extra) {
        ((Extra) extra).setAll(carrier, setter, nameToKey);
        break;
      }
    }
    delegate.inject(traceContext, carrier);
  }
}

static final class One extends Extra {
  String name, value;

  @Override void put(String name, String value) {
    this.name = name;
    this.value = value;
  }

  @Override String get(String name) {
    return name.equals(this.name) ? value : null;
  }

  @Override <C, K> void setAll(C carrier, Setter<C, K> setter, Map<String, K> nameToKey) {
    K key = nameToKey.get(name);
    if (key == null) return;
    setter.put(carrier, key, value);
  }

  @Override public String toString() {
    return "ExtraFieldPropagation{" + name + "=" + value + "}";
  }
}

static final class Many extends Extra {
  final LinkedHashMap<String, String> fields = new LinkedHashMap<>();

  @Override void put(String name, String value) {
    fields.put(name, value);
  }

  @Override String get(String name) {
    return fields.get(name);
  }

  @Override <C, K> void setAll(C carrier, Setter<C, K> setter, Map<String, K> nameToKey) {
    for (Map.Entry<String, String> field : fields.entrySet()) {
      K key = nameToKey.get(field.getKey());
      if (key == null) continue;
      setter.put(carrier, nameToKey.get(field.getKey()), field.getValue());
    }
  }

  @Override public String toString() {
    return "ExtraFieldPropagation" + fields;
  }
}

static final class ExtraFieldExtractor<C, K> implements Extractor<C> {
  final Extractor<C> delegate;
  final Propagation.Getter<C, K> getter;
  final Map<String, K> names;

  ExtraFieldExtractor(Extractor<C> delegate, Getter<C, K> getter, Map<String, K> names) {
    this.delegate = delegate;
    this.getter = getter;
    this.names = names;
  }

  @Override public TraceContextOrSamplingFlags extract(C carrier) {
    TraceContextOrSamplingFlags result = delegate.extract(carrier);

    Extra extra = null;
    for (Map.Entry<String, K> field : names.entrySet()) {
      String maybeValue = getter.get(carrier, field.getValue());
      if (maybeValue == null) continue;
      if (extra == null) {
        extra = new One();
      } else if (extra instanceof One) {
        One one = (One) extra;
        extra = new Many();
        extra.put(one.name, one.value);
      }
      extra.put(field.getKey(), maybeValue);
    }
    if (extra == null) return result;
    return result.toBuilder().addExtra(extra).build();
  }
}