Refactor AsyncAwaitInstrumentationAgent (move responsibilities to RuntimeBytecodeInjector were necessary)

vsilaev · vsilaev · commit f877a0e19640 · 2021-12-22T16:14:52.000+03:00
diff --git a/net.tascalate.async.examples/src/main/java/net/tascalate/async/examples/generator/GeneratorExample.java b/net.tascalate.async.examples/src/main/java/net/tascalate/async/examples/generator/GeneratorExample.java
@@ -41,6 +41,7 @@
 import java.util.concurrent.Future;
 
 import net.tascalate.async.AsyncGenerator;
+import net.tascalate.async.AsyncValue;
 import net.tascalate.async.Sequence;
 import net.tascalate.async.YieldReply;
 import net.tascalate.async.async;
@@ -97,7 +98,7 @@ CompletableFuture<String> mergeStrings(String delimeter) {
     }
     
     @async
-    CompletionStage<String> iterateStringsEx(int z) {
+    AsyncValue<String> iterateStringsEx(int z) {
         z += 2;
         System.out.println(z);
         int x = 3;
diff --git a/net.tascalate.async.tools.javaagent/pom.xml b/net.tascalate.async.tools.javaagent/pom.xml
@@ -127,6 +127,7 @@
 									<excludes>
 										<exclude>META-INF/maven/**/*.*</exclude>
 										<exclude>**/module-info.class</exclude>
+										<exclude>META-INF/MANIFEST.MF</exclude>
 									</excludes>
 								</filter>
 								<filter>
diff --git a/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/AbstractInstrumentationAgent.java b/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/AbstractInstrumentationAgent.java
@@ -44,7 +44,7 @@
 
 public abstract class AbstractInstrumentationAgent {
 
-    private final Logger log = LoggerFactory.getLogger(getClass());
+    protected final Logger log = LoggerFactory.getLogger(getClass());
     
     protected AbstractInstrumentationAgent() {
     }
diff --git a/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/AsyncAwaitInstrumentationAgent.java b/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/AsyncAwaitInstrumentationAgent.java
@@ -24,7 +24,6 @@
  */
 package net.tascalate.async.tools.instrumentation;
 
-import java.io.PrintStream;
 import java.lang.instrument.ClassFileTransformer;
 import java.lang.instrument.Instrumentation;
 import java.util.HashSet;
@@ -57,30 +56,28 @@ public static void premain(String args, Instrumentation instrumentation) throws
         System.setProperty(JavaFlowClassTransformer.class.getName(), "true");
         System.setProperty(AsyncAwaitInstrumentationAgent.class.getName(), "true");
         
-        if (getJdkVersion() < 9) {
+        int jdkVersion = getJdkVersion(); 
+        if (jdkVersion < 9) {
+            if (log.isDebugEnabled()) {
+                log.debug("JDK verion " + jdkVersion + " is less than 9, no lambda instrumentation patch required.");
+            }
             return;
         }
         
         // Need to patch lambda metafactory to support agent instrumentation for lambda classes
         // in JDK 9+
         try {
-            instrumentation.redefineClasses(RuntimeBytecodeInjector.modifyLambdaMetafactory());
-            System.setOut(new PrintStream(System.out, true) {
+            if (!RuntimeBytecodeInjector.isInjectionApplied()) {
+                log.debug("Applying lambda post-processing injection...");
+                instrumentation.redefineClasses(RuntimeBytecodeInjector.modifyLambdaMetafactory());
+                log.debug("Lambda post-processing injection is applied.");
+            } else {
+                log.warn("Lambda post-processing injection was already applied, probably by some other agent.");
+            }
+            RuntimeBytecodeInjector.installTransformer(new RuntimeBytecodeInjector.LambdaClassTransformer() {
                 @Override
-                public void println(Object o) {
-                    if (o instanceof Object[]) {
-                        Object[] params = (Object[])o;
-                        if (params.length == 4 &&
-                            RuntimeBytecodeInjector.isValidCaller(params[0]) &&    
-                            params[1] instanceof Class &&
-                            params[2] instanceof byte[] && 
-                            (params[3] == null || params[3] instanceof byte[])) {
-                            
-                            agent.enhanceLambdaClass((Class<?>)params[1], (byte[])params[2], params);
-                            return;
-                        }
-                    }
-                    super.println(o);
+                public byte[] transform(Class<?> lambdaOwningClass, byte[] lambdaClassBytes) throws Throwable {
+                    return agent.transformLambdaClass(lambdaOwningClass, lambdaClassBytes);
                 }
             });
         } catch (Throwable ex) {
@@ -113,24 +110,26 @@ protected ClassFileTransformer createRetransformableTransformer(String args, Ins
         return new AsyncAwaitClassFileTransformer(continuationsTransformer, instrumentation);
     }
     
-    void enhanceLambdaClass(Class<?> lambdaOwningClass, byte[] input, Object[] output) {
+    byte[] transformLambdaClass(Class<?> lambdaOwningClass, byte[] input) throws Throwable {
         try {
+            String className = null;            
             if (log.isDebugEnabled()) {
-                log.debug("Start transforming lambda class " + ClassEmitters.classNameOf(input) + ", defined in " + lambdaOwningClass);
+                className = ClassEmitters.classNameOf(input);
+                log.debug("Start transforming lambda class " + className + ", defined in " + lambdaOwningClass.getName());
             }
-            byte[] altered = continuationsTransformer.transform(lambdaOwningClass.getClassLoader(), 
-                                                                null, null, // Neither name, nor class 
-                                                                lambdaOwningClass.getProtectionDomain(), 
-                                                                input);
-            if (input != altered && altered != null) {
-                if (log.isDebugEnabled()) {
-                    log.debug("Done transforming lambda class " + ClassEmitters.classNameOf(input) + ", defined in " + lambdaOwningClass);
-                }                
-            }
-            output[3] = null == altered ? input : altered;
+            byte[] output = continuationsTransformer.transform(lambdaOwningClass.getClassLoader(), 
+                                                               null, null, // Neither name, nor class
+                                                               lambdaOwningClass.getProtectionDomain(), 
+                                                               input);
+            if (log.isDebugEnabled()) {
+                log.debug("Done transforming lambda class " + className + 
+                           ", defined in " + lambdaOwningClass.getName() + 
+                           ", modified = " + (input != output && output != null));
+            }                
+            return output;
         } catch (Throwable ex) {
-            ex.printStackTrace();
-            output[3] = input;
+            log.error("Error transforming lambda class defined in " + lambdaOwningClass.getName(), ex);
+            throw ex;
         }
         
     }
diff --git a/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/RuntimeBytecodeInjector.java b/net.tascalate.async.tools.javaagent/src/main/java/net/tascalate/async/tools/instrumentation/RuntimeBytecodeInjector.java
@@ -31,6 +31,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.PrintStream;
+import java.lang.annotation.Documented;
 import java.lang.instrument.ClassDefinition;
 
 import net.tascalate.asmx.ClassReader;
@@ -43,29 +44,22 @@
 class RuntimeBytecodeInjector {
     
     private static final String CLASS = "java.lang.invoke.InnerClassLambdaMetafactory";
-
-    private static ClassDefinition loadClassDefinition(Class<?> clazz) throws IOException {
-        ByteArrayOutputStream out = new ByteArrayOutputStream();
-        byte[] buff = new byte[1024];
-        try (InputStream in = Object.class.getResourceAsStream('/' + clazz.getName().replace('.', '/') + ".class")) {
-            int c = 0;
-            while ((c= in.read(buff)) > 0) {
-                out.write(buff, 0, c);
-            }
-        }
-        out.close();
-        byte[] bytes = out.toByteArray();
-        return new ClassDefinition(clazz, bytes);
-    }
     
-    private static ClassDefinition loadClassDefinition(String className) throws ClassNotFoundException, IOException {
-        return loadClassDefinition(Class.forName(className));
+    static interface LambdaClassTransformer {
+        byte[] transform(Class<?> lambdaOwningClass, byte[] lambdaClassBytes) throws Throwable;
     }
     
-    static boolean isValidCaller(Object o) {
-        return o != null && CLASS.equals(o.getClass().getName());
+    static boolean isInjectionApplied() {
+        Class<?> cls;
+        try {
+            cls = ClassLoader.getSystemClassLoader().loadClass(CLASS);
+        } catch (ClassNotFoundException ex) {
+            throw new RuntimeException(ex);
+        }
+        Documented anno = cls.getAnnotation(Documented.class);
+        return anno != null;
     }
-    
+
     static ClassDefinition modifyLambdaMetafactory() throws ClassNotFoundException, IOException {
         ClassDefinition original = loadClassDefinition(CLASS);
         ClassReader classReader = new ClassReader(new ByteArrayInputStream(original.getDefinitionClassFile()));
@@ -93,14 +87,7 @@ public void visitMethodInsn(int opcode, String owner, String name, String descri
                     super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
                     visitVarInsn(bytesType.getOpcode(ISTORE), bytesVar);
                     
-                    /*
-                    Type stringType = Type.getType(String.class);
-                    visitFieldInsn(GETSTATIC, systemType.getInternalName(), "out", printStreamType.getDescriptor());
-                    visitLdcInsn("HERE WE ARE: ");
-                    visitMethodInsn(INVOKEVIRTUAL, printStreamType.getInternalName(), "print", Type.getMethodDescriptor(Type.VOID_TYPE, stringType), false);
-                    */
-                    
-                    visitInsn(ICONST_4);
+                    visitInsn(ICONST_3);
                     visitTypeInsn(ANEWARRAY, objectType.getInternalName());
                     visitVarInsn(objectsType.getOpcode(ISTORE), params);
                     
@@ -117,30 +104,19 @@ public void visitMethodInsn(int opcode, String owner, String name, String descri
                     visitFieldInsn(GETFIELD, "java/lang/invoke/AbstractValidatingLambdaMetafactory", "targetClass", classType.getDescriptor());
                     visitInsn(AASTORE);
                     
-                    // params[2] = inBytes
+                    // params[2] = inBytes, after call replaced by outBytes
                     visitVarInsn(objectsType.getOpcode(ILOAD), params);
                     visitInsn(ICONST_2);
                     visitVarInsn(bytesType.getOpcode(ILOAD), bytesVar);
                     visitInsn(AASTORE);
                     
-                    // params[3] = inBytes, after call replaced by outBytes
-                    visitVarInsn(objectsType.getOpcode(ILOAD), params);
-                    visitInsn(ICONST_3);
-                    visitVarInsn(bytesType.getOpcode(ILOAD), bytesVar);
-                    visitInsn(AASTORE);
-                    
                     visitFieldInsn(GETSTATIC, systemType.getInternalName(), "out", printStreamType.getDescriptor());
                     visitVarInsn(objectsType.getOpcode(ILOAD), params);
-                    visitMethodInsn(INVOKEVIRTUAL, printStreamType.getInternalName(), "println", Type.getMethodDescriptor(Type.VOID_TYPE, objectType), false);
-                    
-                    /*
-                    visitFieldInsn(GETSTATIC, systemType.getInternalName(), "out", printStreamType.getDescriptor());
-                    visitMethodInsn(INVOKEVIRTUAL, printStreamType.getInternalName(), "println", Type.getMethodDescriptor(Type.VOID_TYPE), false);
-                    */
+                    visitMethodInsn(INVOKEVIRTUAL, printStreamType.getInternalName(), "print", Type.getMethodDescriptor(Type.VOID_TYPE, objectType), false);
 
-                    // get outBytes, params[3]
+                    // get outBytes, params[2]
                     visitVarInsn(objectsType.getOpcode(ILOAD), params);
-                    visitInsn(ICONST_3);
+                    visitInsn(ICONST_2);
                     visitInsn(AALOAD);
                 } else {
                     super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
@@ -149,6 +125,18 @@ public void visitMethodInsn(int opcode, String owner, String name, String descri
         };
 
         ClassVisitor cv = new ClassVisitor(ASM9, classWriter) {
+            @Override
+            public void visit(int version,
+                              int access,
+                              String name,
+                              String signature,
+                              String superName,String[] interfaces) {
+                super.visit(version, access, name, signature, superName, interfaces);
+                Type annoType = Type.getType(Documented.class);
+                visitAnnotation(annoType.getDescriptor(), true).visitEnd();
+            }
+            
+            @Override
             public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
                 MethodVisitor mv = super.visitMethod(access, name, desc, signature, exceptions);
                 if ("spinInnerClass".equals(name) || "generateInnerClass".equals(name)) {
@@ -162,4 +150,54 @@ public MethodVisitor visitMethod(int access, String name, String desc, String si
         classReader.accept(cv, ClassReader.EXPAND_FRAMES);
         return new ClassDefinition(original.getDefinitionClass(), classWriter.toByteArray());
     }
+    
+    static void installTransformer(final LambdaClassTransformer transformer) {
+        System.setOut(new PrintStream(System.out, true) {
+            @Override
+            public void print(Object o) {
+                if (o instanceof Object[]) {
+                    Object[] params = (Object[])o;
+                    if (params.length == 3 &&
+                        RuntimeBytecodeInjector.isValidCaller(params[0]) &&    
+                        params[1] instanceof Class &&
+                        params[2] instanceof byte[]) {
+                        
+                        byte[] inBytes = (byte[])params[2];
+                        byte[] outBytes = inBytes;
+                        try {
+                            outBytes = transformer.transform((Class<?>)params[1], inBytes);
+                        } catch (Throwable ex) {
+                            
+                        }
+                        params[2] = outBytes == null ? inBytes : outBytes;
+                        return;
+                    }
+                }
+                super.println(o);
+            }
+        });        
+    }
+    
+    private static boolean isValidCaller(Object o) {
+        return o != null && CLASS.equals(o.getClass().getName());
+    }
+    
+    private static ClassDefinition loadClassDefinition(Class<?> clazz) throws IOException {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        byte[] buff = new byte[1024];
+        try (InputStream in = Object.class.getResourceAsStream('/' + clazz.getName().replace('.', '/') + ".class")) {
+            int c = 0;
+            while ((c= in.read(buff)) > 0) {
+                out.write(buff, 0, c);
+            }
+        }
+        out.close();
+        byte[] bytes = out.toByteArray();
+        return new ClassDefinition(clazz, bytes);
+    }
+    
+    private static ClassDefinition loadClassDefinition(String className) throws ClassNotFoundException, IOException {
+        return loadClassDefinition(Class.forName(className));
+    }
+
 }

Original file line number	Diff line number	Diff line change
`@@ -44,7 +44,7 @@`
`44`	`44`
`45`	`45`	`public abstract class AbstractInstrumentationAgent {`
`46`	`46`
`47`		`- private final Logger log = LoggerFactory.getLogger(getClass());`
	`47`	`+ protected final Logger log = LoggerFactory.getLogger(getClass());`
`48`	`48`
`49`	`49`	`protected AbstractInstrumentationAgent() {`
`50`	`50`	`}`