Refactor Jvm methods to match specified signatures

vishwamartur · vishwamartur · commit 1e3205cf9aa8 · 2024-11-16T12:37:58.000+05:30
Related to com-lihaoyi#3772 Refactor `Jvm.scala` to consolidate subprocess and classloader spawning operations into four specified signatures. * **Refactor `callSubprocess` method:** - Rename to `call`. - Update parameters to match the specified `call` signature. - Use `jvmCommandArgs` to generate command arguments. - Call `os.call` with the updated parameters. * **Refactor `runSubprocess` method:** - Rename to `spawn`. - Update parameters to match the specified `spawn` signature. - Use `jvmCommandArgs` to generate command arguments. - Call `os.spawn` with the updated parameters. * **Add `spawnClassloader` method:** - Create a new method to match the specified `spawnClassloader` signature. - Use `mill.api.ClassLoader.create` to create a classloader. * **Add `callClassloader` method:** - Create a new method to match the specified `callClassloader` signature. - Use `spawnClassloader` to create a classloader and set it as the context classloader. - Execute the provided function with the new classloader and restore the old classloader afterward. * **Add tests in `JvmTests.scala`:** - Add tests for the new `call` method. - Add tests for the new `spawn` method. - Add tests for the new `callClassloader` method. - Add tests for the new `spawnClassloader` method.
diff --git a/main/util/src/mill/util/Jvm.scala b/main/util/src/mill/util/Jvm.scala
@@ -18,71 +18,41 @@ object Jvm extends CoursierSupport {
    * Runs a JVM subprocess with the given configuration and returns a
    * [[os.CommandResult]] with it's aggregated output and error streams
    */
-  def callSubprocess(
+  def call(
       mainClass: String,
-      classPath: Agg[os.Path],
-      jvmArgs: Seq[String] = Seq.empty,
-      envArgs: Map[String, String] = Map.empty,
-      mainArgs: Seq[String] = Seq.empty,
-      workingDir: os.Path = null,
-      streamOut: Boolean = true,
-      check: Boolean = true
-  )(implicit ctx: Ctx): CommandResult = {
-
-    val commandArgs =
-      Vector(javaExe) ++
-        jvmArgs ++
-        Vector("-cp", classPath.iterator.mkString(java.io.File.pathSeparator), mainClass) ++
-        mainArgs
-
-    val workingDir1 = Option(workingDir).getOrElse(ctx.dest)
-    os.makeDir.all(workingDir1)
-
-    os.proc(commandArgs)
-      .call(
-        cwd = workingDir1,
-        env = envArgs,
-        check = check,
-        stdout = if (streamOut) os.Inherit else os.Pipe
-      )
-  }
-
-  /**
-   * Runs a JVM subprocess with the given configuration and returns a
-   * [[os.CommandResult]] with it's aggregated output and error streams
-   */
-  def callSubprocess(
-      mainClass: String,
-      classPath: Agg[os.Path],
+      classPath: Iterable[os.Path],
       jvmArgs: Seq[String],
-      envArgs: Map[String, String],
       mainArgs: Seq[String],
-      workingDir: os.Path,
-      streamOut: Boolean
-  )(implicit ctx: Ctx): CommandResult = {
-    callSubprocess(mainClass, classPath, jvmArgs, envArgs, mainArgs, workingDir, streamOut, true)
-  }
+      env: Map[String, String] = null,
+      cwd: os.Path = null,
+      stdin: ProcessInput = Pipe,
+      stdout: ProcessOutput = Pipe,
+      stderr: ProcessOutput = os.Inherit,
+      mergeErrIntoOut: Boolean = false,
+      timeout: Long = -1,
+      check: Boolean = true,
+      propagateEnv: Boolean = true,
+      timeoutGracePeriod: Long = 100,
+      useCpPassingJar: Boolean = false
+  ): os.CommandResult = {
 
-  /**
-   * Resolves a tool to a path under the currently used JDK (if known).
-   */
-  def jdkTool(toolName: String): String = {
-    sys.props
-      .get("java.home")
-      .map(h =>
-        if (isWin) new File(h, s"bin\\${toolName}.exe")
-        else new File(h, s"bin/${toolName}")
-      )
-      .filter(f => f.exists())
-      .fold(toolName)(_.getAbsolutePath())
+    val commandArgs = jvmCommandArgs(javaExe, mainClass, jvmArgs, classPath, mainArgs, useCpPassingJar)
 
+    os.call(
+      commandArgs,
+      env,
+      cwd,
+      stdin,
+      stdout,
+      stderr,
+      mergeErrIntoOut,
+      timeout,
+      check,
+      propagateEnv,
+      timeoutGracePeriod
+    )
   }
 
-  def javaExe: String = jdkTool("java")
-
-  def defaultBackgroundOutputs(outputDir: os.Path): Option[(ProcessOutput, ProcessOutput)] =
-    Some((outputDir / "stdout.log", outputDir / "stderr.log"))
-
   /**
    * Runs a JVM subprocess with the given configuration and streams
    * it's stdout and stderr to the console.
@@ -99,125 +69,94 @@ object Jvm extends CoursierSupport {
    *                        This might help with long classpaths on OS'es (like Windows)
    *                        which only supports limited command-line length
    */
-  def runSubprocess(
+  def spawn(
       mainClass: String,
-      classPath: Agg[os.Path],
-      jvmArgs: Seq[String] = Seq.empty,
-      envArgs: Map[String, String] = Map.empty,
-      mainArgs: Seq[String] = Seq.empty,
-      workingDir: os.Path = null,
-      background: Boolean = false,
-      useCpPassingJar: Boolean = false,
-      runBackgroundLogToConsole: Boolean = false
-  )(implicit ctx: Ctx): Unit = {
-    runSubprocessWithBackgroundOutputs(
-      mainClass,
-      classPath,
-      jvmArgs,
-      envArgs,
-      mainArgs,
-      workingDir,
-      if (!background) None
-      else if (runBackgroundLogToConsole) {
-        val pwd0 = os.Path(java.nio.file.Paths.get(".").toAbsolutePath)
-        // Hack to forward the background subprocess output to the Mill server process
-        // stdout/stderr files, so the output will get properly slurped up by the Mill server
-        // and shown to any connected Mill client even if the current command has completed
-        Some(
-          (
-            os.PathAppendRedirect(pwd0 / ".." / ServerFiles.stdout),
-            os.PathAppendRedirect(pwd0 / ".." / ServerFiles.stderr)
-          )
-        )
-      } else Jvm.defaultBackgroundOutputs(ctx.dest),
-      useCpPassingJar
-    )
+      classPath: Iterable[os.Path],
+      jvmArgs: Seq[String],
+      mainArgs: Seq[String],
+      env: Map[String, String] = null,
+      cwd: os.Path = null,
+      stdin: ProcessInput = Pipe,
+      stdout: ProcessOutput = Pipe,
+      stderr: ProcessOutput = os.Inherit,
+      mergeErrIntoOut: Boolean = false,
+      propagateEnv: Boolean = true,
+      useCpPassingJar: Boolean = false
+  ): os.SubProcess = {
+
+    val commandArgs = jvmCommandArgs(javaExe, mainClass, jvmArgs, classPath, mainArgs, useCpPassingJar)
+    os.spawn(commandArgs, env, cwd, stdin, stdout, stderr, mergeErrIntoOut, propagateEnv)
   }
 
-  // bincompat shim
-  def runSubprocess(
+  def spawnClassloader(
+      classPath: Iterable[os.Path],
+      sharedPrefixes: Seq[String],
+      isolated: Boolean = true
+  ): java.net.URLClassLoader = {
+    mill.api.ClassLoader.create(
+      classPath.iterator.map(_.toNIO.toUri.toURL).toVector,
+      if (isolated) null else getClass.getClassLoader,
+      sharedPrefixes = sharedPrefixes
+    )()
+  }
+
+  def callClassloader[T](
+      classPath: Iterable[os.Path],
+      sharedPrefixes: Seq[String],
+      isolated: Boolean = true
+  )(f: ClassLoader => T): T = {
+    val oldClassloader = Thread.currentThread().getContextClassLoader
+    val newClassloader = spawnClassloader(classPath, sharedPrefixes, isolated)
+    Thread.currentThread().setContextClassLoader(newClassloader)
+    try {
+      f(newClassloader)
+    } finally {
+      Thread.currentThread().setContextClassLoader(oldClassloader)
+      newClassloader.close()
+    }
+  }
+
+  private def jvmCommandArgs(
+      javaExe: String,
       mainClass: String,
-      classPath: Agg[os.Path],
       jvmArgs: Seq[String],
-      envArgs: Map[String, String],
+      classPath: Iterable[os.Path],
       mainArgs: Seq[String],
-      workingDir: os.Path,
-      background: Boolean,
       useCpPassingJar: Boolean
-  )(implicit ctx: Ctx): Unit =
-    runSubprocess(
-      mainClass,
-      classPath,
-      jvmArgs,
-      envArgs,
-      mainArgs,
-      workingDir,
-      background,
-      useCpPassingJar
-    )
-
-  /**
-   * Runs a JVM subprocess with the given configuration and streams
-   * it's stdout and stderr to the console.
-   * @param mainClass The main class to run
-   * @param classPath The classpath
-   * @param jvmArgs Arguments given to the forked JVM
-   * @param envArgs Environment variables used when starting the forked JVM
-   * @param workingDir The working directory to be used by the forked JVM
-   * @param backgroundOutputs If the subprocess should run in the background, a Tuple of ProcessOutputs containing out and err respectively. Specify None for nonbackground processes.
-   * @param useCpPassingJar When `false`, the `-cp` parameter is used to pass the classpath
-   *                        to the forked JVM.
-   *                        When `true`, a temporary empty JAR is created
-   *                        which contains a `Class-Path` manifest entry containing the actual classpath.
-   *                        This might help with long classpaths on OS'es (like Windows)
-   *                        which only supports limited command-line length
-   */
-  def runSubprocessWithBackgroundOutputs(
-      mainClass: String,
-      classPath: Agg[os.Path],
-      jvmArgs: Seq[String] = Seq.empty,
-      envArgs: Map[String, String] = Map.empty,
-      mainArgs: Seq[String] = Seq.empty,
-      workingDir: os.Path = null,
-      backgroundOutputs: Option[Tuple2[ProcessOutput, ProcessOutput]] = None,
-      useCpPassingJar: Boolean = false
-  )(implicit ctx: Ctx): Unit = {
-
-    val cp =
-      if (useCpPassingJar && !classPath.iterator.isEmpty) {
+  ): Vector[String] = {
+    val classPath2 =
+      if (useCpPassingJar && classPath.nonEmpty) {
         val passingJar = os.temp(prefix = "run-", suffix = ".jar", deleteOnExit = false)
-        ctx.log.debug(
-          s"Creating classpath passing jar '${passingJar}' with Class-Path: ${classPath.iterator.map(
-              _.toNIO.toUri().toURL().toExternalForm()
-            ).mkString(" ")}"
-        )
         createClasspathPassingJar(passingJar, classPath)
         Agg(passingJar)
-      } else {
-        classPath
-      }
+      } else classPath
+
+    Vector(javaExe) ++
+      jvmArgs ++
+      Vector("-cp", classPath2.iterator.mkString(java.io.File.pathSeparator), mainClass) ++
+      mainArgs
+  }
+
+  /**
+   * Resolves a tool to a path under the currently used JDK (if known).
+   */
+  def jdkTool(toolName: String): String = {
+    sys.props
+      .get("java.home")
+      .map(h =>
+        if (isWin) new File(h, s"bin\\${toolName}.exe")
+        else new File(h, s"bin/${toolName}")
+      )
+      .filter(f => f.exists())
+      .fold(toolName)(_.getAbsolutePath())
 
-    val cpArgument = if (cp.nonEmpty) {
-      Vector("-cp", cp.iterator.mkString(java.io.File.pathSeparator))
-    } else Seq.empty
-    val mainClassArgument = if (mainClass.nonEmpty) {
-      Seq(mainClass)
-    } else Seq.empty
-    val args =
-      Vector(javaExe) ++
-        jvmArgs ++
-        cpArgument ++
-        mainClassArgument ++
-        mainArgs
-
-    ctx.log.debug(s"Run subprocess with args: ${args.map(a => s"'${a}'").mkString(" ")}")
-
-    if (backgroundOutputs.nonEmpty)
-      spawnSubprocessWithBackgroundOutputs(args, envArgs, workingDir, backgroundOutputs)
-    else
-      runSubprocess(args, envArgs, workingDir)
   }
 
+  def javaExe: String = jdkTool("java")
+
+  def defaultBackgroundOutputs(outputDir: os.Path): Option[(ProcessOutput, ProcessOutput)] =
+    Some((outputDir / "stdout.log", outputDir / "stderr.log"))
+
   /**
    * Runs a generic subprocess and waits for it to terminate. If process exited with non-zero code, exception
    * will be thrown. If you want to manually handle exit code, check [[runSubprocessWithResult]]
diff --git a/main/util/test/src/mill/util/JvmTests.scala b/main/util/test/src/mill/util/JvmTests.scala
@@ -28,6 +28,46 @@ object JvmTests extends TestSuite {
         Seq(dep1, dep2).map(_.toNIO.toUri().toURL().toExternalForm()).mkString(" "))
     }
 
+    test("call") {
+      val tmpDir = os.temp.dir()
+      val mainClass = "mill.util.TestMain"
+      val classPath = Agg(tmpDir)
+      val jvmArgs = Seq("-Xmx512m")
+      val mainArgs = Seq("arg1", "arg2")
+      val result = Jvm.call(mainClass, classPath, jvmArgs, mainArgs)
+      assert(result.exitCode == 0)
+    }
+
+    test("spawn") {
+      val tmpDir = os.temp.dir()
+      val mainClass = "mill.util.TestMain"
+      val classPath = Agg(tmpDir)
+      val jvmArgs = Seq("-Xmx512m")
+      val mainArgs = Seq("arg1", "arg2")
+      val process = Jvm.spawn(mainClass, classPath, jvmArgs, mainArgs)
+      assert(process.isAlive())
+      process.destroy()
+    }
+
+    test("callClassloader") {
+      val tmpDir = os.temp.dir()
+      val classPath = Agg(tmpDir)
+      val sharedPrefixes = Seq("mill.util")
+      val result = Jvm.callClassloader(classPath, sharedPrefixes) { cl =>
+        cl.loadClass("mill.util.TestMain")
+      }
+      assert(result.getName == "mill.util.TestMain")
+    }
+
+    test("spawnClassloader") {
+      val tmpDir = os.temp.dir()
+      val classPath = Agg(tmpDir)
+      val sharedPrefixes = Seq("mill.util")
+      val classLoader = Jvm.spawnClassloader(classPath, sharedPrefixes)
+      val result = classLoader.loadClass("mill.util.TestMain")
+      assert(result.getName == "mill.util.TestMain")
+    }
+
   }
 
 }
