enable swapping out of the underlying implementation of a DefaultTrac…

…er after an SDK has been installed.
open-telemetry · Dec 30, 2019 · 31dcc9e · 31dcc9e
1 parent 4d6f4f5
commit 31dcc9e
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 10 deletions.
diff --git a/api/src/main/java/io/opentelemetry/OpenTelemetry.java b/api/src/main/java/io/opentelemetry/OpenTelemetry.java
@@ -23,9 +23,13 @@
 import io.opentelemetry.metrics.Meter;
 import io.opentelemetry.metrics.MeterFactory;
 import io.opentelemetry.metrics.spi.MeterFactoryProvider;
+import io.opentelemetry.trace.DefaultTracer;
 import io.opentelemetry.trace.DefaultTracerFactory;
+import io.opentelemetry.trace.DefaultTracerFactory.TracerKey;
 import io.opentelemetry.trace.Tracer;
 import io.opentelemetry.trace.TracerFactory;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.ServiceLoader;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Logger;
@@ -140,6 +144,14 @@ public static void setTracerFactory(TracerFactory tracerFactory) {
       logger.warning(
           "The global OpenTelemetry TracerFactory instance has already been set. "
               + "Ignoring this assignment");
+      return;
+    }
+
+    // swap out any tracers that have been handed out with the proper ones.
+    Map<TracerKey, DefaultTracer> existingTracers = DefaultTracerFactory.getExistingTracers();
+    for (Entry<TracerKey, DefaultTracer> entry : existingTracers.entrySet()) {
+      TracerKey key = entry.getKey();
+      entry.getValue().setImplementation(tracerFactory.get(key.getName(), key.getVersion()));
     }
   }
 

