Fix NPE when async ES client is nested within Hibernate search (#1531)

Author: eyalkoren

CHANGELOG.asciidoc

@@ -29,6 +29,7 @@ endif::[]
 [float]
 ===== Bug fixes
 * Fix small memory allocation regression introduced with tracestate header {pull}1508[#1508]
+* Fix `NullPointerException` from `WeakConcurrentMap.put` through the Elasticsearch client instrumentation - {pull}1531[1531]
 
 [float]
 ===== Refactors

apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-5_6/src/main/java/co/elastic/apm/agent/es/restclient/v5_6/ElasticsearchClientAsyncInstrumentation.java

@@ -24,13 +24,11 @@
  */
 package co.elastic.apm.agent.es.restclient.v5_6;
 
-import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.es.restclient.ElasticsearchRestClientInstrumentation;
 import co.elastic.apm.agent.es.restclient.ElasticsearchRestClientInstrumentationHelper;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.transaction.Span;
 import co.elastic.apm.agent.sdk.advice.AssignTo;
-import co.elastic.apm.agent.sdk.state.GlobalThreadLocal;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -77,32 +75,37 @@ public ElementMatcher<? super MethodDescription> getMethodMatcher() {
 
 
  public static class ElasticsearchRestClientAsyncAdvice {
- @VisibleForAdvice
- public static final GlobalThreadLocal<Span> spanTls = GlobalThreadLocal.get(ElasticsearchRestClientAsyncAdvice.class, "spanTls");
- @AssignTo.Argument(5)
+
+ @Nullable
+ @AssignTo.Argument(index = 1, value = 5)
  @Advice.OnMethodEnter(suppress = Throwable.class)
- public static ResponseListener onBeforeExecute(@Advice.Argument(0) String method,
-  @Advice.Argument(1) String endpoint,
-  @Advice.Argument(3) @Nullable HttpEntity entity,
-  @Advice.Argument(5) ResponseListener responseListener) {
+ public static Object[] onBeforeExecute(@Advice.Argument(0) String method,
+ @Advice.Argument(1) String endpoint,
+ @Advice.Argument(3) @Nullable HttpEntity entity,
+ @Advice.Argument(5) ResponseListener responseListener) {
 
  ElasticsearchRestClientInstrumentationHelper<HttpEntity, Response, ResponseListener> helper = esClientInstrHelperManager.getForClassLoaderOfClass(Response.class);
  if (helper != null) {
  Span span = helper.createClientSpan(method, endpoint, entity);
- spanTls.set(span);
  if (span != null) {
- return helper.<ResponseListener>wrapResponseListener(responseListener, span);
+ Object[] ret = new Object[2];
+ ret[0] = span;
+ ret[1] = helper.<ResponseListener>wrapResponseListener(responseListener, span);
+ return ret;
  }
  }
- return responseListener;
+ return null;
  }
 
  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
- public static void onAfterExecute(@Advice.Thrown @Nullable Throwable t) {
- final Span span = spanTls.getAndRemove();
- if (span != null) {
- // Deactivate in this thread. Span will be ended and reported by the listener
- span.deactivate();
+ public static void onAfterExecute(@Advice.Thrown @Nullable Throwable t,
+ @Advice.Enter @Nullable Object[] entryArgs) {
+ if (entryArgs != null) {
+ final Span span = (Span) entryArgs[0];
+ if (span != null) {
+ // Deactivate in this thread. Span will be ended and reported by the listener
+ span.deactivate();
+ }
  }
  }
  }

apm-agent-plugins/apm-es-restclient-plugin/apm-es-restclient-plugin-6_4/src/main/java/co/elastic/apm/agent/es/restclient/v6_4/ElasticsearchClientAsyncInstrumentation.java

@@ -24,13 +24,11 @@
  */
 package co.elastic.apm.agent.es.restclient.v6_4;
 
-import co.elastic.apm.agent.bci.VisibleForAdvice;
 import co.elastic.apm.agent.es.restclient.ElasticsearchRestClientInstrumentation;
 import co.elastic.apm.agent.es.restclient.ElasticsearchRestClientInstrumentationHelper;
 import co.elastic.apm.agent.impl.ElasticApmTracer;
 import co.elastic.apm.agent.impl.transaction.Span;
 import co.elastic.apm.agent.sdk.advice.AssignTo;
-import co.elastic.apm.agent.sdk.state.GlobalThreadLocal;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
@@ -71,30 +69,34 @@ public ElementMatcher<? super MethodDescription> getMethodMatcher() {
  }
 
  public static class ElasticsearchRestClientAsyncAdvice {
- @VisibleForAdvice
- public static final GlobalThreadLocal<Span> spanTls = GlobalThreadLocal.get(ElasticsearchRestClientAsyncAdvice.class, "spanTls");
 
- @AssignTo.Argument(1)
+ @Nullable
+ @AssignTo.Argument(index = 1, value = 1)
  @Advice.OnMethodEnter(suppress = Throwable.class)
- public static ResponseListener onBeforeExecute(@Advice.Argument(0) Request request,
-  @Advice.Argument(1) ResponseListener responseListener) {
+ public static Object[] onBeforeExecute(@Advice.Argument(0) Request request,
+ @Advice.Argument(1) ResponseListener responseListener) {
  ElasticsearchRestClientInstrumentationHelper<HttpEntity, Response, ResponseListener> helper = esClientInstrHelperManager.getForClassLoaderOfClass(Request.class);
  if (helper != null) {
  Span span = helper.createClientSpan(request.getMethod(), request.getEndpoint(), request.getEntity());
  if (span != null) {
- spanTls.set(span);
- return helper.<ResponseListener>wrapResponseListener(responseListener, span);
+ Object[] ret = new Object[2];
+ ret[0] = span;
+ ret[1] = helper.<ResponseListener>wrapResponseListener(responseListener, span);
+ return ret;
  }
  }
- return responseListener;
+ return null;
  }
 
  @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
- public static void onAfterExecute(@Advice.Thrown @Nullable Throwable t) {
- final Span span = spanTls.getAndRemove();
- if (span != null) {
- // Deactivate in this thread. Span will be ended and reported by the listener
- span.deactivate();
+ public static void onAfterExecute(@Advice.Thrown @Nullable Throwable t,
+ @Advice.Enter @Nullable Object[] entryArgs) {
+ if (entryArgs != null) {
+ final Span span = (Span) entryArgs[0];
+ if (span != null) {
+ // Deactivate in this thread. Span will be ended and reported by the listener
+ span.deactivate();
+ }
  }
  }
  }