Adding ES instrumentation

dd-java-agent/pom.xml

@@ -107,6 +107,19 @@
             </exclusions>
         </dependency>
 
+        <!-- ES instrumentation -->
+        <dependency>
+            <groupId>io.opentracing.contrib</groupId>
+            <artifactId>opentracing-elasticsearch-client</artifactId>
+            <version>0.0.1</version>
+            <!--<exclusions>-->
+            <!--<exclusion>-->
+                <!--<groupId>com.amazonaws</groupId>-->
+                <!--<artifactId>aws-java-sdk</artifactId>-->
+            <!--</exclusion>-->
+            <!--</exclusions>-->
+        </dependency>
+
         <!-- JUnit tests -->
         <dependency>
             <groupId>io.opentracing</groupId>
@@ -177,6 +190,32 @@
             <scope>test</scope>
         </dependency>
 
+        <!--Elasticsearch tests-->
+        <dependency>
+            <groupId>org.elasticsearch.client</groupId>
+            <artifactId>transport</artifactId>
+            <version>5.4.0</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-api</artifactId>
+            <version>2.7</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-core</artifactId>
+            <version>2.7</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>io.opentracing.contrib</groupId>
+            <artifactId>opentracing-spanmanager</artifactId>
+            <version>0.0.5</version>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
     <build>
         <finalName>${project.artifactId}</finalName>
@@ -195,7 +234,9 @@
                 <artifactId>maven-surefire-plugin</artifactId>
                 <version>2.20</version>
                 <configuration>
-					<argLine>-javaagent:${M2_REPO}/com/datadoghq/dd-java-agent/${project.version}/dd-java-agent-${project.version}.jar</argLine>
+                    <argLine>
+                        -javaagent:${M2_REPO}/com/datadoghq/dd-java-agent/${project.version}/dd-java-agent-${project.version}.jar
+                    </argLine>
                     <includes>
                         <include>*Test.java</include>
                     </includes>
@@ -227,7 +268,8 @@
                                         implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
                                     <manifestEntries>
                                         <Agent-Class>io.opentracing.contrib.agent.AnnotationsTracingAgent</Agent-Class>
-										<Premain-Class>io.opentracing.contrib.agent.AnnotationsTracingAgent</Premain-Class>
+                                        <Premain-Class>io.opentracing.contrib.agent.AnnotationsTracingAgent
+                                        </Premain-Class>
                                         <Can-Redefine-Classes>true</Can-Redefine-Classes>
                                         <Can-Retransform-Classes>true</Can-Retransform-Classes>
                                         <Boot-Class-Path>./${project.artifactId}.jar</Boot-Class-Path>

dd-java-agent/src/main/resources/otarules.btm

@@ -16,8 +16,6 @@ ENDRULE
 # --------------------------------------------------------
 
 
-
-
 # --------------------------------------------------------
 # Instrument AWS client
 RULE AWS SDK Client instrumentation 1
@@ -44,3 +42,57 @@ DO
   $this.requestHandlers.add(0, requestTracerHandler);
 ENDRULE
 # --------------------------------------------------------
+
+
+
+# --------------------------------------------------------
+# Instrument Elasticsearch client
+RULE Elasticsearch Client instrumentation 1
+CLASS org.elasticsearch.client.RestClientBuilder
+METHOD build()
+BIND
+  callback:io.opentracing.contrib.elasticsearch.TracingHttpClientConfigCallback =  new io.opentracing.contrib.elasticsearch.TracingHttpClientConfigCallback();
+AT ENTRY
+IF TRUE
+DO
+  $this.setHttpClientConfigCallback(callback);
+ENDRULE
+
+
+RULE Elasticsearch Client instrumentation 2
+CLASS org.elasticsearch.client.transport.TransportClient
+METHOD doExecute
+HELPER io.opentracing.contrib.agent.OpenTracingHelper
+BIND
+  span:io.opentracing.Span = null;
+#  decorator:io.opentracing.contrib.elasticsearch.SpanDecorator = io.opentracing.contrib.elasticsearch.SpanDecorator;
+AT ENTRY
+IF currentSpan() == null
+DO
+    span = getTracer().buildSpan($2.getClass().getSimpleName())
+        .asChildOf()
+        .withTag(io.opentracing.tag.Tags.SPAN_KIND.getKey(), io.opentracing.tag.Tags.SPAN_KIND_CLIENT)
+        .start();
+#    decorator.onRequest(span);
+    $3 = new io.opentracing.contrib.elasticsearch.TracingResponseListener($3, span);
+ENDRULE
+
+
+RULE Elasticsearch Client instrumentation 3
+CLASS org.elasticsearch.client.transport.TransportClient
+METHOD doExecute
+HELPER io.opentracing.contrib.agent.OpenTracingHelper
+BIND
+  span:io.opentracing.Span = null;
+#  decorator:io.opentracing.contrib.elasticsearch.SpanDecorator = io.opentracing.contrib.elasticsearch.SpanDecorator;
+AT ENTRY
+IF currentSpan() != null
+DO
+    span = getTracer().buildSpan($2.getClass().getSimpleName())
+        .asChildOf(currentSpan())
+        .withTag(io.opentracing.tag.Tags.SPAN_KIND.getKey(), io.opentracing.tag.Tags.SPAN_KIND_CLIENT)
+        .start();
+#    decorator.onRequest(span);
+    $3 = new io.opentracing.contrib.elasticsearch.TracingResponseListener($3, span);
+ENDRULE
+# --------------------------------------------------------

