diff --git a/pythonlib/package.mill b/pythonlib/package.mill index cbc3a6a73f2..8e9f58da976 100644 --- a/pythonlib/package.mill +++ b/pythonlib/package.mill @@ -9,4 +9,5 @@ object `package` extends RootModule with build.MillPublishScalaModule { // we depend on scalalib for re-using some common infrastructure (e.g. License // management of projects), NOT for reusing build logic def moduleDeps = Seq(build.main, build.scalalib) + def testTransitiveDeps = super.testTransitiveDeps() ++ Seq(build.scalalib.backgroundwrapper.testDep()) } diff --git a/pythonlib/src/mill/pythonlib/PythonModule.scala b/pythonlib/src/mill/pythonlib/PythonModule.scala index 219673a3254..e8b1ad57219 100644 --- a/pythonlib/src/mill/pythonlib/PythonModule.scala +++ b/pythonlib/src/mill/pythonlib/PythonModule.scala @@ -100,16 +100,20 @@ trait PythonModule extends PipModule with TaskModule { outer => def runner: Task[PythonModule.Runner] = Task.Anon { new PythonModule.RunnerImpl( command0 = pythonExe().path.toString, - env0 = Map( - "PYTHONPATH" -> transitivePythonPath().map(_.path).mkString(java.io.File.pathSeparator), - "PYTHONPYCACHEPREFIX" -> (T.dest / "cache").toString, - if (Task.log.colored) { "FORCE_COLOR" -> "1" } - else { "NO_COLOR" -> "1" } - ), + env0 = runnerEnvTask(), workingDir0 = Task.workspace ) } + private def runnerEnvTask = Task.Anon { + Map( + "PYTHONPATH" -> transitivePythonPath().map(_.path).mkString(java.io.File.pathSeparator), + "PYTHONPYCACHEPREFIX" -> (T.dest / "cache").toString, + if (Task.log.colored) { "FORCE_COLOR" -> "1" } + else { "NO_COLOR" -> "1" } + ) + } + /** * Run a typechecker on this module. */ @@ -140,6 +144,37 @@ trait PythonModule extends PipModule with TaskModule { outer => ) } + /** + * Run the main python script of this module. + * + * @see [[mainScript]] + */ + def runBackground(args: mill.define.Args) = Task.Command { + val (procUuidPath, procLockfile, procUuid) = mill.scalalib.RunModule.backgroundSetup(T.dest) + + Jvm.runSubprocess( + mainClass = "mill.scalalib.backgroundwrapper.MillBackgroundWrapper", + classPath = mill.scalalib.ZincWorkerModule.backgroundWrapperClasspath().map(_.path).toSeq, + jvmArgs = Nil, + envArgs = runnerEnvTask(), + mainArgs = Seq( + procUuidPath.toString, + procLockfile.toString, + procUuid, + "500", + "", + pythonExe().path.toString, + mainScript().path.toString + ) ++ args.value, + workingDir = T.workspace, + background = true, + useCpPassingJar = false, + runBackgroundLogToConsole = true, + javaHome = mill.scalalib.ZincWorkerModule.javaHome().map(_.path) + ) + () + } + override def defaultCommandName(): String = "run" /** diff --git a/pythonlib/test/resources/run-background/foo/src/foo.py b/pythonlib/test/resources/run-background/foo/src/foo.py new file mode 100644 index 00000000000..d5c468b254b --- /dev/null +++ b/pythonlib/test/resources/run-background/foo/src/foo.py @@ -0,0 +1,12 @@ +import sys +import time +import fcntl + +def main(): + with open(sys.argv[1], 'a+') as file: + fcntl.lockf(file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + time.sleep(9999) + +if __name__ == "__main__": + main() + diff --git a/pythonlib/test/src/mill/pythonlib/RunBackgroundTests.scala b/pythonlib/test/src/mill/pythonlib/RunBackgroundTests.scala new file mode 100644 index 00000000000..289d2e6a0e0 --- /dev/null +++ b/pythonlib/test/src/mill/pythonlib/RunBackgroundTests.scala @@ -0,0 +1,42 @@ +package mill.pythonlib + +import mill.testkit.{TestBaseModule, UnitTester} +import utest.* +import mill._ + +object RunBackgroundTests extends TestSuite { + + object HelloWorldPython extends TestBaseModule { + object foo extends PythonModule { + override def mainScript = T.source(millSourcePath / "src" / "foo.py") + } + } + + val resourcePath = os.Path(sys.env("MILL_TEST_RESOURCE_DIR")) / "run-background" + def tests: Tests = Tests { + test("runBackground") { + val eval = UnitTester(HelloWorldPython, resourcePath) + + val lockedFile = os.temp() + val Right(result) = eval.apply(HelloWorldPython.foo.runBackground(Args(lockedFile))) + val maxSleep = 20000 + val now1 = System.currentTimeMillis() + val lock = mill.main.client.lock.Lock.file(lockedFile.toString()) + + def sleepIfTimeAvailable(error: String) = { + Thread.sleep(100) + if (System.currentTimeMillis() - now1 > maxSleep) throw new Exception(error) + } + + Thread.sleep(1000) // Make sure that the file remains locked even after a significant sleep + + while (lock.probe()) sleepIfTimeAvailable("File never locked by python subprocess") + + os.remove.all(eval.outPath / "foo" / "runBackground.dest") + + while (!lock.probe()) { + sleepIfTimeAvailable("File never unlocked after runBackground.dest removed") + } + } + } +} diff --git a/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/MillBackgroundWrapper.java b/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/MillBackgroundWrapper.java index ebdd96f84d9..c7d9947c884 100644 --- a/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/MillBackgroundWrapper.java +++ b/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/MillBackgroundWrapper.java @@ -2,10 +2,7 @@ import java.io.RandomAccessFile; import java.nio.channels.FileChannel; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.Paths; -import java.nio.file.StandardOpenOption; +import java.nio.file.*; public class MillBackgroundWrapper { public static void main(String[] args) throws Exception { @@ -55,6 +52,27 @@ public static void main(String[] args) throws Exception { // Actually start the Java main method we wanted to run in the background String realMain = args[4]; String[] realArgs = java.util.Arrays.copyOfRange(args, 5, args.length); - Class.forName(realMain).getMethod("main", String[].class).invoke(null, (Object) realArgs); + if (!realMain.equals("")) { + Class.forName(realMain).getMethod("main", String[].class).invoke(null, (Object) realArgs); + } else { + Process subprocess = new ProcessBuilder().command(realArgs).inheritIO().start(); + + Runtime.getRuntime().addShutdownHook(new Thread(() -> { + subprocess.destroy(); + + long now = System.currentTimeMillis(); + + while (subprocess.isAlive() && System.currentTimeMillis() - now < 100) { + try { + Thread.sleep(1); + } catch (InterruptedException e) { + } + if (subprocess.isAlive()) { + subprocess.destroyForcibly(); + } + } + })); + System.exit(subprocess.waitFor()); + } } } diff --git a/scalalib/src/mill/scalalib/RunModule.scala b/scalalib/src/mill/scalalib/RunModule.scala index a371830e1e7..04cd5f41505 100644 --- a/scalalib/src/mill/scalalib/RunModule.scala +++ b/scalalib/src/mill/scalalib/RunModule.scala @@ -166,7 +166,7 @@ trait RunModule extends WithZincWorker { def runBackgroundTask(mainClass: Task[String], args: Task[Args] = Task.Anon(Args())): Task[Unit] = Task.Anon { - val (procUuidPath, procLockfile, procUuid) = backgroundSetup(T.dest) + val (procUuidPath, procLockfile, procUuid) = RunModule.backgroundSetup(T.dest) runner().run( args = Seq( procUuidPath.toString, @@ -207,7 +207,7 @@ trait RunModule extends WithZincWorker { runUseArgsFile: Boolean, backgroundOutputs: Option[Tuple2[ProcessOutput, ProcessOutput]] )(args: String*): Ctx => Result[Unit] = ctx => { - val (procUuidPath, procLockfile, procUuid) = backgroundSetup(taskDest) + val (procUuidPath, procLockfile, procUuid) = RunModule.backgroundSetup(taskDest) try Result.Success( Jvm.runSubprocessWithBackgroundOutputs( "mill.scalalib.backgroundwrapper.MillBackgroundWrapper", @@ -231,16 +231,16 @@ trait RunModule extends WithZincWorker { Result.Failure("subprocess failed") } } +} + +object RunModule { - private[this] def backgroundSetup(dest: os.Path): (Path, Path, String) = { + private[mill] def backgroundSetup(dest: os.Path): (Path, Path, String) = { val procUuid = java.util.UUID.randomUUID().toString val procUuidPath = dest / ".mill-background-process-uuid" val procLockfile = dest / ".mill-background-process-lock" (procUuidPath, procLockfile, procUuid) } -} - -object RunModule { trait Runner { def run( args: os.Shellable,