initial checkin

elasticsearch-highlevel-rest/build.gradle

@@ -0,0 +1,32 @@
+
+// Build.gradle generated for instrumentation module elasticsearch-highlevel-rest
+
+apply plugin: 'java'
+
+dependencies {
+   // https://mvnrepository.com/artifact/org.elasticsearch.client/elasticsearch-rest-high-level-client
+implementation group: 'org.elasticsearch.client', name: 'elasticsearch-rest-high-level-client', version: '7.0.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.elasticsearch-highlevel-rest'
+    attributes 'Implementation-Vendor': 'New Relic Labs'
+    attributes 'Implementation-Vendor-Id': 'com.newrelic.labs'
+    attributes 'Implementation-Version': 1.0
+  }
+}
+
+verifyInstrumentation {
+  // Verifier plugin documentation:
+  // https://github.com/newrelic/newrelic-gradle-verify-instrumentation
+  // Example:
+  // passes 'javax.servlet:servlet-api:[2.2,2.5]'
+  // exclude 'javax.servlet:servlet-api:2.4.public_draft'
+}

elasticsearch-highlevel-rest/src/main/java/com/newrelic/instrumentation/labs/elasticsearch/highlevelrest/ESQueryConverter.java

@@ -0,0 +1,25 @@
+package com.newrelic.instrumentation.labs.elasticsearch.highlevelrest;
+
+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) + "/?";
+	}
+
+}

elasticsearch-highlevel-rest/src/main/java/com/newrelic/instrumentation/labs/elasticsearch/highlevelrest/ElasticSearchUtils.java

