changed entry point for tracking suspends

Author: dhilpipre

Kotlin-Coroutines-Suspends/src/main/java/com/newrelic/instrumentation/kotlin/coroutines/tracing/SuspendClassMatcher.java

@@ -1,49 +1,35 @@
 package com.newrelic.instrumentation.kotlin.coroutines.tracing;
 
-import java.util.ArrayList;
 import java.util.Collection;
 
 import com.newrelic.agent.deps.org.objectweb.asm.ClassReader;
 import com.newrelic.agent.instrumentation.classmatchers.ChildClassMatcher;
 import com.newrelic.agent.instrumentation.classmatchers.ClassMatcher;
-import com.newrelic.agent.instrumentation.classmatchers.ExactClassMatcher;
-import com.newrelic.agent.instrumentation.classmatchers.NotMatcher;
+import com.newrelic.agent.instrumentation.classmatchers.OrClassMatcher;
 
 public class SuspendClassMatcher extends ClassMatcher {
 
-	ChildClassMatcher matcher;
-	NotMatcher nrMatcher;
+	private ClassMatcher orMatcher;
 
 	public SuspendClassMatcher() {
-		matcher = new ChildClassMatcher("kotlin.coroutines.jvm.internal.BaseContinuationImpl",false);
-		nrMatcher = new NotMatcher(new ExactClassMatcher("com.newrelic.instrumentation.kotlin.coroutines.NRWrappedSuspend"));
+		ChildClassMatcher lambdamatcher = new ChildClassMatcher("kotlin.coroutines.jvm.internal.SuspendLambda",true);
+		ChildClassMatcher restrictedlambdamatcher = new ChildClassMatcher("kotlin.coroutines.jvm.internal.RestrictedSuspendLambda",true);
+		orMatcher = OrClassMatcher.getClassMatcher(lambdamatcher,restrictedlambdamatcher);
 	}
 
 	@Override
 	public boolean isMatch(ClassLoader loader, ClassReader cr) {
-		return matcher.isMatch(loader, cr) && nrMatcher.isMatch(loader, cr);
+		return orMatcher.isMatch(loader, cr);
 	}
 
 	@Override
 	public boolean isMatch(Class<?> clazz) {
-		return matcher.isMatch(clazz) && nrMatcher.isMatch(clazz);
+		return orMatcher.isMatch(clazz);
 	}
 
 	@Override
 	public Collection<String> getClassNames() {
-		Collection<String> childClasses = matcher.getClassNames();
-		Collection<String> nrClasses = nrMatcher.getClassNames();
-		if(childClasses == null && nrClasses == null) return null;
-		
-		ArrayList<String> list = new ArrayList<>();
-		if(childClasses != null) {
-			list.addAll(childClasses);
-		}
-		if(nrClasses != null) {
-			list.addAll(nrClasses);
-		}
-		
-		return list;
+		return orMatcher.getClassNames();
 	}
 
 }

Kotlin-Coroutines_1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/NRContinuationWrapper.java

@@ -3,8 +3,6 @@
 import com.newrelic.api.agent.Trace;
 import com.newrelic.api.agent.weaver.Weave;
 import com.newrelic.api.agent.weaver.Weaver;
-import com.newrelic.instrumentation.kotlin.coroutines_14.NRFunction1Wrapper;
-import com.newrelic.instrumentation.kotlin.coroutines_14.NRFunction2Wrapper;
 
 import kotlin.jvm.functions.Function1;
 import kotlin.jvm.functions.Function2;
