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

…4085)

Fixes #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
#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
#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,