Fix coroutines tasks

instrumentation/kotlin-coroutines-1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/NRFunction1SuspendWrapper.java

@@ -2,6 +2,7 @@
 
 import com.newrelic.agent.bridge.AgentBridge;
 import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Trace;
 import kotlin.coroutines.Continuation;
 import kotlin.jvm.functions.Function1;
 
@@ -23,6 +24,7 @@ public NRFunction1SuspendWrapper(Function1<P1,R> d) {
 	}
 
 	@Override
+	@Trace
 	public R invoke(P1 p1) {
 		if(p1 instanceof Continuation) {
 			Continuation<?> cont = (Continuation<?>)p1;

instrumentation/kotlin-coroutines-1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/NRFunction2SuspendWrapper.java

@@ -2,7 +2,8 @@
 
 import com.newrelic.agent.bridge.AgentBridge;
 import com.newrelic.api.agent.NewRelic;
-import com.newrelic.api.agent.Token;
+
+import com.newrelic.api.agent.Trace;
 import kotlin.coroutines.Continuation;
 import kotlin.coroutines.CoroutineContext;
 import kotlin.jvm.functions.Function2;
@@ -15,52 +16,62 @@
 public class NRFunction2SuspendWrapper<S, T, R> implements Function2<S, T, R> {
 
 	private Function2<S, T, R> delegate = null;
+	private String name = null;
+	private String type = null;
 	private static boolean isTransformed = false;
 
-	public NRFunction2SuspendWrapper(Function2<S, T, R> d) {
+	public NRFunction2SuspendWrapper(String nameToUse, String typeToUse, Function2<S, T, R> d) {
 		delegate = d;
-		if(!isTransformed) {
+		this.name = nameToUse;
+		this.type = typeToUse;
+		if (!isTransformed) {
 			isTransformed = true;
 			AgentBridge.instrumentation.retransformUninstrumentedClass(getClass());
 		}
 	}
 
 	@Override
+	@Trace
 	public R invoke(S s, T t) {
-
-		// set name
-		boolean name_set = false;
-		if(s instanceof CoroutineScope) {
-			CoroutineScope scope = (CoroutineScope)s;
-			CoroutineContext ctx = scope.getCoroutineContext();
-			Token token = Utils.getToken(ctx);
-			if(token != null) {
-				token.link();
+		if (name != null) {
+			if (type != null) {
+				NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "Kotlin-Coroutines", "Block", type, name);
+			} else {
+				NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "Kotlin-Coroutines", "Block", "UnknownType", name);
+			}
+		} else {
+			String generatedName = getName(s,t);
+			if (type != null) {
+				NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "Kotlin-Coroutines", "Block", type, generatedName);
+			} else {
+				NewRelic.getAgent().getTracedMethod().setMetricName("Custom", "Kotlin-Coroutines", "Block", "UnknownType", generatedName);
 			}
+		}
+		if (delegate != null) {
+			return delegate.invoke(s, t);
+		}
+		return null;
+	}
+
+	private String getName(S s, T t) {
+		if (s instanceof CoroutineScope) {
+			CoroutineScope scope = (CoroutineScope) s;
 			CoroutineContext context = scope.getCoroutineContext();
 			String coroutineName = Utils.getCoroutineName(context);
-			if(coroutineName != null) {
-				NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Block","SuspendFunction",coroutineName);
-				name_set = true;
+			if (coroutineName != null) {
+				return coroutineName;
 			}
-
 		}
-
-		if(!name_set && t instanceof Continuation) {
-			Continuation<?> cont = (Continuation<?>)t;
 
-			String cont_string = Utils.getContinuationString(cont);
-			if(cont_string != null) {
-				NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Block","SuspendFunction",cont_string);
-			} else {
-				NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Block","SuspendFunction","UnknownSource");
+		if (t instanceof Continuation) {
+			Continuation<?> cont = (Continuation<?>) t;
 
+			String cont_string = Utils.getContinuationString(cont);
+			if (cont_string != null) {
+				return cont_string;
 			}
 		}
-		if(delegate != null) {
-			return delegate.invoke(s, t);
-		}
-		return null;
+		return "UnknownSource";
 	}
 
 }

instrumentation/kotlin-coroutines-1.4/src/main/java/com/newrelic/instrumentation/kotlin/coroutines_14/Utils.java

@@ -84,6 +84,9 @@ public static boolean continueWithScope(CoroutineScope scope) {
 	* coroutineScope can be a Coroutine name or CoroutineScope class name
 	*/
 	public static boolean continueWithScope(String coroutineScope) {
+		if(coroutineScope == null) {
+			return true;
+		}
 		for(Pattern ignoredScope : ignoredScopePatterns) {
 			if(ignoredScope.matcher(coroutineScope).matches()) {
 				return false;

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlin/coroutines/ContinuationKt_Instrumentation.java

@@ -23,7 +23,7 @@ public static <T> void startCoroutine(Function1<? super Continuation<? super T>,
 	@Trace
 	public static <R, T> void startCoroutine(Function2<? super R, ? super Continuation<? super T>, ?> f, R receiver, Continuation<? super T> cont) {
 		if(!(f instanceof NRFunction2SuspendWrapper)) {
-            f = new NRFunction2SuspendWrapper<>(f);
+            f = new NRFunction2SuspendWrapper<>(null, "Coroutine",f);
 		}
 		Weaver.callOriginal();
 	}

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/AbstractCoroutine_Instrumentation.java

@@ -36,11 +36,11 @@ protected void onCancelled(Throwable t, boolean b) {
 
 	@Trace
 	public <R> void start(CoroutineStart start, R receiver, Function2<? super R, ? super Continuation<? super T>, ?> block) {
-		if(!(block instanceof NRFunction2SuspendWrapper)) {
-            block = new NRFunction2SuspendWrapper<>(block);
-		}
 		String ctxName = Utils.getCoroutineName(getContext());
 		String name = ctxName != null ? ctxName : nameString$kotlinx_coroutines_core();
+		if(!(block instanceof NRFunction2SuspendWrapper)) {
+			block = new NRFunction2SuspendWrapper<>(name, "Coroutine", block);
+		}
 		TracedMethod traced = NewRelic.getAgent().getTracedMethod();
 		traced.addCustomAttribute("Coroutine-Name", name);
 		traced.addCustomAttribute("Block", block.toString());

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/AwaitKt_Instrumentation.java

@@ -0,0 +1,40 @@
+package kotlinx.coroutines;
+
+import com.newrelic.api.agent.NewRelic;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import kotlin.Unit;
+import kotlin.coroutines.Continuation;
+
+import java.util.Collection;
+import java.util.List;
+
+@Weave(originalName = "kotlinx.coroutines.AwaitKt")
+public class AwaitKt_Instrumentation {
+
+    @Trace
+    public static <T> Object awaitAll(Deferred<? extends T>[] deferreds, Continuation<? super List<? extends T>> continuation) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Kotlin","Coroutines","AwaitKt","awaitAll");
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public static <T> Object awaitAll(Collection<? extends Deferred<? extends T>> deferreds, Continuation<? super List<? extends T>> continuation) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Kotlin","Coroutines","AwaitKt","awaitAll");
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public static Object joinAll(Job[] jobs, Continuation<? super Unit> continuation) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Kotlin","Coroutines","AwaitKt","joinAll");
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public static Object joinAll(Collection<? extends Job> jobs, Continuation<? super Unit>  continuation) {
+        NewRelic.getAgent().getTracedMethod().setMetricName("Custom","Kotlin","Coroutines","AwaitKt","joinAll");
+        return Weaver.callOriginal();
+    }
+
+}

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/BuildersKt_Instrumentation.java

@@ -29,7 +29,7 @@ public static <T> T runBlocking(CoroutineContext context, Function2<? super Coro
 		NewRelic.getAgent().getTracedMethod().addCustomAttribute("Block", block.toString());
 
 		if (!(block instanceof NRFunction2SuspendWrapper)) {
-            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(block);
+            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(name, "RunBlocking", block);
 		}
         return Weaver.callOriginal();
 	}
@@ -55,7 +55,7 @@ public static <T> Deferred<T> async(CoroutineScope scope, CoroutineContext conte
 			}
 
 			if(!(block instanceof NRFunction2SuspendWrapper)) {
-                block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(block);
+                block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(name, "Async", block);
 			}
 		} else {
 			NewRelic.getAgent().getTransaction().ignore();
@@ -69,7 +69,7 @@ public static <T> Object invoke(CoroutineDispatcher_Instrumentation dispatcher,
 
 		NewRelic.getAgent().getTracedMethod().addCustomAttribute("Continuation", cont.toString());
 		if(!(block instanceof NRFunction2SuspendWrapper)) {
-            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(block);
+            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(null, "Invoke", block);
 		}
 		if(Utils.continueWithContinuation(cont)) {
 			boolean isSuspend = cont instanceof SuspendFunction;
@@ -102,8 +102,7 @@ public static kotlinx.coroutines.Job launch(CoroutineScope scope, CoroutineConte
 			}
 			NewRelic.getAgent().getTracedMethod().addCustomAttribute("Block", block.toString());
 			if (!(block instanceof NRFunction2SuspendWrapper)) {
-                block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super Unit>, ?>) new NRFunction2SuspendWrapper(
-                        block);
+                block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super Unit>, ?>) new NRFunction2SuspendWrapper(name, "Launch", block);
 			} 
 		} else {
 			NewRelic.getAgent().getTransaction().ignore();
@@ -125,7 +124,7 @@ public static <T> Object withContext(CoroutineContext context, Function2<? super
 		}
 
 		if(!(block instanceof NRFunction2SuspendWrapper)) {
-            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(block);
+            block = (NRFunction2SuspendWrapper<? super CoroutineScope, ? super Continuation<? super T>, ?>) new NRFunction2SuspendWrapper(null, "WithContext", block);
 		}
 		if(completion != null && Utils.continueWithContinuation(completion)) {
 			if(!(completion instanceof NRContinuationWrapper)) {

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/CancellableContinuationImpl_Instrumentation.java

@@ -0,0 +1,64 @@
+package kotlinx.coroutines;
+
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import kotlin.Unit;
+import kotlin.coroutines.CoroutineContext;
+import kotlin.jvm.functions.Function1;
+import kotlin.jvm.functions.Function3;
+
+@Weave(originalName = "kotlinx.coroutines.CancellableContinuationImpl")
+public abstract class CancellableContinuationImpl_Instrumentation<T> {
+
+    @Trace
+    public void resumeWith(Object obj) {
+        Weaver.callOriginal();
+    }
+
+    @Trace
+    public java.lang.Object tryResumeWithException(Throwable t) {
+
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public Object tryResume(T t, Object o) {
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public java.lang.Object tryResume(T r, Object obj, Function1<? super Throwable, Unit> function1) {
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public void completeResume(Object object) {
+        Weaver.callOriginal();
+    }
+
+    @Trace
+    public boolean cancel(Throwable t) {
+        return Weaver.callOriginal();
+    }
+
+    @Trace
+    public void invokeOnCancellation(Function1<? super Throwable, Unit> function1) {
+        Weaver.callOriginal();
+    }
+
+    @Trace
+    public void resumeUndispatched(CoroutineDispatcher dispatcher, T t) {
+        Weaver.callOriginal();
+    }
+
+    @Trace
+    public void resumeUndispatchedWithException(CoroutineDispatcher dispatcher, Throwable t) {
+        Weaver.callOriginal();
+    }
+
+    @Trace
+    public void resume(T r, Function1<? super Throwable,Unit> function3) {
+        Weaver.callOriginal();
+    }
+}

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/Deferred_Instrumentation.java

@@ -0,0 +1,16 @@
+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 kotlin.coroutines.Continuation;
+
+@Weave(originalName = "kotlinx.coroutines.Deferred", type = MatchType.Interface)
+public class Deferred_Instrumentation<T> {
+
+    @Trace
+    public Object await(Continuation<? super T>  continuation) {
+        return Weaver.callOriginal();
+    }
+}

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/EventLoopImplBase_Instrumentation.java

@@ -12,6 +12,21 @@
 @Weave(originalName = "kotlinx.coroutines.EventLoopImplBase")
 public abstract class EventLoopImplBase_Instrumentation {
 
+	public void schedule(long nanos, DelayedTask_Instrumentation task) {
+		if(task.token == null) {
+			Token t = NewRelic.getAgent().getTransaction().getToken();
+			if(t != null) {
+				if(!t.isActive()) {
+					t.expire();
+					t = null;
+				} else {
+					task.token = t;
+				}
+			}
+		}
+		Weaver.callOriginal();
+	}
+
 	@Weave(type = MatchType.BaseClass, originalName = "kotlinx.coroutines.EventLoopImplBase$DelayedTask")
 	public static abstract class DelayedTask_Instrumentation {
 
@@ -26,6 +41,14 @@ public DelayedTask_Instrumentation(long nanos) {
 		@NewField
 		protected Token token = null;
 
+		public int scheduleTask(long delay, DelayedTaskQueue_Instrumentation queue, EventLoopImplBase_Instrumentation eventLoop) {
+			if(token == null) {
+				token = NewRelic.getAgent().getTransaction().getToken();
+			}
+
+			return Weaver.callOriginal();
+		}
+
 		public void dispose() {
 			if(token != null) {
 				token.expire();

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/intrinsics/CancellableKt_Instrumentation.java

@@ -57,7 +57,7 @@ public static <R, T> void startCoroutineCancellable(Function2<? super R, ? super
 
 		}
 		if(!(f instanceof NRFunction2SuspendWrapper)) {
-            f = new NRFunction2SuspendWrapper<>(f);
+            f = new NRFunction2SuspendWrapper<>(null, "CoroutineCancellable", f);
 		}
 		Weaver.callOriginal();
 	}

instrumentation/kotlin-coroutines-1.4/src/main/java/kotlinx/coroutines/intrinsics/UndispatchedKt_Instrumentation.java

@@ -54,7 +54,7 @@ public static <R, T> void startCoroutineUndispatched(Function2<? super R, ? supe
 		}
 		traced.addCustomAttribute("Receiver", receiver.getClass().getName());
 		if(!(f instanceof NRFunction2SuspendWrapper)) {
-            f = new NRFunction2SuspendWrapper<>(f);
+            f = new NRFunction2SuspendWrapper<>(null, "CoroutineUndispatched", f);
 		}
 		Weaver.callOriginal();
 	}
@@ -66,7 +66,7 @@ public static <T, R> Object startUndispatchedOrReturn(ScopeCoroutine<? super T>
 		traced.addCustomAttribute("Suspend-Type", "Function2");
 		traced.addCustomAttribute("Receiver", receiver.getClass().getName());
 		if(!(f instanceof NRFunction2SuspendWrapper)) {
-            f = new NRFunction2SuspendWrapper<>(f);
+			f = new NRFunction2SuspendWrapper<>(null, "CoroutineUndispatched", f);
 		}
 		return Weaver.callOriginal();
 	}
@@ -78,7 +78,7 @@ public static <T, R> Object startUndispatchedOrReturnIgnoreTimeout(ScopeCoroutin
 		traced.addCustomAttribute("Suspend-Type", "Function2");
 		traced.addCustomAttribute("Receiver", receiver.getClass().getName());
 		if(!(f instanceof NRFunction2SuspendWrapper)) {
-            f = new NRFunction2SuspendWrapper<>(f);
+			f = new NRFunction2SuspendWrapper<>(null, "CoroutineUndispatched", f);
 		}
 		return Weaver.callOriginal();
 	}