invokeSuspend edge case fix

newrelic-agent/src/test/java/com/newrelic/agent/language/SourceLibraryDetectorTest.java

@@ -48,7 +48,7 @@ public void setup() throws Exception {
     public void testSamplerRun() throws Exception {
         sourceLibraryDetector.run();
         ArgumentCaptor<StatsWork> captor = ArgumentCaptor.forClass(StatsWork.class);
-        verify(ServiceFactory.getStatsService(), times(2)).doStatsWork(captor.capture(), anyString());
+        verify(ServiceFactory.getStatsService(), times(3)).doStatsWork(captor.capture(), anyString());
 
         StatsWork statsWork = captor.getValue();
 

newrelic-weaver/build.gradle

@@ -3,6 +3,7 @@ import com.nr.builder.publish.PublishConfig
 plugins {
     id("maven-publish")
     id("signing")
+    id("kotlin")
 }
 
 configurations {
@@ -35,6 +36,10 @@ dependencies {
 
     // used to test weaving java 1.4 class files
     testImplementation("org.apache.struts:struts-core:1.3.5")
+
+    //testing kotlin coroutines weaving
+    testImplementation("org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.9.21")
+    testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.9.0")
 }
 
 jar {

newrelic-weaver/src/main/java/com/newrelic/weave/MethodCallInlinerAdapter.java

@@ -7,7 +7,9 @@
 
 package com.newrelic.weave;
 
+import com.newrelic.weave.utils.ReturnInsnProcessor;
 import com.newrelic.weave.utils.WeaveUtils;
+import org.objectweb.asm.tree.*;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -16,15 +18,8 @@
 import org.objectweb.asm.commons.LocalVariablesSorter;
 import org.objectweb.asm.commons.MethodRemapper;
 import org.objectweb.asm.commons.Remapper;
-import org.objectweb.asm.tree.MethodNode;
-import org.objectweb.asm.tree.TryCatchBlockNode;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 public abstract class MethodCallInlinerAdapter extends LocalVariablesSorter {
     /**
@@ -113,7 +108,6 @@ public void visitMethodInsn(int opcode, String owner, String name, String desc,
             }
             int access = opcode == Opcodes.INVOKESTATIC ? Opcodes.ACC_STATIC : 0;
             inliner.inliner = new InliningAdapter(api, access, desc, this, mv, inliner.remapper);
-
         }
         inliner.method.accept(inliner.inliner);
     }
@@ -133,6 +127,11 @@ private InlinedMethod getInliner(String owner, String name, String desc) {
             } else {
                 // Copy the MethodNode before modifying the instructions list (which is not thread safe)
                 MethodNode methodNodeCopy = WeaveUtils.copy(method.method);
+                if (shouldClearReturnStacks(name, desc)){
+                    MethodNode result = WeaveUtils.newMethodNode(methodNodeCopy);
+                    methodNodeCopy.accept(new ClearReturnAdapter(owner, methodNodeCopy, result));
+                    methodNodeCopy = result;
+                }
                 methodNodeCopy.instructions.resetLabels();
                 InliningAdapter originalInliner = method.inliner;
 
@@ -234,6 +233,42 @@ protected int newLocalMapping(final Type type) {
         }
     }
 
+    /**
+    Flags method nodes requiring additional return insn processing, which for now is only invokeSuspend.
+    */
+    private boolean shouldClearReturnStacks(String name, String desc) {
+        final String invokeSuspendName = "invokeSuspend";
+        final String invokeSuspendDesc = "(Ljava/lang/Object;)Ljava/lang/Object;";
+        return invokeSuspendName.equals(name) && invokeSuspendDesc.equals(desc);
+    }
+
+    /**
+    This adapter checks a method's return instructions, adding additional POPs prior to return instructions
+    if the return is made with extra (>1) operands on the stack.
+
+    ReturnInsnProcessor.clearReturnStacks is placed in visitEnd() of this adapter, rather than called directly on
+    the source node, for thread safety reasons. Our thread safety strategy is to synchronize on the accept method
+    of the source node, and require that all modifications to bytecode be initiated by an invocation of accept():
+
+    source.accept(new ClearReturnAdapter(owner, source, next)) //conforms to thread-safe model
+
+    ReturnInsnProcessor.clearReturnStacks(owner, source) //does not conform to thread-safe model
+     */
+    class ClearReturnAdapter extends MethodNode {
+        String owner;
+        public ClearReturnAdapter(String owner, MethodNode source, MethodVisitor next) {
+            super(WeaveUtils.ASM_API_LEVEL, source.access, source.name, source.desc, source.signature, source.exceptions.toArray(new String[source.exceptions.size()]));
+            this.mv = next;
+            this.owner = owner;
+        }
+
+        @Override
+        public void visitEnd() {
+            ReturnInsnProcessor.clearReturnStacks(owner, this);
+            accept(mv);
+        }
+    }
+
     static class MergeFrameAdapter extends MethodVisitor {
 
         private final AnalyzerAdapter analyzerAdapter;

newrelic-weaver/src/main/java/com/newrelic/weave/PreparedMatch.java

@@ -297,7 +297,6 @@ public void visitParameter(String name, int access) {
     }
 
     private MethodNode prepare(ClassMatch classMatch, MethodNode matchedMethod, Map<Method, CallOriginalReplacement> callOriginalReplacementMap) {
-
         // remove JSR instructions
         MethodNode prepared = MethodProcessors.removeJSRInstructions(matchedMethod);
 

newrelic-weaver/src/main/java/com/newrelic/weave/utils/ReturnInsnProcessor.java

@@ -0,0 +1,103 @@
+package com.newrelic.weave.utils;
+
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.*;
+import org.objectweb.asm.tree.analysis.*;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class ReturnInsnProcessor {
+
+    private static final int EXPECTED_RETURN_STACK_SIZE = 1;
+
+    /***
+     * Clear the stacks at every A and I-type return instruction of a method, to their EXPECTED_RETURN_STACK_SIZE of 1.
+     * This mutates the given MethodNode and is not thread safe.
+     *
+     * @param owner The owning class of the method
+     * @param mn MethodNode representing the method to be modified
+     */
+    public static void clearReturnStacks(String owner, MethodNode mn) {
+        Map<AbstractInsnNode, Integer> stacks;
+        try {
+            stacks = getReturnStacks(owner, mn);
+        } catch (AnalyzerException e) {
+            //Something went wrong calculating return stacks, leave the method unmodified.
+            return;
+        }
+        if (stacks != null && !stacks.isEmpty()) {
+            int maxLocals = mn.maxLocals;
+            int nextLocalIdx = maxLocals; //create a new local variable
+            for (Map.Entry<AbstractInsnNode, Integer> entry : stacks.entrySet()) {
+                AbstractInsnNode insn = entry.getKey();
+                Integer stackSize = entry.getValue();
+                InsnList clearInsns = insnsBeforeReturn(insn.getOpcode(), stackSize, nextLocalIdx);
+                mn.instructions.insertBefore(insn, clearInsns);
+            }
+            mn.maxLocals = nextLocalIdx + 1;
+        }
+    }
+
+    /*
+    The return insn issue has only been observed happening with reference (A-type) instructions in coroutines.
+    But for safety (and for reuse in future contexts), this method
+    checks the type of the return and pairs it with the appropriate store and load types.
+
+    I-type is also allowed (though it hasn't been observed in the wild before).
+
+    All other return types will return an empty InsnList.
+     */
+
+    private static InsnList insnsBeforeReturn(int opcode, int stackSize, int varIndex){
+        InsnList insns = new InsnList();
+        int store;
+        int load;
+        switch (opcode) {
+            case Opcodes.ARETURN:
+                store = Opcodes.ASTORE;
+                load = Opcodes.ALOAD;
+                break;
+            case Opcodes.IRETURN :
+                store = Opcodes.ISTORE;
+                load = Opcodes.ILOAD;
+                break;
+            default:
+                return insns;
+        }
+        insns.add(new VarInsnNode(store, varIndex));
+        for (int i = stackSize; i > EXPECTED_RETURN_STACK_SIZE; i--){
+            insns.add(new InsnNode(Opcodes.POP));
+        }
+        insns.add(new VarInsnNode(load, varIndex));
+        return insns;
+    }
+
+    /*
+     * Compute the size of the operand stack at each return instruction exceeding the EXPECTED_RETURN_STACK_SIZE of 1.
+     * Only applies to methods returning I (int type) or A (reference type).
+     *
+     * @param owner The owning class
+     * @param method The method to analyze
+     * @return A Map of return instructions having extra operands on the stack, whose keys are return insn instances and values are stack sizes at the insn
+     */
+    private static Map<AbstractInsnNode, Integer> getReturnStacks(String owner, MethodNode method) throws AnalyzerException {
+        BasicInterpreter interpreter = new BasicInterpreter();
+        Analyzer<BasicValue> a = new Analyzer<>(interpreter);
+        Frame<BasicValue>[] frames = a.analyze(owner, method);
+        Map<AbstractInsnNode, Integer> rtStacks = new HashMap<>();
+        for (int j = 0; j < method.instructions.size(); ++j) {
+            AbstractInsnNode insn = method.instructions.get(j);
+            if (insn.getOpcode() == Opcodes.IRETURN || insn.getOpcode() == Opcodes.ARETURN) {
+                Frame<BasicValue> f = frames[j];
+                if (f != null && f.getStackSize() > EXPECTED_RETURN_STACK_SIZE) {
+                    rtStacks.put(insn, f.getStackSize());
+                }
+            }
+        }
+        return rtStacks;
+    }
+
+
+
+}

newrelic-weaver/src/main/java/com/newrelic/weave/utils/WeaveUtils.java

@@ -15,38 +15,20 @@
 import com.newrelic.api.agent.weaver.WeaveWithAnnotation;
 import com.newrelic.api.agent.weaver.Weaver;
 import com.newrelic.weave.MethodKey;
-import org.objectweb.asm.AnnotationVisitor;
-import org.objectweb.asm.ClassReader;
-import org.objectweb.asm.ClassVisitor;
-import org.objectweb.asm.ClassWriter;
-import org.objectweb.asm.MethodVisitor;
-import org.objectweb.asm.Opcodes;
-import org.objectweb.asm.Type;
+import org.objectweb.asm.*;
 import org.objectweb.asm.commons.JSRInlinerAdapter;
 import org.objectweb.asm.commons.Method;
 import org.objectweb.asm.commons.Remapper;
-import org.objectweb.asm.tree.AbstractInsnNode;
-import org.objectweb.asm.tree.AnnotationNode;
-import org.objectweb.asm.tree.ClassNode;
-import org.objectweb.asm.tree.FieldNode;
-import org.objectweb.asm.tree.InnerClassNode;
-import org.objectweb.asm.tree.InsnList;
-import org.objectweb.asm.tree.LabelNode;
-import org.objectweb.asm.tree.MethodInsnNode;
-import org.objectweb.asm.tree.MethodNode;
-import org.objectweb.asm.tree.TypeInsnNode;
+import org.objectweb.asm.tree.*;
+import org.objectweb.asm.util.*;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.lang.annotation.Annotation;
 import java.net.URL;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.regex.Pattern;
 
 /**
@@ -538,6 +520,39 @@ public static ClassNode convertToClassNode(byte[] classBytes) {
         return result;
     }
 
+    /**
+     * Utility method to print human-readable bytecode instructions of a MethodNode.
+     *
+     * @param mn the node to print
+     */
+    public static void printAllInstructions(MethodNode mn) {
+        for (AbstractInsnNode insn : mn.instructions) {
+            System.out.println(stringifyInstruction(insn));
+        }
+    }
+
+    public static String stringifyInstruction(AbstractInsnNode node) {
+        Textifier p = new Textifier(WeaveUtils.ASM_API_LEVEL) {
+            @Override
+            protected void appendLabel(Label l) {
+                if (labelNames == null) {
+                    labelNames = new HashMap<>();
+                }
+                String name = labelNames.get(l);
+                if (name == null) {
+                    name = l.toString();
+                    labelNames.put(l, name);
+                }
+                stringBuilder.append(name);
+            }
+        };
+        TraceMethodVisitor mv = new TraceMethodVisitor(p);
+        node.accept(mv);
+        StringWriter sw = new StringWriter();
+        p.print(new PrintWriter(sw));
+        return sw.toString().replace('\n', ' ');
+    }
+
     /**
      * Converts an ASM {@link ClassNode} to a byte array.
      *