From b6c9952f70ca814832e3479e65bcd6e8be904d52 Mon Sep 17 00:00:00 2001
From: Kate Anderson <kateanderson@newrelic.com>
Date: Tue, 24 Dec 2024 12:58:29 -0800
Subject: [PATCH] Refactor patch to use direct token linking

---
 .../com/nr/agent/instrumentation/Utils.java   | 55 -------------------
 .../main/scala/slick/util/AsyncExecutor.scala | 31 +++++------
 2 files changed, 13 insertions(+), 73 deletions(-)
 delete mode 100644 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/java/com/nr/agent/instrumentation/Utils.java b/instrumentation/slick-2.12_3.2.0/src/main/java/com/nr/agent/instrumentation/Utils.java
deleted file mode 100644
index f3769245cc..0000000000
--- a/instrumentation/slick-2.12_3.2.0/src/main/java/com/nr/agent/instrumentation/Utils.java
+++ /dev/null
@@ -1,55 +0,0 @@
-package com.nr.agent.instrumentation;
-
-import com.newrelic.agent.bridge.AgentBridge;
-import com.newrelic.agent.bridge.Transaction;
-
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.logging.Level;
-
-public final class Utils {
-    private Utils() {}
-
-    public static AgentBridge.TokenAndRefCount getThreadTokenAndRefCount() {
-        AgentBridge.TokenAndRefCount tokenAndRefCount = AgentBridge.activeToken.get();
-        if (tokenAndRefCount == null) {
-            Transaction txn = AgentBridge.getAgent().getTransaction(false);
-            if (txn != null) {
-                tokenAndRefCount = new AgentBridge.TokenAndRefCount(
-                        txn.getToken(),
-                        AgentBridge.getAgent().getTracedMethod(),
-                        new AtomicInteger(1));
-                AgentBridge.activeToken.set(tokenAndRefCount);
-            }
-        }
-        logTokenInfo(tokenAndRefCount, "getThreadTokenAndRefCount");
-        return tokenAndRefCount;
-    }
-
-    public static void setThreadTokenAndRefCount(AgentBridge.TokenAndRefCount tokenAndRefCount) {
-        if (tokenAndRefCount != null && tokenAndRefCount.token != null) {
-            logTokenInfo(tokenAndRefCount, "setThreadTokenAndRefCount");
-            AgentBridge.activeToken.set(tokenAndRefCount);
-            tokenAndRefCount.token.link();
-        }
-    }
-
-    public static void clearThreadTokenAndRefCount(AgentBridge.TokenAndRefCount tokenAndRefCount) {
-        AgentBridge.activeToken.remove();
-        if (tokenAndRefCount != null && tokenAndRefCount.refCount.decrementAndGet() == 0) {
-            if (tokenAndRefCount.token != null) {
-                logTokenInfo(tokenAndRefCount, "clearThreadTokenAndRefCount");
-                tokenAndRefCount.token.expire();
-                tokenAndRefCount.token = null;
-            }
-        }
-    }
-
-    public static void logTokenInfo(AgentBridge.TokenAndRefCount tokenAndRefCount, String message) {
-        if (AgentBridge.getAgent().getLogger().isLoggable(Level.FINEST)) {
-            String tokenMessage = (tokenAndRefCount != null && tokenAndRefCount.token != null)
-                    ? String.format("[%s:%s:%d]", tokenAndRefCount.token, tokenAndRefCount.token.getTransaction(), tokenAndRefCount.refCount.get())
-                    : "[Empty token]";
-            AgentBridge.getAgent().getLogger().log(Level.FINEST, String.format("%s: %s", message, tokenMessage));
-        }
-    }
-}
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
index 26ca0f3f92..a5eea77369 100644
--- 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()
   }
 }