Refactor patch to use direct token linking

newrelic · Dec 24, 2024 · b6c9952 · b6c9952
1 parent b21ab77
commit b6c9952
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 73 deletions.
diff --git a/instrumentation/slick-2.12_3.2.0/src/main/java/com/nr/agent/instrumentation/Utils.java b/instrumentation/slick-2.12_3.2.0/src/main/java/com/nr/agent/instrumentation/Utils.java
diff --git a/instrumentation/slick-2.12_3.2.0/src/main/scala/slick/util/AsyncExecutor.scala b/instrumentation/slick-2.12_3.2.0/src/main/scala/slick/util/AsyncExecutor.scala
@@ -13,7 +13,6 @@ import com.newrelic.api.agent.{Trace, TracedMethod}
 import com.newrelic.api.agent.weaver.Weaver
 import com.newrelic.api.agent.weaver.internal.WeavePackageType
 import com.newrelic.api.agent.weaver.scala.{ScalaMatchType, ScalaWeave}
-import com.nr.agent.instrumentation.Utils
 
 import scala.concurrent.ExecutionContext
 import scala.concurrent.duration._
@@ -74,29 +73,25 @@ class NewRelicExecutionContext(delegatee :ExecutionContext) extends ExecutionCon
 }
 
 class NewRelicRunnable(var runnable :Runnable) extends Runnable {
-  private val tokenAndRefCount = Utils.getThreadTokenAndRefCount()
+  private val token = com.newrelic.api.agent.NewRelic.getAgent.getTransaction.getToken
   @Trace(async = true)
   override def run() {
     try {
-      try {
-        Utils.setThreadTokenAndRefCount(tokenAndRefCount);
-        AgentBridge.currentApiSource.set(WeavePackageType.INTERNAL)
+      token.linkAndExpire()
+      AgentBridge.currentApiSource.set(WeavePackageType.INTERNAL)
 
-        if(AgentBridge.getAgent().startAsyncActivity(runnable)) {
-          val tm = AgentBridge.getAgent().getTransaction().getTracedMethod().asInstanceOf[TracedMethod]
-          tm.setMetricName("ORM", "Slick",  "slickQuery")
-        }
-      } catch {
-        case t: Throwable => {
-          AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
-        }
-        case _ => ;
-      } finally {
-        AgentBridge.currentApiSource.remove()
+      if(AgentBridge.getAgent().startAsyncActivity(runnable)) {
+        val tm = AgentBridge.getAgent().getTransaction().getTracedMethod().asInstanceOf[TracedMethod]
+        tm.setMetricName("ORM", "Slick",  "slickQuery")
+      }
+    } catch {
+      case t: Throwable => {
+        AgentBridge.instrumentation.noticeInstrumentationError(t, Weaver.getImplementationTitle());
       }
-      runnable.run()
+      case _ => ;
     } finally {
-      Utils.clearThreadTokenAndRefCount(tokenAndRefCount);
+      AgentBridge.currentApiSource.remove()
     }
+    runnable.run()
   }
 }