refactor(agent): Refactor Hook.apply and ClassFileTransformer

dividedmind · dividedmind · commit db459a4c9961 · 2026-01-15T04:19:41.000+01:00
Refactor `Hook.apply` to return a `Set&lt;ApplicationAction&gt;` indicating
whether the method was marked for ByteBuddy instrumentation or
instrumented by Javassist. This allows `ClassFileTransformer` to
conditionally apply the `AppMapInstrumented` annotation only when
ByteBuddy instrumentation is actually needed.

Additionally, add improved logging for Javassist instrumentation failures
and guard against excessive logging of these exceptions.
diff --git a/agent/src/main/java/com/appland/appmap/transform/ClassFileTransformer.java b/agent/src/main/java/com/appland/appmap/transform/ClassFileTransformer.java
@@ -189,47 +189,24 @@ private void processClass(CtClass ctClass) {
     }
   }
 
-  private boolean applyHooks(CtBehavior behavior) {
-    boolean traceClass = tracePrefix == null || behavior.getDeclaringClass().getName().startsWith(tracePrefix);
-
+  private Set<Hook.ApplicationAction> applyHooks(CtBehavior behavior, boolean traceClass) {
     try {
       List<HookSite> hookSites = getHookSites(behavior);
 
       if (hookSites == null) {
         if (traceClass) {
           logger.trace("no hook sites");
         }
-        return false;
+        return java.util.Collections.emptySet();
       }
 
-      Hook.apply(behavior, hookSites);
-
-      if (logger.isDebugEnabled()) {
-        for (HookSite hookSite : hookSites) {
-          final Hook hook = hookSite.getHook();
-          String className = behavior.getDeclaringClass().getName();
-          if (tracePrefix != null && !className.startsWith(tracePrefix)) {
-            continue;
-          }
-
-          if (traceClass) {
-            logger.trace("hooked {}.{}{} on ({},{}) with {}",
-                className,
-                behavior.getName(),
-                behavior.getMethodInfo().getDescriptor(),
-                hook.getMethodEvent().getEventString(),
-                hook.getPosition(),
-                hook);
-          }
-        }
-      }
-      return true;
+      return Hook.apply(behavior, hookSites, traceClass);
 
     } catch (NoSourceAvailableException e) {
       Logger.println(e);
     }
 
-    return false;
+    return java.util.Collections.emptySet();
   }
 
   public List<HookSite> getHookSites(CtBehavior behavior) {
@@ -292,7 +269,7 @@ public byte[] transform(ClassLoader loader,
       try {
         ClassPool classPool = AppMapClassPool.get();
         if (traceClass) {
-          logger.debug("className: {}", className);
+          logger.trace("className: {}", className);
         }
 
         ctClass = classPool.makeClass(new ByteArrayInputStream(bytes));
@@ -317,7 +294,8 @@ public byte[] transform(ClassLoader loader,
         return null;
       }
 
-      boolean hookApplied = false;
+      boolean bytecodeModified = false;
+      boolean needsByteBuddy = false;
       for (CtBehavior behavior : ctClass.getDeclaredBehaviors()) {
         if (traceClass) {
           logger.trace("behavior: {}", behavior.getLongName());
@@ -331,24 +309,27 @@ public byte[] transform(ClassLoader loader,
         }
 
         methodsExamined++;
-        if (this.applyHooks(behavior)) {
-          hookApplied = true;
+        Set<Hook.ApplicationAction> actions = this.applyHooks(behavior, traceClass);
+        if (!actions.isEmpty()) {
+          bytecodeModified = true;
           methodsHooked++;
+          if (actions.contains(Hook.ApplicationAction.MARKED)) {
+            needsByteBuddy = true;
+          }
         }
       }
 
-      if (hookApplied) {
-        // One or more of the methods in the the class were marked as needing to
+      if (bytecodeModified) {
+        // One or more of the methods in the class were marked as needing to
         // be instrumented. Mark the class so the bytebuddy transformer will
         // know it needs to be instrumented.
-        ClassFile classFile = ctClass.getClassFile();
-        ConstPool constPool = classFile.getConstPool();
-        Annotation annot = new Annotation(AppMapInstrumented.class.getName(), constPool);
-        AnnotationUtil.setAnnotation(new AnnotatedClass(ctClass), annot);
-
-        if (traceClass) {
-          logger.trace("hooks applied to {}", className);
+        if (needsByteBuddy) {
+          ClassFile classFile = ctClass.getClassFile();
+          ConstPool constPool = classFile.getConstPool();
+          Annotation annot = new Annotation(AppMapInstrumented.class.getName(), constPool);
+          AnnotationUtil.setAnnotation(new AnnotatedClass(ctClass), annot);
         }
+
         if (logger.isDebugEnabled()) {
           packagesHooked.compute(ctClass.getPackageName(), (k, v) -> v == null ? 1 : v + 1);
         }
diff --git a/agent/src/main/java/com/appland/appmap/transform/annotations/Hook.java b/agent/src/main/java/com/appland/appmap/transform/annotations/Hook.java
@@ -2,6 +2,7 @@
 
 import java.util.Arrays;
 import java.util.Comparator;
+import java.util.EnumSet;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -107,25 +108,72 @@ public HookSite prepare(CtBehavior targetBehavior, Map<String, Object> hookConte
     return new HookSite(this, behaviorOrdinal, binding);
   }
 
-  public static void apply(CtBehavior targetBehavior, List<HookSite> hookSites) {
+  public enum ApplicationAction {
+    MARKED,
+    INSTRUMENTED
+  }
+
+  // We only log the first exception to avoid flooding the logs at the debug level.
+  // Note this variable is not thread safe, but this is okay; the worst that can happen is
+  // that we log more than one exception in a multi-threaded scenario.
+  private static boolean exceptionLogged = false;
+
+  public static Set<ApplicationAction> apply(CtBehavior targetBehavior, List<HookSite> hookSites, boolean traceClass) {
+    Set<ApplicationAction> actions = EnumSet.noneOf(ApplicationAction.class);
     MethodInfo methodInfo = targetBehavior.getMethodInfo();
     AnnotationsAttribute attr =
         (AnnotationsAttribute)methodInfo.getAttribute(AnnotationsAttribute.visibleTag);
 
-    // If the behavior is marked as an app method, update the annotation with
-    // the behavior ordinals so the bytebuddy transformer can instrument it.
-    if (attr.getAnnotation(AppMapAppMethod.class.getName()) != null) {
-      setBehaviorOrdinals(targetBehavior, hookSites);
-    }
+    if (attr != null) {
+      // If the behavior is marked as an app method, update the annotation with
+      // the behavior ordinals so the bytebuddy transformer can instrument it.
+      if (attr.getAnnotation(AppMapAppMethod.class.getName()) != null) {
+        setBehaviorOrdinals(targetBehavior, hookSites);
+        actions.add(ApplicationAction.MARKED);
+        if (traceClass) {
+          logger.debug("tracing {}.{}{}",
+              targetBehavior.getDeclaringClass().getName(),
+              targetBehavior.getName(),
+              targetBehavior.getMethodInfo().getDescriptor());
+        }
+      }
 
-    // If it's (also) marked as an agent method, it needs to be instrumented
-    // by javassist.
-    if (attr.getAnnotation(AppMapAgentMethod.class.getName()) != null) {
-      instrument(targetBehavior, hookSites);
+      // If it's (also) marked as an agent method, it needs to be instrumented
+      // by javassist.
+      if (attr.getAnnotation(AppMapAgentMethod.class.getName()) != null) {
+        try {
+          instrument(targetBehavior, hookSites);
+          actions.add(ApplicationAction.INSTRUMENTED);
+          if (traceClass) {
+            String hooks = hookSites.stream()
+                .map(h -> h.getHook().toString())
+                .collect(Collectors.joining(", "));
+            logger.debug("{}.{}{} instrumented with hooks: {}",
+                targetBehavior.getDeclaringClass().getName(),
+                targetBehavior.getName(),
+                targetBehavior.getMethodInfo().getDescriptor(),
+                hooks);
+          }
+        } catch (CannotCompileException | NotFoundException e) {
+          String msg = String.format("failed to instrument %s.%s: %s",
+              targetBehavior.getDeclaringClass().getName(), targetBehavior.getName(), e.getMessage());
+          if (!exceptionLogged) {
+            logger.debug(e, msg);
+            exceptionLogged = true;
+          } else {
+            // Log at trace level after the first one to avoid flooding the debug logs
+            logger.trace(e, msg);
+            logger.debug(msg);
+          }
+        }
+      }
     }
+
+    return actions;
   }
 
-  public static void instrument(CtBehavior targetBehavior, List<HookSite> hookSites) {
+  public static void instrument(CtBehavior targetBehavior, List<HookSite> hookSites)
+      throws CannotCompileException, NotFoundException {
     final CtClass returnType = getReturnType(targetBehavior);
     final Boolean returnsVoid = (returnType == CtClass.voidType);
 
@@ -150,44 +198,36 @@ public static void instrument(CtBehavior targetBehavior, List<HookSite> hookSite
 
     }
 
-    try {
-      String beforeSrcBlock = beforeSrcBlock(uniqueLocks.toString(),
-          invocations[MethodEvent.METHOD_INVOCATION.getIndex()]);
-      logger.trace("{}: beforeSrcBlock:\n{}", targetBehavior::getName, beforeSrcBlock::toString);
-      targetBehavior.insertBefore(
-          beforeSrcBlock);
-
-      String afterSrcBlock = afterSrcBlock(invocations[MethodEvent.METHOD_RETURN.getIndex()]);
-      logger.trace("{}: afterSrcBlock:\n{}", targetBehavior::getName, afterSrcBlock::toString);
-
-      targetBehavior.insertAfter(
-          afterSrcBlock);
-
-      ClassPool cp = AppMapClassPool.get();
-      String exitEarlyCatchSrc = "{com.appland.appmap.process.ThreadLock.current().exit();return;}";
-      if (returnsVoid) {
-        targetBehavior.addCatch(exitEarlyCatchSrc,
-            cp.get("com.appland.appmap.process.ExitEarly"));
-      } else if (!returnType.isPrimitive()) {
-        exitEarlyCatchSrc = "{com.appland.appmap.process.ThreadLock.current().exit();return("
-            + returnType.getName() + ")$e.getReturnValue();}";
-        targetBehavior
-            .addCatch(exitEarlyCatchSrc, cp.get("com.appland.appmap.process.ExitEarly"));
-      }
-      logger.trace("{}: catch1Src:\n{}", targetBehavior::getName, exitEarlyCatchSrc::toString);
-
-      String catchSrcBlock = catchSrcBlock(invocations[MethodEvent.METHOD_EXCEPTION.getIndex()]);
-      targetBehavior.addCatch(
-          catchSrcBlock,
-          cp.get("java.lang.Throwable"));
-      logger.trace("{}: catchSrcBlock:\n{}", targetBehavior::getName, catchSrcBlock::toString);
-
-    } catch (CannotCompileException e) {
-      logger.debug(e, "failed to compile {}.{}", targetBehavior.getDeclaringClass().getName(),
-          targetBehavior.getName());
-    } catch (NotFoundException e) {
-      logger.debug(e);
+    String beforeSrcBlock = beforeSrcBlock(uniqueLocks.toString(),
+        invocations[MethodEvent.METHOD_INVOCATION.getIndex()]);
+    logger.trace("{}: beforeSrcBlock:\n{}", targetBehavior::getName, beforeSrcBlock::toString);
+    targetBehavior.insertBefore(
+        beforeSrcBlock);
+
+    String afterSrcBlock = afterSrcBlock(invocations[MethodEvent.METHOD_RETURN.getIndex()]);
+    logger.trace("{}: afterSrcBlock:\n{}", targetBehavior::getName, afterSrcBlock::toString);
+
+    targetBehavior.insertAfter(
+        afterSrcBlock);
+
+    ClassPool cp = AppMapClassPool.get();
+    String exitEarlyCatchSrc = "{com.appland.appmap.process.ThreadLock.current().exit();return;}";
+    if (returnsVoid) {
+      targetBehavior.addCatch(exitEarlyCatchSrc,
+          cp.get("com.appland.appmap.process.ExitEarly"));
+    } else if (!returnType.isPrimitive()) {
+      exitEarlyCatchSrc = "{com.appland.appmap.process.ThreadLock.current().exit();return("
+          + returnType.getName() + ")$e.getReturnValue();}";
+      targetBehavior
+          .addCatch(exitEarlyCatchSrc, cp.get("com.appland.appmap.process.ExitEarly"));
     }
+    logger.trace("{}: catch1Src:\n{}", targetBehavior::getName, exitEarlyCatchSrc::toString);
+
+    String catchSrcBlock = catchSrcBlock(invocations[MethodEvent.METHOD_EXCEPTION.getIndex()]);
+    targetBehavior.addCatch(
+        catchSrcBlock,
+        cp.get("java.lang.Throwable"));
+    logger.trace("{}: catchSrcBlock:\n{}", targetBehavior::getName, catchSrcBlock::toString);
   }
 
   private static void setBehaviorOrdinals(CtBehavior behavior,
