Merge pull request #100 from DataDog/gpolaert/extra-services

Ability to set extra service information

dd-trace/src/main/java/com/datadoghq/trace/DDTracer.java

@@ -7,13 +7,19 @@ import com.datadoghq.trace.sampling.AllSampler;
 import com.datadoghq.trace.sampling.Sampler;
 import com.datadoghq.trace.writer.LoggingWriter;
 import com.datadoghq.trace.writer.Writer;
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import io.opentracing.ActiveSpan;
 import io.opentracing.ActiveSpanSource;
 import io.opentracing.BaseSpan;
 import io.opentracing.SpanContext;
 import io.opentracing.propagation.Format;
 import io.opentracing.util.ThreadLocalActiveSpanSource;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
 import java.util.concurrent.ThreadLocalRandom;
 import lombok.extern.slf4j.Slf4j;
 
@@ -37,6 +43,7 @@ public class DDTracer extends ThreadLocalActiveSpanSource implements io.opentrac
   private final Map<String, List<AbstractDecorator>> spanContextDecorators = new HashMap<>();
 
   private final CodecRegistry registry;
+  private final Map<String, Service> services = new HashMap<>();
 
   /** Default constructor, trace/spans are logged, no trace/span dropped */
   public DDTracer() {
@@ -137,6 +144,33 @@ public class DDTracer extends ThreadLocalActiveSpanSource implements io.opentrac
     return "DDTracer{" + "writer=" + writer + ", sampler=" + sampler + '}';
   }
 
