Merge pull request #18 from newrelic/undispatched

Fixed "Unable to ignore UndispatchedCoroutine"

Kotlin-Coroutines_1.2/build.gradle

@@ -23,7 +23,7 @@ jar {
 
 verifyInstrumentation {
 	passes 'org.jetbrains.kotlinx:kotlinx-coroutines-core:[1.2.0,1.4.0)'
-	passes 'org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:[1.3.9,1.4.0)'
+	passes 'org.jetbrains.kotlinx:kotlinx-coroutines-core-jvm:[1.2.0,1.4.0)'
 	excludeRegex '.*SNAPSHOT'
 	excludeRegex '.*alpha'
 	excludeRegex '.*-eap-.*'

Kotlin-Coroutines_1.2/src/main/java/com/newrelic/instrumentation/kotlin/coroutines/NRFunction2Wrapper.java

@@ -27,58 +27,69 @@ public NRFunction2Wrapper(Function2<P1, P2, R> d,String n) {
 
 	@SuppressWarnings({ "rawtypes", "unchecked" })
 	@Override
-	@Trace(async=true)
+	@Trace(async=true, excludeFromTransactionTrace = true)
 	public R invoke(P1 p1, P2 p2) {
-		String nameStr = null;
-		boolean linked = false;
-		if(p1 instanceof CoroutineContext) {
-			CoroutineContext ctx = (CoroutineContext)p1;
-			nameStr = Utils.getCoroutineName(ctx);
-			Token token = Utils.getToken(ctx);
-			if(token != null) {
-				token.link();
-				linked = true;
-			}
-		}
-		if(p1 instanceof CoroutineScope) {
-			CoroutineScope scope = (CoroutineScope)p1;
-			nameStr = Utils.getCoroutineName(scope.getCoroutineContext());
-			if (!linked) {
-				Token token = Utils.getToken(scope.getCoroutineContext());
+
+		boolean isUndispatched = p1.getClass().getName().equals("kotlinx.coroutines.UndispatchedCoroutine")  || p2.getClass().getName().equals("kotlinx.coroutines.UndispatchedCoroutine");
+		if (!isUndispatched) {
+			String nameStr = null;
+			boolean linked = false;
+			if (p1 instanceof CoroutineContext) {
+				CoroutineContext ctx = (CoroutineContext) p1;
+				nameStr = Utils.getCoroutineName(ctx);
+				Token token = Utils.getToken(ctx);
 				if (token != null) {
 					token.link();
 					linked = true;
-				} 
+				}
 			}
-		}
-		if(p2 instanceof Continuation) {
-			Continuation continuation = (Continuation)p2;
-			if(nameStr == null) nameStr = Utils.getCoroutineName(continuation.getContext(), continuation);
-			if(nameStr == null || nameStr.equals(Utils.CREATEMETHOD1) || nameStr.equals(Utils.CREATEMETHOD2)) nameStr = name;
-
-			if(!linked) {
-				Token token = Utils.getToken(continuation.getContext());
-				if (token != null) {
-					token.link();
-					linked = true;
-				} 
+			if (p1 instanceof CoroutineScope) {
+				CoroutineScope scope = (CoroutineScope) p1;
+				nameStr = Utils.getCoroutineName(scope.getCoroutineContext());
+				if (!linked) {
+					Token token = Utils.getToken(scope.getCoroutineContext());
+					if (token != null) {
+						token.link();
+						linked = true;
+					}
+				}
 			}
-
-			if (!Utils.ignoreContinuation(name) && !Utils.ignoreContinuation(continuation.getClass(), continuation.getContext())) {
-				NRContinuationWrapper wrapper = new NRContinuationWrapper(continuation, nameStr);
-				p2 = (P2) wrapper;
+			if (p2 instanceof Continuation) {
+				Continuation continuation = (Continuation) p2;
+				if (nameStr == null)
+					nameStr = Utils.getCoroutineName(continuation.getContext(), continuation);
+				if (nameStr == null || nameStr.equals(Utils.CREATEMETHOD1) || nameStr.equals(Utils.CREATEMETHOD2))
+					nameStr = name;
+
+				if (!linked) {
+					Token token = Utils.getToken(continuation.getContext());
+					if (token != null) {
+						token.link();
+						linked = true;
+					}
+				}
+
+				if (!Utils.ignoreContinuation(name)
+						&& !Utils.ignoreContinuation(continuation.getClass(), continuation.getContext())) {
+					NRContinuationWrapper wrapper = new NRContinuationWrapper(continuation, nameStr);
+					p2 = (P2) wrapper;
+				}
 			}
-		}
-		if(nameStr == null) {
-			nameStr = name;
-		}
-		if(nameStr != null) {
-			NewRelic.getAgent().getTracedMethod().setMetricName("Custom","WrappedSuspend",nameStr);
+			if (nameStr == null) {
+				nameStr = name;
+			}
+			if (nameStr != null) {
+				NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "WrappedSuspend", nameStr);
+			} 
 		}
 		if(delegate != null) {
 			return delegate.invoke(p1, p2);
 		}
 		return null;
 	}
+
+	public void markUndispatched() {
+
+	}
 
 }

Kotlin-Coroutines_1.2/src/main/java/com/newrelic/instrumentation/kotlin/coroutines/NRRunnable.java

@@ -0,0 +1,34 @@
+package com.newrelic.instrumentation.kotlin.coroutines;
+
+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 static boolean isTransformed = false;
+
+	public NRRunnable(Runnable r,Token t) {
+		if(!isTransformed) {
+			AgentBridge.instrumentation.retransformUninstrumentedClass(getClass());
+			isTransformed = true;
+		}
+		delegate = r;
+		token = t;
+	}
+
+	@Override
+	@Trace(async = true)
+	public void run() {
+		if(token != null) {
+			token.linkAndExpire();
+			token = null;
+		}
+		if(delegate != null) {
+			delegate.run();
+		}
+	}
+
+}

Kotlin-Coroutines_1.2/src/main/java/com/newrelic/instrumentation/kotlin/coroutines/Utils.java

@@ -3,7 +3,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.StringTokenizer;
-import java.util.logging.Level;
 
 import com.newrelic.agent.config.AgentConfig;
 import com.newrelic.agent.config.AgentConfigListener;
@@ -14,8 +13,10 @@
 
 import kotlin.coroutines.Continuation;
 import kotlin.coroutines.CoroutineContext;
+import kotlin.coroutines.jvm.internal.SuspendLambda;
 import kotlinx.coroutines.AbstractCoroutine;
 import kotlinx.coroutines.CoroutineName;
+import kotlinx.coroutines.DispatchedTask;
 
 public class Utils implements AgentConfigListener {
 
@@ -37,14 +38,31 @@ public class Utils implements AgentConfigListener {
 		loadConfig(config);
 	}
 
+	public static NRRunnable getRunnableWrapper(Runnable r) {
+		if(r instanceof NRRunnable) {
+			return null;
+		}
+		if(r instanceof DispatchedTask) {
+			return null;
+		}
+
+		Token t = NewRelic.getAgent().getTransaction().getToken();
+		if(t != null && t.isActive()) {
+			return new NRRunnable(r, t);
+		} else if(t != null) {
+			t.expire();
+			t = null;
+		}
+		return null;
+	}
+
 	private static void loadConfig(Config config) {
 		String ignores = config.getValue(SUSPENDSIGNORECONFIG);
 		if (ignores != null && !ignores.isEmpty()) {
 			StringTokenizer st = new StringTokenizer(ignores, ",");
 			while (st.hasMoreTokens()) {
 				String token = st.nextToken();
 				if (token != null && !token.isEmpty()) {
-					NewRelic.getAgent().getLogger().log(Level.INFO, "will ignore suspend: {0}", token);
 					ignoredSuspends.add(token);
 				}
 			} 
@@ -55,7 +73,6 @@ private static void loadConfig(Config config) {
 			while (st.hasMoreTokens()) {
 				String token = st.nextToken();
 				if (token != null && !token.isEmpty()) {
-					NewRelic.getAgent().getLogger().log(Level.INFO, "will ignore continuation: {0}", token);
 					ignoredContinuations.add(token);
 				}
 			} 
@@ -66,7 +83,6 @@ private static void loadConfig(Config config) {
 			while (st.hasMoreTokens()) {
 				String token = st.nextToken();
 				if (token != null && !token.isEmpty()) {
-					NewRelic.getAgent().getLogger().log(Level.INFO, "will ignore dispatched continuation: {0}", token);
 					ignoredDispatchs.add(token);
 				}
 			} 
@@ -118,7 +134,7 @@ public static boolean ignoreDispatched(Class<?> dispatched, CoroutineContext con
 	}
 
 	public static boolean ignoreSuspend(Class<?> suspend, CoroutineContext context) {
-
+		
 		String classname = suspend.getName();
 
 		if(ignoredSuspends.contains(classname)) return true;
@@ -183,6 +199,9 @@ public static <T> String getCoroutineName(CoroutineContext context, Continuation
 		if(continuation instanceof AbstractCoroutine) {
 			return ((AbstractCoroutine<T>)continuation).nameString$kotlinx_coroutines_core();
 		}
+		if(continuation instanceof SuspendLambda) {
+			return ((SuspendLambda)continuation).toString();
+		}
 		return getCoroutineName(context,continuation.getClass());
 	}
 

Kotlin-Coroutines_1.2/src/main/java/kotlinx/coroutines/BuildersKt.java

@@ -39,7 +39,7 @@ public static final <T> Deferred<T> async(CoroutineScope scope, CoroutineContext
 		}
 		if(name == null) name = block.getClass().getName();
 		NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Builders","async",name);
-		if(!Utils.ignoreSuspend(block.getClass(),context) && !Utils.ignoreSuspend(block.getClass(), scope.getCoroutineContext())) {
+		if(cStart != CoroutineStart.UNDISPATCHED && !Utils.ignoreSuspend(block.getClass(),context) && !Utils.ignoreSuspend(block.getClass(), scope.getCoroutineContext())) {
 
 			NRCoroutineToken nrContextToken = Utils.setToken(context);
 			if(nrContextToken != null) {
@@ -67,13 +67,16 @@ public static final <T> Object invoke(CoroutineDispatcher dispatcher, Function2<
 
 	@Trace
 	public static final kotlinx.coroutines.Job launch(CoroutineScope scope, CoroutineContext context, CoroutineStart cStart, Function2<? super CoroutineScope, ? super Continuation<? super Unit>, ? extends Object> block) {
+
 		String name = Utils.getCoroutineName(context);
 		if(name == null) {
 			name = Utils.getCoroutineName(scope.getCoroutineContext());
 		}
 		if(name == null) name = block.getClass().getName();
 		NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Builders","launch",name);
-		if(!Utils.ignoreSuspend(block.getClass(), context) && !Utils.ignoreSuspend(block.getClass(), scope.getCoroutineContext())) {
+		boolean check1 = Utils.ignoreSuspend(block.getClass(), context);
+		boolean check2 = Utils.ignoreSuspend(block.getClass(), scope.getCoroutineContext());
+		if(cStart != CoroutineStart.UNDISPATCHED && !check1 && !check2) {
 			NRCoroutineToken nrContextToken = Utils.setToken(context);
 			if(nrContextToken != null) {
 				context = context.plus(nrContextToken);

Kotlin-Coroutines_1.2/src/main/java/kotlinx/coroutines/CoroutineDispatcher.java

@@ -0,0 +1,24 @@
+package kotlinx.coroutines;
+
+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.newrelic.instrumentation.kotlin.coroutines.NRRunnable;
+import com.newrelic.instrumentation.kotlin.coroutines.Utils;
+
+import kotlin.coroutines.CoroutineContext;
+
+@Weave(type = MatchType.BaseClass)
+public abstract class CoroutineDispatcher {
+
+	@Trace
+	public void dispatch(CoroutineContext ctx, Runnable r) {
+		NRRunnable wrapper = Utils.getRunnableWrapper(r);
+		if(wrapper != null) {
+			r = wrapper;
+		}
+
+		Weaver.callOriginal();
+	}
+}