BAEL-1728: add java instrumentation

2018-07-22 00:11:46 +03:00
parent 2fec9da2d5
commit ecb14dd834
9 changed files with 365 additions and 0 deletions
@@ -0,0 +1,70 @@
+package com.baeldung.instrumentation.agent;
+
+import javassist.CannotCompileException;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.CtMethod;
+import javassist.NotFoundException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.lang.instrument.ClassFileTransformer;
+import java.lang.instrument.IllegalClassFormatException;
+import java.security.ProtectionDomain;
+
+public class AtmTransformer implements ClassFileTransformer {
+
+    private static Logger LOGGER = LoggerFactory.getLogger(AtmTransformer.class);
+
+    private static final String WITHDRAW_MONEY_METHOD = "withdrawMoney";
+
+    /** The internal form class name of the class to transform */
+    private String targetClassName;
+    /** The class loader of the class we want to transform */
+    private ClassLoader targetClassLoader;
+
+    public AtmTransformer(String targetClassName, ClassLoader targetClassLoader) {
+        this.targetClassName = targetClassName;
+        this.targetClassLoader = targetClassLoader;
+    }
+
+    @Override
+    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
+            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
+        byte[] byteCode = classfileBuffer;
+
+        String finalTargetClassName = this.targetClassName.replaceAll("\\.", "/"); //replace . with /
+        if (!className.equals(finalTargetClassName)) {
+            return byteCode;
+        }
+
+        if (className.equals(finalTargetClassName) && loader.equals(targetClassLoader)) {
+            LOGGER.info("[Agent] Transforming class MyAtm");
+            try {
+                ClassPool cp = ClassPool.getDefault();
+                CtClass cc = cp.get(targetClassName);
+                CtMethod m = cc.getDeclaredMethod(WITHDRAW_MONEY_METHOD);
+                m.addLocalVariable("startTime", CtClass.longType);
+                m.insertBefore("startTime = System.currentTimeMillis();");
+
+                StringBuilder endBlock = new StringBuilder();
+
+                m.addLocalVariable("endTime", CtClass.longType);
+                m.addLocalVariable("opTime", CtClass.longType);
+                endBlock.append("endTime = System.currentTimeMillis();");
+                endBlock.append("opTime = (endTime-startTime)/1000;");
+
+                endBlock.append("LOGGER.info(\"[Application] Withdrawal operation completed in:\" + opTime + \" seconds!\");");
+
+                m.insertAfter(endBlock.toString());
+
+                byteCode = cc.toBytecode();
+                cc.detach();
+            } catch (NotFoundException | CannotCompileException | IOException e) {
+                LOGGER.error("Exception", e);
+            }
+        }
+        return byteCode;
+    }
+}
@@ -0,0 +1,59 @@
+package com.baeldung.instrumentation.agent;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.instrument.Instrumentation;
+
+public class MyInstrumentationAgent {
+    private static Logger LOGGER = LoggerFactory.getLogger(MyInstrumentationAgent.class);
+
+    public static void premain(String agentArgs, Instrumentation inst) {
+        LOGGER.info("[Agent] In premain method");
+
+        String className = "com.baeldung.instrumentation.application.MyAtm";
+        transformClass(className,inst);
+    }
+
+    public static void agentmain(String agentArgs, Instrumentation inst) {
+        LOGGER.info("[Agent] In agentmain method");
+
+        String className = "com.baeldung.instrumentation.application.MyAtm";
+        transformClass(className,inst);
+    }
+
+    private static void transformClass(String className, Instrumentation instrumentation) {
+        Class<?> targetCls = null;
+        ClassLoader targetClassLoader = null;
+        // see if we can get the class using forName
+        try {
+            targetCls = Class.forName(className);
+            targetClassLoader = targetCls.getClassLoader();
+            transform(targetCls, targetClassLoader, instrumentation);
+            return;
+        } catch (Exception ex) {
+            LOGGER.error("Class [{}] not found with Class.forName");
+        }
+        // otherwise iterate all loaded classes and find what we want
+        for(Class<?> clazz: instrumentation.getAllLoadedClasses()) {
+            if(clazz.getName().equals(className)) {
+                targetCls = clazz;
+                targetClassLoader = targetCls.getClassLoader();
+                transform(targetCls, targetClassLoader, instrumentation);
+                return;
+            }
+        }
+        throw new RuntimeException("Failed to find class [" + className + "]");
+    }
+
+    private static void transform(Class<?> clazz, ClassLoader classLoader, Instrumentation instrumentation) {
+        AtmTransformer dt = new AtmTransformer(clazz.getName(), classLoader);
+        instrumentation.addTransformer(dt, true);
+        try {
+            instrumentation.retransformClasses(clazz);
+        } catch (Exception ex) {
+            throw new RuntimeException("Transform failed for class: [" + clazz.getName() + "]", ex);
+        }
+    }
+
+}
@@ -0,0 +1,46 @@
+package com.baeldung.instrumentation.application;
+
+import com.sun.tools.attach.VirtualMachine;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.util.Optional;
+
+/**
+ * Created by adi on 6/10/18.
+ */
+public class AgentLoader {
+    private static Logger LOGGER = LoggerFactory.getLogger(AgentLoader.class);
+
+    public static void run(String[] args) {
+        String agentFilePath = "/home/adi/Desktop/agent-1.0.0-jar-with-dependencies.jar";
+        String applicationName = "MyAtmApplication";
+
+        //iterate all jvms and get the first one that matches our application name
+        Optional<String> jvmProcessOpt = Optional.ofNullable(VirtualMachine.list()
+          .stream()
+          .filter(jvm -> {
+              LOGGER.info("jvm:{}", jvm.displayName());
+              return jvm.displayName().contains(applicationName);
+          })
+          .findFirst().get().id());
+
+        if(!jvmProcessOpt.isPresent()) {
+            LOGGER.error("Target Application not found");
+            return;
+        }
+        File agentFile = new File(agentFilePath);
+        try {
+            String jvmPid = jvmProcessOpt.get();
+            LOGGER.info("Attaching to target JVM with PID: " + jvmPid);
+            VirtualMachine jvm = VirtualMachine.attach(jvmPid);
+            jvm.loadAgent(agentFile.getAbsolutePath());
+            jvm.detach();
+            LOGGER.info("Attached to target JVM and loaded Java agent successfully");
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+}
@@ -0,0 +1,14 @@
+package com.baeldung.instrumentation.application;
+
+/**
+ * Created by adi on 6/14/18.
+ */
+public class Launcher {
+    public static void main(String[] args) throws Exception {
+        if(args[0].equals("StartMyAtmApplication")) {
+            new MyAtmApplication().run(args);
+        } else if(args[0].equals("LoadAgent")) {
+            new AgentLoader().run(args);
+        }
+    }
+}
@@ -0,0 +1,19 @@
+package com.baeldung.instrumentation.application;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Created by adi on 6/11/18.
+ */
+public class MyAtm {
+    private static Logger LOGGER = LoggerFactory.getLogger(MyAtm.class);
+
+    private static final int account = 10;
+
+    public static void withdrawMoney(int amount) throws InterruptedException {
+        Thread.sleep(2000l); //processing going on here
+        LOGGER.info("[Application] Successful Withdrawal of [{}] units!", amount);
+
+    }
+}
@@ -0,0 +1,19 @@
+package com.baeldung.instrumentation.application;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class MyAtmApplication {
+
+    private static Logger LOGGER = LoggerFactory.getLogger(MyAtmApplication.class);
+
+    public static void run(String[] args) throws Exception {
+        LOGGER.info("[Application] Starting ATM application");
+        MyAtm.withdrawMoney(Integer.parseInt(args[2]));
+
+        Thread.sleep(Long.valueOf(args[1]));
+
+        MyAtm.withdrawMoney(Integer.parseInt(args[3]));
+    }
+
+}