Merge pull request #700 from DataDog/tyler/api-tracer-shutdown

New API: Add shutdown callback to shutdown the Tracer

dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java

@@ -25,6 +25,7 @@ import io.opentracing.SpanContext;
 import io.opentracing.propagation.Format;
 import io.opentracing.propagation.TextMap;
 import java.io.Closeable;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -227,13 +228,7 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
     this.serviceNameMappings = serviceNameMappings;
     this.partialFlushMinSpans = partialFlushMinSpans;
 
-    shutdownCallback =
-        new Thread() {
-          @Override
-          public void run() {
-            close();
-          }
-        };
+    shutdownCallback = new ShutdownHook(this);
     try {
       Runtime.getRuntime().addShutdownHook(shutdownCallback);
     } catch (final IllegalStateException ex) {
@@ -266,8 +261,12 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
 
   @Override
   public void finalize() {
-    Runtime.getRuntime().removeShutdownHook(shutdownCallback);
-    shutdownCallback.run();
+    try {
+      Runtime.getRuntime().removeShutdownHook(shutdownCallback);
+      shutdownCallback.run();
+    } catch (final Exception e) {
+      log.error("Error while finalizing DDTracer.", e);
+    }
   }
 
   /**
@@ -719,4 +718,20 @@ public class DDTracer implements io.opentracing.Tracer, Closeable, datadog.trace
       return context;
     }
   }
+
+  private static class ShutdownHook extends Thread {
+    private final WeakReference<DDTracer> reference;
+
+    private ShutdownHook(final DDTracer tracer) {
+      reference = new WeakReference<>(tracer);
+    }
+
+    @Override
+    public void run() {
+      final DDTracer tracer = reference.get();
+      if (tracer != null) {
+        tracer.close();
+      }
+    }
+  }
 }

dd-trace/src/main/java/datadog/trace/tracer/Span.java

@@ -1,5 +1,7 @@
 package datadog.trace.tracer;
 
+import java.util.Map;
+
 /**
  * A single measurement of time with arbitrary key-value attributes.
  *
@@ -94,6 +96,13 @@ public interface Span {
    */
   void attachThrowable(Throwable throwable);
 
+  /**
+   * Get all the metadata attached to this span.
+   *
+   * @return immutable map of span metadata.
+   */
+  Map<String, Object> getMeta();
+
   /**
    * Get a meta value on a span.
    *
@@ -127,5 +136,6 @@ public interface Span {
    *
    * @param finishTimestampNanoseconds Epoch time in nanoseconds.
    */
+  // FIXME: This should take a Timestamp object instead.
   void finish(long finishTimestampNanoseconds);
 }

dd-trace/src/main/java/datadog/trace/tracer/SpanImpl.java

@@ -13,6 +13,7 @@ import java.io.IOException;
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.math.BigInteger;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -212,8 +213,18 @@ class SpanImpl implements Span {
     }
   }
 
+  @Override
+  public synchronized Map<String, Object> getMeta() {
+    return Collections.unmodifiableMap(meta);
+  }
+
+  /**
+   * The agent expects meta's values to be strings.
+   *
+   * @return a copy of meta with all values converted to strings.
+   */
   @JsonGetter("meta")
