feat(asm): Clean up DelayCertainInsMethodVisitor

Author: PerfectSlayer

dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty9/DelayCertainInsMethodVisitor.java

@@ -2,30 +2,40 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Function;
 import net.bytebuddy.jar.asm.AnnotationVisitor;
 import net.bytebuddy.jar.asm.Handle;
 import net.bytebuddy.jar.asm.Label;
 import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.jar.asm.Opcodes;
 import net.bytebuddy.jar.asm.TypePath;
 
+/**
+ * This method visitor delays the following instructions:
+ *
+ * <ul>
+ *   <li>Local variable instruction: {@code ALOAD},</li>
+ *   <li>field instructions: {@code GETSTATIC}, {@code GETFIELD}</li>
+ *   <li>method instructions: {@code INVOKEVIRTUAL}</li>
+ * </ul>
+ *
+ * They can be queried using {@link #transferVisitations()} and manually commited using {@link #commitVisitations(List)}.
+ */
 public class DelayCertainInsMethodVisitor extends MethodVisitor {
-  private final List<Function> heldVisitations = new ArrayList();
+  private final List<Runnable> heldVisitations = new ArrayList<>();
 
   public DelayCertainInsMethodVisitor(int api, MethodVisitor methodVisitor) {
     super(api, methodVisitor);
   }
 
-  public void commitVisitations(List<Function> heldVisitations) {
-    for (Function fun : heldVisitations) {
-      fun.apply(null);
+  public void commitVisitations(List<Runnable> heldVisitations) {
+    for (Runnable r : heldVisitations) {
+      r.run();
     }
     heldVisitations.clear();
   }
 
-  public List<Function> transferVisitations() {
-    ArrayList<Function> copy = new ArrayList<>(this.heldVisitations);
+  public List<Runnable> transferVisitations() {
+    ArrayList<Runnable> copy = new ArrayList<>(this.heldVisitations);
     this.heldVisitations.clear();
     return copy;
   }
@@ -207,7 +217,7 @@ public void visitEnd() {
     super.visitEnd();
   }
 
-  public class InvokeDynamicInsn implements Function {
+  public class InvokeDynamicInsn implements Runnable {
     public final String name;
     public final String descriptor;
     public final Handle bootstrapMethodHandle;
@@ -225,13 +235,12 @@ public InvokeDynamicInsn(
     }
 
     @Override
-    public Object apply(Object input) {
+    public void run() {
       mv.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments);
-      return null;
     }
   }
 
-  public class VirtualMethodInsn implements Function {
+  public class VirtualMethodInsn implements Runnable {
     public final int opcode;
     public final String owner;
     public final String name;
@@ -248,13 +257,12 @@ public VirtualMethodInsn(
     }
 
     @Override
-    public Object apply(Object input) {
+    public void run() {
       mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
-      return null;
     }
   }
 
-  public class GetStaticFieldInsn implements Function {
+  public class GetStaticFieldInsn implements Runnable {
     public final int opcode;
     public final String owner;
     public final String name;
@@ -268,13 +276,12 @@ public GetStaticFieldInsn(int opcode, String owner, String name, String descript
     }
 
     @Override
-    public Object apply(Object input) {
+    public void run() {
       mv.visitFieldInsn(opcode, owner, name, descriptor);
-      return null;
     }
   }
 
-  public class GetFieldInsn implements Function {
+  public class GetFieldInsn implements Runnable {
     public final int opcode;
     public final String owner;
     public final String name;
@@ -288,13 +295,12 @@ public GetFieldInsn(int opcode, String owner, String name, String descriptor) {
     }
 
     @Override
-    public Object apply(Object input) {
+    public void run() {
       mv.visitFieldInsn(opcode, owner, name, descriptor);
-      return null;
     }
   }
 
-  public class ALoadVarInsn implements Function {
+  public class ALoadVarInsn implements Runnable {
     public final int opcode;
     public final int varIndex;
 
@@ -304,9 +310,8 @@ public ALoadVarInsn(int opcode, int varIndex) {
     }
 
     @Override
-    public Object apply(Object input) {
+    public void run() {
       mv.visitVarInsn(opcode, varIndex);
-      return null;
     }
   }
 }

dd-java-agent/instrumentation/jetty-common/src/main/java/datadog/trace/instrumentation/jetty9/HandleVisitor.java

@@ -141,7 +141,7 @@ public void visitMethodInsn(
         && descriptor.equals("(Lorg/eclipse/jetty/server/HttpChannel;)V")) {
       debug("handle bytecode found");
       DelayCertainInsMethodVisitor mv = delayVisitorDelegate();
-      List<Function> savedVisitations = mv.transferVisitations();
+      List<Runnable> savedVisitations = mv.transferVisitations();
       /*
        * Saved visitations should be for:
        *
@@ -156,8 +156,9 @@ public void visitMethodInsn(
         return;
       }
 
+      // Declare label to insert after Server.handle() call
       Label afterHandle = new Label();
-
+      // Inject blocking helper call
       super.visitVarInsn(Opcodes.ALOAD, 0);
       super.visitMethodInsn(
           Opcodes.INVOKEVIRTUAL,
@@ -181,10 +182,12 @@ public void visitMethodInsn(
               + Type.getDescriptor(Context.class)
               + ")Z",
           false);
+      // Inject jump to after Server.handle() call if blocked
       super.visitJumpInsn(Opcodes.IFNE, afterHandle);
-
+      // Inject getServer() and Server.handle() calls
       mv.commitVisitations(savedVisitations);
       super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
+      // Inject label after Server.handle() call to jump here when blocked
       super.visitLabel(afterHandle);
       super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
       debug("handle bytecode injected");
@@ -200,7 +203,7 @@ public void visitMethodInsn(
                 "(Ljakarta/servlet/DispatcherType;Lorg/eclipse/jetty/server/HttpChannel$Dispatchable;)V"))) {
 
       DelayCertainInsMethodVisitor mv = delayVisitorDelegate();
-      List<Function> savedVisitations = mv.transferVisitations();
+      List<Runnable> savedVisitations = mv.transferVisitations();
 
       // check that we've queued up what we're supposed to
       if (!checkDispatchMethodState(savedVisitations)) {
@@ -237,7 +240,7 @@ public void visitMethodInsn(
       super.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
       // dispatch with a Dispatchable created from JettyBlockingHelper::block
       // first set up the first two arguments to dispatch (this and DispatcherType.REQUEST)
-      List<Function> loadThisAndEnum = new ArrayList<>(savedVisitations.subList(0, 2));
+      List<Runnable> loadThisAndEnum = new ArrayList<>(savedVisitations.subList(0, 2));
       mv.commitVisitations(loadThisAndEnum);
       // set up the arguments to the method underlying the lambda (Request, Response,
       // RequestBlockingAction, AgentSpan)
@@ -309,7 +312,7 @@ public void visitMethodInsn(
     super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
   }
 
-  private boolean checkDispatchMethodState(final List<Function> savedVisitations) {
+  private boolean checkDispatchMethodState(final List<Runnable> savedVisitations) {
     if (savedVisitations.size() != 4) {
       return false;
     }
@@ -329,7 +332,7 @@ private boolean checkDispatchMethodState(final List<Function> savedVisitations)
       return false;
     }
 
-    final Function last = savedVisitations.get(3);
+    final Runnable last = savedVisitations.get(3);
     // jetty < 11.16.0
     //  this.dispatch(DispatcherType.REQUEST, () -> { ... });
     if (last instanceof DelayCertainInsMethodVisitor.InvokeDynamicInsn) {