Minimize thrown IllegalStateException when profiling (#1931)

Sometimes we see a thrown IllegalStateException when CPU profiling,
where the underlying cause is a ConcurrentModificationException in
the underlying GraalVM internals.

This attempts to minimize the chances of this exception being thrown
by starting proiling after the evaluator and its context have been
created.

pkl-cli/src/main/kotlin/org/pkl/cli/CliCommandRunner.kt

@@ -59,7 +59,7 @@ constructor(
   override fun doRun() {
     val builder = evaluatorBuilder()
     try {
-      withCpuProfiler(profilerOptions) { evalCmd(builder) }
+      withCpuProfiler(profilerOptions) { startProfiling -> evalCmd(builder, startProfiling) }
     } finally {
       Closeables.closeQuietly(builder.moduleKeyFactories)
       Closeables.closeQuietly(builder.resourceReaders)
@@ -69,8 +69,9 @@ constructor(
     }
   }
 
-  private fun evalCmd(builder: EvaluatorBuilder) {
+  private fun evalCmd(builder: EvaluatorBuilder, startProfiling: () -> Unit) {
     val evaluator = builder.build()
+    startProfiling()
     evaluator.use {
       evaluator.evaluateCommand(
         uri(options.normalizedSourceModules.first()),

pkl-cli/src/main/kotlin/org/pkl/cli/CliEvaluator.kt

@@ -100,11 +100,11 @@ constructor(
   override fun doRun() {
     val builder = evaluatorBuilder()
     try {
-      withCpuProfiler(options.profilerOptions) {
+      withCpuProfiler(options.profilerOptions) { startProfiling ->
         if (options.multipleFileOutputPath != null) {
-          writeMultipleFileOutput(builder)
+          writeMultipleFileOutput(builder, startProfiling)
         } else {
-          writeOutput(builder)
+          writeOutput(builder, startProfiling)
         }
       }
     } finally {
@@ -157,8 +157,9 @@ constructor(
   }
 
   /** Renders each module's `output.bytes`, writing it to the specified output file. */
-  private fun writeOutput(builder: EvaluatorBuilder) {
+  private fun writeOutput(builder: EvaluatorBuilder, startProfiling: () -> Unit) {
     val evaluator = builder.setOutputFormat(options.outputFormat).build()
+    startProfiling()
     evaluator.use { ev ->
       val outputFiles = fileOutputPaths
       if (outputFiles != null) {
@@ -232,10 +233,11 @@ constructor(
    * Renders each module's `output.files`, writing each entry as a file into the specified output
    * directory.
    */
-  private fun writeMultipleFileOutput(builder: EvaluatorBuilder) {
+  private fun writeMultipleFileOutput(builder: EvaluatorBuilder, startProfiling: () -> Unit) {
     val outputDirs = directoryOutputPaths!!
     val writtenFiles = mutableMapOf<Path, OutputFile>()
     builder.setOutputFormat(options.outputFormat).build().use { evaluator ->
+      startProfiling()
       for ((moduleUri, outputDir) in outputDirs) {
         if (outputDir.exists() && !outputDir.isDirectory()) {
           throw CliException("Output path `$outputDir` exists and is not a directory.")

pkl-cli/src/main/kotlin/org/pkl/cli/CliTestRunner.kt

@@ -43,7 +43,7 @@ constructor(
   override fun doRun() {
     val builder = evaluatorBuilder()
     try {
-      withCpuProfiler(profilerOptions) { evalTest(builder) }
+      withCpuProfiler(profilerOptions) { startProfiling -> evalTest(builder, startProfiling) }
     } finally {
       Closeables.closeQuietly(builder.moduleKeyFactories)
       Closeables.closeQuietly(builder.resourceReaders)
@@ -53,7 +53,7 @@ constructor(
     }
   }
 
-  private fun evalTest(builder: EvaluatorBuilder) {
+  private fun evalTest(builder: EvaluatorBuilder, startProfiling: () -> Unit) {
     val sources =
       options.normalizedSourceModules.ifEmpty { project?.tests?.map { it.toUri() } }
         ?:
@@ -68,6 +68,7 @@ constructor(
         )
 
     val evaluator = builder.build()
+    startProfiling()
     evaluator.use {
       var failed = false
       var isExampleWrittenFailure = true

pkl-cli/src/main/kotlin/org/pkl/cli/utils.kt

@@ -23,15 +23,24 @@ import org.pkl.profiler.CpuProfilerInstrument
 /**
  * Runs [body], sampling it with the CPU profiler if [options] enables CPU profiling.
  *
- * Profiler output is written when [body] finishes, whether normally or exceptionally.
+ * Sampling does not begin until [body] calls the `startProfiling` function it is passed. Call it
+ * after the evaluator has been built: starting Truffle's sampling thread while the main thread is
+ * still creating the engine context can race inside Truffle's thread-local handshake setup.
+ *
+ * Profiler output is written when [body] finishes, whether normally or exceptionally. If
+ * `startProfiling` was never called, no output is written.
  */
-internal inline fun <T> withCpuProfiler(options: ProfilerOptions, body: () -> T): T {
-  if (!options.cpu.isEnabled) return body()
+internal inline fun <T> withCpuProfiler(
+  options: ProfilerOptions,
+  body: (startProfiling: () -> Unit) -> T,
+): T {
+  if (!options.cpu().isEnabled) {
+    return body {}
+  }
   val profiler =
     VmUtils.getEngine().instruments[CpuProfilerInstrument.ID]?.lookup(CpuProfiler::class.java)!!
-  profiler.start(options.cpu)
   try {
-    return body()
+    return body { profiler.start(options.cpu) }
   } finally {
     profiler.stop()
   }