dd-java-agent/src/test/java/com/datadoghq/trace/instrument/ElasticSearchInstrumentationTest.java

@@ -0,0 +1,116 @@
+package com.datadoghq.trace.instrument;
+
+import io.opentracing.Tracer;
+import io.opentracing.contrib.spanmanager.DefaultSpanManager;
+import io.opentracing.mock.MockTracer;
+import io.opentracing.util.GlobalTracer;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.index.IndexResponse;
+import org.elasticsearch.client.transport.TransportClient;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.transport.InetSocketTransportAddress;
+import org.elasticsearch.node.InternalSettingsPreparer;
+import org.elasticsearch.node.Node;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.transport.Netty3Plugin;
+import org.elasticsearch.transport.client.PreBuiltTransportClient;
+import org.junit.*;
+
+import java.io.IOException;
+import java.net.InetAddress;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Date;
+
+import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat;
+import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
+
+/**
+ * Created by gpolaert on 5/23/17.
+ */
+public class ElasticSearchInstrumentationTest {
+
+
+    private static final int HTTP_PORT = 9205;
+    private static final String HTTP_TRANSPORT_PORT = "9305";
+    private static final String ES_WORKING_DIR = "target/es";
+    private static String clusterName = "cluster-name";
+    private static Node node;
+    private static MockTracer tracer;
+
+    @BeforeClass
+    public static void startElasticsearch() throws Exception {
+
+        tracer = new MockTracer();
+        GlobalTracer.register(tracer);
+
+
+        Settings settings = Settings.builder()
+                .put("path.home", ES_WORKING_DIR)
+                .put("path.data", ES_WORKING_DIR + "/data")
+                .put("path.logs", ES_WORKING_DIR + "/logs")
+                .put("transport.type", "netty3")
+                .put("http.type", "netty3")
+                .put("cluster.name", clusterName)
+                .put("http.port", HTTP_PORT)
+                .put("transport.tcp.port", HTTP_TRANSPORT_PORT)
+                .put("network.host", "127.0.0.1")
+                .build();
+        Collection plugins = Collections.singletonList(Netty3Plugin.class);
+        node = new PluginConfigurableNode(settings, plugins);
+        node.start();
+
+    }
+
+    @AfterClass
+    public static void stopElasticsearch() throws Exception {
+        node.close();
+    }
+
+    @Before
+    public void setUp() {
+        DefaultSpanManager.getInstance().activate(tracer.buildSpan("parent").start());
+    }
+
+
+    @After
+    public void tearDown() {
+        DefaultSpanManager.getInstance().current().close();
+    }
+
+    @Test
+    public void Test() throws IOException {
+
+        Settings settings = Settings.builder()
+                .put("cluster.name", clusterName).build();
+
+
+        TransportClient client = new PreBuiltTransportClient(settings)
+                .addTransportAddress(new InetSocketTransportAddress(InetAddress.getByName("localhost"),
+                        Integer.parseInt(HTTP_TRANSPORT_PORT)));
+
+        IndexRequest indexRequest = new IndexRequest("twitter").type("tweet").id("1").
+                source(jsonBuilder()
+                        .startObject()
+                        .field("user", "kimchy")
+                        .field("postDate", new Date())
+                        .field("message", "trying out Elasticsearch")
+                        .endObject()
+                );
+
+        IndexResponse indexResponse = client.index(indexRequest).actionGet();
+
+
+        client.close();
+        DefaultSpanManager.getInstance().current().close();
+        assertThat(tracer.finishedSpans().size()).isEqualTo(2);
+    }
+
+
+    private static class PluginConfigurableNode extends Node {
+        public PluginConfigurableNode(Settings settings, Collection<Class<? extends Plugin>> classpathPlugins) {
+            super(InternalSettingsPreparer.prepareEnvironment(settings, null), classpathPlugins);
+        }
+    }
+
+}