Fix performance drop drastically while profiling a Ballerina program

bvm/ballerina-profiler/build.gradle

@@ -24,12 +24,14 @@ dependencies {
     implementation "com.google.code.gson:gson:${project.gsonVersion}"
     implementation "commons-io:commons-io:${project.commonsIoVersion}"
     implementation project(':identifier-util')
+    implementation project(':ballerina-runtime')
 }
 
 version = 1.0
 
 jar {
     dependsOn(':identifier-util:jar')
+    dependsOn(':ballerina-runtime:jar')
     from(sourceSets.main.output)
     from(sourceSets.main.java) {
         include "**/*.java"
@@ -43,11 +45,6 @@ jar {
     duplicatesStrategy = DuplicatesStrategy.EXCLUDE
 
     manifest {
-        attributes 'Implementation-Title': 'Foobar',
-                'Implementation-Version': '',
-                'Built-By': System.getProperty('user.name'),
-                'Built-Date': new Date(),
-                'Built-JDK': System.getProperty('java.version'),
-                'Main-Class': 'io.ballerina.runtime.profiler.Main'
+        attributes 'Main-Class': 'io.ballerina.runtime.profiler.Main'
     }
 }

bvm/ballerina-profiler/spotbugs-exclude.xml

@@ -20,4 +20,8 @@
         <Class name="io.ballerina.runtime.profiler.Profiler"/>
         <Bug pattern="DP_CREATE_CLASSLOADER_INSIDE_DO_PRIVILEGED"/>
     </Match>
+    <Match>
+        <Class name="io.ballerina.runtime.profiler.ui.JsonParser$Data"/>
+        <Bug pattern="URF_UNREAD_FIELD"/>
+    </Match>
 </FindBugsFilter>

bvm/ballerina-profiler/src/main/java/io/ballerina/runtime/profiler/Profiler.java

@@ -31,25 +31,23 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.PrintWriter;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
-import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.jar.JarFile;
 import java.util.stream.Collectors;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 
 import static io.ballerina.runtime.profiler.util.Constants.CPU_PRE_JSON;
+import static io.ballerina.runtime.profiler.util.Constants.CURRENT_DIR_KEY;
 import static io.ballerina.runtime.profiler.util.Constants.OUT_STREAM;
+import static io.ballerina.runtime.profiler.util.Constants.PERFORMANCE_JSON;
 
 /**
  * This class is used to as the driver class of the Ballerina profiler.
@@ -59,7 +57,6 @@
 public class Profiler {
 
     private final long profilerStartTime;
-    private String skipFunctionString = null;
     private String balJarArgs = null;
     private String balJarName = null;
     private String targetDir = null;
@@ -72,10 +69,12 @@ public class Profiler {
     private int balFunctionCount = 0;
     private int moduleCount = 0;
     private final ProfilerMethodWrapper profilerMethodWrapper;
+    private final String currentDir;
 
     public Profiler(long profilerStartTime) {
         this.profilerStartTime = profilerStartTime;
         this.profilerMethodWrapper = new ProfilerMethodWrapper();
+        this.currentDir = System.getenv(CURRENT_DIR_KEY);
     }
 
     private void addShutdownHookAndCleanup() {
@@ -88,13 +87,12 @@ private void addShutdownHookAndCleanup() {
                 OUT_STREAM.printf("%s[6/6] Generating output...%s%n", Constants.ANSI_CYAN, Constants.ANSI_RESET);
                 JsonParser jsonParser = new JsonParser();
                 HttpServer httpServer = new HttpServer();
-                jsonParser.initializeCPUParser(skipFunctionString);
-                deleteFileIfExists("usedPathsList.txt");
-                deleteFileIfExists(CPU_PRE_JSON);
+                String cpuFilePath = Paths.get(currentDir, CPU_PRE_JSON).toString();
+                jsonParser.initializeCPUParser(cpuFilePath);
+                deleteFileIfExists(cpuFilePath);
                 OUT_STREAM.printf(" ○ Execution time: %d seconds %n", profilerTotalTime / 1000);
-                deleteTempData();
                 httpServer.initializeHTMLExport(this.sourceRoot);
-                deleteFileIfExists("performance_report.json");
+                deleteFileIfExists(PERFORMANCE_JSON);
                 OUT_STREAM.println("--------------------------------------------------------------------------------");
             } catch (IOException e) {
                 throw new ProfilerException("Error occurred while generating the output", e);
@@ -116,18 +114,6 @@ private void deleteFileIfExists(String filePath) {
         }
     }
 
-    private void deleteTempData() {
-        String filePrefix = "jartmp";
-        File[] files = new File(System.getProperty("user.dir")).listFiles();
-        if (files != null) {
-            for (File file : files) {
-                if (file.getName().startsWith(filePrefix)) {
-                    FileUtils.deleteQuietly(file);
-                }
-            }
-        }
-    }
-
     private void printHeader() {
         OUT_STREAM.printf("%n%s================================================================================%s",
                 Constants.ANSI_GRAY, Constants.ANSI_RESET);
@@ -154,10 +140,6 @@ private void handleProfilerArguments(String[] args) {
                     this.balJarArgs = extractBalJarArgs(args[i + 1]);
                     addToUsedArgs(args, usedArgs, i);
                 }
-                case "--skip" -> {
-                    this.skipFunctionString = extractSkipFunctionString(args[i + 1]);
-                    addToUsedArgs(args, usedArgs, i);
-                }
                 case "--target" -> {
                     this.targetDir = args[i + 1];
                     addToUsedArgs(args, usedArgs, i);
@@ -189,13 +171,6 @@ private String extractBalJarArgs(String value) {
         return value.substring(1, value.length() - 1);
     }
 
-    private String extractSkipFunctionString(String value) {
-        if (value == null || !value.matches("\\[.*\\]")) {
-            throw new ProfilerException("Invalid skip functions found: " + value);
-        }
-        return value.substring(1, value.length() - 1);
-    }
-
     private void handleUnrecognizedArgument(String argument, List<String> usedArgs) {
         if (!usedArgs.contains(argument)) {
             throw new ProfilerException("Unrecognized argument found: " + argument);
@@ -242,16 +217,15 @@ private void initializeProfiling() throws ProfilerException {
                     new File(balJarName).toURI().toURL()}));
             ProfilerClassLoader profilerClassLoader = new ProfilerClassLoader(new URLClassLoader(new URL[]{
                     new File(balJarName).toURI().toURL()}));
-            Set<String> usedPaths = new HashSet<>();
             for (String className : classNames) {
-                if (mainClassPackage == null) {
+                if (mainClassPackage == null || className.contains("$gen$")) {
                     continue;
                 }
                 if (className.startsWith(mainClassPackage.split("/")[0]) || utilPaths.contains(className)) {
                     try (InputStream inputStream = jarFile.getInputStream(jarFile.getJarEntry(className))) {
-                        byte[] code = profilerMethodWrapper.modifyMethods(inputStream);
+                        String sourceClassName = className.replace(Constants.CLASS_SUFFIX, "");
+                        byte[] code = profilerMethodWrapper.modifyMethods(inputStream, sourceClassName);
                         profilerClassLoader.loadClass(code);
-                        usedPaths.add(className.replace(Constants.CLASS_SUFFIX, "").replace("/", "."));
                         profilerMethodWrapper.printCode(className, code, getFileNameWithoutExtension(balJarName));
                     }
                 }
@@ -260,9 +234,6 @@ private void initializeProfiling() throws ProfilerException {
                 }
             }
             OUT_STREAM.printf(" ○ Instrumented module count: %d%n", moduleCount);
-            try (PrintWriter printWriter = new PrintWriter("usedPathsList.txt", StandardCharsets.UTF_8)) {
-                printWriter.println(String.join(", ", usedPaths));
-            }
             OUT_STREAM.printf(" ○ Instrumented function count: %d%n", balFunctionCount);
             modifyJar();
         } catch (Throwable throwable) {

bvm/ballerina-profiler/src/main/java/io/ballerina/runtime/profiler/codegen/ProfilerClassVisitor.java

@@ -32,19 +32,20 @@
  */
 public class ProfilerClassVisitor extends ClassVisitor {
 
-    public ProfilerClassVisitor(ClassVisitor classVisitor) {
+    private final String className;
+
+    public ProfilerClassVisitor(String className, ClassVisitor classVisitor) {
         super(Opcodes.ASM9, classVisitor);
+        this.className = className;
     }
 
     @Override
     public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
         MethodVisitor methodVisitor = super.visitMethod(access, name, desc, signature, exceptions);
-        if (!name.startsWith("$") && desc.startsWith(STRAND_ARG)) {
+        if (desc.startsWith(STRAND_ARG) && !name.contains("$gen$")) {
             Main.incrementBalFunctionCount();
+            return new StrandCheckAdapter(className, access, methodVisitor, name, desc, (access & Opcodes.ACC_STATIC));
         }
-        if (desc.startsWith(STRAND_ARG)) {
-            return new StrandCheckAdapter(access, methodVisitor, name, desc, (access & Opcodes.ACC_STATIC));
-        }
-        return new NonStrandCheckAdapter(access, methodVisitor, name, desc);
+        return methodVisitor;
     }
 }

bvm/ballerina-profiler/src/main/java/io/ballerina/runtime/profiler/codegen/ProfilerMethodWrapper.java

@@ -34,13 +34,16 @@
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.jar.Attributes;
 import java.util.jar.Manifest;
 
+import static io.ballerina.runtime.profiler.util.Constants.CURRENT_DIR_KEY;
 import static io.ballerina.runtime.profiler.util.Constants.ERROR_STREAM;
 import static io.ballerina.runtime.profiler.util.Constants.OUT_STREAM;
+import static io.ballerina.runtime.profiler.util.Constants.USER_DIR;
 
 /**
  * This class is used as the method wrapper for the Ballerina profiler.
@@ -57,12 +60,13 @@ public void invokeMethods(String debugArg) throws IOException, InterruptedExcept
         if (debugArg != null) {
             commands.add(debugArg);
         }
-        commands.add(Constants.TEMP_JAR_FILE_NAME);
+        commands.add(Paths.get(System.getProperty(USER_DIR), Constants.TEMP_JAR_FILE_NAME).toString());
         if (balJarArgs != null) {
             commands.add(balJarArgs);
         }
         ProcessBuilder processBuilder = new ProcessBuilder(commands);
-        processBuilder.redirectErrorStream(true);
+        processBuilder.inheritIO();
+        processBuilder.directory(new File(System.getenv(CURRENT_DIR_KEY)));
         Process process = processBuilder.start();
         OUT_STREAM.printf(Constants.ANSI_CYAN + "[5/6] Running executable..." + Constants.ANSI_RESET + "%n");
         try (BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream(),
@@ -84,13 +88,13 @@ public String mainClassFinder(URLClassLoader manifestClassLoader) {
         }
     }
 
-    public byte[] modifyMethods(InputStream inputStream) {
+    public byte[] modifyMethods(InputStream inputStream, String className) {
         byte[] code;
         try {
             ClassReader reader = new ClassReader(inputStream);
             ClassWriter classWriter = new ProfilerClassWriter(reader, ClassWriter.COMPUTE_MAXS |
                     ClassWriter.COMPUTE_FRAMES);
-            ClassVisitor change = new ProfilerClassVisitor(classWriter);
+            ClassVisitor change = new ProfilerClassVisitor(className, classWriter);
             reader.accept(change, ClassReader.EXPAND_FRAMES);
             code = classWriter.toByteArray();
             return code;

bvm/ballerina-profiler/src/main/java/io/ballerina/runtime/profiler/codegen/StrandCheckAdapter.java

@@ -18,33 +18,48 @@
 
 package io.ballerina.runtime.profiler.codegen;
 
+import io.ballerina.runtime.internal.scheduling.Strand;
+import io.ballerina.runtime.profiler.runtime.Data;
 import org.objectweb.asm.Label;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.Type;
 import org.objectweb.asm.commons.AdviceAdapter;
 
+import static io.ballerina.runtime.profiler.util.Constants.DATA_CLASS;
 import static io.ballerina.runtime.profiler.util.Constants.GET_INSTANCE_DESCRIPTOR;
 import static io.ballerina.runtime.profiler.util.Constants.PROFILE_ANALYZER;
-import static io.ballerina.runtime.profiler.util.Constants.STRAND_CLASS;
+import static io.ballerina.runtime.profiler.util.Constants.START_PROFILE_DESCRIPTOR;
+import static io.ballerina.runtime.profiler.util.Constants.STOP_PROFILE_DESCRIPTOR;
 
 /**
  * This class is used as the advice adapter for the Ballerina profiler.
- * This class only manages the functions that contain the strand parameter.
+ * This will wrap the ballerina methods with
+ * {@link io.ballerina.runtime.profiler.runtime.ProfileAnalyzer#start(Strand, String, String)} and
+ * {@link io.ballerina.runtime.profiler.runtime.ProfileAnalyzer#stop(Strand, Data)} methods.
  *
  * @since 2201.8.0
  */
 public class StrandCheckAdapter extends AdviceAdapter {
 
+    private final String className;
+    private final String methodName;
+    private final Type dataType;
     Label tryStart = new Label();
     int load;
+    int stackKeyIndex;
 
-    public StrandCheckAdapter(int access, MethodVisitor mv, String methodName, String description, int load) {
+    public StrandCheckAdapter(String className, int access, MethodVisitor mv, String methodName,
+                              String description, int load) {
         super(Opcodes.ASM9, mv, access, methodName, description);
+        this.className = className;
+        this.methodName = methodName;
         if (load == 0) {
             this.load = 1;
         } else {
             this.load = 0;
         }
+        this.dataType = Type.getObjectType(DATA_CLASS);
     }
 
     // It adds a label to the try block of the wrapped method.
@@ -57,10 +72,13 @@ public void visitCode() {
     // It retrieves the profiler instance, gets the strand id and starts the profiling.
     @Override
     protected void onMethodEnter() {
+        this.stackKeyIndex = this.newLocal(dataType);
         mv.visitMethodInsn(INVOKESTATIC, PROFILE_ANALYZER, "getInstance", GET_INSTANCE_DESCRIPTOR, false);
         mv.visitVarInsn(ALOAD, load);
-        mv.visitMethodInsn(INVOKEVIRTUAL, STRAND_CLASS, "getId", "()I", false);
-        mv.visitMethodInsn(INVOKEVIRTUAL, PROFILE_ANALYZER, "start", "(I)V", false);
+        mv.visitLdcInsn(className);
+        mv.visitLdcInsn(methodName);
+        mv.visitMethodInsn(INVOKEVIRTUAL, PROFILE_ANALYZER, "start", START_PROFILE_DESCRIPTOR, false);
+        mv.visitVarInsn(ASTORE, this.stackKeyIndex);
     }
 
     // If the exit is not due to an exception, it calls the onFinally method.
@@ -88,12 +106,7 @@ public void visitMaxs(int maxStack, int maxLocals) {
     private void onFinally() {
         mv.visitMethodInsn(INVOKESTATIC, PROFILE_ANALYZER, "getInstance", GET_INSTANCE_DESCRIPTOR, false);
         mv.visitVarInsn(ALOAD, load);
-        mv.visitMethodInsn(INVOKEVIRTUAL, STRAND_CLASS,
-                "getState", "()Lio/ballerina/runtime/internal/scheduling/State;", false);
-        mv.visitMethodInsn(INVOKEVIRTUAL, "io/ballerina/runtime/internal/scheduling/State",
-                "toString", "()Ljava/lang/String;", false);
-        mv.visitVarInsn(ALOAD, load);
-        mv.visitMethodInsn(INVOKEVIRTUAL, STRAND_CLASS, "getId", "()I", false);
-        mv.visitMethodInsn(INVOKEVIRTUAL, PROFILE_ANALYZER, "stop", "(Ljava/lang/String;I)V", false);
+        mv.visitVarInsn(ALOAD, this.stackKeyIndex);
+        mv.visitMethodInsn(INVOKEVIRTUAL, PROFILE_ANALYZER, "stop", STOP_PROFILE_DESCRIPTOR, false);
     }
 }