support hook parameters @returned and @Thrown

promagent-dist/promagent-api/src/main/java/io/promagent/annotations/Returned.java

@@ -0,0 +1,45 @@
+// Copyright 2017 The Promagent Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package io.promagent.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * A parameter annotated with @Returned can be used in a Hook's @After method to capture the return value of the instrumented method.
+ *
+ * Example: In order to instrument the following method:
+ *
+ * <pre>
+ *     int sum(int a, int b) {...}
+ * </pre>
+ *
+ * A Hook could use an @After method like this:
+ *
+ * <pre>
+ *    {@literal @}After(method = "sum")
+ *     void after(int a, int b, @Returned int sum) {...}
+ * </pre>
+ *
+ * The parameter annotated with @Returned is optional, if the hook does not use the return value, the parameter can be omitted.
+ * <p/>
+ * If the instrumented method terminates exceptionally, the type's default value is assigned to the parameter,
+ * i.e. {@code 0} for numeric types and {@code null} for reference types.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.PARAMETER)
+public @interface Returned {}

promagent-dist/promagent-api/src/main/java/io/promagent/annotations/Thrown.java

@@ -0,0 +1,47 @@
+// Copyright 2017 The Promagent Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package io.promagent.annotations;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * A parameter annotated with @Thrown can be used in a Hook's @After method to capture an Exception thrown in an instrumented method.
+ *
+ * Example: In order to instrument the following method:
+ *
+ * <pre>
+ *     int div(int a, int b) {
+ *         return a / b;
+ *     }
+ * </pre>
+ *
+ * A Hook could use an @After method like this:
+ *
+ * <pre>
+ *    {@literal @}After(method = "div")
+ *     void after(int a, int b, @Returned int result, @Thrown Throwable exception) {...}
+ * </pre>
+ *
+ * In case everything goes well, the {@code result} will be the return value of {@code div()}, and {@code excption} will be {@code null}.
+ * If {@code b} is {@code 0}, the {@code result} will be {@code 0}, and {@code exception} will be an {@link ArithmeticException} (division by zero).
+ * <p/>
+ * The parameter annotated with @Thrown is optional, if the hook does not use the exception, the parameter can be omitted.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.PARAMETER)
+public @interface Thrown {}

promagent-dist/promagent-internal/src/main/java/io/promagent/internal/Delegator.java

@@ -17,6 +17,8 @@
 import io.promagent.agent.ClassLoaderCache;
 import io.promagent.annotations.After;
 import io.promagent.annotations.Before;
+import io.promagent.annotations.Returned;
+import io.promagent.annotations.Thrown;
 import io.promagent.hookcontext.MetricsStore;
 
 import java.lang.annotation.Annotation;
@@ -46,7 +48,7 @@ private Delegator(SortedSet<HookMetadata> hookMetadata, MetricsStore metricsStor
         this.threadLocal = ThreadLocal.withInitial(HashMap::new);
     }
 
