From 06874a5c755f6f81c068a858b186d61e47e14638 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Thu, 2 Feb 2023 21:16:31 +0100 Subject: [PATCH 1/8] draft implementation --- .../apm/agent/concurrent/JavaConcurrent.java | 32 +++++++++++++++++++ .../agent/concurrent/ScopeManagementTest.java | 17 ++++++++++ 2 files changed, 49 insertions(+) diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java index fa083d0a20..c84bffff37 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java @@ -35,7 +35,9 @@ import java.util.List; import java.util.Set; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ForkJoinTask; +import java.util.function.BiConsumer; // Not strictly necessary as AbstractJavaConcurrentInstrumentation returns an empty collection for pluginClassLoaderRootPackages // but this signals the intent that this class must not be loaded from the IndyBootstrapClassLoader so that the state in this class applies globally @@ -163,6 +165,36 @@ public static ForkJoinTask withContext(@Nullable ForkJoinTask task, Tr return task; } + @Nullable + public static CompletableFuture withContext(CompletableFuture completableFuture, Tracer tracer){ + if (shouldAvoidContextPropagation(completableFuture)) { + return completableFuture; + } + needsContext.set(Boolean.FALSE); + AbstractSpan active = tracer.getActive(); + if (active == null) { + return completableFuture; + } + + CompletableFuture wrapped = new CompletableFuture(); + wrapped.whenComplete(new BiConsumer() { + @Override + public void accept(T value, @Nullable Throwable thrown) { + try { + active.activate(); + if(thrown != null){ + completableFuture.completeExceptionally(thrown); + } else { + completableFuture.complete(value); + } + } finally { + active.deactivate(); + } + } + }); + return wrapped; + } + public static void doFinally(@Nullable Throwable thrown, @Nullable Object contextObject) { needsContext.set(Boolean.TRUE); if (thrown != null && contextObject != null) { diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java index ac78d46be2..d954e787ac 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java @@ -26,6 +26,8 @@ import org.junit.jupiter.api.Test; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -142,4 +144,19 @@ void testSpanAndContextCallableActivationInDifferentThread() throws Exception { assertThat(tracer.getActive()).isNull(); } + + @Test + void testCompletableFutureWrapping() throws ExecutionException, InterruptedException { + final Transaction transaction = tracer.startRootTransaction(null).activate(); + + CompletableFuture completableFuture = CompletableFuture.runAsync(() -> { + assertThat(tracer.currentTransaction()).isNotNull(); + }, Executors.newSingleThreadExecutor()); + + CompletableFuture wrapped = JavaConcurrent.withContext(completableFuture, tracer); + + wrapped.get(); + + transaction.deactivate().end(); + } } From 8726e0d10c031c280c7fd66aa88736cb1acbb094 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Thu, 2 Feb 2023 21:19:13 +0100 Subject: [PATCH 2/8] fix parent pom version --- apm-agent-java-17/apm-agent-spring-6/pom.xml | 2 +- apm-agent-java-17/pom.xml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/apm-agent-java-17/apm-agent-spring-6/pom.xml b/apm-agent-java-17/apm-agent-spring-6/pom.xml index b88f261b27..b45de09af3 100755 --- a/apm-agent-java-17/apm-agent-spring-6/pom.xml +++ b/apm-agent-java-17/apm-agent-spring-6/pom.xml @@ -5,7 +5,7 @@ co.elastic.apm apm-agent-java-17 - 1.36.0-SNAPSHOT + 1.36.1-SNAPSHOT apm-agent-spring-6 diff --git a/apm-agent-java-17/pom.xml b/apm-agent-java-17/pom.xml index ae8f968c8a..fac25d2d83 100644 --- a/apm-agent-java-17/pom.xml +++ b/apm-agent-java-17/pom.xml @@ -5,7 +5,7 @@ apm-agent-parent co.elastic.apm - 1.36.0-SNAPSHOT + 1.36.1-SNAPSHOT apm-agent-java-17 From 2a017cd6c914e9dd7b9b886669b1eb7b33123cf4 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Wed, 8 Feb 2023 12:10:39 +0100 Subject: [PATCH 3/8] wip-broken --- ...asticsearchClientAsyncInstrumentation.java | 2 +- ...asticsearchClientAsyncInstrumentation.java | 2 +- .../apm-es-restclient-plugin-7_x/pom.xml | 2 - .../v7_x/RestClientInstrumentation.java | 79 ++++++++++++++++ .../agent/esrestclient/v7_x/package-info.java | 22 +++++ ...ic.apm.agent.sdk.ElasticApmInstrumentation | 1 + ...searchRestClientInstrumentationHelper.java | 10 ++- .../esrestclient/ResponseListenerWrapper.java | 90 +++++++++++++++++-- .../apm/agent/concurrent/JavaConcurrent.java | 31 +------ .../agent/concurrent/ScopeManagementTest.java | 16 ---- apm-agent/pom.xml | 5 ++ 11 files changed, 199 insertions(+), 61 deletions(-) create mode 100644 apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java create mode 100644 apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/package-info.java create mode 100644 apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-5_6/src/main/java/co/elastic/apm/agent/esrestclient/v5_6/ElasticsearchClientAsyncInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-5_6/src/main/java/co/elastic/apm/agent/esrestclient/v5_6/ElasticsearchClientAsyncInstrumentation.java index 892cc37197..10b9b54fa5 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-5_6/src/main/java/co/elastic/apm/agent/esrestclient/v5_6/ElasticsearchClientAsyncInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-5_6/src/main/java/co/elastic/apm/agent/esrestclient/v5_6/ElasticsearchClientAsyncInstrumentation.java @@ -79,7 +79,7 @@ public static Object[] onBeforeExecute(@Advice.Argument(0) String method, if (span != null) { Object[] ret = new Object[2]; ret[0] = span; - ret[1] = helper.wrapResponseListener(responseListener, span); + ret[1] = helper.wrapClientResponseListener(responseListener, span); return ret; } return null; diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java index fcc2b093b3..1e00191c11 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java @@ -69,7 +69,7 @@ public static Object[] onBeforeExecute(@Advice.Argument(0) Request request, if (span != null) { Object[] ret = new Object[2]; ret[0] = span; - ret[1] = helper.wrapResponseListener(responseListener, span); + ret[1] = helper.wrapClientResponseListener(responseListener, span); return ret; } return null; diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml index 2daaff06fd..aa3105c6fc 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml @@ -27,13 +27,11 @@ org.elasticsearch.client elasticsearch-rest-client ${version.elasticsearch} - test org.elasticsearch.client elasticsearch-rest-high-level-client ${version.elasticsearch} - test org.apache.logging.log4j diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java new file mode 100644 index 0000000000..68ee3406ef --- /dev/null +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java @@ -0,0 +1,79 @@ +/* + * Licensed to Elasticsearch B.V. under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch B.V. licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package co.elastic.apm.agent.esrestclient.v7_x; + +import co.elastic.apm.agent.esrestclient.ElasticsearchRestClientInstrumentation; +import co.elastic.apm.agent.esrestclient.ElasticsearchRestClientInstrumentationHelper; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument; +import net.bytebuddy.description.method.MethodDescription; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +import javax.annotation.Nullable; + +import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.returns; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import static net.bytebuddy.matcher.ElementMatchers.takesArguments; + +public class RestClientInstrumentation extends ElasticsearchRestClientInstrumentation { + + @Override + public ElementMatcher getTypeMatcher() { + return named("org.elasticsearch.client.RestClient"); + } + + @Override + public ElementMatcher getMethodMatcher() { + return named("performRequestAsync").and(isPublic()) + .and(takesArguments(2)) + .and(takesArgument(0, named("org.elasticsearch.client.Request"))) + .and(takesArgument(1, named("org.elasticsearch.client.ResponseListener"))) + .and(returns(named("org.elasticsearch.client.Cancellable"))); + } + + @Override + public String getAdviceClassName() { + return getClass().getName() + "$ElasticsearchRestClientAsyncAdvice"; + } + + public static class ElasticsearchRestClientAsyncAdvice { + +// private static final ElasticsearchRestClientInstrumentationHelper helper = ElasticsearchRestClientInstrumentationHelper.get(); + + @Nullable + @Advice.AssignReturned.ToArguments(@ToArgument(value = 1, typing = DYNAMIC)) + @Advice.OnMethodEnter(suppress = Throwable.class, inline = false) + public static Object onEnter(@Advice.Argument(1) @Nullable Object listenerArg) { + +// AbstractSpan activeContext = tracer.getActive(); +// +// ResponseListener listener = (ResponseListener) listenerArg; +// if (listener != null && activeContext != null) { +// listener = helper.wrapContextPropagationContextListener(listener, activeContext); +// } + + return listenerArg; + } + + } +} diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/package-info.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/package-info.java new file mode 100644 index 0000000000..83e80ff931 --- /dev/null +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/package-info.java @@ -0,0 +1,22 @@ +/* + * Licensed to Elasticsearch B.V. under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch B.V. licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +@NonnullApi +package co.elastic.apm.agent.esrestclient.v7_x; + +import co.elastic.apm.agent.sdk.NonnullApi; diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation new file mode 100644 index 0000000000..c572b1c253 --- /dev/null +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation @@ -0,0 +1 @@ +#co.elastic.apm.agent.esrestclient.v7_x.RestClientInstrumentation diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java index 7e093f333c..b071200946 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java @@ -74,7 +74,7 @@ private ElasticsearchRestClientInstrumentationHelper(ElasticApmTracer tracer) { private class ResponseListenerAllocator implements Allocator { @Override public ResponseListenerWrapper createInstance() { - return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this); + return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this, ElasticsearchRestClientInstrumentationHelper.this.tracer); } } @@ -162,8 +162,12 @@ public void finishClientSpan(@Nullable Response response, Span span, @Nullable T } } - public ResponseListener wrapResponseListener(ResponseListener listener, Span span) { - return responseListenerObjectPool.createInstance().with(listener, span); + public ResponseListener wrapClientResponseListener(ResponseListener listener, Span span) { + return responseListenerObjectPool.createInstance().withClientSpan(listener, span); + } + + public ResponseListener wrapContextPropagationContextListener(ResponseListener listener, AbstractSpan activeContext) { + return responseListenerObjectPool.createInstance().withContextPropagation(listener, activeContext); } void recycle(ResponseListenerWrapper listenerWrapper) { diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java index 8bc0636323..de31f44fb3 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java @@ -18,6 +18,8 @@ */ package co.elastic.apm.agent.esrestclient; +import co.elastic.apm.agent.impl.Tracer; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; import co.elastic.apm.agent.impl.transaction.Span; import co.elastic.apm.agent.objectpool.Recyclable; import org.elasticsearch.client.Response; @@ -28,24 +30,69 @@ public class ResponseListenerWrapper implements ResponseListener, Recyclable { private final ElasticsearchRestClientInstrumentationHelper helper; + private final Tracer tracer; + @Nullable private ResponseListener delegate; + + /** + * When {@literal true}, the context object is a client span that needs to be ended on completion, when {@literal false} + * it means that only context-propagation (aka activation/deactivation) is required. + */ + private boolean isClientSpan; + @Nullable - private volatile Span span; + private volatile AbstractSpan context; - ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper) { + ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper, Tracer tracer) { this.helper = helper; + this.tracer = tracer; } - ResponseListenerWrapper with(ResponseListener delegate, Span span) { + ResponseListenerWrapper withClientSpan(ResponseListener delegate, Span span) { // Order is important due to visibility - write to span last on this (initiating) thread this.delegate = delegate; - this.span = span; + this.isClientSpan = true; + this.context = span; + return this; + } + + ResponseListenerWrapper withContextPropagation(ResponseListener delegate, AbstractSpan context) { + // Order is important due to visibility - write to span last on this (initiating) thread + this.delegate = delegate; + this.isClientSpan = false; + this.context = context; return this; } @Override public void onSuccess(Response response) { + if (isClientSpan) { + onSuccessClient(response); + } else { + onSuccessContextPropagation(response); + } + } + + private void onSuccessContextPropagation(Response response) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onSuccess(response); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + + private void onSuccessClient(Response response){ try { finishClientSpan(response, null); } finally { @@ -58,6 +105,14 @@ public void onSuccess(Response response) { @Override public void onFailure(Exception exception) { + if (isClientSpan) { + onFailureClient(exception); + } else { + onFailureContextPropagation(exception); + } + } + + private void onFailureClient(Exception exception) { try { finishClientSpan(null, exception); } finally { @@ -68,17 +123,36 @@ public void onFailure(Exception exception) { } } + private void onFailureContextPropagation(Exception exception) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onFailure(exception); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + private void finishClientSpan(@Nullable Response response, @Nullable Throwable throwable) { // First read volatile span to ensure visibility on executing thread - Span localSpan = span; - if (localSpan != null) { - helper.finishClientSpan(response, localSpan, throwable); + AbstractSpan localSpan = context; + if (localSpan instanceof Span) { + helper.finishClientSpan(response, (Span) localSpan, throwable); } } @Override public void resetState() { delegate = null; - span = null; + context = null; + isClientSpan = false; } } diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java index c84bffff37..9a6b7db48e 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java @@ -165,36 +165,6 @@ public static ForkJoinTask withContext(@Nullable ForkJoinTask task, Tr return task; } - @Nullable - public static CompletableFuture withContext(CompletableFuture completableFuture, Tracer tracer){ - if (shouldAvoidContextPropagation(completableFuture)) { - return completableFuture; - } - needsContext.set(Boolean.FALSE); - AbstractSpan active = tracer.getActive(); - if (active == null) { - return completableFuture; - } - - CompletableFuture wrapped = new CompletableFuture(); - wrapped.whenComplete(new BiConsumer() { - @Override - public void accept(T value, @Nullable Throwable thrown) { - try { - active.activate(); - if(thrown != null){ - completableFuture.completeExceptionally(thrown); - } else { - completableFuture.complete(value); - } - } finally { - active.deactivate(); - } - } - }); - return wrapped; - } - public static void doFinally(@Nullable Throwable thrown, @Nullable Object contextObject) { needsContext.set(Boolean.TRUE); if (thrown != null && contextObject != null) { @@ -285,4 +255,5 @@ public V call() throws Exception { return delegate.call(); } } + } diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java index d954e787ac..813c0a7c01 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java @@ -26,8 +26,6 @@ import org.junit.jupiter.api.Test; import java.util.concurrent.Callable; -import java.util.concurrent.CompletableFuture; -import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -145,18 +143,4 @@ void testSpanAndContextCallableActivationInDifferentThread() throws Exception { assertThat(tracer.getActive()).isNull(); } - @Test - void testCompletableFutureWrapping() throws ExecutionException, InterruptedException { - final Transaction transaction = tracer.startRootTransaction(null).activate(); - - CompletableFuture completableFuture = CompletableFuture.runAsync(() -> { - assertThat(tracer.currentTransaction()).isNotNull(); - }, Executors.newSingleThreadExecutor()); - - CompletableFuture wrapped = JavaConcurrent.withContext(completableFuture, tracer); - - wrapped.get(); - - transaction.deactivate().end(); - } } diff --git a/apm-agent/pom.xml b/apm-agent/pom.xml index c5378b2dba..5ddf25ed80 100644 --- a/apm-agent/pom.xml +++ b/apm-agent/pom.xml @@ -80,6 +80,11 @@ apm-es-restclient-plugin-6_4 ${project.version} + + ${project.groupId} + apm-es-restclient-plugin-7_x + ${project.version} + ${project.groupId} apm-grails-plugin From 0bc2c3fa871d1a0d3ba9a81ba9aa3ce20a4de040 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Wed, 8 Feb 2023 14:34:09 +0100 Subject: [PATCH 4/8] still broken --- ...asticsearchClientAsyncInstrumentation.java | 9 +- .../v7_x/RestClientInstrumentation.java | 5 +- ...ic.apm.agent.sdk.ElasticApmInstrumentation | 1 - ...searchRestClientInstrumentationHelper.java | 9 +- .../esrestclient/ResponseListenerWrapper.java | 88 +--------- .../ResponseListenerWrapper2.java | 159 ++++++++++++++++++ 6 files changed, 179 insertions(+), 92 deletions(-) create mode 100644 apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java index 1e00191c11..ecefa86231 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/esrestclient/v6_4/ElasticsearchClientAsyncInstrumentation.java @@ -38,10 +38,6 @@ public class ElasticsearchClientAsyncInstrumentation extends ElasticsearchRestClientInstrumentation { - @Override - public String getAdviceClassName() { - return getClass().getName() + "$ElasticsearchRestClientAsyncAdvice"; - } @Override public ElementMatcher getTypeMatcher() { @@ -56,6 +52,11 @@ public ElementMatcher getMethodMatcher() { .and(takesArgument(1, named("org.elasticsearch.client.ResponseListener")))); } + @Override + public String getAdviceClassName() { + return getClass().getName() + "$ElasticsearchRestClientAsyncAdvice"; + } + public static class ElasticsearchRestClientAsyncAdvice { private static final ElasticsearchRestClientInstrumentationHelper helper = ElasticsearchRestClientInstrumentationHelper.get(); diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java index 68ee3406ef..f5f2e75621 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java @@ -35,6 +35,7 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +// co.elastic.apm.agent.esrestclient.v7_x.RestClientInstrumentation public class RestClientInstrumentation extends ElasticsearchRestClientInstrumentation { @Override @@ -53,10 +54,10 @@ public ElementMatcher getMethodMatcher() { @Override public String getAdviceClassName() { - return getClass().getName() + "$ElasticsearchRestClientAsyncAdvice"; + return getClass().getName() + "$RestClientAsyncAdvice"; } - public static class ElasticsearchRestClientAsyncAdvice { + public static class RestClientAsyncAdvice { // private static final ElasticsearchRestClientInstrumentationHelper helper = ElasticsearchRestClientInstrumentationHelper.get(); diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation index c572b1c253..e69de29bb2 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation @@ -1 +0,0 @@ -#co.elastic.apm.agent.esrestclient.v7_x.RestClientInstrumentation diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java index b071200946..969d4d4548 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java @@ -74,7 +74,8 @@ private ElasticsearchRestClientInstrumentationHelper(ElasticApmTracer tracer) { private class ResponseListenerAllocator implements Allocator { @Override public ResponseListenerWrapper createInstance() { - return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this, ElasticsearchRestClientInstrumentationHelper.this.tracer); + return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this); +// return new ResponseListenerWrapper2(ElasticsearchRestClientInstrumentationHelper.this, ElasticsearchRestClientInstrumentationHelper.this.tracer); } } @@ -166,9 +167,9 @@ public ResponseListener wrapClientResponseListener(ResponseListener listener, Sp return responseListenerObjectPool.createInstance().withClientSpan(listener, span); } - public ResponseListener wrapContextPropagationContextListener(ResponseListener listener, AbstractSpan activeContext) { - return responseListenerObjectPool.createInstance().withContextPropagation(listener, activeContext); - } +// public ResponseListener wrapContextPropagationContextListener(ResponseListener listener, AbstractSpan activeContext) { +// return responseListenerObjectPool.createInstance().withContextPropagation(listener, activeContext); +// } void recycle(ResponseListenerWrapper listenerWrapper) { responseListenerObjectPool.recycle(listenerWrapper); diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java index de31f44fb3..735c578c2f 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java @@ -18,8 +18,6 @@ */ package co.elastic.apm.agent.esrestclient; -import co.elastic.apm.agent.impl.Tracer; -import co.elastic.apm.agent.impl.transaction.AbstractSpan; import co.elastic.apm.agent.impl.transaction.Span; import co.elastic.apm.agent.objectpool.Recyclable; import org.elasticsearch.client.Response; @@ -30,69 +28,24 @@ public class ResponseListenerWrapper implements ResponseListener, Recyclable { private final ElasticsearchRestClientInstrumentationHelper helper; - private final Tracer tracer; - @Nullable private ResponseListener delegate; - - /** - * When {@literal true}, the context object is a client span that needs to be ended on completion, when {@literal false} - * it means that only context-propagation (aka activation/deactivation) is required. - */ - private boolean isClientSpan; - @Nullable - private volatile AbstractSpan context; + private volatile Span span; - ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper, Tracer tracer) { + ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper) { this.helper = helper; - this.tracer = tracer; } ResponseListenerWrapper withClientSpan(ResponseListener delegate, Span span) { // Order is important due to visibility - write to span last on this (initiating) thread this.delegate = delegate; - this.isClientSpan = true; - this.context = span; - return this; - } - - ResponseListenerWrapper withContextPropagation(ResponseListener delegate, AbstractSpan context) { - // Order is important due to visibility - write to span last on this (initiating) thread - this.delegate = delegate; - this.isClientSpan = false; - this.context = context; + this.span = span; return this; } @Override public void onSuccess(Response response) { - if (isClientSpan) { - onSuccessClient(response); - } else { - onSuccessContextPropagation(response); - } - } - - private void onSuccessContextPropagation(Response response) { - AbstractSpan localContext = context; - boolean activate = localContext != null && tracer.getActive() != localContext; - try { - if (activate) { - localContext.activate(); - } - if (delegate != null) { - delegate.onSuccess(response); - } - } finally { - if (activate) { - localContext.deactivate(); - } - helper.recycle(this); - } - } - - private void onSuccessClient(Response response){ try { finishClientSpan(response, null); } finally { @@ -105,14 +58,6 @@ private void onSuccessClient(Response response){ @Override public void onFailure(Exception exception) { - if (isClientSpan) { - onFailureClient(exception); - } else { - onFailureContextPropagation(exception); - } - } - - private void onFailureClient(Exception exception) { try { finishClientSpan(null, exception); } finally { @@ -123,36 +68,17 @@ private void onFailureClient(Exception exception) { } } - private void onFailureContextPropagation(Exception exception) { - AbstractSpan localContext = context; - boolean activate = localContext != null && tracer.getActive() != localContext; - try { - if (activate) { - localContext.activate(); - } - if (delegate != null) { - delegate.onFailure(exception); - } - } finally { - if (activate) { - localContext.deactivate(); - } - helper.recycle(this); - } - } - private void finishClientSpan(@Nullable Response response, @Nullable Throwable throwable) { // First read volatile span to ensure visibility on executing thread - AbstractSpan localSpan = context; - if (localSpan instanceof Span) { - helper.finishClientSpan(response, (Span) localSpan, throwable); + Span localSpan = span; + if (localSpan != null) { + helper.finishClientSpan(response, localSpan, throwable); } } @Override public void resetState() { delegate = null; - context = null; - isClientSpan = false; + span = null; } } diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java new file mode 100644 index 0000000000..d30343f3b9 --- /dev/null +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java @@ -0,0 +1,159 @@ +/* + * Licensed to Elasticsearch B.V. under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch B.V. licenses this file to you under + * the Apache License, Version 2.0 (the "License"); you may + * not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package co.elastic.apm.agent.esrestclient; + +import co.elastic.apm.agent.impl.Tracer; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; +import co.elastic.apm.agent.impl.transaction.Span; +import co.elastic.apm.agent.objectpool.Recyclable; +import org.elasticsearch.client.Response; +import org.elasticsearch.client.ResponseListener; + +import javax.annotation.Nullable; + +public class ResponseListenerWrapper2 extends ResponseListenerWrapper implements ResponseListener, Recyclable { + + private final ElasticsearchRestClientInstrumentationHelper helper; + private final Tracer tracer; + + @Nullable + private ResponseListener delegate; + + /** + * When {@literal true}, the context object is a client span that needs to be ended on completion, when {@literal false} + * it means that only context-propagation (aka activation/deactivation) is required. + */ + private boolean isClientSpan; + + @Nullable + private volatile AbstractSpan context; + + ResponseListenerWrapper2(ElasticsearchRestClientInstrumentationHelper helper, Tracer tracer) { + super(helper); + this.helper = helper; + this.tracer = tracer; + } + + ResponseListenerWrapper withClientSpan(ResponseListener delegate, Span span) { + // Order is important due to visibility - write to span last on this (initiating) thread + this.delegate = delegate; + this.isClientSpan = true; + this.context = span; + return this; + } + + ResponseListenerWrapper2 withContextPropagation(ResponseListener delegate, AbstractSpan context) { + // Order is important due to visibility - write to span last on this (initiating) thread + this.delegate = delegate; + this.isClientSpan = false; + this.context = context; + return this; + } + + @Override + public void onSuccess(Response response) { + if (isClientSpan) { + onSuccessClient(response); + } else { + onSuccessContextPropagation(response); + } + } + + private void onSuccessContextPropagation(Response response) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onSuccess(response); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + + private void onSuccessClient(Response response){ + try { + finishClientSpan(response, null); + } finally { + if (delegate != null) { + delegate.onSuccess(response); + } + helper.recycle(this); + } + } + + @Override + public void onFailure(Exception exception) { + if (isClientSpan) { + onFailureClient(exception); + } else { + onFailureContextPropagation(exception); + } + } + + private void onFailureClient(Exception exception) { + try { + finishClientSpan(null, exception); + } finally { + if (delegate != null) { + delegate.onFailure(exception); + } + helper.recycle(this); + } + } + + private void onFailureContextPropagation(Exception exception) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onFailure(exception); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + + private void finishClientSpan(@Nullable Response response, @Nullable Throwable throwable) { + // First read volatile span to ensure visibility on executing thread + AbstractSpan localSpan = context; + if (localSpan instanceof Span) { + helper.finishClientSpan(response, (Span) localSpan, throwable); + } + } + + @Override + public void resetState() { + delegate = null; + context = null; + isClientSpan = false; + } +} From 5715397e8748a1caed005c656548fd1e3095feb7 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Wed, 8 Feb 2023 17:28:55 +0100 Subject: [PATCH 5/8] final working version + test --- .../apm-es-restclient-plugin-7_x/pom.xml | 2 + .../v7_x/RestClientInstrumentation.java | 21 +-- ...ic.apm.agent.sdk.ElasticApmInstrumentation | 1 + ...sticsearchRestClientInstrumentationIT.java | 58 ++++++- ...searchRestClientInstrumentationHelper.java | 9 +- .../esrestclient/ResponseListenerWrapper.java | 88 +++++++++- .../ResponseListenerWrapper2.java | 159 ------------------ 7 files changed, 155 insertions(+), 183 deletions(-) delete mode 100644 apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml index aa3105c6fc..13735abbfb 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/pom.xml @@ -27,11 +27,13 @@ org.elasticsearch.client elasticsearch-rest-client ${version.elasticsearch} + provided org.elasticsearch.client elasticsearch-rest-high-level-client ${version.elasticsearch} + provided org.apache.logging.log4j diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java index f5f2e75621..cff8551126 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java @@ -20,11 +20,13 @@ import co.elastic.apm.agent.esrestclient.ElasticsearchRestClientInstrumentation; import co.elastic.apm.agent.esrestclient.ElasticsearchRestClientInstrumentationHelper; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; import net.bytebuddy.asm.Advice; import net.bytebuddy.asm.Advice.AssignReturned.ToArguments.ToArgument; import net.bytebuddy.description.method.MethodDescription; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; +import org.elasticsearch.client.ResponseListener; import javax.annotation.Nullable; @@ -35,7 +37,6 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; -// co.elastic.apm.agent.esrestclient.v7_x.RestClientInstrumentation public class RestClientInstrumentation extends ElasticsearchRestClientInstrumentation { @Override @@ -59,21 +60,21 @@ public String getAdviceClassName() { public static class RestClientAsyncAdvice { -// private static final ElasticsearchRestClientInstrumentationHelper helper = ElasticsearchRestClientInstrumentationHelper.get(); + private static final ElasticsearchRestClientInstrumentationHelper helper = ElasticsearchRestClientInstrumentationHelper.get(); @Nullable @Advice.AssignReturned.ToArguments(@ToArgument(value = 1, typing = DYNAMIC)) @Advice.OnMethodEnter(suppress = Throwable.class, inline = false) - public static Object onEnter(@Advice.Argument(1) @Nullable Object listenerArg) { + public static ResponseListener onEnter(@Advice.Argument(1) @Nullable ResponseListener listenerArg) { -// AbstractSpan activeContext = tracer.getActive(); -// -// ResponseListener listener = (ResponseListener) listenerArg; -// if (listener != null && activeContext != null) { -// listener = helper.wrapContextPropagationContextListener(listener, activeContext); -// } + AbstractSpan activeContext = tracer.getActive(); - return listenerArg; + ResponseListener listener = listenerArg; + if (listener != null && activeContext != null) { + listener = helper.wrapContextPropagationContextListener(listener, activeContext); + } + + return listener; } } diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation index e69de29bb2..3832826a36 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/resources/META-INF/services/co.elastic.apm.agent.sdk.ElasticApmInstrumentation @@ -0,0 +1 @@ +co.elastic.apm.agent.esrestclient.v7_x.RestClientInstrumentation diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java index 6ef2e5cbfe..16dd565e8d 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java @@ -19,8 +19,10 @@ package co.elastic.apm.agent.esrestclient.v7_x; import co.elastic.apm.agent.esrestclient.v6_4.AbstractEs6_4ClientInstrumentationTest; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; import co.elastic.apm.agent.impl.transaction.Outcome; import co.elastic.apm.agent.impl.transaction.Span; +import co.elastic.apm.agent.impl.transaction.Transaction; import org.apache.http.HttpHost; import org.apache.http.auth.AuthScope; import org.apache.http.auth.UsernamePasswordCredentials; @@ -44,7 +46,10 @@ import org.elasticsearch.action.update.UpdateRequest; import org.elasticsearch.action.update.UpdateResponse; import org.elasticsearch.client.Cancellable; +import org.elasticsearch.client.Request; import org.elasticsearch.client.RequestOptions; +import org.elasticsearch.client.Response; +import org.elasticsearch.client.ResponseListener; import org.elasticsearch.client.RestClient; import org.elasticsearch.client.RestClientBuilder; import org.elasticsearch.client.RestHighLevelClient; @@ -62,7 +67,11 @@ import org.junit.runners.Parameterized; import java.io.IOException; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicReference; import static org.assertj.core.api.Assertions.assertThat; import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder; @@ -72,6 +81,8 @@ public class ElasticsearchRestClientInstrumentationIT extends AbstractEs6_4Clien private static final String ELASTICSEARCH_CONTAINER_VERSION = "docker.elastic.co/elasticsearch/elasticsearch:7.11.0"; + private static RestClientBuilder clientBuilder; + public ElasticsearchRestClientInstrumentationIT(boolean async) { this.async = async; } @@ -83,10 +94,10 @@ public static void startElasticsearchContainerAndClient() throws IOException { CredentialsProvider credentialsProvider = new BasicCredentialsProvider(); credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(USER_NAME, PASSWORD)); - RestClientBuilder builder = RestClient.builder(HttpHost.create(container.getHttpHostAddress())) + clientBuilder = RestClient.builder(HttpHost.create(container.getHttpHostAddress())) .setHttpClientConfigCallback(httpClientBuilder -> httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider)); - client = new RestHighLevelClient(builder); + client = new RestHighLevelClient(clientBuilder); client.indices().create(new CreateIndexRequest(INDEX), RequestOptions.DEFAULT); reporter.reset(); @@ -137,6 +148,49 @@ public void onFailure(Exception e) { deleteDocument(); } + @Test + public void testRestClientAsyncContextPropagation() throws InterruptedException, ExecutionException, TimeoutException { + if(!async) { + // test only relevant when testing async stuff + return; + } + // start without an active transaction + tracer.getActive().deactivate().end(); + assertThat(tracer.getActive()).isNull(); + + reporter.reset(); + + AtomicReference> observedActive = new AtomicReference<>(); + CountDownLatch endLatch = new CountDownLatch(1); + + Transaction transaction = startTestRootTransaction("test-root-transaction"); + + RestClient restClient = clientBuilder.build(); + + Request request = new Request("GET", "/"); + + restClient.performRequestAsync(request, new ResponseListener() { + @Override + public void onSuccess(Response response) { + observedActive.set(tracer.getActive()); + endLatch.countDown(); + } + + @Override + public void onFailure(Exception exception) { + observedActive.set(tracer.getActive()); + endLatch.countDown(); + } + }); + + endLatch.await(1, TimeUnit.SECONDS); + + assertThat(observedActive.get()) + .isSameAs(transaction); + + transaction.deactivate().end(); + } + @Override protected void verifyMultiSearchTemplateSpanContent(Span span) { validateDbContextContent(span, "{\"index\":[\"my-index\"],\"types\":[],\"search_type\":\"query_then_fetch\",\"ccs_minimize_roundtrips\":true}\n" + diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java index 969d4d4548..b071200946 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ElasticsearchRestClientInstrumentationHelper.java @@ -74,8 +74,7 @@ private ElasticsearchRestClientInstrumentationHelper(ElasticApmTracer tracer) { private class ResponseListenerAllocator implements Allocator { @Override public ResponseListenerWrapper createInstance() { - return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this); -// return new ResponseListenerWrapper2(ElasticsearchRestClientInstrumentationHelper.this, ElasticsearchRestClientInstrumentationHelper.this.tracer); + return new ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper.this, ElasticsearchRestClientInstrumentationHelper.this.tracer); } } @@ -167,9 +166,9 @@ public ResponseListener wrapClientResponseListener(ResponseListener listener, Sp return responseListenerObjectPool.createInstance().withClientSpan(listener, span); } -// public ResponseListener wrapContextPropagationContextListener(ResponseListener listener, AbstractSpan activeContext) { -// return responseListenerObjectPool.createInstance().withContextPropagation(listener, activeContext); -// } + public ResponseListener wrapContextPropagationContextListener(ResponseListener listener, AbstractSpan activeContext) { + return responseListenerObjectPool.createInstance().withContextPropagation(listener, activeContext); + } void recycle(ResponseListenerWrapper listenerWrapper) { responseListenerObjectPool.recycle(listenerWrapper); diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java index 735c578c2f..de31f44fb3 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper.java @@ -18,6 +18,8 @@ */ package co.elastic.apm.agent.esrestclient; +import co.elastic.apm.agent.impl.Tracer; +import co.elastic.apm.agent.impl.transaction.AbstractSpan; import co.elastic.apm.agent.impl.transaction.Span; import co.elastic.apm.agent.objectpool.Recyclable; import org.elasticsearch.client.Response; @@ -28,24 +30,69 @@ public class ResponseListenerWrapper implements ResponseListener, Recyclable { private final ElasticsearchRestClientInstrumentationHelper helper; + private final Tracer tracer; + @Nullable private ResponseListener delegate; + + /** + * When {@literal true}, the context object is a client span that needs to be ended on completion, when {@literal false} + * it means that only context-propagation (aka activation/deactivation) is required. + */ + private boolean isClientSpan; + @Nullable - private volatile Span span; + private volatile AbstractSpan context; - ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper) { + ResponseListenerWrapper(ElasticsearchRestClientInstrumentationHelper helper, Tracer tracer) { this.helper = helper; + this.tracer = tracer; } ResponseListenerWrapper withClientSpan(ResponseListener delegate, Span span) { // Order is important due to visibility - write to span last on this (initiating) thread this.delegate = delegate; - this.span = span; + this.isClientSpan = true; + this.context = span; + return this; + } + + ResponseListenerWrapper withContextPropagation(ResponseListener delegate, AbstractSpan context) { + // Order is important due to visibility - write to span last on this (initiating) thread + this.delegate = delegate; + this.isClientSpan = false; + this.context = context; return this; } @Override public void onSuccess(Response response) { + if (isClientSpan) { + onSuccessClient(response); + } else { + onSuccessContextPropagation(response); + } + } + + private void onSuccessContextPropagation(Response response) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onSuccess(response); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + + private void onSuccessClient(Response response){ try { finishClientSpan(response, null); } finally { @@ -58,6 +105,14 @@ public void onSuccess(Response response) { @Override public void onFailure(Exception exception) { + if (isClientSpan) { + onFailureClient(exception); + } else { + onFailureContextPropagation(exception); + } + } + + private void onFailureClient(Exception exception) { try { finishClientSpan(null, exception); } finally { @@ -68,17 +123,36 @@ public void onFailure(Exception exception) { } } + private void onFailureContextPropagation(Exception exception) { + AbstractSpan localContext = context; + boolean activate = localContext != null && tracer.getActive() != localContext; + try { + if (activate) { + localContext.activate(); + } + if (delegate != null) { + delegate.onFailure(exception); + } + } finally { + if (activate) { + localContext.deactivate(); + } + helper.recycle(this); + } + } + private void finishClientSpan(@Nullable Response response, @Nullable Throwable throwable) { // First read volatile span to ensure visibility on executing thread - Span localSpan = span; - if (localSpan != null) { - helper.finishClientSpan(response, localSpan, throwable); + AbstractSpan localSpan = context; + if (localSpan instanceof Span) { + helper.finishClientSpan(response, (Span) localSpan, throwable); } } @Override public void resetState() { delegate = null; - span = null; + context = null; + isClientSpan = false; } } diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java deleted file mode 100644 index d30343f3b9..0000000000 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-common/src/main/java/co/elastic/apm/agent/esrestclient/ResponseListenerWrapper2.java +++ /dev/null @@ -1,159 +0,0 @@ -/* - * Licensed to Elasticsearch B.V. under one or more contributor - * license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright - * ownership. Elasticsearch B.V. licenses this file to you under - * the Apache License, Version 2.0 (the "License"); you may - * not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package co.elastic.apm.agent.esrestclient; - -import co.elastic.apm.agent.impl.Tracer; -import co.elastic.apm.agent.impl.transaction.AbstractSpan; -import co.elastic.apm.agent.impl.transaction.Span; -import co.elastic.apm.agent.objectpool.Recyclable; -import org.elasticsearch.client.Response; -import org.elasticsearch.client.ResponseListener; - -import javax.annotation.Nullable; - -public class ResponseListenerWrapper2 extends ResponseListenerWrapper implements ResponseListener, Recyclable { - - private final ElasticsearchRestClientInstrumentationHelper helper; - private final Tracer tracer; - - @Nullable - private ResponseListener delegate; - - /** - * When {@literal true}, the context object is a client span that needs to be ended on completion, when {@literal false} - * it means that only context-propagation (aka activation/deactivation) is required. - */ - private boolean isClientSpan; - - @Nullable - private volatile AbstractSpan context; - - ResponseListenerWrapper2(ElasticsearchRestClientInstrumentationHelper helper, Tracer tracer) { - super(helper); - this.helper = helper; - this.tracer = tracer; - } - - ResponseListenerWrapper withClientSpan(ResponseListener delegate, Span span) { - // Order is important due to visibility - write to span last on this (initiating) thread - this.delegate = delegate; - this.isClientSpan = true; - this.context = span; - return this; - } - - ResponseListenerWrapper2 withContextPropagation(ResponseListener delegate, AbstractSpan context) { - // Order is important due to visibility - write to span last on this (initiating) thread - this.delegate = delegate; - this.isClientSpan = false; - this.context = context; - return this; - } - - @Override - public void onSuccess(Response response) { - if (isClientSpan) { - onSuccessClient(response); - } else { - onSuccessContextPropagation(response); - } - } - - private void onSuccessContextPropagation(Response response) { - AbstractSpan localContext = context; - boolean activate = localContext != null && tracer.getActive() != localContext; - try { - if (activate) { - localContext.activate(); - } - if (delegate != null) { - delegate.onSuccess(response); - } - } finally { - if (activate) { - localContext.deactivate(); - } - helper.recycle(this); - } - } - - private void onSuccessClient(Response response){ - try { - finishClientSpan(response, null); - } finally { - if (delegate != null) { - delegate.onSuccess(response); - } - helper.recycle(this); - } - } - - @Override - public void onFailure(Exception exception) { - if (isClientSpan) { - onFailureClient(exception); - } else { - onFailureContextPropagation(exception); - } - } - - private void onFailureClient(Exception exception) { - try { - finishClientSpan(null, exception); - } finally { - if (delegate != null) { - delegate.onFailure(exception); - } - helper.recycle(this); - } - } - - private void onFailureContextPropagation(Exception exception) { - AbstractSpan localContext = context; - boolean activate = localContext != null && tracer.getActive() != localContext; - try { - if (activate) { - localContext.activate(); - } - if (delegate != null) { - delegate.onFailure(exception); - } - } finally { - if (activate) { - localContext.deactivate(); - } - helper.recycle(this); - } - } - - private void finishClientSpan(@Nullable Response response, @Nullable Throwable throwable) { - // First read volatile span to ensure visibility on executing thread - AbstractSpan localSpan = context; - if (localSpan instanceof Span) { - helper.finishClientSpan(response, (Span) localSpan, throwable); - } - } - - @Override - public void resetState() { - delegate = null; - context = null; - isClientSpan = false; - } -} From 31fc142e0adf56b8ef18ac830d5accb7f81200a9 Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Wed, 8 Feb 2023 17:54:06 +0100 Subject: [PATCH 6/8] simplify a bit what we can --- .../esrestclient/v7_x/RestClientInstrumentation.java | 3 +++ .../ElasticsearchRestClientInstrumentationIT.java | 12 ++++-------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java index cff8551126..36e7125be2 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/main/java/co/elastic/apm/agent/esrestclient/v7_x/RestClientInstrumentation.java @@ -37,6 +37,9 @@ import static net.bytebuddy.matcher.ElementMatchers.takesArgument; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +/** + * Instruments {@link org.elasticsearch.client.RestClient#performRequestAsync(org.elasticsearch.client.Request, org.elasticsearch.client.ResponseListener)} + */ public class RestClientInstrumentation extends ElasticsearchRestClientInstrumentation { @Override diff --git a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java index 16dd565e8d..64fb1f8d27 100644 --- a/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java +++ b/apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-7_x/src/test/java/co/elastic/apm/agent/esrestclient/v7_x/ElasticsearchRestClientInstrumentationIT.java @@ -154,17 +154,15 @@ public void testRestClientAsyncContextPropagation() throws InterruptedException, // test only relevant when testing async stuff return; } - // start without an active transaction - tracer.getActive().deactivate().end(); - assertThat(tracer.getActive()).isNull(); + + AbstractSpan active = tracer.getActive(); + assertThat(active).isInstanceOf(Transaction.class); reporter.reset(); AtomicReference> observedActive = new AtomicReference<>(); CountDownLatch endLatch = new CountDownLatch(1); - Transaction transaction = startTestRootTransaction("test-root-transaction"); - RestClient restClient = clientBuilder.build(); Request request = new Request("GET", "/"); @@ -186,9 +184,7 @@ public void onFailure(Exception exception) { endLatch.await(1, TimeUnit.SECONDS); assertThat(observedActive.get()) - .isSameAs(transaction); - - transaction.deactivate().end(); + .isSameAs(active); } @Override From cfd570c6305e0a1fdf936ca11262b8791a41e6ee Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Wed, 8 Feb 2023 17:57:20 +0100 Subject: [PATCH 7/8] cleanup --- .../java/co/elastic/apm/agent/concurrent/JavaConcurrent.java | 2 -- .../co/elastic/apm/agent/concurrent/ScopeManagementTest.java | 1 - 2 files changed, 3 deletions(-) diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java index 9a6b7db48e..6518cfd2fa 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/main/java/co/elastic/apm/agent/concurrent/JavaConcurrent.java @@ -35,9 +35,7 @@ import java.util.List; import java.util.Set; import java.util.concurrent.Callable; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.ForkJoinTask; -import java.util.function.BiConsumer; // Not strictly necessary as AbstractJavaConcurrentInstrumentation returns an empty collection for pluginClassLoaderRootPackages // but this signals the intent that this class must not be loaded from the IndyBootstrapClassLoader so that the state in this class applies globally diff --git a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java index 813c0a7c01..ac78d46be2 100644 --- a/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java +++ b/apm-agent-plugins/apm-java-concurrent-plugin/src/test/java/co/elastic/apm/agent/concurrent/ScopeManagementTest.java @@ -142,5 +142,4 @@ void testSpanAndContextCallableActivationInDifferentThread() throws Exception { assertThat(tracer.getActive()).isNull(); } - } From fef0896d93df20bd5d1f2de25d0e87e67fcd736e Mon Sep 17 00:00:00 2001 From: Sylvain Juge Date: Tue, 14 Feb 2023 10:03:14 +0100 Subject: [PATCH 8/8] update changelog --- CHANGELOG.asciidoc | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.asciidoc b/CHANGELOG.asciidoc index 79eead5859..9ee657fe7a 100644 --- a/CHANGELOG.asciidoc +++ b/CHANGELOG.asciidoc @@ -35,6 +35,7 @@ communication - {pull}2996[#2996] * Fixed used instrumentations printed on shutdown {pull}3001[#3001] * Prevent potential connection leak on network failure - {pull}2869[#2869] * Fix for inferred spans where the parent id was also a child id - {pull}2686[#2686] +* Fix context propagation for async 7.x and 8.x Elasticsearch clients - {pull}3015[#3015] [[release-notes-1.x]] === Java Agent version 1.x