@@ -14,19 +12,11 @@ public abstract class ContinuationKt {
 
 	@Trace
 	public static <T> void startCoroutine(Function1<? super Continuation<? super T>, ? extends Object> f, Continuation<? super T> cont ) {
-		if(!(f instanceof NRFunction1Wrapper)) {
-			NRFunction1Wrapper<? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction1Wrapper<>(f);
-			f = wrapper;
-		}
 		Weaver.callOriginal();
 	}
 
 	@Trace
 	public static <R, T> void startCoroutine(Function2<? super R, ? super Continuation<? super T>, ? extends Object> f, R receiver, Continuation<? super T> cont) {
-		if(!(f instanceof NRFunction2Wrapper)) {
-			NRFunction2Wrapper<? super R, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper<>(f);
-			f = wrapper;
-		}
 		Weaver.callOriginal();
 	}
 

Kotlin-Coroutines_1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/NRFunction1Wrapper.java

@@ -1,14 +1,13 @@
 package kotlinx.coroutines;
 
 import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Token;
 import com.newrelic.api.agent.Trace;
 import com.newrelic.api.agent.TracedMethod;
 import com.newrelic.api.agent.TransactionNamePriority;
 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_14.NRFunction1Wrapper;
-import com.newrelic.instrumentation.kotlin.coroutines_14.NRFunction2Wrapper;
 import com.newrelic.instrumentation.kotlin.coroutines_14.Utils;
 
 import kotlin.coroutines.Continuation;
@@ -39,10 +38,6 @@ protected void onCancelled(Throwable t, boolean b) {
 
 	@Trace
 	public void start(CoroutineStart start, Function1<? super Continuation<? super T>, ? extends Object> block) {
-		if(!(block instanceof NRFunction1Wrapper)) {
-			NRFunction1Wrapper<? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction1Wrapper<>(block);
-			block = wrapper;
-		}
 		String ctxName = Utils.getCoroutineName(getContext());
 		String name = ctxName != null ? ctxName : nameString$kotlinx_coroutines_core();
 		TracedMethod traced = NewRelic.getAgent().getTracedMethod();
@@ -60,10 +55,6 @@ public void start(CoroutineStart start, Function1<? super Continuation<? super T
 
 	@Trace
 	public <R> void start(CoroutineStart start, R receiver, Function2<? super R, ? super Continuation<? super T>, ? extends Object> block) {
-		if(!(block instanceof NRFunction2Wrapper)) {
-			NRFunction2Wrapper<? super R, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper<>(block);
-			block = wrapper;
-		}
 		String ctxName = Utils.getCoroutineName(getContext());
 		String name = ctxName != null ? ctxName : nameString$kotlinx_coroutines_core();
 		TracedMethod traced = NewRelic.getAgent().getTracedMethod();
@@ -80,4 +71,12 @@ public <R> void start(CoroutineStart start, R receiver, Function2<? super R, ? s
 		Weaver.callOriginal();
 	}
 
+	@Trace(async = true)
+	public void resumeWith(Object result) {
+		Token token = Utils.getToken(getContext());
+		if(token != null) {
+			token.link();
+		}
+		Weaver.callOriginal();
+	}
 }

Kotlin-Coroutines_1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/NRFunction2Wrapper.java

@@ -5,18 +5,14 @@
 import com.newrelic.api.agent.Trace;
 import com.newrelic.api.agent.weaver.Weave;
 import com.newrelic.api.agent.weaver.Weaver;
-import com.newrelic.instrumentation.kotlin.coroutines_14.NRContinuationWrapper;
 import com.newrelic.instrumentation.kotlin.coroutines_14.NRCoroutineToken;
-import com.newrelic.instrumentation.kotlin.coroutines_14.NRFunction2Wrapper;
 import com.newrelic.instrumentation.kotlin.coroutines_14.Utils;
 
 import kotlin.Unit;
 import kotlin.coroutines.Continuation;
 import kotlin.coroutines.CoroutineContext;
-import kotlin.coroutines.jvm.internal.SuspendFunction;
 import kotlin.jvm.functions.Function2;
 
-@SuppressWarnings({ "unchecked", "rawtypes" })
 @Weave
 public class BuildersKt {
 
@@ -39,10 +35,6 @@ public static final <T> T runBlocking(CoroutineContext context, Function2<? supe
 		}
 		NewRelic.getAgent().getTracedMethod().addCustomAttribute("Block", block.toString());
 
-		if (!(block instanceof NRFunction2Wrapper)) {
-			NRFunction2Wrapper<? super CoroutineScope, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper(block);
-			block = wrapper;
-		} 
 		T t = Weaver.callOriginal();
 		return t;
 	}
@@ -69,10 +61,6 @@ public static final <T> Deferred<T> async(CoroutineScope scope, CoroutineContext
 					context = context.plus(nrContextToken);
 				}
 			}
-			if(!(block instanceof NRFunction2Wrapper)) {
-				NRFunction2Wrapper<? super CoroutineScope, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper(block);
-				block = wrapper;
-			}
 		} else {
 			NewRelic.getAgent().getTransaction().ignore();
 		}
@@ -83,18 +71,6 @@ public static final <T> Deferred<T> async(CoroutineScope scope, CoroutineContext
 	public static final <T> Object invoke(CoroutineDispatcher dispatcher, Function2<? super CoroutineScope, ? super Continuation<? super T>, ? extends Object> block, Continuation<? super T> c) {
 
 		NewRelic.getAgent().getTracedMethod().addCustomAttribute("Continuation", c.toString());
-		if(!(block instanceof NRFunction2Wrapper)) {
-			NRFunction2Wrapper<? super CoroutineScope, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper(block);
-			block = wrapper;
-		}
-		if(c != null && !Utils.ignoreContinuation(c.toString())) {
-			boolean isSuspend = c instanceof SuspendFunction;
-			if(!isSuspend) {
-				String cont_string = Utils.getContinuationString(c);
-				NRContinuationWrapper wrapper = new NRContinuationWrapper<>(c, cont_string);
-				c = wrapper;
-			}
-		}
 		Object t = Weaver.callOriginal();
 		return t;
 	}
@@ -126,11 +102,6 @@ public static final kotlinx.coroutines.Job launch(CoroutineScope scope, Coroutin
 					context = context.plus(nrContextToken);
 				}
 			}
-			if (!(block instanceof NRFunction2Wrapper)) {
-				NRFunction2Wrapper<? super CoroutineScope, ? super Continuation<? super Unit>, ? extends Object> wrapper = new NRFunction2Wrapper(
-						block);
-				block = wrapper;
-			} 
 		} else {
 			NewRelic.getAgent().getTransaction().ignore();
 		}
@@ -159,17 +130,6 @@ public static final <T> Object withContext(CoroutineContext context,Function2<?
 				context = context.plus(nrContextToken);
 			}
 		}
-		if(!(block instanceof NRFunction2Wrapper)) {
-			NRFunction2Wrapper<? super CoroutineScope, ? super Continuation<? super T>, ? extends Object> wrapper = new NRFunction2Wrapper(block);
-			block = wrapper;
-		}
-		if(completion != null && !Utils.ignoreContinuation(completion.toString())) {
-			if(!(completion instanceof NRContinuationWrapper)) {
-				String cont_string = Utils.getContinuationString(completion);
-				NRContinuationWrapper wrapper = new NRContinuationWrapper<>(completion, cont_string);
-				completion = wrapper;
-			}
-		}
 		return Weaver.callOriginal();
 	}
 }