Use abstract factory pattern to allow subclassing core Jaeger types (#…

…509)

Signed-off-by: Isaac Hier <ihier@uber.com>

jaeger-core/src/main/java/io/jaegertracing/Configuration.java

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2018, The Jaeger Authors
  * Copyright (c) 2016, Uber Technologies, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
@@ -213,7 +214,7 @@ public JaegerTracer.Builder getTracerBuilder() {
     Metrics metrics = new Metrics(metricsFactory);
     Reporter reporter = reporterConfig.getReporter(metrics);
     Sampler sampler = samplerConfig.createSampler(serviceName, metrics);
-    JaegerTracer.Builder builder = new JaegerTracer.Builder(serviceName)
+    JaegerTracer.Builder builder = createTracerBuilder(serviceName)
         .withSampler(sampler)
         .withReporter(reporter)
         .withMetrics(metrics)
@@ -222,6 +223,10 @@ public JaegerTracer.Builder getTracerBuilder() {
     return builder;
   }
 
+  protected JaegerTracer.Builder createTracerBuilder(String serviceName) {
+    return new JaegerTracer.Builder(serviceName);
+  }
+
   public synchronized JaegerTracer getTracer() {
     if (tracer != null) {
       return tracer;

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerObjectFactory.java

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2018, The Jaeger Authors
+ *
+ * Licensed 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
+ *
+ * http://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 io.jaegertracing.internal;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Implements abstract factory pattern for creating spans, span contexts, and span builders. This
+ * pattern allows subclasses of JaegerSpan, JaegerSpanContext, and JaegerTracer.SpanBuilder to be
+ * used consistently in the trace instead of the base class.
+ *
+ * <p>Example usage:</p>
+ *
+ * <pre>{@code
+ *     public class CustomObjectFactory extends JaegerObjectFactory {
+ *         \@Override
+ *         public JaegerSpan createSpan(...) {
+ *           return new CustomSpan(...);
+ *         }
+ *
+ *         // Override other methods...
+ *     }
+ * }</pre>
+ */
+public class JaegerObjectFactory {
+  public JaegerSpan createSpan(
+      JaegerTracer tracer,
+      String operationName,
+      JaegerSpanContext context,
+      long startTimeMicroseconds,
+      long startTimeNanoTicks,
+      boolean computeDurationViaNanoTicks,
+      Map<String, Object> tags,
+      List<Reference> references) {
+    return new JaegerSpan(
+        tracer,
+        operationName,
+        context,
+        startTimeMicroseconds,
+        startTimeNanoTicks,
+        computeDurationViaNanoTicks,
+        tags,
+        references);
+  }
+
+  public JaegerSpanContext createSpanContext(
+      long traceId,
+      long spanId,
+      long parentId,
+      byte flags,
+      Map<String, String> baggage,
+      String debugId) {
+    return new JaegerSpanContext(traceId, spanId, parentId, flags, baggage, debugId, this);
+  }
+
+  public JaegerTracer.SpanBuilder createSpanBuilder(JaegerTracer tracer, String operationName) {
+    return tracer.new SpanBuilder(operationName);
+  }
+}

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerSpan.java

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2018, The Jaeger Authors
  * Copyright (c) 2016, Uber Technologies, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
@@ -46,7 +47,7 @@ public class JaegerSpan implements Span {
   private List<LogData> logs;
   private boolean finished = false; // to prevent the same span from getting reported multiple times
 
-  JaegerSpan(
+  protected JaegerSpan(
       JaegerTracer tracer,
       String operationName,
       JaegerSpanContext context,

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerSpanContext.java

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2018, The Jaeger Authors
  * Copyright (c) 2016, Uber Technologies, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
@@ -30,18 +31,27 @@ public class JaegerSpanContext implements SpanContext {
   private final byte flags;
   private final Map<String, String> baggage;
   private final String debugId;
+  private final JaegerObjectFactory objectFactory;
 
   public JaegerSpanContext(long traceId, long spanId, long parentId, byte flags) {
-    this(traceId, spanId, parentId, flags, Collections.<String, String>emptyMap(), null);
+    this(
+        traceId,
+        spanId,
+        parentId,
+        flags,
+        Collections.<String, String>emptyMap(),
+        null, // debugId
+        new JaegerObjectFactory());
   }
 
-  JaegerSpanContext(
+  protected JaegerSpanContext(
       long traceId,
       long spanId,
       long parentId,
       byte flags,
       Map<String, String> baggage,
-      String debugId) {
+      String debugId,
+      JaegerObjectFactory objectFactory) {
     if (baggage == null) {
       throw new NullPointerException();
     }
@@ -51,6 +61,7 @@ public JaegerSpanContext(long traceId, long spanId, long parentId, byte flags) {
     this.flags = flags;
     this.baggage = baggage;
     this.debugId = debugId;
+    this.objectFactory = objectFactory;
   }
 
   @Override
@@ -102,15 +113,15 @@ public JaegerSpanContext withBaggageItem(String key, String val) {
     } else {
       newBaggage.put(key, val);
     }
-    return new JaegerSpanContext(traceId, spanId, parentId, flags, newBaggage, debugId);
+    return objectFactory.createSpanContext(traceId, spanId, parentId, flags, newBaggage, debugId);
   }
 
   public JaegerSpanContext withBaggage(Map<String, String> newBaggage) {
-    return new JaegerSpanContext(traceId, spanId, parentId, flags, newBaggage, debugId);
+    return objectFactory.createSpanContext(traceId, spanId, parentId, flags, newBaggage, debugId);
   }
 
   public JaegerSpanContext withFlags(byte flags) {
-    return new JaegerSpanContext(traceId, spanId, parentId, flags, baggage, debugId);
+    return objectFactory.createSpanContext(traceId, spanId, parentId, flags, baggage, debugId);
   }
 
   /**
@@ -127,21 +138,12 @@ boolean isDebugIdContainerOnly() {
   }
 
   /**
-   * Create a new dummy JaegerSpanContext as a container for debugId string. This is used when
-   * "jaeger-debug-id" header is passed in the request headers and forces the trace to be sampled as
-   * debug trace, and the value of header recorded as a span tag to serve as a searchable
-   * correlation ID.
-   *
-   * @param debugId arbitrary string used as correlation ID
-   *
-   * @return new dummy JaegerSpanContext that serves as a container for debugId only.
+   * debugId is used when "jaeger-debug-id" header is passed in the request headers and forces the
+   * trace to be sampled as debug trace, and the value of header recorded as a span tag to serve as
+   * a searchable correlation ID.
    *
    * @see Constants#DEBUG_ID_HEADER_KEY
    */
-  public static JaegerSpanContext withDebugId(String debugId) {
-    return new JaegerSpanContext(0, 0, 0, (byte) 0, Collections.<String, String>emptyMap(), debugId);
-  }
-
   String getDebugId() {
     return debugId;
   }

jaeger-core/src/main/java/io/jaegertracing/internal/JaegerTracer.java

@@ -1,4 +1,5 @@
 /*
+ * Copyright (c) 2018, The Jaeger Authors
  * Copyright (c) 2016, Uber Technologies, Inc
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
@@ -74,32 +75,24 @@ public class JaegerTracer implements Tracer, Closeable {
   private final ScopeManager scopeManager;
   private final BaggageSetter baggageSetter;
   private final boolean expandExceptionLogs;
-
-  private JaegerTracer(
-      String serviceName,
-      Reporter reporter,
-      Sampler sampler,
-      PropagationRegistry registry,
-      Clock clock,
-      Metrics metrics,
-      Map<String, Object> tags,
-      boolean zipkinSharedRpcSpan,
-      ScopeManager scopeManager,
-      BaggageRestrictionManager baggageRestrictionManager,
-      boolean expandExceptionLogs) {
-    this.serviceName = serviceName;
-    this.reporter = reporter;
-    this.sampler = sampler;
-    this.registry = registry;
-    this.clock = clock;
-    this.metrics = metrics;
-    this.zipkinSharedRpcSpan = zipkinSharedRpcSpan;
-    this.scopeManager = scopeManager;
-    this.baggageSetter = new BaggageSetter(baggageRestrictionManager, metrics);
-    this.expandExceptionLogs = expandExceptionLogs;
+  private final JaegerObjectFactory objectFactory;
+
+  protected JaegerTracer(JaegerTracer.Builder builder) {
+    this.serviceName = builder.serviceName;
+    this.reporter = builder.reporter;
+    this.sampler = builder.sampler;
+    this.registry = builder.registry;
+    this.clock = builder.clock;
+    this.metrics = builder.metrics;
+    this.zipkinSharedRpcSpan = builder.zipkinSharedRpcSpan;
+    this.scopeManager = builder.scopeManager;
+    this.baggageSetter = new BaggageSetter(builder.baggageRestrictionManager, metrics);
+    this.expandExceptionLogs = builder.expandExceptionLogs;
+    this.objectFactory = builder.objectFactory;
 
     this.version = loadVersion();
 
+    Map<String, Object> tags = new HashMap<String, Object>(builder.tags);
     tags.put(Constants.JAEGER_CLIENT_VERSION_TAG_KEY, this.version);
     if (tags.get(Constants.TRACER_HOSTNAME_TAG_KEY) == null) {
       String hostname = getHostName();
@@ -184,7 +177,7 @@ public Span activeSpan() {
 
   @Override
   public JaegerTracer.SpanBuilder buildSpan(String operationName) {
-    return new SpanBuilder(operationName);
+    return objectFactory.createSpanBuilder(this, operationName);
   }
 
   @Override
@@ -223,10 +216,11 @@ public class SpanBuilder implements Tracer.SpanBuilder {
      * a collection of references.
      */
     private List<Reference> references = Collections.emptyList();
+
     private final Map<String, Object> tags = new HashMap<String, Object>();
     private boolean ignoreActiveSpan = false;
 
-    SpanBuilder(String operationName) {
+    protected SpanBuilder(String operationName) {
       this.operationName = operationName;
     }
 
@@ -315,7 +309,13 @@ private JaegerSpanContext createNewContext(String debugId) {
         }
       }
 
-      return new JaegerSpanContext(id, id, 0, flags);
+      return getObjectFactory().createSpanContext(
+          id,
+          id,
+          0,
+          flags,
+          Collections.<String, String>emptyMap(),
+          debugId);
     }
 
     private Map<String, String> createChildBaggage() {
@@ -354,7 +354,7 @@ private JaegerSpanContext createChildContext() {
         }
       }
 
-      return new JaegerSpanContext(
+      return getObjectFactory().createSpanContext(
           preferredReference.getTraceId(),
           Utils.uniqueId(),
           preferredReference.getSpanId(),
@@ -429,7 +429,7 @@ public JaegerSpan start() {
         }
       }
 
-      JaegerSpan jaegerSpan = new JaegerSpan(
+      JaegerSpan jaegerSpan = getObjectFactory().createSpan(
               JaegerTracer.this,
               operationName,
               context,
@@ -462,12 +462,16 @@ public JaegerTracer.SpanBuilder ignoreActiveSpan() {
     public JaegerSpan startManual() {
       return start();
     }
+
+    private JaegerObjectFactory getObjectFactory() {
+      return JaegerTracer.this.objectFactory;
+    }
   }
 
   /**
    * Builds a {@link JaegerTracer} with options.
    */
-  public static final class Builder {
+  public static class Builder {
     private Sampler sampler;
     private Reporter reporter;
     private final PropagationRegistry registry = new PropagationRegistry();
@@ -479,13 +483,28 @@ public static final class Builder {
     private ScopeManager scopeManager = new ThreadLocalScopeManager();
     private BaggageRestrictionManager baggageRestrictionManager = new DefaultBaggageRestrictionManager();
     private boolean expandExceptionLogs;
+    private final JaegerObjectFactory objectFactory;
 
     public Builder(String serviceName) {
+      this(serviceName, new JaegerObjectFactory());
+    }
+
+    protected Builder(String serviceName, JaegerObjectFactory objectFactory) {
       this.serviceName = checkValidServiceName(serviceName);
-      TextMapCodec textMapCodec = new TextMapCodec(false);
+      this.objectFactory = objectFactory;
+
+      TextMapCodec textMapCodec =
+          TextMapCodec.builder()
+              .withUrlEncoding(false)
+              .withObjectFactory(this.objectFactory)
+              .build();
       this.registerInjector(Format.Builtin.TEXT_MAP, textMapCodec);
       this.registerExtractor(Format.Builtin.TEXT_MAP, textMapCodec);
-      TextMapCodec httpCodec = new TextMapCodec(true);
+      TextMapCodec httpCodec =
+          TextMapCodec.builder()
+              .withUrlEncoding(true)
+              .withObjectFactory(this.objectFactory)
+              .build();
       this.registerInjector(Format.Builtin.HTTP_HEADERS, httpCodec);
       this.registerExtractor(Format.Builtin.HTTP_HEADERS, httpCodec);
       // TODO binary codec not implemented
@@ -590,8 +609,11 @@ public JaegerTracer build() {
             .withMetrics(metrics)
             .build();
       }
-      return new JaegerTracer(serviceName, reporter, sampler, registry, clock, metrics, tags,
-          zipkinSharedRpcSpan, scopeManager, baggageRestrictionManager, expandExceptionLogs);
+      return createTracer();
+    }
+
+    protected JaegerTracer createTracer() {
+      return new JaegerTracer(this);
     }
 
     public static String checkValidServiceName(String serviceName) {