Fix elasticsearch latestDepTest (#9066)

open-telemetry · Jul 28, 2023 · 1773cb2 · 1773cb2
1 parent dc5d76a
commit 1773cb2
Show file tree

Hide file tree

Showing 4 changed files with 112 additions and 3 deletions.
diff --git a/.../instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java b/.../instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java
@@ -5,7 +5,7 @@
 
 package io.opentelemetry.javaagent.instrumentation.elasticsearch.apiclient;
 
-import static java.util.Collections.singletonList;
+import static java.util.Arrays.asList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
@@ -20,6 +20,7 @@ public ElasticsearchApiClientInstrumentationModule() {
 
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
-    return singletonList(new ApiClientInstrumentation());
+    return asList(
+        new RestClientTransportInstrumentation(), new RestClientHttpClientInstrumentation());
   }
 }
diff --git a/...n/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/EndpointId.java b/...n/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/EndpointId.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.elasticsearch.apiclient;
+
+import static io.opentelemetry.context.ContextKey.named;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.ContextKey;
+import javax.annotation.Nullable;
+
+public final class EndpointId {
+
+  private static final ContextKey<String> KEY = named("elasticsearch-api-client-endpoint-id");
+
+  public static Context storeInContext(Context context, String endpointId) {
+    return context.with(KEY, endpointId);
+  }
+
+  @Nullable
+  public static String get(Context context) {
+    return context.get(KEY);
+  }
+
+  private EndpointId() {}
+}
diff --git a/...avaagent/instrumentation/elasticsearch/apiclient/RestClientHttpClientInstrumentation.java b/...avaagent/instrumentation/elasticsearch/apiclient/RestClientHttpClientInstrumentation.java
@@ -0,0 +1,79 @@
+/*
+ * 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.namedOneOf;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+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;
+
+// starting from 8.9
+public class RestClientHttpClientInstrumentation implements TypeInstrumentation {
+
+  @Override
+  public ElementMatcher<TypeDescription> typeMatcher() {
+    return named("co.elastic.clients.transport.rest_client.RestClientHttpClient");
+  }
+
+  @Override
+  public void transform(TypeTransformer transformer) {
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(namedOneOf("performRequest", "performRequestAsync"))
+            .and(takesArgument(0, String.class)),
+        this.getClass().getName() + "$PerformRequestAdvice");
+    transformer.applyAdviceToMethod(
+        isMethod()
+            .and(named("createRestRequest"))
+            .and(returns(named("org.elasticsearch.client.Request"))),
+        this.getClass().getName() + "$CreateRestRequestAdvice");
+  }
+
+  @SuppressWarnings("unused")
+  public static class PerformRequestAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static Scope onEnter(@Advice.Argument(0) String endpointId) {
+      return EndpointId.storeInContext(Context.current(), endpointId).makeCurrent();
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(@Advice.Enter Scope scope) {
+      if (scope != null) {
+        scope.close();
+      }
+    }
+  }
+
+  @SuppressWarnings("unused")
+  public static class CreateRestRequestAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void onExit(@Advice.Return Request request) {
+      String endpointId = EndpointId.get(Context.current());
+      if (endpointId == null) {
+        return;
+      }
+      if (endpointId.startsWith("es/") && endpointId.length() > 3) {
+        endpointId = endpointId.substring(3);
+      }
+      VirtualField.find(Request.class, ElasticsearchEndpointDefinition.class)
+          .set(request, ElasticsearchEndpointMap.get(endpointId));
+    }
+  }
+}
diff --git a/...h/apiclient/ApiClientInstrumentation.java → ...t/RestClientTransportInstrumentation.java b/...h/apiclient/ApiClientInstrumentation.java → ...t/RestClientTransportInstrumentation.java
@@ -20,7 +20,8 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import org.elasticsearch.client.Request;
 
-public class ApiClientInstrumentation implements TypeInstrumentation {
+// up to 8.8 (included)
+public class RestClientTransportInstrumentation implements TypeInstrumentation {
 
   @Override
   public ElementMatcher<TypeDescription> typeMatcher() {