From 4457d4e98f91501bb7914cbb29e440a857972fee Mon Sep 17 00:00:00 2001
From: eyalkoren <eyal.koren@elastic.co>
Date: Tue, 18 Dec 2018 15:45:14 +0200
Subject: [PATCH] Move wrapRunnable to Tracer and avoid trace ref recycling

---
 .../co/elastic/apm/agent/impl/ElasticApmTracer.java  |  4 ++--
 .../apm/agent/impl/InScopeRunnableWrapper.java       | 12 +++---------
 .../apm/agent/servlet/AsyncInstrumentation.java      |  4 +---
 .../servlet/helper/AsyncContextAdviceHelperImpl.java |  5 -----
 4 files changed, 6 insertions(+), 19 deletions(-)

diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java
index 21a5f3f2666..dc1676c1a0e 100644
--- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java
+++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/ElasticApmTracer.java
@@ -304,8 +304,8 @@ public void recycle(ErrorCapture error) {
         errorPool.recycle(error);
     }
 
-    public InScopeRunnableWrapper allocateRunnableWrapper() {
-        return runnableWrapperObjectPool.createInstance();
+    public Runnable wrapRunnable(Runnable delegate, AbstractSpan<?> span) {
+        return runnableWrapperObjectPool.createInstance().wrap(delegate, span);
     }
 
     public void recycle(InScopeRunnableWrapper wrapper) {
diff --git a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/InScopeRunnableWrapper.java b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/InScopeRunnableWrapper.java
index a98f6cd2187..0613938e187 100644
--- a/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/InScopeRunnableWrapper.java
+++ b/apm-agent-core/src/main/java/co/elastic/apm/agent/impl/InScopeRunnableWrapper.java
@@ -36,19 +36,15 @@ public class InScopeRunnableWrapper implements Runnable, Recyclable {
     private Runnable delegate;
     @Nullable
     private AbstractSpan<?> span;
-    @Nullable
-    private ElasticApmTracer tracer;
+
+    private final ElasticApmTracer tracer;
 
     InScopeRunnableWrapper(ElasticApmTracer tracer) {
         this.tracer = tracer;
     }
 
-    public InScopeRunnableWrapper wrap(Runnable delegate) {
+    public InScopeRunnableWrapper wrap(Runnable delegate, AbstractSpan<?> span) {
         this.delegate = delegate;
-        return this;
-    }
-
-    public InScopeRunnableWrapper withScope(AbstractSpan<?> span) {
         this.span = span;
         return this;
     }
@@ -75,7 +71,6 @@ public void run() {
                 if (span != null) {
                     span.deactivate();
                 }
-                //noinspection ConstantConditions
                 tracer.recycle(this);
             } catch (Throwable throwable) {
                 try {
@@ -91,6 +86,5 @@ public void run() {
     public void resetState() {
         delegate = null;
         span = null;
-        tracer = null;
     }
 }
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java
index 005a9337349..ed50f404acc 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/AsyncInstrumentation.java
@@ -81,8 +81,6 @@ public Collection<String> getInstrumentationGroupNames() {
 
     public interface AsyncContextAdviceHelper<T> {
         void onExitStartAsync(T asyncContext);
-
-        Runnable wrapRunnable(Runnable delegate, Transaction transaction);
     }
 
     public static class StartAsyncInstrumentation extends AsyncInstrumentation {
@@ -174,7 +172,7 @@ private static void onEnterAsyncContextStart(@Advice.Argument(value = 0, readOnl
                 if (tracer != null) {
                     final Transaction transaction = tracer.currentTransaction();
                     if (transaction != null && runnable != null && asyncHelper != null) {
-                        runnable = asyncHelper.getForClassLoaderOfClass(AsyncContext.class).wrapRunnable(runnable, transaction);
+                        runnable = tracer.wrapRunnable(runnable, transaction);
                     }
                 }
             }
diff --git a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java
index 18061aca122..e817f264f58 100644
--- a/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java
+++ b/apm-agent-plugins/apm-servlet-plugin/src/main/java/co/elastic/apm/agent/servlet/helper/AsyncContextAdviceHelperImpl.java
@@ -60,9 +60,4 @@ public void onExitStartAsync(AsyncContext asyncContext) {
                 asyncContext.getRequest(), asyncContext.getResponse());
         }
     }
-
-    @Override
-    public Runnable wrapRunnable(Runnable delegate, Transaction transaction) {
-        return tracer.allocateRunnableWrapper().wrap(delegate).withScope(transaction);
-    }
 }