Fix maxStack for computeFrames (#12054)

Fix maxStack for computeFrames

increase method maxStack before running computeFrames analysis
to avoid stack error related to previous instrumentation fromm probe
on the same method.
Change the classGeneratedFailed test using large tablelookup
to generate incorrect class instead of incorrect frame
which is now caught by computeFrames analysis.
Fix off-by-1 the newVar id.

fix oversizing max stack

set a probable max stack of 16.
detect if an analysis failed because of insufficient stack, and double
the stack up to 2* the number of instructions.
move computeFrames to ASMHelper and add test

fix newVar

Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>

Author: jpbempel

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/ASMHelper.java

@@ -19,7 +19,9 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.IdentityHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.StringJoiner;
 import java.util.stream.Collectors;
 import org.objectweb.asm.ClassWriter;
@@ -38,12 +40,21 @@
 import org.objectweb.asm.tree.MethodInsnNode;
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.TypeInsnNode;
+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 org.objectweb.asm.util.Printer;
 import org.objectweb.asm.util.Textifier;
 import org.objectweb.asm.util.TraceClassVisitor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /** Helper class for bytecode generation */
 public class ASMHelper {
+  private static final Logger LOGGER = LoggerFactory.getLogger(ASMHelper.class);
+
   public static final Type INT_TYPE = new Type(org.objectweb.asm.Type.INT_TYPE);
   public static final Type OBJECT_TYPE = new Type(Types.OBJECT_TYPE);
   public static final Type STRING_TYPE = new Type(Types.STRING_TYPE);
@@ -428,6 +439,65 @@ public static List<Integer> getLineNumbers(MethodNode methodNode) {
     return lines;
   }
 
+  /**
+   * Returns a map for all instructions what is the current frame with value on the stacks Helps to
+   * know at exit points of a method what remains on the stack
+   */
+  protected static Map<AbstractInsnNode, Frame<BasicValue>> computeFrames(
+      String owner, MethodNode methodNode) {
+    // Initial guess used for methodNode.maxStack when retrying computeFrames() below; small
+    // enough to be cheap, large enough to cover the vast majority of probe-generated methods on
+    // the first attempt.
+    final int INITIAL_MAX_STACK_GUESS = 16;
+    // Probes already applied to this method may have inserted bytecode that requires more
+    // stack than originally declared, without updating methodNode.maxStack (only maxLocals is
+    // tracked, see newVar()). The ASM analyzer sizes its Frame objects from the declared
+    // maxStack (each Frame costs maxLocals+maxStack slots), so a stale/too-small value makes
+    // the analysis fail with "Insufficient maximum stack size" even though the class itself is
+    // valid. Rather than guessing an upper bound from the instruction count (which blows up
+    // Frame memory quadratically for large methods), retry with a growing maxStack only when
+    // the Analyzer itself reports it as insufficient - this converges to close to the actual
+    // stack depth the method needs. The real maxStack is recomputed again from scratch by the
+    // ClassWriter's COMPUTE_FRAMES pass when the class is finally written.
+    Map<AbstractInsnNode, Frame<BasicValue>> frames = new IdentityHashMap<>();
+    int attemptMaxStack = Math.max(methodNode.maxStack, INITIAL_MAX_STACK_GUESS);
+    int hardLimitMaxStack = methodNode.instructions.size() * 2;
+    while (true) {
+      methodNode.maxStack = attemptMaxStack;
+      try {
+        Frame<BasicValue>[] frameArray =
+            new Analyzer<>(new BasicInterpreter()).analyze(owner, methodNode);
+        AbstractInsnNode current = methodNode.instructions.getFirst();
+        int idx = 0;
+        while (current != null) {
+          frames.put(current, frameArray[idx++]);
+          current = current.getNext();
+        }
+        break;
+      } catch (AnalyzerException ex) {
+        if (isInsufficientMaxStack(ex) && attemptMaxStack < hardLimitMaxStack) {
+          attemptMaxStack *= 2;
+          continue;
+        }
+        LOGGER.debug(
+            "Failed to analyze method[{}::{}{}] instructions",
+            owner,
+            methodNode.name,
+            methodNode.desc,
+            ex);
+        // when running tests, fails if an exception is thrown during analysis
+        // Gradle is running tests with assertions enbabled
+        assert false;
+        break;
+      }
+    }
+    return frames;
+  }
+
+  private static boolean isInsufficientMaxStack(AnalyzerException ex) {
+    return ex.getMessage() != null && ex.getMessage().contains("Insufficient maximum stack size");
+  }
+
   /** Wraps ASM's {@link org.objectweb.asm.Type} with associated generic types */
   public static class Type {
     private final org.objectweb.asm.Type mainType;

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/CapturedContextInstrumenter.java

@@ -1,5 +1,6 @@
 package com.datadog.debugger.instrumentation;
 
+import static com.datadog.debugger.instrumentation.ASMHelper.computeFrames;
 import static com.datadog.debugger.instrumentation.ASMHelper.getStatic;
 import static com.datadog.debugger.instrumentation.ASMHelper.invokeConstructor;
 import static com.datadog.debugger.instrumentation.ASMHelper.invokeStatic;

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/ExceptionInstrumenter.java

@@ -25,7 +25,8 @@ public InstrumentationResult.Status instrument() {
     // a try/catch that create a subscobe and even level method local vars are not accessible
     // in the catch clause for capture
     hoistedLocalVars = initAndHoistLocalVars(methodNode);
-    Map<AbstractInsnNode, Frame<BasicValue>> frames = computeFrames(classNode.name, methodNode);
+    Map<AbstractInsnNode, Frame<BasicValue>> frames =
+        ASMHelper.computeFrames(classNode.name, methodNode);
     processInstructions(frames); // fill returnHandlerLabel
     addFinallyHandler(methodEnterLabel, returnHandlerLabel);
     installFinallyBlocks();

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/Instrumenter.java

@@ -13,7 +13,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
-import java.util.IdentityHashMap;
 import java.util.List;
 import java.util.Map;
 import org.objectweb.asm.Opcodes;
@@ -29,17 +28,11 @@
 import org.objectweb.asm.tree.MethodNode;
 import org.objectweb.asm.tree.TryCatchBlockNode;
 import org.objectweb.asm.tree.TypeInsnNode;
-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 org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 /** Common class for generating instrumentation */
 public abstract class Instrumenter {
-  private static final Logger LOGGER = LoggerFactory.getLogger(Instrumenter.class);
   protected static final String CONSTRUCTOR_NAME = "<init>";
   protected static final String PROBEID_TAG_NAME = "debugger.probeid";
 
@@ -193,29 +186,6 @@ protected InsnList stackCleanupInsnList(
     return result;
   }
 
-  protected static Map<AbstractInsnNode, Frame<BasicValue>> computeFrames(
-      String owner, MethodNode methodNode) {
-    Map<AbstractInsnNode, Frame<BasicValue>> frames = new IdentityHashMap<>();
-    try {
-      Frame<BasicValue>[] frameArray =
-          new Analyzer<>(new BasicInterpreter()).analyze(owner, methodNode);
-      AbstractInsnNode current = methodNode.instructions.getFirst();
-      int idx = 0;
-      while (current != null) {
-        frames.put(current, frameArray[idx++]);
-        current = current.getNext();
-      }
-    } catch (AnalyzerException ex) {
-      LOGGER.debug(
-          "Failed to analyze method[{}::{}{}] instructions",
-          owner,
-          methodNode.name,
-          methodNode.desc,
-          ex);
-    }
-    return frames;
-  }
-
   protected AbstractInsnNode processInstruction(
       AbstractInsnNode node, Map<AbstractInsnNode, Frame<BasicValue>> frames) {
     switch (node.getOpcode()) {
@@ -285,13 +255,11 @@ protected void pushTags(InsnList insnList, ProbeDefinition.Tag[] tags) {
   }
 
   protected int newVar(Type type) {
-    int varId = methodNode.maxLocals + 1;
-    methodNode.maxLocals += type.getSize();
-    return varId;
+    return newVar(type.getSize());
   }
 
   protected int newVar(int size) {
-    int varId = methodNode.maxLocals + 1;
+    int varId = methodNode.maxLocals;
     methodNode.maxLocals += size;
     return varId;
   }

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/MetricInstrumenter.java

@@ -112,7 +112,7 @@ public InstrumentationResult.Status instrument() {
       case EXIT:
         {
           Map<AbstractInsnNode, Frame<BasicValue>> frames =
-              computeFrames(classNode.name, methodNode);
+              ASMHelper.computeFrames(classNode.name, methodNode);
           processInstructions(frames);
           addFinallyHandler(returnHandlerLabel);
           installFinallyBlocks();

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/instrumentation/SpanInstrumenter.java

@@ -42,7 +42,8 @@ public InstrumentationResult.Status instrument() {
       return addRangeSpan(classFileLines);
     }
     spanVar = newVar(DEBUGGER_SPAN_TYPE);
-    Map<AbstractInsnNode, Frame<BasicValue>> frames = computeFrames(classNode.name, methodNode);
+    Map<AbstractInsnNode, Frame<BasicValue>> frames =
+        ASMHelper.computeFrames(classNode.name, methodNode);
     processInstructions(frames);
     LabelNode initSpanLabel = new LabelNode();
     InsnList insnList = createSpan(initSpanLabel);

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/agent/DebuggerTransformerTest.java

@@ -55,7 +55,10 @@
 import org.junit.jupiter.api.condition.DisabledIf;
 import org.mockito.ArgumentCaptor;
 import org.objectweb.asm.Opcodes;
-import org.objectweb.asm.tree.VarInsnNode;
+import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.InsnNode;
+import org.objectweb.asm.tree.LabelNode;
+import org.objectweb.asm.tree.TableSwitchInsnNode;
 
 public class DebuggerTransformerTest {
   private static final String LANGUAGE = "java";
@@ -350,23 +353,23 @@ public void classGenerationFailed() {
             .startsWith(
                 "Instrumentation failed for "
                     + CLASS_NAME
-                    + ": java.lang.ArrayIndexOutOfBoundsException:"));
+                    + ": org.objectweb.asm.MethodTooLargeException:"));
     assertTrue(
         strCaptor
             .getAllValues()
             .get(1)
             .startsWith(
                 "Instrumentation failed for "
                     + CLASS_NAME
-                    + ": java.lang.ArrayIndexOutOfBoundsException:"));
+                    + ": org.objectweb.asm.MethodTooLargeException:"));
     assertTrue(
         strCaptor
             .getAllValues()
             .get(2)
             .startsWith(
                 "Instrumentation failed for "
                     + CLASS_NAME
-                    + ": java.lang.ArrayIndexOutOfBoundsException:"));
+                    + ": org.objectweb.asm.MethodTooLargeException:"));
   }
 
   @Test
@@ -443,11 +446,19 @@ public MockProbe(ProbeId probeId, Where where) {
     @Override
     public InstrumentationResult.Status instrument(
         MethodInfo methodInfo, List<DiagnosticMessage> diagnostics, List<Integer> probeIndices) {
-      methodInfo
-          .getMethodNode()
-          .instructions
-          .insert(
-              new VarInsnNode(Opcodes.ASTORE, methodInfo.getMethodNode().localVariables.size()));
+      // Build a single TABLESWITCH with enough case entries (4 bytes each) to push the
+      // method's bytecode past the 65535-byte limit, while keeping the ASM tree
+      // instruction *count* tiny so that dataflow analysis (computeFrames) stays cheap
+      // and succeeds; the failure should only manifest when ASM serializes the method.
+      LabelNode target = new LabelNode();
+      int high = 20_000;
+      LabelNode[] labels = new LabelNode[high + 1];
+      Arrays.fill(labels, target);
+      InsnList list = new InsnList();
+      list.add(new InsnNode(Opcodes.ICONST_0));
+      list.add(new TableSwitchInsnNode(0, high, target, labels));
+      list.add(target);
+      methodInfo.getMethodNode().instructions.insert(list);
       return InstrumentationResult.Status.INSTALLED;
     }
 

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/instrumentation/ASMHelperTest.java

@@ -11,12 +11,18 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
+import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.Type;
+import org.objectweb.asm.tree.AbstractInsnNode;
+import org.objectweb.asm.tree.InsnNode;
 import org.objectweb.asm.tree.LabelNode;
 import org.objectweb.asm.tree.LineNumberNode;
 import org.objectweb.asm.tree.LocalVariableNode;
 import org.objectweb.asm.tree.MethodNode;
+import org.objectweb.asm.tree.analysis.BasicValue;
+import org.objectweb.asm.tree.analysis.Frame;
 
 public class ASMHelperTest {
 
@@ -187,4 +193,22 @@ void getLineNumbers() {
     assertEquals(2, lineNumbers.get(1).intValue());
     assertEquals(82, lineNumbers.get(2).intValue());
   }
+
+  @Test
+  void computeFramesGrowsMaxStackWhenInsufficient() {
+    MethodNode methodNode = new MethodNode(Opcodes.ACC_STATIC, "deepLongStack", "()V", null, null);
+    methodNode.maxStack = 1; // deliberately too small for the actual stack depth needed below
+    int longCount = 20; // stack depth of 20, well above the initial guess of 16
+    for (int i = 0; i < longCount; i++) {
+      methodNode.instructions.add(new InsnNode(Opcodes.LCONST_0));
+    }
+    for (int i = 0; i < longCount; i++) {
+      methodNode.instructions.add(new InsnNode(Opcodes.POP2));
+    }
+    methodNode.instructions.add(new InsnNode(Opcodes.RETURN));
+    Map<AbstractInsnNode, Frame<BasicValue>> frames =
+        ASMHelper.computeFrames("com/datadog/debugger/FakeOwner", methodNode);
+    assertEquals(methodNode.instructions.size(), frames.size());
+    assertTrue(methodNode.maxStack >= longCount);
+  }
 }