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,6 +7,7 @@
 
 package com.newrelic.weave;
 
+import com.newrelic.weave.utils.ReturnInsnProcessor;
 import com.newrelic.weave.utils.WeaveUtils;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
@@ -19,12 +20,12 @@
 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.ArrayList;
 
 public abstract class MethodCallInlinerAdapter extends LocalVariablesSorter {
     /**
@@ -113,7 +114,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 +133,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 +239,58 @@ protected int newLocalMapping(final Type type) {
         }
     }
 
+    /**
+     * Flags method nodes requiring additional return insn processing, which for now is only invokeSuspend.
+     * For future reference, if other code surfaces a bytecode verification failure (such as an ArrayIndexOutOfBoundsException),
+     * modify this guard to process additional methods beyond invokeSuspend.
+     */
+    private boolean shouldClearReturnStacks(String name, String desc) {
+        if (clearReturnStacksDisabled()) {
+            return false;
+        }
+        final String invokeSuspendName = "invokeSuspend";
+        final String invokeSuspendDesc = "(Ljava/lang/Object;)Ljava/lang/Object;";
+        return invokeSuspendName.equals(name) && invokeSuspendDesc.equals(desc);
+    }
+
+    /**
+     * Feature flag to disable return stack processing.
+     * To use this feature flag, set -Dnewrelic.config.class_transformer.clear_return_stacks=false at JVM startup.
+     */
+    private boolean clearReturnStacksDisabled() {
+        String enabled = System.getProperty("newrelic.config.class_transformer.clear_return_stacks", "true");
+        return enabled.equalsIgnoreCase("false");
+    }
+
+    /**
+     * 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.
+     * <p>
+     * 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():
+     * <p>
+     * source.accept(new ClearReturnAdapter(owner, source, next)) //conforms to thread-safe model
+     * <p>
+     * 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,118 @@
+/*
+ *
+ *  * Copyright 2025 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.weave.utils;
+
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.AbstractInsnNode;
+import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.InsnNode;
+import org.objectweb.asm.tree.MethodNode;
+import org.objectweb.asm.tree.VarInsnNode;
+import org.objectweb.asm.tree.analysis.Analyzer;
+import org.objectweb.asm.tree.analysis.AnalyzerException;
+import org.objectweb.asm.tree.analysis.BasicInterpreter;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
+
+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.
+            //This is a silent failure that indicates that the source (client) bytecode is invalid.
+            //This should never happen!
+            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;
+    }
+
+}