-  synchronized Map<String, String> getMeta() {
+  synchronized Map<String, String> getMetaString() {
     final Map<String, String> result = new HashMap<>(meta.size());
     for (final Map.Entry<String, Object> entry : meta.entrySet()) {
       result.put(entry.getKey(), String.valueOf(entry.getValue()));
@@ -264,6 +275,7 @@ class SpanImpl implements Span {
     }
   }
 
+  // FIXME: This should take a Timestamp object instead.
   @Override
   public synchronized void finish(final long finishTimestampNanoseconds) {
     if (isFinished()) {

dd-trace/src/main/java/datadog/trace/tracer/Tracer.java

@@ -5,7 +5,7 @@ import datadog.trace.tracer.sampling.AllSampler;
 import datadog.trace.tracer.sampling.Sampler;
 import datadog.trace.tracer.writer.Writer;
 import java.io.Closeable;
-import java.io.IOException;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -29,6 +29,12 @@ public class Tracer implements Closeable {
   /** Interceptors to be called on certain trace and span events */
   private final List<Interceptor> interceptors;
 
+  /**
+   * JVM shutdown callback, keeping a reference to it to remove this if Tracer gets destroyed
+   * earlier
+   */
+  private final Thread shutdownCallback;
+
   @Builder
   private Tracer(
       final Config config,
@@ -45,6 +51,13 @@ public class Tracer implements Closeable {
         interceptors != null
             ? Collections.unmodifiableList(new ArrayList<>(interceptors))
             : Collections.<Interceptor>emptyList();
+
+    shutdownCallback = new ShutdownHook(this);
+    try {
+      Runtime.getRuntime().addShutdownHook(shutdownCallback);
+    } catch (final IllegalStateException ex) {
+      // The JVM is already shutting down.
+    }
   }
 
   /** @return {@link Writer} used by this tracer */
@@ -124,7 +137,35 @@ public class Tracer implements Closeable {
   }
 
   @Override
-  public void close() throws IOException {
+  public void finalize() {
+    try {
+      Runtime.getRuntime().removeShutdownHook(shutdownCallback);
+      shutdownCallback.run();
+    } catch (final Exception e) {
+      log.error("Error while finalizing Tracer.", e);
+    }
+  }
+
+  @Override
+  public void close() {
+    // FIXME: Handle the possibility of close being called more than once or not at all.
+    // FIXME: Depends on order of execution between finalize, GC, and the shutdown hook.
     writer.close();
   }
+
+  private static class ShutdownHook extends Thread {
+    private final WeakReference<Tracer> reference;
+
+    private ShutdownHook(final Tracer tracer) {
+      reference = new WeakReference<>(tracer);
+    }
+
+    @Override
+    public void run() {
+      final Tracer tracer = reference.get();
+      if (tracer != null) {
+        tracer.close();
+      }
+    }
+  }
 }

dd-trace/src/test/groovy/datadog/trace/tracer/JsonSpan.groovy

@@ -43,7 +43,7 @@ class JsonSpan {
   @JsonCreator
   JsonSpan() {}
 
-  JsonSpan(Span span) {
+  JsonSpan(SpanImpl span) {
     traceId = new BigInteger(span.getContext().getTraceId())
     parentId = new BigInteger(span.getContext().getParentId())
     spanId = new BigInteger(span.getContext().getSpanId())
@@ -58,6 +58,6 @@ class JsonSpan {
 
     error = span.isErrored() ? 1 : 0
 
-    meta = span.getMeta()
+    meta = span.getMetaString()
   }
 }

dd-trace/src/test/groovy/datadog/trace/tracer/SpanImplTest.groovy

@@ -129,7 +129,8 @@ class SpanImplTest extends Specification {
     span.setMeta("boolean.key", true)
 
     then:
-    span.getMeta() == ["number.key": "123", "string.key": "meta string", "boolean.key": "true"]
+    span.getMeta() == ["number.key": 123, "string.key": "meta string", "boolean.key": true]
+    span.getMetaString() == ["number.key": "123", "string.key": "meta string", "boolean.key": "true"]
   }
 
   def "test meta setter on finished span for #key"() {
@@ -227,7 +228,7 @@ class SpanImplTest extends Specification {
 
     when: "finish/finalize span"
     span."$method"(*methodArgs)
-    
+
     then: "interceptors called"
     interceptors.reverseEach({ interceptor ->
       then:

dd-trace/src/test/groovy/datadog/trace/tracer/TracerTest.groovy

@@ -45,6 +45,26 @@ class TracerTest extends Specification {
     0 * _  // don't allow any other interaction
   }
 
+  def "finalize closes writer"() {
+    setup:
+    def writer = Mock(Writer)
+    def tracer = Tracer.builder().writer(writer).build()
+
+    when:
+    tracer.finalize()
+
+    then: "closed writer"
+    1 * writer.close()
+    0 * _  // don't allow any other interaction
+
+    when:
+    tracer.finalize()
+
+    then: "thrown error swallowed"
+    1 * writer.close() >> { throw new Exception("test error") }
+    0 * _  // don't allow any other interaction
+  }
+
   def "test create current timestamp"() {
     setup:
     def tracer = Tracer.builder().config(config).build()