diff --git a/api/src/main/java/io/opentelemetry/trace/DefaultTracer.java b/api/src/main/java/io/opentelemetry/trace/DefaultTracer.java
@@ -24,6 +24,7 @@
 import io.opentelemetry.trace.propagation.HttpTraceContext;
 import io.opentelemetry.trace.unsafe.ContextUtils;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import javax.annotation.concurrent.ThreadSafe;
 
 /**
@@ -33,10 +34,13 @@
  */
 @ThreadSafe
 public final class DefaultTracer implements Tracer {
+
   private static final DefaultTracer INSTANCE = new DefaultTracer();
   private static final BinaryFormat<SpanContext> BINARY_FORMAT = new BinaryTraceContext();
   private static final HttpTextFormat<SpanContext> HTTP_TEXT_FORMAT = new HttpTraceContext();
 
+  private final AtomicReference<Tracer> implementation = new AtomicReference<>();
+
   /**
    * Returns a {@code Tracer} singleton that is the default implementations for {@link Tracer}.
    *
@@ -47,35 +51,65 @@ public static Tracer getInstance() {
     return INSTANCE;
   }
 
+  /**
+   * Replace the default implementation with a real one. This can only be done once.
+   *
+   * @param implementation The Tracer implementation that should be used for real.
+   */
+  public void setImplementation(Tracer implementation) {
+    this.implementation.compareAndSet(null, implementation);
+  }
+
   @Override
   public Span getCurrentSpan() {
-    return ContextUtils.getValue();
+    if (useDefault()) {
+      return ContextUtils.getValue();
+    }
+    return implementation.get().getCurrentSpan();
   }
 
   @Override
+  @SuppressWarnings("MustBeClosedChecker")
   public Scope withSpan(Span span) {
-    return ContextUtils.withSpan(span);
+    if (useDefault()) {
+      return ContextUtils.withSpan(span);
+    }
+    return implementation.get().withSpan(span);
   }
 
   @Override
   public Span.Builder spanBuilder(String spanName) {
-    return NoopSpanBuilder.create(this, spanName);
+    if (useDefault()) {
+      return NoopSpanBuilder.create(this, spanName);
+    }
+    return implementation.get().spanBuilder(spanName);
+  }
+
+  private boolean useDefault() {
+    return implementation.get() == null;
   }
 
   @Override
   public BinaryFormat<SpanContext> getBinaryFormat() {
-    return BINARY_FORMAT;
+    if (useDefault()) {
+      return BINARY_FORMAT;
+    }
+    return implementation.get().getBinaryFormat();
   }
 
   @Override
   public HttpTextFormat<SpanContext> getHttpTextFormat() {
-    return HTTP_TEXT_FORMAT;
+    if (useDefault()) {
+      return HTTP_TEXT_FORMAT;
+    }
+    return implementation.get().getHttpTextFormat();
   }
 
-  private DefaultTracer() {}
+  DefaultTracer() {}
 
   // Noop implementation of Span.Builder.
   private static final class NoopSpanBuilder implements Span.Builder {
+
     static NoopSpanBuilder create(Tracer tracer, String spanName) {
       return new NoopSpanBuilder(tracer, spanName);
     }

diff --git a/api/src/main/java/io/opentelemetry/trace/DefaultTracerFactory.java b/api/src/main/java/io/opentelemetry/trace/DefaultTracerFactory.java
@@ -16,21 +16,61 @@
 
 package io.opentelemetry.trace;
 
+import com.google.auto.value.AutoValue;
+import java.util.HashMap;
+import java.util.Map;
+import javax.annotation.Nullable;
+
 public class DefaultTracerFactory implements TracerFactory {
 
-  private static final TracerFactory instance = new DefaultTracerFactory();
+  private final Map<TracerKey, DefaultTracer> tracerRegistry = new HashMap<>();
+
+  private static final DefaultTracerFactory instance = new DefaultTracerFactory();
 
   public static TracerFactory getInstance() {
     return instance;
   }
 
+  /**
+   * Get a map of all the default tracers that were created via this factory.
+   *
+   * @return A Map of {@link TracerKey} to {@link DefaultTracer}
+   */
+  public static Map<TracerKey, DefaultTracer> getExistingTracers() {
+    synchronized (instance.tracerRegistry) {
+      return new HashMap<>(instance.tracerRegistry);
+    }
+  }
+
   @Override
   public Tracer get(String instrumentationName) {
     return get(instrumentationName, null);
   }
 
   @Override
   public Tracer get(String instrumentationName, String instrumentationVersion) {
-    return DefaultTracer.getInstance();
+    synchronized (instance.tracerRegistry) {
+      TracerKey key = TracerKey.makeKey(instrumentationName, instrumentationVersion);
+      DefaultTracer result = tracerRegistry.get(key);
+      if (result != null) {
+        return result;
+      }
+      DefaultTracer defaultTracer = new DefaultTracer();
+      tracerRegistry.put(key, defaultTracer);
+      return defaultTracer;
+    }
+  }
+
+  @AutoValue
+  public abstract static class TracerKey {
+
+    public static TracerKey makeKey(String name, String version) {
+      return new AutoValue_DefaultTracerFactory_TracerKey(name, version);
+    }
+
+    public abstract String getName();
+
+    @Nullable
+    public abstract String getVersion();
   }
 }
diff --git a/api/src/test/java/io/opentelemetry/trace/DefaultTracerTest.java b/api/src/test/java/io/opentelemetry/trace/DefaultTracerTest.java
@@ -17,22 +17,28 @@
 package io.opentelemetry.trace;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.context.propagation.BinaryFormat;
+import io.opentelemetry.context.propagation.HttpTextFormat;
 import io.opentelemetry.trace.propagation.BinaryTraceContext;
 import io.opentelemetry.trace.propagation.HttpTraceContext;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
 
 /** Unit tests for {@link DefaultTracer}. */
-@RunWith(JUnit4.class)
+@RunWith(MockitoJUnitRunner.class)
 // Need to suppress warnings for MustBeClosed because Android 14 does not support
 // try-with-resources.
 @SuppressWarnings("MustBeClosedChecker")
 public class DefaultTracerTest {
+
   private static final Tracer defaultTracer = DefaultTracer.getInstance();
   private static final String SPAN_NAME = "MySpanName";
   private static final byte[] firstBytes =
@@ -46,11 +52,35 @@ public class DefaultTracerTest {
 
   @Rule public final ExpectedException thrown = ExpectedException.none();
 
+  @Mock HttpTextFormat<SpanContext> newHttpTextFormat;
+  @Mock BinaryFormat<SpanContext> newBinaryFormat;
+
   @Test
   public void defaultGetCurrentSpan() {
     assertThat(defaultTracer.getCurrentSpan()).isInstanceOf(DefaultSpan.class);
   }
 
+  @Test
+  public void testImplementationSwap() {
+    DefaultTracer testTracer = new DefaultTracer();
+    assertThat(testTracer.getCurrentSpan()).isInstanceOf(DefaultSpan.class);
+
+    Tracer replacementTracer = mock(Tracer.class);
+    Span newSpan = mock(Span.class);
+    Span.Builder newBuilder = mock(Span.Builder.class);
+
+    when(replacementTracer.getCurrentSpan()).thenReturn(newSpan);
+    when(replacementTracer.spanBuilder("mySpan")).thenReturn(newBuilder);
+    when(replacementTracer.getHttpTextFormat()).thenReturn(newHttpTextFormat);
+    when(replacementTracer.getBinaryFormat()).thenReturn(newBinaryFormat);
+
+    testTracer.setImplementation(replacementTracer);
+    assertThat(testTracer.getCurrentSpan()).isSameInstanceAs(newSpan);
+    assertThat(testTracer.spanBuilder("mySpan")).isSameInstanceAs(newBuilder);
+    assertThat(testTracer.getHttpTextFormat()).isSameInstanceAs(newHttpTextFormat);
+    assertThat(testTracer.getBinaryFormat()).isSameInstanceAs(newBinaryFormat);
+  }
+
   @Test
   public void getCurrentSpan_WithSpan() {
     assertThat(defaultTracer.getCurrentSpan()).isInstanceOf(DefaultSpan.class);