Instrumentation for Elasticsearch 8+

instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent-unit-tests/build.gradle.kts

@@ -0,0 +1,8 @@
+plugins {
+  id("otel.java-conventions")
+}
+
+dependencies {
+  testImplementation(project(":instrumentation:elasticsearch:elasticsearch-rest-common:javaagent"))
+  testImplementation(project(":instrumentation:elasticsearch:elasticsearch-api-client-7.16:javaagent"))
+}

instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent-unit-tests/src/test/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/rest/ElasticsearchEndpointMapTest.java

@@ -0,0 +1,123 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.elasticsearch.rest;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.opentelemetry.javaagent.instrumentation.elasticsearch.apiclient.ElasticsearchEndpointMap;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import org.junit.jupiter.api.Test;
+
+public class ElasticsearchEndpointMapTest {
+
+  private static final Set<String> SEARCH_ENDPOINTS =
+      new HashSet<>(
+          Arrays.asList(
+              "search",
+              "async_search.submit",
+              "msearch",
+              "eql.search",
+              "terms_enum",
+              "search_template",
+              "msearch_template",
+              "render_search_template"));
+
+  private static List<String> getPathParts(String route) {
+    List<String> pathParts = new ArrayList<>();
+    String routeFragment = route;
+    int paramStartIndex = routeFragment.indexOf('{');
+    while (paramStartIndex >= 0) {
+      int paramEndIndex = routeFragment.indexOf('}');
+      if (paramEndIndex < 0 || paramEndIndex <= paramStartIndex + 1) {
+        throw new IllegalStateException("Invalid route syntax!");
+      }
+      pathParts.add(routeFragment.substring(paramStartIndex + 1, paramEndIndex));
+
+      int nextIdx = paramEndIndex + 1;
+      if (nextIdx >= routeFragment.length()) {
+        break;
+      }
+
+      routeFragment = routeFragment.substring(nextIdx);
+      paramStartIndex = routeFragment.indexOf('{');
+    }
+    return pathParts;
+  }
+
+  @Test
+  public void testIsSearchEndpoint() {
+    for (ElasticsearchEndpointDefinition esEndpointDefinition :
+        ElasticsearchEndpointMap.getAllEndpoints()) {
+      String endpointId = esEndpointDefinition.getEndpointName();
+      assertEquals(SEARCH_ENDPOINTS.contains(endpointId), esEndpointDefinition.isSearchEndpoint());
+    }
+  }
+
+  @Test
+  public void testProcessPathParts() {
+    for (ElasticsearchEndpointDefinition esEndpointDefinition :
+        ElasticsearchEndpointMap.getAllEndpoints()) {
+      for (String route :
+          esEndpointDefinition.getRoutes().stream()
+              .map(ElasticsearchEndpointDefinition.Route::getName)
+              .collect(Collectors.toList())) {
+        List<String> pathParts = getPathParts(route);
+        String resolvedRoute = route.replace("{", "").replace("}", "");
+        Map<String, String> observedParams = new HashMap<>();
+        esEndpointDefinition.processPathParts(resolvedRoute, (k, v) -> observedParams.put(k, v));
+
+        Map<String, String> expectedMap = new HashMap<>();
+        pathParts.forEach(part -> expectedMap.put(part, part));
+
+        assertEquals(expectedMap, observedParams);
+      }
+    }
+  }
+
+  @Test
+  public void testSearchEndpoint() {
+    ElasticsearchEndpointDefinition esEndpoint = ElasticsearchEndpointMap.get("search");
+    Map<String, String> observedParams = new HashMap<>();
+    esEndpoint.processPathParts(
+        "/test-index-1,test-index-2/_search", (k, v) -> observedParams.put(k, v));
+
+    assertEquals("test-index-1,test-index-2", observedParams.get("index"));
+  }
+
+  @Test
+  public void testBuildRegexPattern() {
+    Pattern pattern =
+        ElasticsearchEndpointDefinition.EndpointPattern.buildRegexPattern(
+            "/_nodes/{node_id}/shutdown");
+    assertEquals("^/_nodes/(?<node0id>[^/]+)/shutdown$", pattern.pattern());
+
+    pattern =
+        ElasticsearchEndpointDefinition.EndpointPattern.buildRegexPattern(
+            "/_snapshot/{repository}/{snapshot}/_mount");
+    assertEquals("^/_snapshot/(?<repository>[^/]+)/(?<snapshot>[^/]+)/_mount$", pattern.pattern());
+
+    pattern =
+        ElasticsearchEndpointDefinition.EndpointPattern.buildRegexPattern(
+            "/_security/profile/_suggest");
+    assertEquals("^/_security/profile/_suggest$", pattern.pattern());
+
+    pattern =
+        ElasticsearchEndpointDefinition.EndpointPattern.buildRegexPattern(
+            "/_application/search_application/{name}");
+    assertEquals("^/_application/search_application/(?<name>[^/]+)$", pattern.pattern());
+
+    pattern = ElasticsearchEndpointDefinition.EndpointPattern.buildRegexPattern("/");
+    assertEquals("^/$", pattern.pattern());
+  }
+}

instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/build.gradle.kts

@@ -0,0 +1,31 @@
+plugins {
+  id("otel.javaagent-instrumentation")
+}
+
+muzzle {
+  pass {
+    group.set("co.elastic.clients")
+    module.set("elasticsearch-java")
+    versions.set("[7.16,)")
+    assertInverse.set(true)
+  }
+}
+
+dependencies {
+  library("co.elastic.clients:elasticsearch-java:7.16.0")
+
+  implementation(project(":instrumentation:elasticsearch:elasticsearch-rest-common:javaagent"))
+
+  testInstrumentation(project(":instrumentation:elasticsearch:elasticsearch-rest-7.0:javaagent"))
+  testInstrumentation(project(":instrumentation:apache-httpclient:apache-httpclient-4.0:javaagent"))
+  testInstrumentation(project(":instrumentation:apache-httpasyncclient-4.1:javaagent"))
+
+  testImplementation("com.fasterxml.jackson.core:jackson-databind:2.14.2")
+  testImplementation("org.testcontainers:elasticsearch")
+}
+
+tasks {
+  test {
+    usesService(gradle.sharedServices.registrations["testcontainersBuildService"].service)
+  }
+}

instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ApiClientInstrumentation.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.elasticsearch.apiclient;
+
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import co.elastic.clients.transport.Endpoint;
+import io.opentelemetry.instrumentation.api.util.VirtualField;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
+import io.opentelemetry.javaagent.instrumentation.elasticsearch.rest.ElasticsearchEndpointDefinition;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.elasticsearch.client.Request;
+
+public class ApiClientInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("co.elastic.clients.transport.rest_client.RestClientTransport");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("prepareLowLevelRequest"))
+            .and(takesArgument(1, named("co.elastic.clients.transport.Endpoint")))
+            .and(returns(named("org.elasticsearch.client.Request"))),
+        this.getClass().getName() + "$RestClientTransportAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class RestClientTransportAdvice {
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onPrepareLowLevelRequest(
+        @Advice.Argument(1) Endpoint<?, ?, ?> endpoint, @Advice.Return Request request) {
+      VirtualField<Request, ElasticsearchEndpointDefinition> virtualField =
+          VirtualField.find(Request.class, ElasticsearchEndpointDefinition.class);
+      String endpointId = endpoint.id();
+      if (endpointId.startsWith("es/") && endpointId.length() > 3) {
+        endpointId = endpointId.substring(3);
+      }
+      virtualField.set(request, ElasticsearchEndpointMap.get(endpointId));
+    }
+  }
+}

instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java

@@ -0,0 +1,25 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.elasticsearch.apiclient;
+
+import static java.util.Collections.singletonList;
+
+import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import java.util.List;
+
+@AutoService(InstrumentationModule.class)
+public class ElasticsearchApiClientInstrumentationModule extends InstrumentationModule {
+  public ElasticsearchApiClientInstrumentationModule() {
+    super("elasticsearch-api-client-7.16", "elasticsearch");
+  }
+
+  @Override
+  public List<TypeInstrumentation> typeInstrumentations() {
+    return singletonList(new ApiClientInstrumentation());
+  }
+}