@@ -0,0 +1,243 @@
+package com.newrelic.instrumentation.labs.elasticsearch.highlevelrest;
+
+import java.util.List;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.delete.DeleteRequest;
+import org.elasticsearch.action.get.GetRequest;
+import org.elasticsearch.action.get.MultiGetRequest;
+import org.elasticsearch.action.get.MultiGetRequest.Item;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.action.search.MultiSearchRequest;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.update.UpdateRequest;
+
+public class ElasticSearchUtils {
+
+	public static Holder getAll(ActionRequest request) {
+		String operation = getOperation(request);
+		String collection = getCollection(operation, request);
+		String query = getQuery(operation,request);
+
+		return new Holder(query, collection, operation);
+	}
+
+	public static String getOperation(ActionRequest request) {
+		String type = "Unknown";
+		String classname = request.getClass().getSimpleName();
+		if(classname.endsWith("Request")) {
+			int index = classname.indexOf("Request");
+			if(index > -1) {
+				type = classname.substring(0, index);
+			}
+		}
+		return type;
+	}
+
+
+	public static String getQuery(String operation,ActionRequest request) {
+		String result = "";
+
+		if(operation.equalsIgnoreCase("get")) {
+			GetRequest getRequest = (GetRequest)request;
+			String index = getRequest.index();
+			String id = getRequest.id();
+
+			return  "get "+index+"/"+id;
+		}
+
+		if(operation.equalsIgnoreCase("update")) {
+			UpdateRequest updateRequest = (UpdateRequest)request;
+			String index = updateRequest.index();
+			String id = updateRequest.id();
+			return  "update "+index+"/"+id;
+		}
+
+		if(operation.equalsIgnoreCase("search")) {
+			SearchRequest searchRequest = (SearchRequest)request;
+			String tmp = getSearchCollection(searchRequest);
+			return "search "+tmp;
+		}
+
+		if(operation.equalsIgnoreCase("index")) {
+			IndexRequest indexRequest = (IndexRequest)request;
+			String index = indexRequest.index();
+			String id = indexRequest.id();
+
+			return "index " +index+"/"+id;
+		}
+
+		if(operation.equalsIgnoreCase("delete")) {
+			DeleteRequest deleteRequest = (DeleteRequest)request;
+			String index = deleteRequest.index();
+			String id = deleteRequest.id();
+
+			return "delete " +index+"/"+id;
+		}
+
+		if(operation.equalsIgnoreCase("multiget")) {
+			MultiGetRequest multiGetRequest = (MultiGetRequest)request;
+			List<Item> items = multiGetRequest.getItems();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<items.size();i++) {
+				Item item = items.get(i);
+				String index = item.index();
+				String type = item.type();
+				String id = item.id();
+				sb.append(index+"/"+type+"/"+id);
+				if(i < items.size()-1) {
+					sb.append(',');
+				}
+			}
+			return "multiget " + sb.toString();
+		}
+
+		if(operation.equalsIgnoreCase("multisearch")) {
+			MultiSearchRequest multiSearchRequest = (MultiSearchRequest)request;
+			List<SearchRequest> searches = multiSearchRequest.requests(); 
+			int size = searches.size();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<size;i++) {
+				SearchRequest searchRequest = searches.get(i);
+				String s = getSearchCollection(searchRequest);
+				sb.append(s);
+				if(i < size -1) {
+					sb.append(',');
+				}
+			}
+			return "multisearch "+sb.toString();
+		}
+
+		return result;
+	}
+
+	public static String getCollection(ActionRequest request) {
+		String result = "Unknown";
+		if(GetRequest.class.isInstance(request)) {
+			GetRequest getRequest = (GetRequest)request;
+			result = getRequest.index();
+		} else if(UpdateRequest.class.isInstance(request)) {
+			UpdateRequest updateRequest = (UpdateRequest)request;
+			result = updateRequest.index();
+		} else if(SearchRequest.class.isInstance(request)) {
+			SearchRequest searchRequest = (SearchRequest)request;
+			result = getSearchCollection(searchRequest);
+		} else if(IndexRequest.class.isInstance(request)) {
+			IndexRequest indexRequest = (IndexRequest)request;
+			result = indexRequest.index();
+		} else if(DeleteRequest.class.isInstance(request)) {
+			DeleteRequest deleteRequest = (DeleteRequest)request;
+			result = deleteRequest.index();
+		} else if(MultiGetRequest.class.isInstance(request)) {
+			MultiGetRequest multiGetRequest = (MultiGetRequest)request;
+			List<Item> items = multiGetRequest.getItems();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<items.size();i++) {
+				Item item = items.get(i);
+				sb.append(item.index()+"-"+item.type());
+				if(i < items.size()-1) {
+					sb.append(',');
+				}
+			}
+			result = sb.toString();
+			if(result.endsWith(",")) {
+				result = result.substring(0, result.length()-1);
+			}
+		} else if(MultiSearchRequest.class.isInstance(request)) {
+			MultiSearchRequest multiSearchRequest = (MultiSearchRequest)request;
+			List<SearchRequest> searches = multiSearchRequest.requests(); 
+			int size = searches.size();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<size;i++) {
+				SearchRequest searchRequest = searches.get(i);
+				String s = getSearchCollection(searchRequest);
+				sb.append(s);
+				if(i < size -1) {
+					sb.append(',');
+				}
+			}
+			result = sb.toString();
+		}
+		return result;
+	}
+
+	public static String getCollection(String operation, ActionRequest request) {
+		String result = "Unknown";
+		if(operation.equalsIgnoreCase("get")) {
+			GetRequest getRequest = (GetRequest)request;
+			result = getRequest.index();
+		} else if(operation.equalsIgnoreCase("update")) {
+			UpdateRequest updateRequest = (UpdateRequest)request;
+			result = updateRequest.index();
+		} else if(operation.equalsIgnoreCase("search")) {
+			SearchRequest searchRequest = (SearchRequest)request;
+			result = getSearchCollection(searchRequest);
+		} else if(operation.equalsIgnoreCase("index")) {
+			IndexRequest indexRequest = (IndexRequest)request;
+			result = indexRequest.index();
+		} else if(operation.equalsIgnoreCase("delete")) {
+			DeleteRequest deleteRequest = (DeleteRequest)request;
+			result = deleteRequest.index();
+		} else if(operation.equalsIgnoreCase("multiget")) {
+			MultiGetRequest multiGetRequest = (MultiGetRequest)request;
+			List<Item> items = multiGetRequest.getItems();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<items.size();i++) {
+				Item item = items.get(i);
+				sb.append(item.index()+"-"+item.index());
+				if(i < items.size()-1) {
+					sb.append(',');
+				}
+			}
+			result = sb.toString();
+			if(result.endsWith(",")) {
+				result = result.substring(0, result.length()-1);
+			}
+		} else if(operation.equalsIgnoreCase("multisearch")) {
+			MultiSearchRequest multiSearchRequest = (MultiSearchRequest)request;
+			List<SearchRequest> searches = multiSearchRequest.requests(); 
+			int size = searches.size();
+			StringBuffer sb = new StringBuffer();
+			for(int i=0;i<size;i++) {
+				SearchRequest searchRequest = searches.get(i);
+				String s = getSearchCollection(searchRequest);
+				sb.append(s);
+				if(i < size -1) {
+					sb.append(',');
+				}
+			}
+			result = sb.toString();
+		}
+		return result;
+	}
+
+	private static String getSearchCollection(SearchRequest request) {
+		String[] indices = request.indices();
+		StringBuffer sb = new StringBuffer();
+		for(int i=0;i<indices.length;i++) {
+			sb.append(indices[i]);
+			if(i < indices.length-1) {
+				sb.append(',');
+			}
+		}
+		if(sb.length() != 0) {
+			sb.append('-');
+		}
+		return sb.toString();
+	}
+
+	public static class Holder {
+		public String rawQuery;
+		public String collection;
+		public String operation;
+
+		public Holder(String rawQuery, String collection, String operation) {
+			super();
+			this.rawQuery = rawQuery;
+			this.collection = collection;
+			this.operation = operation;
+		}
+
+
+	}
+}