+  /**
+   * Register additional information about a service. Service additional information are a Datadog
+   * feature only. Services are reported through a specific Datadog endpoint.
+   *
+   * @param service additional service information
+   */
+  public void addServiceInfo(final Service service) {
+    services.put(service.getName(), service);
+    // Update the writer
+    try {
+      // We don't bother to send multiple times the list of services at this time
+      writer.writeServices(services);
+    } catch (final Throwable ex) {
+      log.warn("Failed to report additional service information, reason: {}", ex.getMessage());
+    }
+  }
+
+  /**
+   * Return the list of additional service information registered
+   *
+   * @return the list of additional service information
+   */
+  @JsonIgnore
+  public Map<String, Service> getServiceInfo() {
+    return services;
+  }
+
   private static class CodecRegistry {
 
     private final Map<Format<?>, Codec<?>> codecs = new HashMap<>();

dd-trace/src/main/java/com/datadoghq/trace/Service.java

@@ -0,0 +1,66 @@
+package com.datadoghq.trace;
+
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonValue;
+
+public class Service {
+
+  private final String name;
+  private final String appName;
+  private final Service.AppType appType;
+
+  public Service(final String name, final String appName, final AppType appType) {
+    this.name = name;
+    this.appName = appName;
+    this.appType = appType;
+  }
+
+  @JsonIgnore
+  public String getName() {
+    return name;
+  }
+
+  @JsonProperty("app")
+  public String getAppName() {
+    return appName;
+  }
+
+  @JsonProperty("app_type")
+  public AppType getAppType() {
+    return appType;
+  }
+
+  @Override
+  public String toString() {
+    return "Service{"
+        + "name='"
+        + name
+        + '\''
+        + ", appName='"
+        + appName
+        + '\''
+        + ", appType="
+        + appType
+        + '}';
+  }
+
+  public enum AppType {
+    WEB("web"),
+    DB("db"),
+    CUSTOM("custom"),
+    CACHE("cache"),
+    WORKER("worker");
+
+    private final String type;
+
+    AppType(final String type) {
+      this.type = type;
+    }
+
+    @JsonValue
+    public String toString() {
+      return type;
+    }
+  }
+}

dd-trace/src/main/java/com/datadoghq/trace/writer/DDAgentWriter.java

@@ -1,8 +1,10 @@
 package com.datadoghq.trace.writer;
 
 import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.Service;
 import com.google.auto.service.AutoService;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -80,6 +82,28 @@ public class DDAgentWriter implements Writer {
     queueFullReported = false;
   }
 
+  /* (non-Javadoc)
+   * @see com.datadoghq.trace.Writer#writeServices(java.util.List)
+   */
+  @Override
+  public void writeServices(final Map<String, Service> services) {
+
+    final Runnable task =
+        new Runnable() {
+          @Override
+          public void run() {
+            //SEND the payload to the agent
+            log.debug("Async writer about to write {} services", services.size());
+            if (api.sendServices(services)) {
+              log.debug("Async writer just sent  {} services", services.size());
+            } else {
+              log.warn("Failed for Async writer to send {} services", services.size());
+            }
+          }
+        };
+    executor.submit(task);
+  }
+
   /* (non-Javadoc)
    * @see com.datadoghq.trace.writer.Writer#start()
    */

dd-trace/src/main/java/com/datadoghq/trace/writer/DDApi.java

@@ -2,6 +2,7 @@ package com.datadoghq.trace.writer;
 
 import com.datadoghq.trace.DDBaseSpan;
 import com.datadoghq.trace.DDTraceInfo;
+import com.datadoghq.trace.Service;
 import com.fasterxml.jackson.core.JsonFactory;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -10,6 +11,7 @@ import java.io.OutputStreamWriter;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.util.List;
+import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
 
 /** The API pointing to a DD agent */
@@ -17,14 +19,17 @@ import lombok.extern.slf4j.Slf4j;
 public class DDApi {
 
   private static final String TRACES_ENDPOINT = "/v0.3/traces";
+  private static final String SERVICES_ENDPOINT = "/v0.3/services";
 
   private final String tracesEndpoint;
+  private final String servicesEndpoint;
 
   private final ObjectMapper objectMapper = new ObjectMapper();
   private final JsonFactory jsonFactory = objectMapper.getFactory();
 
   public DDApi(final String host, final int port) {
     this.tracesEndpoint = "http://" + host + ":" + port + TRACES_ENDPOINT;
+    this.servicesEndpoint = "http://" + host + ":" + port + SERVICES_ENDPOINT;
   }
 
   /**
@@ -34,7 +39,7 @@ public class DDApi {
    * @return the staus code returned
    */
   public boolean sendTraces(final List<List<DDBaseSpan<?>>> traces) {
-    final int status = callPUT(traces);
+    final int status = callPUT(tracesEndpoint, traces);
     if (status == 200) {
       log.debug("Succesfully sent {} traces to the DD agent.", traces.size());
       return true;
@@ -44,6 +49,26 @@ public class DDApi {
     }
   }
 
+  /**
+   * Send service extra information to the services endpoint
+   *
+   * @param services the services to be sent
+   */
+  public boolean sendServices(final Map<String, Service> services) {
+    if (services == null) {
+      return true;
+    }
+    final int status = callPUT(servicesEndpoint, services);
+    if (status == 200) {
+      log.debug("Succesfully sent {} services to the DD agent.", services.size());
+      return true;
+    } else {
+      log.warn(
+          "Error while sending {} services to the DD agent. Status: {}", services.size(), status);
+      return false;
+    }
+  }
+
   /**
    * PUT to an endpoint the provided JSON content
    *
@@ -51,10 +76,10 @@ public class DDApi {
    * @param content
    * @return the status code
    */
-  private int callPUT(final Object content) {
+  private int callPUT(final String endpoint, final Object content) {
     HttpURLConnection httpCon = null;
     try {
-      httpCon = getHttpURLConnection();
+      httpCon = getHttpURLConnection(endpoint);
     } catch (final Exception e) {
       log.warn("Error thrown before PUT call to the DD agent.", e);
       return -1;
@@ -82,9 +107,9 @@ public class DDApi {
     }
   }
 
-  private HttpURLConnection getHttpURLConnection() throws IOException {
+  private HttpURLConnection getHttpURLConnection(final String endpoint) throws IOException {
     final HttpURLConnection httpCon;
-    final URL url = new URL(tracesEndpoint);
+    final URL url = new URL(endpoint);
     httpCon = (HttpURLConnection) url.openConnection();
     httpCon.setDoOutput(true);
     httpCon.setRequestMethod("PUT");

dd-trace/src/main/java/com/datadoghq/trace/writer/ListWriter.java

@@ -1,7 +1,9 @@
 package com.datadoghq.trace.writer;
 
 import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.Service;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CopyOnWriteArrayList;
 
 /** List writer used by tests mostly */
@@ -20,6 +22,11 @@ public class ListWriter extends CopyOnWriteArrayList<List<DDBaseSpan<?>>> implem
     add(trace);
   }
 
+  @Override
+  public void writeServices(final Map<String, Service> services) {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public void start() {
     clear();

dd-trace/src/main/java/com/datadoghq/trace/writer/LoggingWriter.java

@@ -1,8 +1,10 @@
 package com.datadoghq.trace.writer;
 
 import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.Service;
 import com.google.auto.service.AutoService;
 import java.util.List;
+import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
 
 @Slf4j
@@ -14,6 +16,11 @@ public class LoggingWriter implements Writer {
     log.info("write(trace): {}", trace);
   }
 
+  @Override
+  public void writeServices(final Map<String, Service> services) {
+    log.info("additional service information: {}", services.values());
+  }
+
   @Override
   public void close() {
     log.info("close()");

dd-trace/src/main/java/com/datadoghq/trace/writer/Writer.java

@@ -1,7 +1,9 @@
 package com.datadoghq.trace.writer;
 
 import com.datadoghq.trace.DDBaseSpan;
+import com.datadoghq.trace.Service;
 import java.util.List;
+import java.util.Map;
 
 /** A writer is responsible to send collected spans to some place */
 public interface Writer {
@@ -13,6 +15,13 @@ public interface Writer {
    */
   void write(List<DDBaseSpan<?>> trace);
 
+  /**
+   * Report additional service information to the endpoint
+   *
+   * @param services a list of extra information about services
+   */
+  void writeServices(Map<String, Service> services);
+
   /** Start the writer */
   void start();
 

dd-trace/src/test/groovy/com/datadoghq/trace/ServiceTest.groovy

@@ -0,0 +1,67 @@
+package com.datadoghq.trace
+
+import com.datadoghq.trace.sampling.AllSampler
+import com.datadoghq.trace.writer.DDAgentWriter
+import spock.lang.Specification
+
+import static org.mockito.ArgumentMatchers.any
+import static org.mockito.Mockito.*
+
+class ServiceTest extends Specification {
+
+
+  def "getter and setter"() {
+
+    setup:
+    def service = new Service("api-intake", "kafka", Service.AppType.CUSTOM)
+
+    expect:
+    service.getName() == "api-intake"
+    service.getAppName() == "kafka"
+    service.getAppType() == Service.AppType.CUSTOM
+
+  }
+
+  def "enum"() {
+
+    expect:
+    Service.AppType.values().size() == 5
+    Service.AppType.DB.toString() == "db"
+    Service.AppType.WEB.toString() == "web"
+    Service.AppType.CUSTOM.toString() == "custom"
+    Service.AppType.WORKER.toString() == "worker"
+    Service.AppType.CACHE.toString() == "cache"
+
+  }
+
+  def "add extra info about a specific service"() {
+
+    setup:
+    def tracer = new DDTracer()
+    def service = new Service("api-intake", "kafka", Service.AppType.CUSTOM)
+
+    when:
+    tracer.addServiceInfo(service)
+
+    then:
+    tracer.getServiceInfo().size() == 1
+    tracer.getServiceInfo().get("api-intake") == service
+
+  }
+
+  def "add a extra info is reported to the writer"() {
+
+    setup:
+    def writer = spy(new DDAgentWriter())
+    def tracer = new DDTracer(writer, new AllSampler())
+
+
+    when:
+    tracer.addServiceInfo(new Service("api-intake", "kafka", Service.AppType.CUSTOM))
+
+    then:
+    verify(writer, times(1)).writeServices(any(Map))
+
+  }
+
+}

dd-trace/src/test/groovy/com/datadoghq/trace/writer/DDApiTest.groovy

@@ -1,5 +1,6 @@
 package com.datadoghq.trace.writer
 
+import com.datadoghq.trace.Service
 import com.datadoghq.trace.SpanFactory
 import com.fasterxml.jackson.core.type.TypeReference
 import com.fasterxml.jackson.databind.ObjectMapper
@@ -76,7 +77,7 @@ class DDApiTest extends Specification {
     requestHeaders.get().get("Datadog-Meta-Lang") == "java"
     requestHeaders.get().get("Datadog-Meta-Lang-Version") == System.getProperty("java.version", "unknown")
     requestHeaders.get().get("Datadog-Meta-Tracer-Version") == "unknown"
-    convert(requestBody.get()) == expectedRequestBody
+    convertList(requestBody.get()) == expectedRequestBody
 
     cleanup:
     agent.close()
@@ -113,7 +114,88 @@ class DDApiTest extends Specification {
     ])]
   }
 
-  static List<TreeMap<String, Object>> convert(String json) {
+  // Services endpoint
+  def "sending an empty map of services returns no errors"() {
+    setup:
+    def agent = ratpack {
+      handlers {
+        put("v0.3/services") {
+          response.status(200).send()
+        }
+      }
+    }
+    def client = new DDApi("localhost", agent.address.port)
+
+    expect:
+    client.sendServices()
+
+    cleanup:
+    agent.close()
+  }
+
+  def "non-200 response results in false returned for services endpoint"() {
+    setup:
+    def agent = ratpack {
+      handlers {
+        put("v0.3/services") {
+          response.status(404).send()
+        }
+      }
+    }
+    def client = new DDApi("localhost", agent.address.port)
+
+    expect:
+    !client.sendServices([:])
+
+    cleanup:
+    agent.close()
+  }
+
+  def "services content is sent as JSON"() {
+    setup:
+    def requestContentType = new AtomicReference<MediaType>()
+    def requestHeaders = new AtomicReference<Headers>()
+    def requestBody = new AtomicReference<String>()
+    def agent = ratpack {
+      handlers {
+        put("v0.3/services") {
+          requestContentType.set(request.contentType)
+          requestHeaders.set(request.headers)
+          request.body.then {
+            requestBody.set(it.text)
+            response.send()
+          }
+        }
+      }
+    }
+    def client = new DDApi("localhost", agent.address.port)
+
+    expect:
+    client.sendServices(services)
+    requestContentType.get().type == APPLICATION_JSON
+    requestHeaders.get().get("Datadog-Meta-Lang") == "java"
+    requestHeaders.get().get("Datadog-Meta-Lang-Version") == System.getProperty("java.version", "unknown")
+    requestHeaders.get().get("Datadog-Meta-Tracer-Version") == "unknown"
+    convertMap(requestBody.get()) == expectedRequestBody
+
+    cleanup:
+    agent.close()
+
+    // Populate thread info dynamically as it is different when run via gradle vs idea.
+    where:
+    services                                                                          | expectedRequestBody
+    [:]                                                                               | [:]
+    ["service-name": new Service("service-name", "app-name", Service.AppType.CUSTOM)] | ["service-name": new TreeMap<>([
+      "app"     : "app-name",
+      "app_type": "custom"])
+    ]
+  }
+
+  static List<TreeMap<String, Object>> convertList(String json) {
     return mapper.readValue(json, new TypeReference<List<TreeMap<String, Object>>>() {})
   }
+
+  static TreeMap<String, Object> convertMap(String json) {
+    return mapper.readValue(json, new TypeReference<TreeMap<String, Object>>() {})
+  }
 }

dd-trace/src/test/java/ExampleWithLoggingWriter.java

@@ -1,30 +1,28 @@
 import com.datadoghq.trace.DDTracer;
+import com.datadoghq.trace.Service;
 import com.datadoghq.trace.sampling.AllSampler;
 import com.datadoghq.trace.writer.LoggingWriter;
 import io.opentracing.Span;
 
 public class ExampleWithLoggingWriter {
 
-  public static void main(String[] args) throws Exception {
+  public static void main(final String[] args) throws Exception {
 
-    DDTracer tracer = new DDTracer(new LoggingWriter(), new AllSampler());
+    final DDTracer tracer = new DDTracer(new LoggingWriter(), new AllSampler());
+    tracer.addServiceInfo(new Service("api-intake", "spark", Service.AppType.CACHE));
 
-    Span parent =
-        tracer
-            .buildSpan("hello-world")
-            .withServiceName("service-name")
-            .withSpanType("web")
-            .startManual();
+    final Span parent =
+        tracer.buildSpan("fetch.backend").withServiceName("api-intake").startManual();
 
-    parent.setBaggageItem("a-baggage", "value");
+    parent.setBaggageItem("scope-id", "a-1337");
 
     Thread.sleep(100);
 
-    Span child =
+    final Span child =
         tracer
-            .buildSpan("hello-world")
+            .buildSpan("delete.resource")
             .asChildOf(parent)
-            .withResourceName("resource-name")
+            .withResourceName("delete")
             .startManual();
 
     Thread.sleep(100);
@@ -34,5 +32,6 @@ public class ExampleWithLoggingWriter {
     Thread.sleep(100);
 
     parent.finish();
+    tracer.close();
   }
 }

dd-trace/src/test/resources/logback.xml

@@ -9,7 +9,7 @@
     </layout>
   </appender>
 
-  <root level="info">
+  <root level="debug">
     <appender-ref ref="console"/>
   </root>