-    static void init(SortedSet<HookMetadata> hookMetadata, MetricsStore metricsStore, ClassLoaderCache classLoaderCache) {
+    public static void init(SortedSet<HookMetadata> hookMetadata, MetricsStore metricsStore, ClassLoaderCache classLoaderCache) {
         instance = new Delegator(hookMetadata, metricsStore, classLoaderCache);
     }
 
@@ -68,33 +70,21 @@ private List<HookInstance> doBefore(Object that, Method interceptedMethod, Objec
     /**
      * Should be called from the Advice's @OnMethodExit method. First parameter is the list of hooks returned by before()
      */
-    public static void after(List<HookInstance> hookInstances, Method interceptedMethod, Object[] args) {
-        instance.doAfter(hookInstances, interceptedMethod, args);
+    public static void after(List<HookInstance> hookInstances, Method interceptedMethod, Object[] args, Object returned, Throwable thrown) {
+        instance.doAfter(hookInstances, interceptedMethod, args, returned, thrown);
     }
 
-    private void doAfter(List<HookInstance> hookInstances, Method interceptedMethod, Object[] args) {
+    private void doAfter(List<HookInstance> hookInstances, Method interceptedMethod, Object[] args, Object returned, Throwable thrown) {
         if (hookInstances != null) {
             for (HookInstance hookInstance : hookInstances) {
-                invokeAfter(hookInstance.getInstance(), interceptedMethod, args);
+                invokeAfter(hookInstance.getInstance(), interceptedMethod, args, returned, thrown);
                 if (!hookInstance.isRecursiveCall()) {
                     threadLocal.get().remove(hookInstance.getInstance().getClass());
                 }
             }
         }
     }
 
-    /**
-     * Take an existing hook instance from a thread local or create a new one.
-     * Hook classes must satisfy the following criteria:
-     * <ul>
-     *     <li>that.getClass() is assignable to the value of the Hook's instruments annotation
-     *     <li>The name of the instrumented method and the number of arguments match.
-     * </ul>
-     * The result may still contain hooks that don't match. This happens if the Hook method differs
-     * only in the argument types of the intercepted method. However, these Hooks will
-     * be ignored when calling {@link #invokeBefore(Object, Method, Object...)}
-     * and {@link #invokeAfter(Object, Method, Object...)}, so it's ok to include them here.
-     */
     private List<HookInstance> loadFromThreadLocalOrCreate(Object that, Method interceptedMethod) {
         return hookMetadata.stream()
                 .filter(hook -> classOrInterfaceMatches(that.getClass(), hook))
@@ -179,11 +169,12 @@ private static List<String> getMethodNames(Annotation annotation) throws HookExc
 
     // TODO: We could extend this to find the "closest" match, like in Java method calls.
     private static boolean parameterTypesMatch(Method hookMethod, Method interceptedMethod) {
-        if (hookMethod.getParameterCount() != interceptedMethod.getParameterCount()) {
+        List<Class<?>> hookParameterTypes = stripReturnedAndThrown(hookMethod);
+        if (hookParameterTypes.size() != interceptedMethod.getParameterCount()) {
             return false;
         }
-        for (int i = 0; i < hookMethod.getParameterCount(); i++) {
-            Class<?> hookParam = hookMethod.getParameterTypes()[i];
+        for (int i = 0; i < hookParameterTypes.size(); i++) {
+            Class<?> hookParam = hookParameterTypes.get(i);
             Class<?> interceptedParam = interceptedMethod.getParameterTypes()[i];
             if (!hookParam.equals(interceptedParam)) {
                 return false;
@@ -192,6 +183,23 @@ private static boolean parameterTypesMatch(Method hookMethod, Method intercepted
         return true;
     }
 
+    private static List<Class<?>> stripReturnedAndThrown(Method hookMethod) {
+        Class<?>[] allTypes = hookMethod.getParameterTypes();
+        Annotation[][] annotations = hookMethod.getParameterAnnotations();
+        if (allTypes.length != annotations.length) {
+            throw new HookException("Method.getParameterAnnotations() returned an unexpected value. This is a bug in promagent.");
+        }
+        List<Class<?>> result = new ArrayList<>();
+        for (int i=0; i<allTypes.length; i++) {
+            if (Arrays.stream(annotations[i])
+                    .map(Annotation::annotationType)
+                    .noneMatch(a -> Returned.class.equals(a) || Thrown.class.equals(a))) {
+                result.add(allTypes[i]);
+            }
+        }
+        return result;
+    }
+
     private HookInstance loadFromTheadLocalOrCreate(Class<?> hookClass) {
         Object existingHookInstance = threadLocal.get().get(hookClass);
         if (existingHookInstance != null) {
@@ -213,18 +221,18 @@ private HookInstance loadFromTheadLocalOrCreate(Class<?> hookClass) {
     /**
      * Invoke the matching Hook methods annotated with @Before
      */
-    private static void invokeBefore(Object hookInstance, Method interceptedMethod, Object... args) throws HookException {
-        invoke(Before.class, hookInstance, interceptedMethod, args);
+    private static void invokeBefore(Object hookInstance, Method interceptedMethod, Object[] args) throws HookException {
+        invoke(Before.class, hookInstance, interceptedMethod, args, null, null);
     }
 
     /**
      * Invoke the matching Hook methods annotated with @After
      */
-    private static void invokeAfter(Object hookInstance, Method interceptedMethod, Object... args) throws HookException {
-        invoke(After.class, hookInstance, interceptedMethod, args);
+    private static void invokeAfter(Object hookInstance, Method interceptedMethod, Object[] args, Object returned, Throwable thrown) throws HookException {
+        invoke(After.class, hookInstance, interceptedMethod, args, returned, thrown);
     }
 
-    private static void invoke(Class<? extends Annotation> annotation, Object hookInstance, Method interceptedMethod, Object... args) throws HookException {
+    private static void invoke(Class<? extends Annotation> annotation, Object hookInstance, Method interceptedMethod, Object[] args, Object returned, Throwable thrown) throws HookException {
         if (args.length != interceptedMethod.getParameterCount()) {
             throw new IllegalArgumentException("Number of provided arguments is " + args.length + ", but interceptedMethod expects " + interceptedMethod.getParameterCount() + " argument(s).");
         }
@@ -233,10 +241,30 @@ private static void invoke(Class<? extends Annotation> annotation, Object hookIn
                 if (!method.isAccessible()) {
                     method.setAccessible(true);
                 }
-                method.invoke(hookInstance, args);
+                method.invoke(hookInstance, addReturnedAndThrownArgs(method, args, returned, thrown));
             } catch (Exception e) {
                 throw new HookException("Failed to call " + method.getName() + "() on " + hookInstance.getClass().getSimpleName() + ": " + e.getMessage(), e);
             }
         }
     }
+
+    private static Object[] addReturnedAndThrownArgs(Method hookMethod, Object[] args, Object returned, Throwable thrown) {
+        Annotation[][] annotations = hookMethod.getParameterAnnotations();
+        List<Object> result = new ArrayList<>();
+        int arg = 0;
+        for (Annotation[] annotation : annotations) {
+            if (Arrays.stream(annotation)
+                    .map(Annotation::annotationType)
+                    .anyMatch(Returned.class::equals)) {
+                result.add(returned);
+            } else if (Arrays.stream(annotation)
+                    .map(Annotation::annotationType)
+                    .anyMatch(Thrown.class::equals)) {
+                result.add(thrown);
+            } else {
+                result.add(args[arg++]);
+            }
+        }
+        return result.toArray();
+    }
 }

promagent-dist/promagent-internal/src/main/java/io/promagent/internal/HookMetadataParser.java

@@ -16,6 +16,8 @@
 
 import io.promagent.annotations.After;
 import io.promagent.annotations.Before;
+import io.promagent.annotations.Returned;
+import io.promagent.annotations.Thrown;
 import io.promagent.internal.HookMetadata.MethodSignature;
 import net.bytebuddy.jar.asm.*;
 import org.apache.commons.io.IOUtils;
@@ -41,11 +43,11 @@
  * classes that are not available in the agent's premain phase. We need to parse the metadata without instantiating
  * the hook classes.
  */
-class HookMetadataParser {
+public class HookMetadataParser {
 
     private final SortedSet<Path> hookJars;
 
-    HookMetadataParser(Collection<Path> hookJars) {
+    public HookMetadataParser(Collection<Path> hookJars) {
         this.hookJars = Collections.unmodifiableSortedSet(new TreeSet<>(hookJars));
     }
 
@@ -85,7 +87,7 @@ SortedSet<HookMetadata> parse() throws IOException, ClassNotFoundException {
      *
      * The classNameFilter is used to parse only specific classes from the JAR files.
      */
-    SortedSet<HookMetadata> parse(Predicate<String> classNameFilter) throws IOException, ClassNotFoundException {
+    public SortedSet<HookMetadata> parse(Predicate<String> classNameFilter) throws IOException, ClassNotFoundException {
         SortedSet<HookMetadata> result = new TreeSet<>();
         for (String className : listAllJavaClasses(hookJars, classNameFilter)) {
             byte[] binaryRepresentation = readBinaryRepresentation(className);
@@ -103,8 +105,19 @@ public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
 
                 @Override
                 public MethodVisitor visitMethod(int i, String method, String desc, String signature, String[] strings) {
-                    MethodSignatureBuilder builder = hookMetadata.newMethodSignature(Arrays.stream(Type.getArgumentTypes(desc)).map(Type::getClassName).collect(Collectors.toList()));
+                    List<String> parameterTypes = Arrays.stream(Type.getArgumentTypes(desc))
+                            .map(Type::getClassName)
+                            .collect(Collectors.toList());
+                    MethodSignatureBuilder builder = hookMetadata.newMethodSignature(parameterTypes);
                     return new MethodVisitor(Opcodes.ASM5, super.visitMethod(i, method, desc, signature, strings)) {
+                        @Override
+                        public AnnotationVisitor visitParameterAnnotation(int parameter, String desc, boolean visible) {
+                            if (visible && typeEquals(desc, Returned.class, Thrown.class)) {
+                                builder.markReturnedOrThrown(parameter);
+                            }
+                            return super.visitParameterAnnotation(parameter, desc, visible);
+                        }
+
                         @Override
                         public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
                             if (visible && typeEquals(desc, Before.class, After.class)) {
@@ -214,24 +227,43 @@ private boolean isComplete() {
      */
     private static class MethodSignatureBuilder {
 
+        private static class ParameterType {
+            final String type;
+            boolean isReturnedOrThrown; // method parameters annotated with @Returned or @Thrown will be ignored.
+
+            private ParameterType(String type) {
+                this.type = type;
+                this.isReturnedOrThrown = false;
+            }
+        }
+
         SortedSet<String> methodNames = new TreeSet<>();
-        List<String> parameterTypes = new ArrayList<>();
+        List<ParameterType> parameterTypes = new ArrayList<>();
 
         private void addMethodName(String methodName) {
             methodNames.add(methodName);
         }
 
         private void addParameterType(String parameterType) {
-            parameterTypes.add(parameterType);
+            parameterTypes.add(new ParameterType(parameterType));
         }
 
         private SortedSet<MethodSignature> build() {
+            List<String> strippedParameterTypes = parameterTypes.stream()
+                    .filter(p -> ! p.isReturnedOrThrown)
+                    .map(p -> p.type)
+                    .collect(Collectors.toList());
             SortedSet<MethodSignature> result = new TreeSet<>();
             for (String methodName : methodNames) {
-                result.add(new MethodSignature(methodName, parameterTypes));
+                result.add(new MethodSignature(methodName, strippedParameterTypes));
             }
             return result;
         }
+
+        public void markReturnedOrThrown(int parameter) {
+            // We know that parameter is a valid index in parameterTypes.
+            parameterTypes.get(parameter).isReturnedOrThrown = true;
+        }
     }
 
     /**

promagent-dist/promagent-internal/src/main/java/io/promagent/internal/PromagentAdvice.java

@@ -15,6 +15,7 @@
 package io.promagent.internal;
 
 import io.promagent.agent.ClassLoaderCache;
+import net.bytebuddy.implementation.bytecode.assign.Assigner;
 
 import java.lang.reflect.Method;
 import java.util.List;
@@ -49,21 +50,23 @@ public static List<Object> before(
         }
     }
 
-    @OnMethodExit
+    @OnMethodExit(onThrowable = Throwable.class)
     public static void after(
             @Enter List<Object> hooks,
             @This Object that,
             @Origin Method method,
-            @AllArguments Object[] args
+            @AllArguments Object[] args,
+            @Return(typing = Assigner.Typing.DYNAMIC) Object returned, // support void == null and int == Integer
+            @Thrown Throwable thrown
     ) {
         try {
             // The following code is equivalent to:
             //     Delegator.after(hooks, method, args);
             // However, the Delegator class will not be available in the context of the instrumented method,
             // so we must use our agent class loader to load the Delegator class and do the call via reflection.
             Class<?> delegator = ClassLoaderCache.getInstance().currentClassLoader().loadClass("io.promagent.internal.Delegator");
-            Method afterMethod = delegator.getMethod("after", List.class, Method.class, Object[].class);
-            afterMethod.invoke(null, hooks, method, args);
+            Method afterMethod = delegator.getMethod("after", List.class, Method.class, Object[].class, Object.class, Throwable.class);
+            afterMethod.invoke(null, hooks, method, args, returned, thrown);
         } catch (Exception e) {
             System.err.println("Error executing Prometheus hook on " + that.getClass().getSimpleName());
             e.printStackTrace();