Skip to content

Commit 1e3205c

Browse files
vishwamarturvishwamartur
committed
Refactor Jvm methods to match specified signatures
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.
1 parent f23475c commit 1e3205c

File tree

2 files changed

+143
-164
lines changed

2 files changed

+143
-164
lines changed

‎main/util/src/mill/util/Jvm.scala

Lines changed: 103 additions & 164 deletions
Original file line numberDiff line numberDiff line change
@@ -18,71 +18,41 @@ object Jvm extends CoursierSupport {
1818
* Runs a JVM subprocess with the given configuration and returns a
1919
* [[os.CommandResult]] with it's aggregated output and error streams
2020
*/
21-
def callSubprocess(
21+
def call(
2222
mainClass: String,
23-
classPath: Agg[os.Path],
24-
jvmArgs: Seq[String] = Seq.empty,
25-
envArgs: Map[String, String] = Map.empty,
26-
mainArgs: Seq[String] = Seq.empty,
27-
workingDir: os.Path = null,
28-
streamOut: Boolean = true,
29-
check: Boolean = true
30-
)(implicit ctx: Ctx): CommandResult = {
31-
32-
val commandArgs =
33-
Vector(javaExe) ++
34-
jvmArgs ++
35-
Vector("-cp", classPath.iterator.mkString(java.io.File.pathSeparator), mainClass) ++
36-
mainArgs
37-
38-
val workingDir1 = Option(workingDir).getOrElse(ctx.dest)
39-
os.makeDir.all(workingDir1)
40-
41-
os.proc(commandArgs)
42-
.call(
43-
cwd = workingDir1,
44-
env = envArgs,
45-
check = check,
46-
stdout = if (streamOut) os.Inherit else os.Pipe
47-
)
48-
}
49-
50-
/**
51-
* Runs a JVM subprocess with the given configuration and returns a
52-
* [[os.CommandResult]] with it's aggregated output and error streams
53-
*/
54-
def callSubprocess(
55-
mainClass: String,
56-
classPath: Agg[os.Path],
23+
classPath: Iterable[os.Path],
5724
jvmArgs: Seq[String],
58-
envArgs: Map[String, String],
5925
mainArgs: Seq[String],
60-
workingDir: os.Path,
61-
streamOut: Boolean
62-
)(implicit ctx: Ctx): CommandResult = {
63-
callSubprocess(mainClass, classPath, jvmArgs, envArgs, mainArgs, workingDir, streamOut, true)
64-
}
26+
env: Map[String, String] = null,
27+
cwd: os.Path = null,
28+
stdin: ProcessInput = Pipe,
29+
stdout: ProcessOutput = Pipe,
30+
stderr: ProcessOutput = os.Inherit,
31+
mergeErrIntoOut: Boolean = false,
32+
timeout: Long = -1,
33+
check: Boolean = true,
34+
propagateEnv: Boolean = true,
35+
timeoutGracePeriod: Long = 100,
36+
useCpPassingJar: Boolean = false
37+
): os.CommandResult = {
6538

66-
/**
67-
* Resolves a tool to a path under the currently used JDK (if known).
68-
*/
69-
def jdkTool(toolName: String): String = {
70-
sys.props
71-
.get("java.home")
72-
.map(h =>
73-
if (isWin) new File(h, s"bin\\${toolName}.exe")
74-
else new File(h, s"bin/${toolName}")
75-
)
76-
.filter(f => f.exists())
77-
.fold(toolName)(_.getAbsolutePath())
39+
val commandArgs = jvmCommandArgs(javaExe, mainClass, jvmArgs, classPath, mainArgs, useCpPassingJar)
7840

41+
os.call(
42+
commandArgs,
43+
env,
44+
cwd,
45+
stdin,
46+
stdout,
47+
stderr,
48+
mergeErrIntoOut,
49+
timeout,
50+
check,
51+
propagateEnv,
52+
timeoutGracePeriod
53+
)
7954
}
8055

81-
def javaExe: String = jdkTool("java")
82-
83-
def defaultBackgroundOutputs(outputDir: os.Path): Option[(ProcessOutput, ProcessOutput)] =
84-
Some((outputDir / "stdout.log", outputDir / "stderr.log"))
85-
8656
/**
8757
* Runs a JVM subprocess with the given configuration and streams
8858
* it's stdout and stderr to the console.
@@ -99,125 +69,94 @@ object Jvm extends CoursierSupport {
9969
* This might help with long classpaths on OS'es (like Windows)
10070
* which only supports limited command-line length
10171
*/
102-
def runSubprocess(
72+
def spawn(
10373
mainClass: String,
104-
classPath: Agg[os.Path],
105-
jvmArgs: Seq[String] = Seq.empty,
106-
envArgs: Map[String, String] = Map.empty,
107-
mainArgs: Seq[String] = Seq.empty,
108-
workingDir: os.Path = null,
109-
background: Boolean = false,
110-
useCpPassingJar: Boolean = false,
111-
runBackgroundLogToConsole: Boolean = false
112-
)(implicit ctx: Ctx): Unit = {
113-
runSubprocessWithBackgroundOutputs(
114-
mainClass,
115-
classPath,
116-
jvmArgs,
117-
envArgs,
118-
mainArgs,
119-
workingDir,
120-
if (!background) None
121-
else if (runBackgroundLogToConsole) {
122-
val pwd0 = os.Path(java.nio.file.Paths.get(".").toAbsolutePath)
123-
// Hack to forward the background subprocess output to the Mill server process
124-
// stdout/stderr files, so the output will get properly slurped up by the Mill server
125-
// and shown to any connected Mill client even if the current command has completed
126-
Some(
127-
(
128-
os.PathAppendRedirect(pwd0 / ".." / ServerFiles.stdout),
129-
os.PathAppendRedirect(pwd0 / ".." / ServerFiles.stderr)
130-
)
131-
)
132-
} else Jvm.defaultBackgroundOutputs(ctx.dest),
133-
useCpPassingJar
134-
)
74+
classPath: Iterable[os.Path],
75+
jvmArgs: Seq[String],
76+
mainArgs: Seq[String],
77+
env: Map[String, String] = null,
78+
cwd: os.Path = null,
79+
stdin: ProcessInput = Pipe,
80+
stdout: ProcessOutput = Pipe,
81+
stderr: ProcessOutput = os.Inherit,
82+
mergeErrIntoOut: Boolean = false,
83+
propagateEnv: Boolean = true,
84+
useCpPassingJar: Boolean = false
85+
): os.SubProcess = {
86+
87+
val commandArgs = jvmCommandArgs(javaExe, mainClass, jvmArgs, classPath, mainArgs, useCpPassingJar)
88+
os.spawn(commandArgs, env, cwd, stdin, stdout, stderr, mergeErrIntoOut, propagateEnv)
13589
}
13690

137-
// bincompat shim
138-
def runSubprocess(
91+
def spawnClassloader(
92+
classPath: Iterable[os.Path],
93+
sharedPrefixes: Seq[String],
94+
isolated: Boolean = true
95+
): java.net.URLClassLoader = {
96+
mill.api.ClassLoader.create(
97+
classPath.iterator.map(_.toNIO.toUri.toURL).toVector,
98+
if (isolated) null else getClass.getClassLoader,
99+
sharedPrefixes = sharedPrefixes
100+
)()
101+
}
102+
103+
def callClassloader[T](
104+
classPath: Iterable[os.Path],
105+
sharedPrefixes: Seq[String],
106+
isolated: Boolean = true
107+
)(f: ClassLoader => T): T = {
108+
val oldClassloader = Thread.currentThread().getContextClassLoader
109+
val newClassloader = spawnClassloader(classPath, sharedPrefixes, isolated)
110+
Thread.currentThread().setContextClassLoader(newClassloader)
111+
try {
112+
f(newClassloader)
113+
} finally {
114+
Thread.currentThread().setContextClassLoader(oldClassloader)
115+
newClassloader.close()
116+
}
117+
}
118+
119+
private def jvmCommandArgs(
120+
javaExe: String,
139121
mainClass: String,
140-
classPath: Agg[os.Path],
141122
jvmArgs: Seq[String],
142-
envArgs: Map[String, String],
123+
classPath: Iterable[os.Path],
143124
mainArgs: Seq[String],
144-
workingDir: os.Path,
145-
background: Boolean,
146125
useCpPassingJar: Boolean
147-
)(implicit ctx: Ctx): Unit =
148-
runSubprocess(
149-
mainClass,
150-
classPath,
151-
jvmArgs,
152-
envArgs,
153-
mainArgs,
154-
workingDir,
155-
background,
156-
useCpPassingJar
157-
)
158-
159-
/**
160-
* Runs a JVM subprocess with the given configuration and streams
161-
* it's stdout and stderr to the console.
162-
* @param mainClass The main class to run
163-
* @param classPath The classpath
164-
* @param jvmArgs Arguments given to the forked JVM
165-
* @param envArgs Environment variables used when starting the forked JVM
166-
* @param workingDir The working directory to be used by the forked JVM
167-
* @param backgroundOutputs If the subprocess should run in the background, a Tuple of ProcessOutputs containing out and err respectively. Specify None for nonbackground processes.
168-
* @param useCpPassingJar When `false`, the `-cp` parameter is used to pass the classpath
169-
* to the forked JVM.
170-
* When `true`, a temporary empty JAR is created
171-
* which contains a `Class-Path` manifest entry containing the actual classpath.
172-
* This might help with long classpaths on OS'es (like Windows)
173-
* which only supports limited command-line length
174-
*/
175-
def runSubprocessWithBackgroundOutputs(
176-
mainClass: String,
177-
classPath: Agg[os.Path],
178-
jvmArgs: Seq[String] = Seq.empty,
179-
envArgs: Map[String, String] = Map.empty,
180-
mainArgs: Seq[String] = Seq.empty,
181-
workingDir: os.Path = null,
182-
backgroundOutputs: Option[Tuple2[ProcessOutput, ProcessOutput]] = None,
183-
useCpPassingJar: Boolean = false
184-
)(implicit ctx: Ctx): Unit = {
185-
186-
val cp =
187-
if (useCpPassingJar && !classPath.iterator.isEmpty) {
126+
): Vector[String] = {
127+
val classPath2 =
128+
if (useCpPassingJar && classPath.nonEmpty) {
188129
val passingJar = os.temp(prefix = "run-", suffix = ".jar", deleteOnExit = false)
189-
ctx.log.debug(
190-
s"Creating classpath passing jar '${passingJar}' with Class-Path: ${classPath.iterator.map(
191-
_.toNIO.toUri().toURL().toExternalForm()
192-
).mkString(" ")}"
193-
)
194130
createClasspathPassingJar(passingJar, classPath)
195131
Agg(passingJar)
196-
} else {
197-
classPath
198-
}
132+
} else classPath
133+
134+
Vector(javaExe) ++
135+
jvmArgs ++
136+
Vector("-cp", classPath2.iterator.mkString(java.io.File.pathSeparator), mainClass) ++
137+
mainArgs
138+
}
139+
140+
/**
141+
* Resolves a tool to a path under the currently used JDK (if known).
142+
*/
143+
def jdkTool(toolName: String): String = {
144+
sys.props
145+
.get("java.home")
146+
.map(h =>
147+
if (isWin) new File(h, s"bin\\${toolName}.exe")
148+
else new File(h, s"bin/${toolName}")
149+
)
150+
.filter(f => f.exists())
151+
.fold(toolName)(_.getAbsolutePath())
199152

200-
val cpArgument = if (cp.nonEmpty) {
201-
Vector("-cp", cp.iterator.mkString(java.io.File.pathSeparator))
202-
} else Seq.empty
203-
val mainClassArgument = if (mainClass.nonEmpty) {
204-
Seq(mainClass)
205-
} else Seq.empty
206-
val args =
207-
Vector(javaExe) ++
208-
jvmArgs ++
209-
cpArgument ++
210-
mainClassArgument ++
211-
mainArgs
212-
213-
ctx.log.debug(s"Run subprocess with args: ${args.map(a => s"'${a}'").mkString(" ")}")
214-
215-
if (backgroundOutputs.nonEmpty)
216-
spawnSubprocessWithBackgroundOutputs(args, envArgs, workingDir, backgroundOutputs)
217-
else
218-
runSubprocess(args, envArgs, workingDir)
219153
}
220154

155+
def javaExe: String = jdkTool("java")
156+
157+
def defaultBackgroundOutputs(outputDir: os.Path): Option[(ProcessOutput, ProcessOutput)] =
158+
Some((outputDir / "stdout.log", outputDir / "stderr.log"))
159+
221160
/**
222161
* Runs a generic subprocess and waits for it to terminate. If process exited with non-zero code, exception
223162
* will be thrown. If you want to manually handle exit code, check [[runSubprocessWithResult]]

‎main/util/test/src/mill/util/JvmTests.scala

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,46 @@ object JvmTests extends TestSuite {
2828
Seq(dep1, dep2).map(_.toNIO.toUri().toURL().toExternalForm()).mkString(" "))
2929
}
3030

31+
test("call") {
32+
val tmpDir = os.temp.dir()
33+
val mainClass = "mill.util.TestMain"
34+
val classPath = Agg(tmpDir)
35+
val jvmArgs = Seq("-Xmx512m")
36+
val mainArgs = Seq("arg1", "arg2")
37+
val result = Jvm.call(mainClass, classPath, jvmArgs, mainArgs)
38+
assert(result.exitCode == 0)
39+
}
40+
41+
test("spawn") {
42+
val tmpDir = os.temp.dir()
43+
val mainClass = "mill.util.TestMain"
44+
val classPath = Agg(tmpDir)
45+
val jvmArgs = Seq("-Xmx512m")
46+
val mainArgs = Seq("arg1", "arg2")
47+
val process = Jvm.spawn(mainClass, classPath, jvmArgs, mainArgs)
48+
assert(process.isAlive())
49+
process.destroy()
50+
}
51+
52+
test("callClassloader") {
53+
val tmpDir = os.temp.dir()
54+
val classPath = Agg(tmpDir)
55+
val sharedPrefixes = Seq("mill.util")
56+
val result = Jvm.callClassloader(classPath, sharedPrefixes) { cl =>
57+
cl.loadClass("mill.util.TestMain")
58+
}
59+
assert(result.getName == "mill.util.TestMain")
60+
}
61+
62+
test("spawnClassloader") {
63+
val tmpDir = os.temp.dir()
64+
val classPath = Agg(tmpDir)
65+
val sharedPrefixes = Seq("mill.util")
66+
val classLoader = Jvm.spawnClassloader(classPath, sharedPrefixes)
67+
val result = classLoader.loadClass("mill.util.TestMain")
68+
assert(result.getName == "mill.util.TestMain")
69+
}
70+
3171
}
3272

3373
}

0 commit comments

Comments
 (0)