Generalize .runBackground to work with all kinds of subprocesses (c…

…om-lihaoyi#4085)

Fixes com-lihaoyi#657

We re-use `MillBackgroundWrapper` and extend it to optionally spawn
subprocess instead of calling the in-process main method. There's a bit
of overhead in the JVM wrapper, but we need some kind of process wrapper
to manage the mutex and termination even when the Mill process
terminates, and eventually when
com-lihaoyi#4007 lands we can use that to
provide lighter-weight wrappers. We add a shutdown hook and copy the
termination grace-period logic from OS-Lib to try and gently kill the
wrapped process when necessary

Integrated this into `PythonModule#runBackground` and added a unit test
to verify the asynchronous launch, background existence (by checking
that a file lock is taken) and termination on deletion. We can integrate
this into `TypescriptModule` as well but punting that for later

The subprocess APIs are a mess but leaving that to fix in 0.13.0
com-lihaoyi#3772

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())
 }

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",
+        "<subprocess>",
+        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"
 
   /**

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()
+

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")
+      }
+    }
+  }
+}

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("<subprocess>")) {
+      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());
+    }
   }
 }

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,