From 1773cb2e2fb3b3d2c9131011c513288a4faacbe2 Mon Sep 17 00:00:00 2001 From: Mateusz Rzeszutek Date: Fri, 28 Jul 2023 09:26:14 +0200 Subject: [PATCH] Fix elasticsearch latestDepTest (#9066) --- ...csearchApiClientInstrumentationModule.java | 5 +- .../elasticsearch/apiclient/EndpointId.java | 28 +++++++ .../RestClientHttpClientInstrumentation.java | 79 +++++++++++++++++++ ...> RestClientTransportInstrumentation.java} | 3 +- 4 files changed, 112 insertions(+), 3 deletions(-) create mode 100644 instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/EndpointId.java create mode 100644 instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientHttpClientInstrumentation.java rename instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/{ApiClientInstrumentation.java => RestClientTransportInstrumentation.java} (95%) diff --git a/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java index b8e2cb866b..31ab3e3b84 100644 --- a/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ElasticsearchApiClientInstrumentationModule.java +++ b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/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 class ElasticsearchApiClientInstrumentationModule extends Instrumentation @Override public List typeInstrumentations() { - return singletonList(new ApiClientInstrumentation()); + return asList( + new RestClientTransportInstrumentation(), new RestClientHttpClientInstrumentation()); } } diff --git a/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/EndpointId.java b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/EndpointId.java new file mode 100644 index 0000000000..96f1468067 --- /dev/null +++ b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/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 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/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientHttpClientInstrumentation.java b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientHttpClientInstrumentation.java new file mode 100644 index 0000000000..b219dc7550 --- /dev/null +++ b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/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 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/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ApiClientInstrumentation.java b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientTransportInstrumentation.java similarity index 95% rename from instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ApiClientInstrumentation.java rename to instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientTransportInstrumentation.java index 9388a92d03..abd731fb16 100644 --- a/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/ApiClientInstrumentation.java +++ b/instrumentation/elasticsearch/elasticsearch-api-client-7.16/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/elasticsearch/apiclient/RestClientTransportInstrumentation.java @@ -20,7 +20,8 @@ import net.bytebuddy.description.type.TypeDescription; 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 typeMatcher() {