add JavaModule#runBackground, to allow a task to kick off processes to run in the background that only die when the task is re-run

3 files changed, 108 insertions, 0 deletions

diff --git a/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/BackgroundWrapper.java b/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/BackgroundWrapper.java
new file mode 100644
index 00000000..64719b53
--- /dev/null
+++ b/scalalib/backgroundwrapper/src/mill/scalalib/backgroundwrapper/BackgroundWrapper.java
@@ -0,0 +1,40 @@
+package mill.scalalib.backgroundwrapper;
+
+public class BackgroundWrapper {
+    public static void main(String[] args) throws Exception{
+        String watched = args[0];
+        String tombstone = args[1];
+        String expected = args[2];
+        Thread watcher = new Thread(new Runnable() {
+            @Override
+            public void run() {
+                while (true) {
+                    try{
+                        Thread.sleep(10);
+                        String token = new String(
+                            java.nio.file.Files.readAllBytes(java.nio.file.Paths.get(watched))
+                        );
+                        if (!token.equals(expected)) {
+                            new java.io.File(tombstone).createNewFile();
+                            System.exit(0);
+                        }
+                    }catch(Exception e){
+                        try {
+                            new java.io.File(tombstone).createNewFile();
+                        }catch(Exception e2){}
+                        System.exit(0);
+                    }
+
+                }
+            }
+        });
+        watcher.setDaemon(true);
+        watcher.start();
+        String realMain = args[3];
+        String[] realArgs = new String[args.length - 4];
+        for(int i = 0; i < args.length-4; i++){
+            realArgs[i] = args[i+4];
+        }
+        Class.forName(realMain).getMethod("main", String[].class).invoke(null, (Object)realArgs);
+    }
+}
diff --git a/scalalib/src/mill/scalalib/JavaModule.scala b/scalalib/src/mill/scalalib/JavaModule.scala
index 3907b6a6..87dfb1b9 100644
--- a/scalalib/src/mill/scalalib/JavaModule.scala
+++ b/scalalib/src/mill/scalalib/JavaModule.scala
@@ -274,6 +274,67 @@ trait JavaModule extends mill.Module with TaskModule { outer =>
     }
   }
 
+  private[this] def backgroundSetup(dest: Path) = {
+    val token = java.util.UUID.randomUUID().toString
+    val procId = dest / ".mill-background-process-id"
+    val procTombstone = dest / ".mill-background-process-tombstone"
+    // The backgrounded subprocesses poll the procId file, and kill themselves
+    // when the procId file is deleted. This deletion happens immediately before
+    // the body of these commands run, but we cannot be sure the subprocess has
+    // had time to notice.
+    //
+    // To make sure we wait for the previous subprocess to
+    // die, we make the subprocess write a tombstone file out when it kills
+    // itself due to procId being deleted, and we wait a short time on task-start
+    // to see if such a tombstone appears. If a tombstone appears, we can be sure
+    // the subprocess has killed itself, and can continue. If a tombstone doesn't
+    // appear in a short amount of time, we assume the subprocess exited or was
+    // killed via some other means, and continue anyway.
+    val start = System.currentTimeMillis()
+    while({
+      if (exists(procTombstone)) {
+        Thread.sleep(10)
+        rm(procTombstone)
+        true
+      } else {
+        Thread.sleep(10)
+        System.currentTimeMillis() - start < 100
+      }
+    })()
+
+    write(procId, token)
+    write(procTombstone, token)
+    (procId, procTombstone, token)
+  }
+  def runBackground(args: String*) = T.command{
+    val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
+    try Result.Success(Jvm.interactiveSubprocess(
+      "mill.scalalib.backgroundwrapper.BackgroundWrapper",
+      (runClasspath() ++ scalaWorker.backgroundWrapperClasspath()).map(_.path),
+      forkArgs(),
+      forkEnv(),
+      Seq(procId.toString, procTombstone.toString, token, finalMainClass()) ++ args,
+      workingDir = ammonite.ops.pwd,
+      background = true
+    )) catch { case e: InteractiveShelloutException =>
+       Result.Failure("subprocess failed")
+    }
+  }
+
+  def runMainBackground(mainClass: String, args: String*) = T.command{
+    val (procId, procTombstone, token) = backgroundSetup(T.ctx().dest)
+    try Result.Success(Jvm.interactiveSubprocess(
+      "mill.scalalib.backgroundwrapper.BackgroundWrapper",
+      (runClasspath() ++ scalaWorker.backgroundWrapperClasspath()).map(_.path),
+      forkArgs(),
+      forkEnv(),
+      Seq(procId.toString, procTombstone.toString, token, mainClass) ++ args,
+      workingDir = ammonite.ops.pwd,
+      background = true
+    )) catch { case e: InteractiveShelloutException =>
+      Result.Failure("subprocess failed")
+    }
+  }
 
   def runMainLocal(mainClass: String, args: String*) = T.command {
     Jvm.runLocal(
diff --git a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
index 6865c541..3d1a195a 100644
--- a/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
+++ b/scalalib/src/mill/scalalib/ScalaWorkerApi.scala
@@ -29,6 +29,13 @@ trait ScalaWorkerModule extends mill.Module{
     mill.modules.Util.millProjectModule("MILL_SCALA_LIB", "mill-scalalib", repositories)
   }
 
+  def backgroundWrapperClasspath = T{
+    mill.modules.Util.millProjectModule(
+      "MILL_BACKGROUNDWRAPPER", "mill-scalalib-backgroundwrapper",
+      repositories, artifactSuffix = ""
+    )
+  }
+
   def worker: Worker[ScalaWorkerApi] = T.worker{
     val cl = mill.util.ClassLoader.create(
       classpath().map(_.path.toNIO.toUri.toURL).toVector,