Merge pull request #9 from newrelic-experimental/elasticrestclient710

Elasticrestclient710

elasticsearch-7.3/build.gradle

@@ -4,7 +4,8 @@
 apply plugin: 'java'
 
 dependencies {
-   implementation 'org.elasticsearch.client:transport:7.3.0'
+implementation 'org.elasticsearch.client:transport:7.3.0'
+implementation 'org.elasticsearch.client:elasticsearch-rest-client:7.10.0'
 
    // New Relic Labs Java Agent dependencies
    implementation 'com.newrelic.agent.java:newrelic-agent:6.4.0'
@@ -27,4 +28,5 @@ verifyInstrumentation {
 	excludeRegex '.*alpha.*'
 	excludeRegex '.*beta.*'
 	excludeRegex '.*rc[1-9].*'
-}
+}
+

esrestclient-7.10/build.gradle

@@ -0,0 +1,35 @@
+
+// Build.gradle generated for instrumentation module esrestclient-7.10
+
+apply plugin: 'java'
+
+dependencies {
+   // Declare a dependency on each JAR you want to instrument
+   // Example:
+   // implementation 'javax.servlet:servlet-api:2.5'
+
+   implementation 'org.elasticsearch.client:elasticsearch-rest-client:7.10.0'
+
+   // New Relic Labs Java Agent dependencies
+   implementation 'com.newrelic.agent.java:newrelic-agent:6.4.0'
+   implementation 'com.newrelic.agent.java:newrelic-api:6.4.0'
+   implementation fileTree(include: ['*.jar'], dir: '../libs')
+}
+
+jar {
+  manifest {
+    attributes 'Implementation-Title': 'com.newrelic.instrumentation.labs.esrestclient-7.10'
+    attributes 'Implementation-Vendor': 'New Relic Labs'
+    attributes 'Implementation-Vendor-Id': 'com.newrelic.labs'
+    attributes 'Implementation-Version': 1.0
+  }
+}
+
+
+verifyInstrumentation {
+	passes 'org.elasticsearch.client:elasticsearch-rest-client:[7.10.0,)'
+	excludeRegex '.*SNAPSHOT'
+	excludeRegex '.*alpha.*'
+	excludeRegex '.*beta.*'
+	excludeRegex '.*rc[1-9].*'
+}

esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/ESQueryConverter.java

@@ -0,0 +1,25 @@
+package com.newrelic.instrumentation.labs.esrestclient;
+
+import com.newrelic.api.agent.QueryConverter;
+
+public class ESQueryConverter implements QueryConverter<String> {
+
+	@Override
+	public String toRawQueryString(String rawQuery) {
+		return rawQuery;
+	}
+
+	@Override
+	/**
+	 * Always obfuscate the id
+	 */
+	public String toObfuscatedQueryString(String rawQuery) {
+		int index = rawQuery.lastIndexOf('/');
+		if(index < 1) {
+			return rawQuery;
+		}
+
+		return rawQuery.substring(0, index) + "/?";
+	}
+
+}

esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/NRHolder.java

@@ -0,0 +1,42 @@
+package com.newrelic.instrumentation.labs.esrestclient;
+
+import java.util.HashMap;
+
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Token;
+
+public class NRHolder {
+
+	private static HashMap<Integer, Token> tokenCache = new HashMap<Integer, Token>();
+	private static HashMap<Integer, Segment> segmentCache = new HashMap<Integer, Segment>();
+	private static HashMap<Integer, DatastoreParameters> paramsCache = new HashMap<Integer, DatastoreParameters>();
+
+	public static void putToken(Integer hash, Token t) {
+		tokenCache.put(hash, t);
+	}
+
+	public static void putSegment(Integer hash, Segment s) {
+		segmentCache.put(hash, s);
+	}
+
+	public static void putParams(Integer hash, DatastoreParameters p) {
+		paramsCache.put(hash, p);
+	}
+
+	public static Token getToken(Integer hash) {
+		Token t = tokenCache.remove(hash);
+		return t;
+	}
+
+	public static Segment getSegment(Integer hash) {
+		Segment s = segmentCache.remove(hash);
+		return s;
+	}
+
+	public static DatastoreParameters getParams(Integer hash) {
+		DatastoreParameters p = paramsCache.remove(hash);
+		return p;
+	}
+
+}

esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/RestClient_Instrumentation.java

@@ -0,0 +1,92 @@
+package com.newrelic.instrumentation.labs.esrestclient;
+
+import org.elasticsearch.client.Cancellable;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseListener;
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+
+
+@Weave(originalName = "org.elasticsearch.client.RestClient")
+public abstract class RestClient_Instrumentation {
+
+  @Trace(dispatcher = true)
+  public Response performRequest(Request request) {
+    // Start a segment for the synchronous request
+    Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequest",
+        request.getMethod());
+
+    try {
+      Utils.logRequestAttributes(request);
+
+      String endPoint = request.getEndpoint();
+
+
+      if (endPoint != null && !endPoint.isEmpty()) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequest", endPoint);
+      } else {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequest");
+      }
+      // Log the request as a database call
+
+      segment.reportAsExternal(DatastoreParameters.product("Elasticsearch")
+          .collection(Utils.extractIndex(request.getEndpoint())).operation(request.getMethod())
+          .noInstance().noDatabaseName().slowQuery(endPoint, new ESQueryConverter()).build());
+
+      // Call the original method
+      return Weaver.callOriginal();
+    } finally {
+      // End the segment
+      segment.end();
+    }
+  }
+
+  @Trace(dispatcher = true)
+  public Cancellable performRequestAsync(Request request, ResponseListener responseListener) {
+    // Start a segment for the asynchronous request
+    Segment segment = NewRelic.getAgent().getTransaction().startSegment("ElasticsearchRequestAsync",
+        request.getMethod());
+
+    try {
+      // Add custom attributes for better traceability
+      Utils.logRequestAttributes(request);
+
+      String endPoint = request.getEndpoint();
+
+
+      if (endPoint != null && !endPoint.isEmpty()) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequestAsync", endPoint);
+      } else {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "ES", "RestClient",
+            getClass().getSimpleName(), "performRequestAsync");
+      }
+
+      DatastoreParameters params = DatastoreParameters.product("Elasticsearch")
+          .collection(Utils.extractIndex(request.getEndpoint())).operation(request.getMethod())
+          .noInstance().noDatabaseName().slowQuery(endPoint, new ESQueryConverter()).build();
+
+      // Log the request as a database call
+      int hash = responseListener.hashCode();
+
+      NRHolder.putParams(hash, params);
+      NRHolder.putSegment(hash, segment);
+      NRHolder.putToken(hash, NewRelic.getAgent().getTransaction().getToken());
+
+
+      // Call the original method with the wrapped listener
+      return Weaver.callOriginal();
+    } catch (Exception e) {
+      // Ensure segment is ended in case of an exception
+      segment.end();
+      throw e;
+    }
+  }
+}

esrestclient-7.10/src/main/java/com/newrelic/instrumentation/labs/esrestclient/Utils.java

@@ -0,0 +1,71 @@
+package com.newrelic.instrumentation.labs.esrestclient;
+
+
+import java.util.HashMap;
+import java.util.Map;
+import org.elasticsearch.client.Request;
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+
+public class Utils {
+
+  public static void addAttribute(Map<String, Object> attributes, String key, Object value) {
+    if (key != null && !key.isEmpty() && value != null) {
+      attributes.put(key, value);
+    }
+  }
+
+
+
+  public static void addRequestAttributes(Map<String, Object> attributes, Object request) {
+    if (request instanceof Request) {
+      Request req = (Request) request;
+      // Add attributes directly from the Request object
+      addAttribute(attributes, "Request-Method", req.getMethod());
+      addAttribute(attributes, "Request-Endpoint", req.getEndpoint());
+
+
+      // Add parameters if present
+      if (!req.getParameters().isEmpty()) {
+        addAttribute(attributes, "Request-Parameters", req.getParameters().toString());
+      }
+
+      // Add entity if present
+      if (req.getEntity() != null) {
+        addAttribute(attributes, "Request-Entity", req.getEntity().toString());
+      }
+
+      // Add options
+      addAttribute(attributes, "Request-Options", req.getOptions().toString());
+    }
+
+    // System.out.println(attributes);
+  }
+
+
+  public static DatastoreParameters createDatastoreParams(String operation, String index,
+      String query) {
+    return DatastoreParameters.product("Elasticsearch").collection(index).operation(operation)
+        .instance("localhost", 9200) // Adjust host and port as necessary
+        .noDatabaseName().slowQuery(query, new ESQueryConverter()).build();
+  }
+
+  public static void logRequestAttributes(Object request) {
+    Map<String, Object> attributes = new HashMap<>();
+    addRequestAttributes(attributes, request);
+    NewRelic.getAgent().getTracedMethod().addCustomAttributes(attributes);
+  }
+
+  public static String extractIndex(String input) {
+    if (input == null || input.isEmpty()) {
+      return "noindex"; // or throw an exception based on your use case
+    }
+
+    int underscoreIndex = input.indexOf('_');
+    if (underscoreIndex == -1) {
+      return "noindex"; // or handle the case where there's no underscore
+    }
+
+    return input.substring(0, underscoreIndex);
+  }
+}

esrestclient-7.10/src/main/java/org/elasticsearch/client/ResponseListener.java

@@ -0,0 +1,59 @@
+package org.elasticsearch.client;
+
+
+
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import com.newrelic.instrumentation.labs.esrestclient.NRHolder;
+
+@Weave(type = MatchType.Interface)
+public abstract class ResponseListener {
+
+
+  @Trace(async = true)
+  public void onSuccess(Response response) {
+    Weaver.callOriginal();
+    Token token = NRHolder.getToken(hashCode());
+    if (token != null) {
+      token.linkAndExpire();
+      token = null;
+    }
+    Segment segment = NRHolder.getSegment(hashCode());
+    DatastoreParameters params = NRHolder.getParams(hashCode());
+    if (segment != null) {
+      if (params != null) {
+        segment.reportAsExternal(params);
+      }
+      segment.end();
+    } else if (params != null) {
+      NewRelic.getAgent().getTracedMethod().reportAsExternal(params);
+    }
+  }
+
+  @Trace(async = true)
+  public void onFailure(Exception e) {
+    NewRelic.noticeError(e);
+    Weaver.callOriginal();
+    Token token = NRHolder.getToken(hashCode());
+    if (token != null) {
+      token.linkAndExpire();
+      token = null;
+    }
+    Segment segment = NRHolder.getSegment(hashCode());
+    DatastoreParameters params = NRHolder.getParams(hashCode());
+    if (segment != null) {
+      if (params != null) {
+        segment.reportAsExternal(params);
+      }
+      segment.end();
+    } else if (params != null) {
+      NewRelic.getAgent().getTracedMethod().reportAsExternal(params);
+    }
+  }
+}

settings.gradle

@@ -6,3 +6,4 @@ include 'elasticsearch-6.x'
 include 'elasticsearch-7.0'
 include 'elasticsearch-7.3'
 include 'elasticsearch-7.14'
+include 'esrestclient-7.10'