Allows subtypes of ITRemote to override the CurrentTraceContext (#1138)

brave-tests/src/main/java/brave/test/ITRemote.java

@@ -14,14 +14,18 @@
 package brave.test;
 
 import brave.Tracing;
-import brave.propagation.B3Propagation;
 import brave.baggage.BaggageField;
 import brave.baggage.BaggagePropagation;
+import brave.propagation.B3Propagation;
+import brave.propagation.CurrentTraceContext;
 import brave.propagation.Propagation;
 import brave.propagation.SamplingFlags;
 import brave.propagation.StrictCurrentTraceContext;
+import brave.propagation.StrictScopeDecorator;
 import brave.propagation.TraceContext;
 import brave.sampler.Sampler;
+import java.io.Closeable;
+import java.io.IOException;
 import org.junit.After;
 import org.junit.Rule;
 import org.junit.rules.DisableOnDebug;
@@ -69,14 +73,32 @@ protected TraceContext newTraceContext(SamplingFlags flags) {
     return tracing.propagationFactory().decorate(result);
   }
 
-  // field because this allows subclasses to initialize a field Tracing
-  protected final StrictCurrentTraceContext currentTraceContext = new StrictCurrentTraceContext();
+  protected final CurrentTraceContext currentTraceContext;
+  protected final Propagation.Factory propagationFactory;
+  protected Tracing tracing; // mutable for test-specific configuration!
 
-  protected final Propagation.Factory propagationFactory =
-    BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
-      .addRemoteField(BAGGAGE_FIELD).build();
+  final Closeable checkForLeakedScopes; // internal to this type
 
-  protected Tracing tracing = tracingBuilder(Sampler.ALWAYS_SAMPLE).build();
+  /** Subclass to override the builder. The result will have {@link StrictScopeDecorator} added */
+  protected CurrentTraceContext.Builder currentTraceContextBuilder() {
+    return StrictCurrentTraceContext.newBuilder();
+  }
+
+  protected ITRemote() {
+    CurrentTraceContext.Builder currentTraceContextBuilder = currentTraceContextBuilder();
+    if (currentTraceContextBuilder instanceof StrictCurrentTraceContext.Builder) {
+      currentTraceContext = currentTraceContextBuilder.build();
+      checkForLeakedScopes = (Closeable) currentTraceContext;
+    } else {
+      StrictScopeDecorator strictScopeDecorator = StrictScopeDecorator.create();
+      currentTraceContext = currentTraceContextBuilder
+        .addScopeDecorator(strictScopeDecorator).build();
+      checkForLeakedScopes = strictScopeDecorator;
+    }
+    propagationFactory = BaggagePropagation.newFactoryBuilder(B3Propagation.FACTORY)
+      .addRemoteField(BAGGAGE_FIELD).build();
+    tracing = tracingBuilder(Sampler.ALWAYS_SAMPLE).build();
+  }
 
   protected Tracing.Builder tracingBuilder(Sampler sampler) {
     return Tracing.newBuilder()
@@ -109,9 +131,9 @@ protected Tracing.Builder tracingBuilder(Sampler sampler) {
     checkForLeakedScopes();
   }
 
-  /** Override to control scope leak enforcement. */
-  protected void checkForLeakedScopes() {
-    currentTraceContext.close();
+  /** Override to disable scope leak enforcement. */
+  protected void checkForLeakedScopes() throws IOException {
+    checkForLeakedScopes.close();
   }
 
   // Assertions below here can eventually move to a new type

instrumentation/http-tests/pom.xml

@@ -29,6 +29,8 @@
     <main.basedir>${project.basedir}/../..</main.basedir>
     <main.java.version>1.8</main.java.version>
     <main.signature.artifact>java18</main.signature.artifact>
+    <!-- disable errorprone warning that "return this" is not synchronized! -->
+    <errorprone.args>-Xep:UnsynchronizedOverridesSynchronized:OFF</errorprone.args>
   </properties>
 
   <dependencies>