Adds ScopeDecorator and backports log correlation to use it (#683)

brave-tests/src/test/java/brave/propagation/DefaultCurrentTraceContextTest.java → brave-tests/src/test/java/brave/propagation/ThreadLocalCurrentTraceContextTest.java

@@ -4,7 +4,7 @@
 import org.junit.Before;
 import org.junit.Test;
 
-public class DefaultCurrentTraceContextTest extends CurrentTraceContextTest {
+public class ThreadLocalCurrentTraceContextTest extends CurrentTraceContextTest {
 
   @Override protected CurrentTraceContext newCurrentTraceContext() {
     return CurrentTraceContext.Default.create();

brave/README.md

@@ -512,7 +512,7 @@ needed to support the current span. It also exposes utilities which you
 can use to decorate executors.
 
 ```java
-CurrentTraceContext currentTraceContext = CurrentTraceContext.Default.create();
+CurrentTraceContext currentTraceContext = ThreadLocalCurrentTraceContext.create();
 tracing = Tracing.newBuilder()
                  .currentTraceContext(currentTraceContext)
                  ...
@@ -666,6 +666,28 @@ class MyFilter extends Filter {
 }
 ```
 
+### Customizing in-process propagation (Ex. log correlation)
+
+The `CurrentTraceContext` makes a trace context visible, such that spans
+aren't accidentally added to the wrong trace. It also accepts hooks like
+log correlation.
+
+For example, if you use Log4J 2, you can copy trace IDs to your logging
+context with [our decorator](../context/log4j2/README.md):
+```java
+tracing = Tracing.newBuilder()
+    .currentTraceContext(ThreadLocalCurrentTraceContext.newBuilder()
+       .addScopeDecorator(ThreadContextScopeDecorator.create())
+       .build()
+    )
+    ...
+    .build();
+```
+
+Besides logging, other tools are available. [StrictScopeDecorator](src/main/java/brave/propagation/StrictScopeDecorator.java) can
+help find out when code is not closing scopes properly. This can be
+useful when writing or diagnosing custom instrumentation.
+
 ## Disabling Tracing
 
 If you are in a situation where you need to turn off tracing at runtime,
@@ -687,15 +709,18 @@ When writing unit tests, there are a few tricks that will make bugs
 easier to find:
 
 * Report spans into a concurrent queue, so you can read them in tests
-* Use `StrictCurrentTraceContext` to reveal subtle propagation bugs
+* Use `StrictScopeDecorator` to reveal subtle thread-related propagation bugs
 * Unconditionally cleanup `Tracing.current()`, to prevent leaks
 
 Here's an example setup for your unit test fixture:
 ```java
 ConcurrentLinkedDeque<Span> spans = new ConcurrentLinkedDeque<>();
 
 Tracing tracing = Tracing.newBuilder()
-                 .currentTraceContext(new StrictCurrentTraceContext())
+                 .currentTraceContext(ThreadLocalCurrentTraceContext.newBuilder()
+                   .addScopeDecorator(StrictScopeDecorator.create())
+                   .build()
+                 )
                  .spanReporter(spans::add)
                  .build();
 

brave/src/main/java/brave/internal/propagation/CorrelationFieldScopeDecorator.java

@@ -0,0 +1,91 @@
+package brave.internal.propagation;
+
+import brave.internal.HexCodec;
+import brave.internal.Nullable;
+import brave.propagation.CurrentTraceContext.Scope;
+import brave.propagation.CurrentTraceContext.ScopeDecorator;
+import brave.propagation.TraceContext;
+
+import static brave.internal.HexCodec.lowerHexEqualsTraceId;
+import static brave.internal.HexCodec.lowerHexEqualsUnsignedLong;
+
+/**
+ * Adds correlation properties "traceId", "parentId" and "spanId" when a {@link
+ * brave.Tracer#currentSpan() span is current}.
+ */
+public abstract class CorrelationFieldScopeDecorator implements ScopeDecorator {
+
+  /**
+   * When the input is not null "traceId", "parentId" and "spanId" correlation properties are saved
+   * off and replaced with those of the current span. When the input is null, these properties are
+   * removed. Either way, "traceId", "parentId" and "spanId" properties are restored on {@linkplain
+   * Scope#close()}.
+   */
+  @Override public Scope decorateScope(@Nullable TraceContext currentSpan, Scope scope) {
+    String previousTraceId = get("traceId");
+    String previousSpanId = get("spanId");
+    String previousParentId = get("parentId");
+
+    if (currentSpan != null) {
+      maybeReplaceTraceContext(currentSpan, previousTraceId, previousParentId, previousSpanId);
+    } else {
+      remove("traceId");
+      remove("parentId");
+      remove("spanId");
+    }
+
+    class CorrelationFieldCurrentTraceContextScope implements Scope {
+      @Override public void close() {
+        scope.close();
+        replace("traceId", previousTraceId);
+        replace("parentId", previousParentId);
+        replace("spanId", previousSpanId);
+      }
+    }
+    return new CorrelationFieldCurrentTraceContextScope();
+  }
+
+  /**
+   * Idempotently sets correlation properties to hex representation of trace identifiers in this
+   * context.
+   */
+  void maybeReplaceTraceContext(
+      TraceContext currentSpan,
+      String previousTraceId,
+      @Nullable String previousParentId,
+      String previousSpanId
+  ) {
+    boolean sameTraceId = lowerHexEqualsTraceId(previousTraceId, currentSpan);
+    if (!sameTraceId) put("traceId", currentSpan.traceIdString());
+
+    long parentId = currentSpan.parentIdAsLong();
+    if (parentId == 0L) {
+      remove("parentId");
+    } else {
+      boolean sameParentId = lowerHexEqualsUnsignedLong(previousParentId, parentId);
+      if (!sameParentId) put("parentId", HexCodec.toLowerHex(parentId));
+    }
+
+    boolean sameSpanId = lowerHexEqualsUnsignedLong(previousSpanId, currentSpan.spanId());
+    if (!sameSpanId) put("spanId", HexCodec.toLowerHex(currentSpan.spanId()));
+  }
+
+  /**
+   * Returns the correlation property of the specified name iff it is a string, or null otherwise.
+   */
+  protected abstract @Nullable String get(String key);
+
+  /** Replaces the correlation property of the specified name */
+  protected abstract void put(String key, String value);
+
+  /** Removes the correlation property of the specified name */
+  protected abstract void remove(String key);
+
+  final void replace(String key, @Nullable String value) {
+    if (value != null) {
+      put(key, value);
+    } else {
+      remove(key);
+    }
+  }
+}

brave/src/main/java/brave/propagation/CurrentTraceContext.java

@@ -4,6 +4,9 @@
 import brave.Tracing;
 import brave.internal.Nullable;
 import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -21,23 +24,81 @@
  * com.google.inject.servlet.RequestScoper and com.github.kristofa.brave.CurrentSpan
  */
 public abstract class CurrentTraceContext {
+
+  /** Implementations of this allow standardized configuration, for example scope decoration. */
+  public abstract static class Builder<B extends Builder<B>> {
+    ArrayList<ScopeDecorator> scopeDecorators = new ArrayList<>();
+
+    /** Implementations call decorators in order to add features like log correlation to a scope. */
+    public B addScopeDecorator(ScopeDecorator scopeDecorator) {
+      if (scopeDecorator == null) throw new NullPointerException("scopeDecorator == null");
+      this.scopeDecorators.add(scopeDecorator);
+      return (B) this;
+    }
+
+    public abstract CurrentTraceContext build();
+  }
+
   /** Returns the current span in scope or null if there isn't one. */
   public abstract @Nullable TraceContext get();
 
   /**
-   * Sets the current span in scope until the returned object is closed. It is a programming
-   * error to drop or never close the result. Using try-with-resources is preferred for this reason.
+   * Sets the current span in scope until the returned object is closed. It is a programming error
+   * to drop or never close the result. Using try-with-resources is preferred for this reason.
    *
    * @param currentSpan span to place into scope or null to clear the scope
    */
   public abstract Scope newScope(@Nullable TraceContext currentSpan);
 
+  final List<ScopeDecorator> scopeDecorators;
+
+  public interface Factory {
+    CurrentTraceContext create(List<ScopeDecorator> scopeDecorators);
+  }
+
+  protected CurrentTraceContext() {
+    this.scopeDecorators = Collections.emptyList();
+  }
+
+  protected CurrentTraceContext(Builder<?> builder) {
+    this.scopeDecorators = new ArrayList<>(builder.scopeDecorators);
+  }
+
+  /**
+   * When implementing {@linkplain #newScope(TraceContext)}, decorate the result before returning
+   * it.
+   *
+   * <p>Ex.
+   * <pre>{@code
+   *   @Override public Scope newScope(@Nullable TraceContext currentSpan) {
+   *     final TraceContext previous = local.get();
+   *     local.set(currentSpan);
+   *     class ThreadLocalScope implements Scope {
+   *       @Override public void close() {
+   *         local.set(previous);
+   *       }
+   *     }
+   *     Scope result = new ThreadLocalScope();
+   *     // ensure scope hooks are attached to the result
+   *     return decorateScope(currentSpan, result);
+   *   }
+   * }</pre>
+   */
+  protected Scope decorateScope(@Nullable TraceContext currentSpan, Scope scope) {
+    int length = scopeDecorators.size();
+    for (int i = 0; i < length; i++) {
+      scope = scopeDecorators.get(i).decorateScope(currentSpan, scope);
+    }
+    return scope;
+  }
+
   /**
    * Like {@link #newScope(TraceContext)}, except returns {@link Scope#NOOP} if the given context is
    * already in scope. This can reduce overhead when scoping callbacks. However, this will not apply
    * any changes, notably in {@link TraceContext#extra()}. As such, it should be used carefully and
    * only in conditions where redundancy is possible and the intent is primarily to facilitate
-   * {@link Tracer#currentSpan}. Most often, this is used to eliminate redundant scopes by wrappers.
+   * {@link Tracer#currentSpan}. Most often, this is used to eliminate redundant scopes by
+   * wrappers.
    *
    * <p>For example, RxJava includes hooks to wrap types that represent an asynchronous functional
    * composition. For example, {@code flowable.parallel().flatMap(Y).sequential()} Assembly hooks
@@ -69,7 +130,8 @@ public Scope maybeScope(@Nullable TraceContext currentSpan) {
   /** A span remains in the scope it was bound to until close is called. */
   public interface Scope extends Closeable {
     /**
-     * Returned when {@link CurrentTraceContext#maybeScope(TraceContext)} detected scope redundancy.
+     * Returned when {@link CurrentTraceContext#maybeScope(TraceContext)} detected scope
+     * redundancy.
      */
     Scope NOOP = new Scope() {
       @Override public void close() {
@@ -84,6 +146,14 @@ public interface Scope extends Closeable {
     @Override void close();
   }
 
+  /**
+   * Use this to add features such as thread checks or log correlation fields when a scope is
+   * created or closed.
+   */
+  public interface ScopeDecorator {
+    Scope decorateScope(@Nullable TraceContext currentSpan, Scope scope);
+  }
+
   /**
    * Default implementation which is backed by a static thread local.
    *
@@ -103,16 +173,13 @@ public interface Scope extends Closeable {
    * <p>If you want a different behavior, use a different subtype of {@link CurrentTraceContext},
    * possibly your own, or raise an issue and explain what your use case is.
    */
-  public static final class Default extends CurrentTraceContext {
-    static final ThreadLocal<TraceContext> DEFAULT = new ThreadLocal<>();
+  public static final class Default extends ThreadLocalCurrentTraceContext {
     // Inheritable as Brave 3's ThreadLocalServerClientAndLocalSpanState was inheritable
     static final InheritableThreadLocal<TraceContext> INHERITABLE = new InheritableThreadLocal<>();
 
-    final ThreadLocal<TraceContext> local;
-
     /** Uses a non-inheritable static thread local */
     public static CurrentTraceContext create() {
-      return new Default(DEFAULT);
+      return new ThreadLocalCurrentTraceContext(new Builder(), DEFAULT);
     }
 
     /**
@@ -125,27 +192,11 @@ public static CurrentTraceContext create() {
      * 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;
+      return new Default();
     }
 
-    @Override public TraceContext get() {
-      return local.get();
-    }
-
-    @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();
+    Default() {
+      super(new Builder(), INHERITABLE);
     }
   }
 

brave/src/main/java/brave/propagation/StrictCurrentTraceContext.java

@@ -1,53 +1,21 @@
 package brave.propagation;
 
-import brave.internal.Nullable;
-
 /**
  * Useful when developing instrumentation as state is enforced more strictly.
  *
  * <p>For example, it is instance scoped as opposed to static scoped, not inheritable and throws an
  * exception if a scope is closed on a different thread that it was opened on.
  *
- * @see CurrentTraceContext.Default
+ * @deprecated use {@linkplain StrictScopeDecorator}. This will be removed in Brave v6.
  */
-public final class StrictCurrentTraceContext extends CurrentTraceContext {
-  // intentionally not inheritable to ensure instrumentation propagation doesn't accidentally work
-  // intentionally not static to make explicit when instrumentation need per thread semantics
-  final ThreadLocal<TraceContext> local = new ThreadLocal<>();
-
-  @Override public TraceContext get() {
-    return local.get();
-  }
-
-  /** Identifies problems by throwing assertion errors when a scope is closed on a different thread. */
-  @Override public Scope newScope(@Nullable TraceContext currentSpan) {
-    TraceContext previous = local.get();
-    local.set(currentSpan);
-    return new StrictScope(previous, new Error(String.format("Thread %s opened scope for %s here:",
-        Thread.currentThread().getName(), currentSpan)));
-  }
-
-  class StrictScope implements Scope {
-    final TraceContext previous;
-    final Throwable caller;
-    final long threadId = Thread.currentThread().getId();
-
-    StrictScope(TraceContext previous, Throwable caller) {
-      this.previous = previous;
-      this.caller = caller;
-    }
-
-    @Override public void close() {
-      if (Thread.currentThread().getId() != threadId) {
-        throw new IllegalStateException(
-            "scope closed in a different thread: " + Thread.currentThread().getName(),
-            caller);
-      }
-      local.set(previous);
-    }
-
-    @Override public String toString() {
-      return caller.toString();
-    }
+@Deprecated
+public final class StrictCurrentTraceContext extends ThreadLocalCurrentTraceContext {
+  static final CurrentTraceContext.Builder SCOPE_DECORATING_BUILDER =
+      ThreadLocalCurrentTraceContext.newBuilder().addScopeDecorator(new StrictScopeDecorator());
+
+  public StrictCurrentTraceContext() { // Preserve historical public ctor
+    // intentionally not inheritable to ensure instrumentation propagation doesn't accidentally work
+    // intentionally not static to make explicit when instrumentation need per thread semantics
+    super(SCOPE_DECORATING_BUILDER, new ThreadLocal<>());
   }
 }