elasticsearch-highlevel-rest/src/main/java/com/newrelic/instrumentation/labs/elasticsearch/highlevelrest/NRActionListenerWrapper.java

@@ -0,0 +1,52 @@
+package com.newrelic.instrumentation.labs.elasticsearch.highlevelrest;
+
+import org.elasticsearch.action.ActionListener;
+
+import com.newrelic.api.agent.DatastoreParameters;
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Segment;
+import com.newrelic.instrumentation.labs.elasticsearch.highlevelrest.ElasticSearchUtils.Holder;
+
+public class NRActionListenerWrapper<Response> implements ActionListener<Response> {
+
+	private ActionListener<Response> delegate = null;
+	private Segment segment = null;
+	private Holder holder = null;
+
+	public NRActionListenerWrapper(ActionListener<Response> d, Holder h) {
+		delegate = d;
+		holder = h;
+		initialize();
+	}
+
+	private void initialize() {
+		DatastoreParameters params = DatastoreParameters.product("ElasticSearch").collection(holder.collection).operation(holder.collection).build();
+		segment = NewRelic.getAgent().getTransaction().startSegment("HighLevelRestClient/"+holder.collection);
+		segment.reportAsExternal(params);
+	}
+
+	@Override
+	public void onResponse(Response response) {
+		if(segment != null) {
+			segment.end();
+			segment = null;
+		}
+		if(delegate != null) {
+			delegate.onResponse(response);
+		}
+	}
+
+	@Override
+	public void onFailure(Exception e) {
+		NewRelic.noticeError(e);
+		if(segment != null) {
+			segment.end();
+			segment = null;
+		}
+		if(delegate != null) {
+			delegate.onFailure(e);
+		}
+
+	}
+
+}

elasticsearch-highlevel-rest/src/main/java/com/newrelic/instrumentation/labs/elasticsearch/highlevelrest/NRHolder.java

@@ -0,0 +1,42 @@
+package com.newrelic.instrumentation.labs.elasticsearch.highlevelrest;
+
+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;
+	}
+
+}