Scheduler changes

newrelic-experimental · Feb 28, 2024 · 58a1981 · 58a1981
1 parent dbb197a
commit 58a1981
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 108 deletions.
diff --git a/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/HeaderUtils.java b/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/HeaderUtils.java
diff --git a/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRCallable.java b/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRCallable.java
@@ -3,30 +3,33 @@
 import java.util.concurrent.Callable;
 
 import com.newrelic.agent.bridge.AgentBridge;
+import com.newrelic.api.agent.Token;
 import com.newrelic.api.agent.Trace;
 
 public class NRCallable<V> implements Callable<V> {
 
 	private Callable<V> delegate = null;
 
-	private NRMuleHeaders headers = null;
+	private Token token = null;
 
 	private static boolean isTransformed = false;
 
-	public NRCallable(Callable<V> c, NRMuleHeaders h) {
+	public NRCallable(Callable<V> c, Token h) {
 		delegate = c;
-		headers = h;
+		token = h;
 		if(!isTransformed) {
 			isTransformed = true;
 			AgentBridge.instrumentation.retransformUninstrumentedClass(getClass());
 		}
 	}
 
 	@Override
-	@Trace(dispatcher = true)
+	@Trace(async = true)
 	public V call() throws Exception {
-		HeaderUtils.acceptHeaders(headers);
-		headers = null;
+		if(token != null) {
+			token.linkAndExpire();
+			token = null;
+		}
 		if(delegate != null) {
 			return delegate.call();
 		}

diff --git a/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRMuleHeaders.java b/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRMuleHeaders.java
diff --git a/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRRunnable.java b/Mule-Scheduler/src/main/java/com/nr/instrumentation/mule/scheduler/NRRunnable.java
@@ -1,30 +1,32 @@
 package com.nr.instrumentation.mule.scheduler;
 
 import com.newrelic.agent.bridge.AgentBridge;
+import com.newrelic.api.agent.Token;
 import com.newrelic.api.agent.Trace;
 
 public class NRRunnable implements Runnable {
 
 	private Runnable delegate = null;
+	private Token token = null;
 
-	private NRMuleHeaders headers = null;
-
 	private static boolean isTransformed = false;
 
-	public NRRunnable(Runnable d, NRMuleHeaders h) {
+	public NRRunnable(Runnable d, Token h) {
 		delegate = d;
-		headers = h;
+		token = h;
 		if(!isTransformed) {
 			isTransformed = true;
 			AgentBridge.instrumentation.retransformUninstrumentedClass(getClass());
 		}
 	}
 
 	@Override
-	@Trace(dispatcher=true)
+	@Trace(async=true)
 	public void run() {
-		HeaderUtils.acceptHeaders(headers);
-		headers = null;
+		if(token != null) {
+			token.linkAndExpire();
+			token = null;
+		}
 		if(delegate != null) {
 			delegate.run();
 		}

diff --git a/Mule-Scheduler/src/main/java/org/mule/service/scheduler/internal/DefaultScheduler.java b/Mule-Scheduler/src/main/java/org/mule/service/scheduler/internal/DefaultScheduler.java
@@ -5,37 +5,45 @@
 import java.util.concurrent.TimeUnit;
 
 import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Token;
+import com.newrelic.api.agent.Trace;
 import com.newrelic.api.agent.weaver.MatchType;
 import com.newrelic.api.agent.weaver.Weave;
 import com.newrelic.api.agent.weaver.Weaver;
 import com.nr.instrumentation.mule.scheduler.NRCallable;
-import com.nr.instrumentation.mule.scheduler.NRMuleHeaders;
 import com.nr.instrumentation.mule.scheduler.NRRunnable;
 
 @Weave(type=MatchType.BaseClass)
 public abstract class DefaultScheduler {
 
+	@Trace(dispatcher = true)
 	public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
 		if(!(command instanceof NRRunnable)) {
-			NRMuleHeaders headers = new NRMuleHeaders();
-			NewRelic.getAgent().getTransaction().insertDistributedTraceHeaders(headers);
-			if(!headers.isEmpty()) {
-				NRRunnable wrapper = new NRRunnable(command,headers);
+			Token token = NewRelic.getAgent().getTransaction().getToken();
+			if(token != null && token.isActive()) {
+				NRRunnable wrapper = new NRRunnable(command,token);
 				command = wrapper;
+			} else if(token != null) {
+				token.expire();
+				token = null;
 			}
 		}
 		return Weaver.callOriginal();
 	}
 
+	@Trace(dispatcher = true)
 	public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
 		if(!(callable instanceof NRCallable)) {
-			NRMuleHeaders headers = new NRMuleHeaders();
-			NewRelic.getAgent().getTransaction().insertDistributedTraceHeaders(headers);
-			if(!headers.isEmpty()) {
-				NRCallable<V> wrapper = new NRCallable<V>(callable, headers);
+			Token token = NewRelic.getAgent().getTransaction().getToken();
+			if(token != null && token.isActive()) {
+				NRCallable<V> wrapper = new NRCallable<V>(callable, token);
 				callable = wrapper;
+			} else if(token != null) {
+				token.expire();
+				token = null;
 			}
 		}
 		return Weaver.callOriginal();
 	}
+
 }