Skip to content

Commit

Permalink
Refactor Jvm methods to match specified signatures
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
vishwamartur committed Nov 16, 2024
1 parent f23475c commit 1e3205c
Show file tree
Hide file tree
Showing 2 changed files with 143 additions and 164 deletions.
267 changes: 103 additions & 164 deletions main/util/src/mill/util/Jvm.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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]]
Expand Down
40 changes: 40 additions & 0 deletions main/util/test/src/mill/util/JvmTests.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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")
}

}

}

0 comments on commit 1e3205c

Please sign in to comment.