Profile-Guided Optimization for Native Executables (#1665)

* Switch GraalVM edition

* add task to compile binary with instrumentation

* Switch back to handle instrumented native outside Gradle

* Ignore profile dir

* Add script to generate PGO binary

* Remove unnecessary gradle task

* use new command to build binaries

* try using Powershell for Windows

* Add logs

* call gradle binary directly from script

* Remove pipe to stdout since it's stuck on Windows

* Use $.sync for Windows

* Add comments describing the process to generate executable

* One more try to get Windows working

* Fix gradle command

* try sync method again

* debug - try a few steps to see which one is blocking

* debug - remove OSes from matrix

* add back await

* debug - direct calls to gradle

* try calling -v

* forgot node CLI

* calling 2 tasks at once

* set task dependency

* Try building before running script

* try running gradle directly

* try setting empty prefix

* try using plain console

* try ignoring stdin

* piping straight to stdout and stderr

* put everything back

* Add CHANGELOG for PGO

.github/workflows/tests-deployments.yml

@@ -14,7 +14,7 @@ jobs:
       fail-fast: false
       matrix:
         node:
-          - "21"
+          - "22"
           - "20"
           - "18"
         java:
@@ -160,8 +160,8 @@ jobs:
       matrix:
         os:
           - ubuntu-latest
-          - windows-latest
           - macos-14
+          - windows-latest
         include:
           - os: windows-latest
             ASSET_NAME: apex-ast-serializer.exe
@@ -187,16 +187,24 @@ jobs:
         working-directory: ./packages/apex-ast-serializer
     steps:
       - uses: actions/[email protected]
+      - name: Setup Node.js
+        uses: actions/[email protected]
+        with:
+          node-version: 18
       - name: Set up JDK
-        uses: graalvm/setup-graalvm@v1
+        uses: graalvm/setup-graalvm@v1.2.4.1
         with:
-          java-version: "17"
-          distribution: "graalvm-community" # See 'Options' for all available distributions
+          java-version: "23"
+          distribution: "graalvm"
           github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Setup pnpm
+        uses: pnpm/action-setup@v3
+      - name: Install dependencies
+        run: pnpm install --frozen-lockfile
       - name: Grant execute permission for gradlew
         run: chmod +x gradlew
-      - name: Build native executable
-        run: ./gradlew clean :parser:buildNative
+      - name: Build native instrumented executable
+        run: pnpm nx run apex-ast-serializer:build:native
       - name: Set version
         run: echo "version=$(cat ../prettier-plugin-apex/package.json | jq -r '.version')" >> $GITHUB_ENV
         if: env.OS != 'windows-latest'

.gitignore

@@ -12,6 +12,7 @@ out
 .vscode/settings.json
 packages/prettier-plugin-apex/tests/local_only
 packages/prettier-plugin-apex/vendor
+packages/apex-ast-serializer/parser/src/pgo-profiles
 
 .nx
 .pnpm-store

.tool-versions

@@ -1,2 +1,2 @@
 nodejs 18.20.4
-java graalvm-22.3.1+java11
+java oracle-graalvm-22

CHANGELOG.md

@@ -1,3 +1,7 @@
+# Unreleased
+
+- Improve native executable performance by utilizing Profile-Guided Optimization.
+
 # 2.1.5
 
 - Fix Windows application not being able to get output from parser when `DEBUG` environment variable is set ([issue](https://github.com/dangmai/prettier-plugin-apex/issues/1513)).

packages/apex-ast-serializer/README.md

@@ -23,7 +23,7 @@ The package will be built and installed directly under `packages/prettier-plugin
 For native image, run:
 
 ```bash
-./gradlew buildNative
+pnpm nx run apex-ast-serializer:build:native
 ```
 
 ## Running

packages/apex-ast-serializer/parser/build.gradle

@@ -136,16 +136,6 @@ graalvmNative {
   }
 }
 
-// This is the task that builds native executable for the host platform.
-// Run it: `gradle buildNative`
-// It contains multiple steps in order to achive this goal:
-// - It runs the entire test suite, and uses native-image-agent
-// to figure out what reflections/proxies are used as part of that run.
-// - Then it uses native-image to compile the final artifacts, using the information
-// in the previous step (which is stored in `build/native/agent-output/`)
-// - The final artifact is produced and stored in `build/native/nativeCompile`
-task buildNative {
-  dependsOn 'nativeInstrumentedTest'
-  dependsOn 'nativeCompile'
-  tasks.findByName('nativeCompile').mustRunAfter 'nativeInstrumentedTest'
+nativeCompile {
+  mustRunAfter nativeInstrumentedTest
 }

packages/apex-ast-serializer/project.json

@@ -26,7 +26,7 @@
       "outputs": ["{projectRoot}/parser/build"],
       "options": {
         "cwd": "packages/apex-ast-serializer",
-        "command": "node tools/run-gradle.mjs clean :parser:buildNative"
+        "command": "node tools/build-native-binary.mjs"
       }
     },
     "test": {

packages/apex-ast-serializer/tools/build-native-binary.mjs

@@ -0,0 +1,98 @@
+#!/usr/bin/env zx
+
+// This is the script that builds native executable for the host platform.
+// It contains multiple steps in order to achive this goal:
+// - It runs the entire test suite, and uses native-image-agent
+// to figure out what reflections/proxies are used as part of that run.
+// - Then it uses native-image to compile an instrumented binary, using the information
+// in the previous step (which is stored in `build/native/agent-output/`)
+// - This artifact is produced and stored in `build/native/nativeCompile`
+// - Then it runs the instrumented binary with all the test classes,
+// in order to produce profiles that will be used later to optimize the binary.
+// - Then it merges all the profiles into a single one.
+// - Finally, it runs the nativeCompile task again, but this time with the merged profile,
+// and produces the final artifact.
+
+import { join } from "path";
+import { $, fs, usePowerShell } from "zx";
+
+$.verbose = false;
+
+let gradle = "./gradlew";
+if (process.platform === "win32") {
+  usePowerShell();
+  gradle += ".bat";
+}
+
+async function getFilesWithSuffix(rootDir, suffix) {
+  let result = [];
+
+  async function walkDir(dir) {
+    const files = await fs.readdir(dir, { withFileTypes: true });
+    for (const file of files) {
+      const fullPath = join(dir, file.name);
+      if (file.isDirectory()) {
+        await walkDir(fullPath);
+      } else if (file.name.endsWith(suffix)) {
+        result.push(fullPath);
+      }
+    }
+  }
+
+  await walkDir(rootDir);
+  return result;
+}
+console.log("Running nativeCompile with PGO instrumentation");
+await $`${gradle} :parser:nativeInstrumentedTest :parser:nativeCompile --pgo-instrument`.stdio(
+  "ignore",
+  process.stdout,
+  process.stderr,
+);
+const classFiles = await getFilesWithSuffix(
+  "./parser/build/resources/test",
+  ".cls",
+);
+await fs.remove("./parser/src/pgo-profiles/main");
+await fs.ensureDir("./parser/src/pgo-profiles/main");
+
+let i = 1;
+for (const classFile of classFiles) {
+  if (classFile.includes("__snapshots__")) {
+    continue;
+  }
+  console.log(`Processing ${classFile}`);
+  const isAnonymous = classFile.includes("anonymous");
+
+  await $({
+    input: await fs.readFile(classFile),
+  })`./parser/build/native/nativeCompile/apex-ast-serializer-instrumented -f json -i ${isAnonymous ? "--anonymous" : ""}`.stdio(
+    "ignore",
+  );
+
+  await fs.move(
+    "./default.iprof",
+    `./parser/src/pgo-profiles/main/${i}.iprof`,
+    { overwrite: true },
+  );
+  i++;
+}
+console.log("Merging profiles");
+await $`native-image-configure merge-pgo-profiles --input-dir=./parser/src/pgo-profiles/main --output-file=merged_profile.iprof`.stdio(
+  "ignore",
+  process.stdout,
+  process.stderr,
+);
+await fs.remove("./parser/src/pgo-profiles/main");
+await fs.ensureDir("./parser/src/pgo-profiles/main");
+await fs.move(
+  "./merged_profile.iprof",
+  `./parser/src/pgo-profiles/main/default.iprof`,
+  { overwrite: true },
+);
+
+console.log("Running nativeCompile for final artifact");
+await $`${gradle} :parser:nativeCompile`.stdio(
+  "ignore",
+  process.stdout,
+  process.stderr,
+);

packages/apex-ast-serializer/tools/run-gradle.mjs

@@ -1,5 +1,6 @@
 #!/usr/bin/env zx
 import { $, usePowerShell } from "zx";
+$.verbose = false;
 
 // This script runs Gradle with the given arguments in a cross-platform way.
 // It's meant to be called from nx run-commands.
@@ -11,4 +12,4 @@ if (process.platform === "win32") {
   command += ".bat";
 }
 const args = process.argv.slice(2);
-await $`${command} ${args}`;
+await $`${command} ${args}`.stdio("ignore", process.stdout, process.stderr);