Configurable recipe classloader for PrepareRecipe

Author: jkschneider

rewrite-core/src/main/java/org/openrewrite/rpc/RewriteRpc.java

@@ -67,6 +67,8 @@ public class RewriteRpc {
     private final AtomicReference<PrintStream> log = new AtomicReference<>();
     private final AtomicReference<TraceGetObject> traceGetObject = new AtomicReference<>(
             new TraceGetObject(false, false));
+    private final AtomicReference<ClassLoader> recipeClassLoader = new AtomicReference<>(
+            RewriteRpc.class.getClassLoader());
 
     final PreparedRecipeCache preparedRecipes = new PreparedRecipeCache();
 
@@ -142,7 +144,7 @@ protected Object handle(Void noParams) {
                 }
             }
         });
-        jsonRpc.rpc("PrepareRecipe", new PrepareRecipe.Handler(preparedRecipes));
+        jsonRpc.rpc("PrepareRecipe", new PrepareRecipe.Handler(preparedRecipes, recipeClassLoader));
         jsonRpc.rpc("Print", new Print.Handler(this::getObject));
 
         jsonRpc.bind();
@@ -169,6 +171,11 @@ public RewriteRpc log(@Nullable PrintStream logFile) {
         return this;
     }
 
+    public RewriteRpc recipeClassLoader(ClassLoader recipeClassLoader) {
+        this.recipeClassLoader.set(recipeClassLoader);
+        return this;
+    }
+
     public void shutdown() {
         //noinspection ConstantValue
         if (log.get() != null) {

rewrite-core/src/main/java/org/openrewrite/rpc/request/PrepareRecipe.java

@@ -25,6 +25,7 @@
 import org.openrewrite.rpc.internal.PreparedRecipeCache;
 
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 import static java.util.Collections.emptyList;
 
@@ -36,10 +37,11 @@ public class PrepareRecipe implements RpcRequest {
     @RequiredArgsConstructor
     public static class Handler extends JsonRpcMethod<PrepareRecipe> {
         private final PreparedRecipeCache preparedRecipes;
+        private final AtomicReference<ClassLoader> recipeClassLoader;
 
         @Override
         protected Object handle(PrepareRecipe request) throws Exception {
-            Recipe recipe = new RecipeLoader(null).load(request.getId(), request.getOptions());
+            Recipe recipe = new RecipeLoader(recipeClassLoader.get()).load(request.getId(), request.getOptions());
             String instanceId = SnowflakeId.generateId();
             preparedRecipes.getInstantiated().put(instanceId, recipe);
             return new PrepareRecipeResponse(

rewrite-javascript/src/main/java/org/openrewrite/javascript/rpc/JavaScriptRewriteRpc.java

@@ -107,9 +107,9 @@ public static class Builder implements Supplier<JavaScriptRewriteRpc> {
         private Duration timeout = Duration.ofSeconds(30);
         private boolean traceRpcMessages;
         private @Nullable Integer inspectBrk;
-        private boolean profiler;
         private @Nullable Integer maxHeapSize;
         private @Nullable Path workingDirectory;
+        private ClassLoader recipeClassLoader = JavaScriptRewriteRpc.class.getClassLoader();
 
         public Builder marketplace(Environment marketplace) {
             this.marketplace = marketplace;
@@ -181,26 +181,6 @@ public Builder inspectBrk() {
             return inspectBrk(9229);
         }
 
-        /**
-         * Enable V8 CPU profiling for performance analysis. When enabled, the process will
-         * generate a CPU profile that can be analyzed to identify performance bottlenecks.
-         * The profile is saved in Chrome DevTools format (.cpuprofile) on shutdown.
-         * <p>
-         * Profiling uses the V8 Inspector API and works with both npx and direct node execution modes.
-         * The profile file can be loaded in Chrome DevTools for analysis.
-         *
-         * @param profiler Whether to enable profiling
-         * @return This builder
-         */
-        public Builder profiler(boolean profiler) {
-            this.profiler = profiler;
-            return this;
-        }
-
-        public Builder profiler() {
-            return profiler(true);
-        }
-
         /**
          * Set the maximum heap size for the Node.js process in megabytes.
          * Default V8 heap size is approximately 1.5-2 GB on 64-bit systems.
@@ -227,6 +207,11 @@ public Builder workingDirectory(@Nullable Path workingDirectory) {
             return this;
         }
 
+        public Builder recipeClassLoader(ClassLoader recipeClassLoader) {
+            this.recipeClassLoader = recipeClassLoader;
+            return this;
+        }
+
         @Override
         public JavaScriptRewriteRpc get() {
             Stream<@Nullable String> cmd;
@@ -245,7 +230,6 @@ public JavaScriptRewriteRpc get() {
                         "node",
                         "--enable-source-maps",
                         "--inspect-brk=" + inspectBrk,
-                        profiler ? "--prof" : null,
                         maxHeapSize != null ? "--max-old-space-size=" + maxHeapSize : null,
                         serverJs.toAbsolutePath().normalize().toString(),
                         log == null ? null : "--log-file=" + log.toAbsolutePath().normalize(),
@@ -262,8 +246,7 @@ public JavaScriptRewriteRpc get() {
                         log == null ? null : "--log-file=" + log.toAbsolutePath().normalize(),
                         metricsCsv == null ? null : "--metrics-csv=" + metricsCsv.toAbsolutePath().normalize(),
                         traceRpcMessages ? "--trace-rpc-messages" : null,
-                        recipeInstallDir == null ? null : "--recipe-install-dir=" + recipeInstallDir.toAbsolutePath().normalize(),
-                        profiler ? "--profile" : null
+                        recipeInstallDir == null ? null : "--recipe-install-dir=" + recipeInstallDir.toAbsolutePath().normalize()
                 );
             }
 
@@ -293,7 +276,8 @@ public JavaScriptRewriteRpc get() {
                         String.join(" ", cmdArr), process.environment())
                         .livenessCheck(process::getLivenessCheck)
                         .timeout(timeout)
-                        .log(log == null ? null : new PrintStream(Files.newOutputStream(log, StandardOpenOption.APPEND, StandardOpenOption.CREATE)));
+                        .log(log == null ? null : new PrintStream(Files.newOutputStream(log, StandardOpenOption.APPEND, StandardOpenOption.CREATE)))
+                        .recipeClassLoader(recipeClassLoader);
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }