WIP keep mill server alive if you Ctrl-C during --watch (#339)

* wip

* Clean up more resources in the Mill client after every command

* catch and ignore SIGINT in Mill server to make it survive Ctrl-C on the client

11 files changed, 159 insertions, 87 deletions

diff --git a/build.sc b/build.sc
index 5f77873a..79b2d46a 100755
--- a/build.sc
+++ b/build.sc
@@ -243,10 +243,10 @@ def launcherScript(shellJvmArgs: Seq[String],
 
       s"""case "$$1" in
          |  -i | --interactive )
-         |    ${java("mill.Main")}
+         |    ${java("mill.MillMain")}
          |    ;;
          |  *)
-         |    ${java("mill.main.client.Main")}
+         |    ${java("mill.main.client.MillClientMain")}
          |    ;;
          |esac""".stripMargin
     },
diff --git a/main/client/src/mill/main/client/Lock.java b/main/client/src/mill/main/client/Lock.java
index 890a352b..6e5f18b0 100644
--- a/main/client/src/mill/main/client/Lock.java
+++ b/main/client/src/mill/main/client/Lock.java
@@ -1,7 +1,8 @@
 package mill.main.client;
-public abstract class Lock{
+public abstract class Lock implements AutoCloseable{
     abstract public Locked lock() throws Exception;
     abstract public Locked tryLock() throws Exception;
+
     public void await() throws Exception{
         lock().release();
     }
diff --git a/main/client/src/mill/main/client/Locks.java b/main/client/src/mill/main/client/Locks.java
index 2843973d..64259293 100644
--- a/main/client/src/mill/main/client/Locks.java
+++ b/main/client/src/mill/main/client/Locks.java
@@ -4,7 +4,7 @@ import java.io.RandomAccessFile;
 import java.nio.channels.FileChannel;
 import java.util.concurrent.locks.ReentrantLock;
 
-public class Locks{
+public class Locks implements AutoCloseable{
     public Lock processLock;
     public Lock serverLock;
     public Lock clientLock;
@@ -24,6 +24,13 @@ public class Locks{
             this.clientLock = new MemoryLock();
         }};
     }
+
+    @Override
+    public void close() throws Exception {
+        processLock.close();
+        serverLock.close();
+        clientLock.close();
+    }
 }
 class FileLocked implements Locked{
     private java.nio.channels.FileLock lock;
@@ -61,6 +68,12 @@ class FileLock extends Lock{
             return true;
         }
     }
+
+    @Override
+    public void close() throws Exception {
+        raf.close();
+        chan.close();
+    }
 }
 class MemoryLocked implements Locked{
     java.util.concurrent.locks.Lock l;
@@ -86,4 +99,9 @@ class MemoryLock extends Lock{
         if (innerLock.tryLock()) return new MemoryLocked(innerLock);
         else return null;
     }
+
+    @Override
+    public void close() throws Exception {
+        innerLock.unlock();
+    }
 }
diff --git a/main/client/src/mill/main/client/Main.java b/main/client/src/mill/main/client/MillClientMain.java
index 98445d3c..17a043f6 100644
--- a/main/client/src/mill/main/client/Main.java
+++ b/main/client/src/mill/main/client/MillClientMain.java
@@ -9,10 +9,10 @@ import java.net.URL;
 import java.nio.channels.FileChannel;
 import java.util.*;
 
-public class Main {
+public class MillClientMain {
     static void initServer(String lockBase, boolean setJnaNoSys) throws IOException,URISyntaxException{
         ArrayList<String> selfJars = new ArrayList<String>();
-        ClassLoader current = Main.class.getClassLoader();
+        ClassLoader current = MillClientMain.class.getClassLoader();
         while(current != null){
             if (current instanceof java.net.URLClassLoader) {
                 URL[] urls = ((java.net.URLClassLoader) current).getURLs();
@@ -38,7 +38,7 @@ public class Main {
         }
         l.add("-cp");
         l.add(String.join(File.pathSeparator, selfJars));
-        l.add("mill.main.ServerMain");
+        l.add("mill.main.MillServerMain");
         l.add(lockBase);
 
         new java.lang.ProcessBuilder()
@@ -48,6 +48,9 @@ public class Main {
                 .start();
     }
     public static void main(String[] args) throws Exception{
+        System.exit(main0(args));
+    }
+    public static int main0(String[] args) throws Exception{
         boolean setJnaNoSys = System.getProperty("jna.nosys") == null;
         Map<String, String> env = System.getenv();
         if (setJnaNoSys) {
@@ -58,33 +61,35 @@ public class Main {
             index += 1;
             String lockBase = "out/mill-worker-" + index;
             new java.io.File(lockBase).mkdirs();
-            RandomAccessFile lockFile = new RandomAccessFile(lockBase + "/clientLock", "rw");
-            FileChannel channel = lockFile.getChannel();
-            java.nio.channels.FileLock tryLock = channel.tryLock();
-            if (tryLock == null) {
-                lockFile.close();
-                channel.close();
-            } else {
-                int exitCode = Main.run(
-                        lockBase,
-                        new Runnable() {
-                            @Override
-                            public void run() {
-                                try{
-                                    initServer(lockBase, setJnaNoSys);
-                                }catch(Exception e){
-                                    throw new RuntimeException(e);
+
+            try(RandomAccessFile lockFile = new RandomAccessFile(lockBase + "/clientLock", "rw");
+                FileChannel channel = lockFile.getChannel();
+                java.nio.channels.FileLock tryLock = channel.tryLock();
+                Locks locks = Locks.files(lockBase)){
+                if (tryLock != null) {
+                    int exitCode = MillClientMain.run(
+                            lockBase,
+                            new Runnable() {
+                                @Override
+                                public void run() {
+                                    try{
+                                        initServer(lockBase, setJnaNoSys);
+                                    }catch(Exception e){
+                                        throw new RuntimeException(e);
+                                    }
                                 }
-                            }
-                        },
-                        Locks.files(lockBase),
-                        System.in,
-                        System.out,
-                        System.err,
-                        args,
-                        env
-                );
-                System.exit(exitCode);
+                            },
+                            locks,
+                            System.in,
+                            System.out,
+                            System.err,
+                            args,
+                            env
+                    );
+                    return exitCode;
+                }
+            } finally{
+
             }
         }
         throw new Exception("Reached max process limit: " + 5);
@@ -99,12 +104,12 @@ public class Main {
                           String[] args,
                           Map<String, String> env) throws Exception{
 
-        FileOutputStream f = new FileOutputStream(lockBase + "/run");
-        f.write(System.console() != null ? 1 : 0);
-        Util.writeString(f, System.getProperty("MILL_VERSION"));
-        Util.writeArgs(args, f);
-        Util.writeMap(env, f);
-        f.close();
+        try(FileOutputStream f = new FileOutputStream(lockBase + "/run")){
+            f.write(System.console() != null ? 1 : 0);
+            Util.writeString(f, System.getProperty("MILL_VERSION"));
+            Util.writeArgs(args, f);
+            Util.writeMap(env, f);
+        }
 
         boolean serverInit = false;
         if (locks.processLock.probe()) {
@@ -120,6 +125,7 @@ public class Main {
         Socket ioSocket = null;
 
         long retryStart = System.currentTimeMillis();
+
         while(ioSocket == null && System.currentTimeMillis() - retryStart < 1000){
             try{
                 ioSocket = Util.isWindows?
@@ -146,16 +152,12 @@ public class Main {
 
         locks.serverLock.await();
 
-        try{
-            return Integer.parseInt(
-                new BufferedReader(
-                    new InputStreamReader(
-                        new FileInputStream(lockBase + "/exitCode")
-                    )
-                ).readLine()
-            );
+        try(FileInputStream fos = new FileInputStream(lockBase + "/exitCode")){
+            return Integer.parseInt(new BufferedReader(new InputStreamReader(fos)).readLine());
         } catch(Throwable e){
             return 1;
+        } finally{
+            ioSocket.close();
         }
     }
 }
diff --git a/main/src/mill/Main.scala b/main/src/mill/MillMain.scala
index 5573a325..ad1bd39d 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/MillMain.scala
@@ -9,7 +9,7 @@ import io.github.retronym.java9rtexport.Export
 import mill.eval.Evaluator
 import mill.util.DummyInputStream
 
-object Main {
+object MillMain {
 
   def main(args: Array[String]): Unit = {
     val as = args match {
@@ -23,7 +23,8 @@ object Main {
       System.in,
       System.out,
       System.err,
-      System.getenv().asScala.toMap
+      System.getenv().asScala.toMap,
+      b => ()
     )
     System.exit(if(result) 0 else 1)
   }
@@ -34,7 +35,8 @@ object Main {
             stdin: InputStream,
             stdout: PrintStream,
             stderr: PrintStream,
-            env: Map[String, String]): (Boolean, Option[Evaluator.State]) = {
+            env: Map[String, String],
+            setIdle: Boolean => Unit): (Boolean, Option[Evaluator.State]) = {
     import ammonite.main.Cli
 
     val removed = Set("predef-code", "no-home-predef")
@@ -103,7 +105,8 @@ object Main {
             config.copy(colored = Some(mainInteractive)),
             stdout, stderr, stdin,
             stateCache,
-            env
+            env,
+            setIdle
           )
 
           if (mill.main.client.Util.isJava9OrAbove) {
diff --git a/main/src/mill/main/MainRunner.scala b/main/src/mill/main/MainRunner.scala
index a289db5f..bf0d5901 100644
--- a/main/src/mill/main/MainRunner.scala
+++ b/main/src/mill/main/MainRunner.scala
@@ -22,7 +22,8 @@ class MainRunner(val config: ammonite.main.Cli.Config,
                  errPrintStream: PrintStream,
                  stdIn: InputStream,
                  stateCache0: Option[Evaluator.State] = None,
-                 env : Map[String, String])
+                 env : Map[String, String],
+                 setIdle: Boolean => Unit)
   extends ammonite.MainRunner(
     config, outprintStream, errPrintStream,
     stdIn, outprintStream, errPrintStream
@@ -35,8 +36,9 @@ class MainRunner(val config: ammonite.main.Cli.Config,
     def statAll() = watched.forall{ case (file, lastMTime) =>
       Interpreter.pathSignature(file) == lastMTime
     }
-
+    setIdle(true)
     while(statAll()) Thread.sleep(100)
+    setIdle(false)
   }
 
   /**
diff --git a/main/src/mill/main/Server.scala b/main/src/mill/main/MillServerMain.scala
index 07703bed..092d958c 100644
--- a/main/src/mill/main/Server.scala
+++ b/main/src/mill/main/MillServerMain.scala
@@ -3,14 +3,16 @@ package mill.main
 import java.io._
 import java.net.Socket
 
-import mill.Main
+import mill.MillMain
+
 import scala.collection.JavaConverters._
 import org.scalasbt.ipcsocket._
 import mill.main.client._
 import mill.eval.Evaluator
 import mill.util.DummyInputStream
+import sun.misc.{Signal, SignalHandler}
 
-trait ServerMain[T]{
+trait MillServerMain[T]{
   var stateCache = Option.empty[T]
   def main0(args: Array[String],
             stateCache: Option[T],
@@ -18,11 +20,22 @@ trait ServerMain[T]{
             stdin: InputStream,
             stdout: PrintStream,
             stderr: PrintStream,
-            env : Map[String, String]): (Boolean, Option[T])
+            env : Map[String, String],
+            setIdle: Boolean => Unit): (Boolean, Option[T])
 }
 
-object ServerMain extends mill.main.ServerMain[Evaluator.State]{
+object MillServerMain extends mill.main.MillServerMain[Evaluator.State]{
   def main(args0: Array[String]): Unit = {
+    // Disable SIGINT interrupt signal in the Mill server.
+    //
+    // This gets passed through from the client to server whenever the user
+    // hits `Ctrl-C`, which by default kills the server, which defeats the purpose
+    // of running a background server. Furthermore, the background server already
+    // can detect when the Mill client goes away, which is necessary to handle
+    // the case when a Mill client that did *not* spawn the server gets `CTRL-C`ed
+    Signal.handle(new Signal("INT"), new SignalHandler () {
+      def handle(sig: Signal) = {} // do nothing
+    })
     new Server(
       args0(0),
       this,
@@ -37,20 +50,24 @@ object ServerMain extends mill.main.ServerMain[Evaluator.State]{
             stdin: InputStream,
             stdout: PrintStream,
             stderr: PrintStream,
-            env : Map[String, String]) = Main.main0(
-    args,
-    stateCache,
-    mainInteractive,
-    DummyInputStream,
-    stdout,
-    stderr,
-    env
-  )
+            env : Map[String, String],
+            setIdle: Boolean => Unit) = {
+    MillMain.main0(
+      args,
+      stateCache,
+      mainInteractive,
+      DummyInputStream,
+      stdout,
+      stderr,
+      env,
+      setIdle
+    )
+  }
 }
 
 
 class Server[T](lockBase: String,
-                sm: ServerMain[T],
+                sm: MillServerMain[T],
                 interruptServer: () => Unit,
                 acceptTimeout: Int,
                 locks: Locks) {
@@ -71,6 +88,7 @@ class Server[T](lockBase: String,
           }
 
           val sockOpt = Server.interruptWith(
+            "MillSocketTimeoutInterruptThread",
             acceptTimeout,
             socketClose(),
             serverSocket.accept()
@@ -111,7 +129,8 @@ class Server[T](lockBase: String,
     val env = Util.parseMap(argStream)
     argStream.close()
 
-    var done = false
+    @volatile var done = false
+    @volatile var idle = false
     val t = new Thread(() =>
 
       try {
@@ -122,7 +141,8 @@ class Server[T](lockBase: String,
           socketIn,
           stdout,
           stderr,
-          env.asScala.toMap
+          env.asScala.toMap,
+          idle = _
         )
 
         sm.stateCache = newStateCache
@@ -134,18 +154,18 @@ class Server[T](lockBase: String,
         sm.stateCache = sc
       } finally{
         done = true
-      }
+        idle = true
+      },
+      "MillServerActionRunner"
     )
-
     t.start()
     // We cannot simply use Lock#await here, because the filesystem doesn't
     // realize the clientLock/serverLock are held by different threads in the
     // two processes and gives a spurious deadlock error
-    while(!done && !locks.clientLock.probe()) {
-      Thread.sleep(3)
-    }
+    while(!done && !locks.clientLock.probe()) Thread.sleep(3)
+
+    if (!idle) interruptServer()
 
-    if (!done) interruptServer()
 
     t.interrupt()
     t.stop()
@@ -176,17 +196,22 @@ object Server{
     }
 
   }
-  def interruptWith[T](millis: Int, close: => Unit, t: => T): Option[T] = {
+  def interruptWith[T](threadName: String, millis: Int, close: => Unit, t: => T): Option[T] = {
     @volatile var interrupt = true
     @volatile var interrupted = false
-    new Thread(() => {
-      Thread.sleep(millis)
-      if (interrupt) {
-        interrupted = true
-        close
-      }
-    }).start()
+    val thread = new Thread(
+      () => {
+        try Thread.sleep(millis)
+        catch{ case t: InterruptedException => /* Do Nothing */ }
+        if (interrupt) {
+          interrupted = true
+          close
+        }
+      },
+      threadName
+    )
 
+    thread.start()
     try {
       val res =
         try Some(t)
@@ -196,6 +221,7 @@ object Server{
       else res
 
     } finally {
+      thread.interrupt()
       interrupt = false
     }
   }
diff --git a/main/test/src/mill/main/ClientServerTests.scala b/main/test/src/mill/main/ClientServerTests.scala
index 7ed826af..2462f650 100644
--- a/main/test/src/mill/main/ClientServerTests.scala
+++ b/main/test/src/mill/main/ClientServerTests.scala
@@ -6,14 +6,15 @@ import mill.main.client.{Util, Locks}
 
 import scala.collection.JavaConverters._
 import utest._
-class EchoServer extends ServerMain[Int]{
+class EchoServer extends MillServerMain[Int]{
   def main0(args: Array[String],
             stateCache: Option[Int],
             mainInteractive: Boolean,
             stdin: InputStream,
             stdout: PrintStream,
             stderr: PrintStream,
-            env: Map[String, String]) = {
+            env: Map[String, String],
+            setIdle: Boolean => Unit) = {
 
     val reader = new BufferedReader(new InputStreamReader(stdin))
     val str = reader.readLine()
@@ -60,7 +61,7 @@ object ClientServerTests extends TestSuite{
                   (env : Map[String, String], args: Array[String]) = {
     val (in, out, err) = initStreams()
     Server.lockBlock(locks.clientLock){
-      mill.main.client.Main.run(
+      mill.main.client.MillClientMain.run(
         tmpDir.toString,
         () => spawnEchoServer(tmpDir, locks),
         locks,
diff --git a/main/test/src/mill/util/ScriptTestSuite.scala b/main/test/src/mill/util/ScriptTestSuite.scala
index e9f31ce3..a23b34fe 100644
--- a/main/test/src/mill/util/ScriptTestSuite.scala
+++ b/main/test/src/mill/util/ScriptTestSuite.scala
@@ -16,7 +16,8 @@ abstract class ScriptTestSuite(fork: Boolean) extends TestSuite{
   val stdIn = new ByteArrayInputStream(Array())
   lazy val runner = new mill.main.MainRunner(
     ammonite.main.Cli.Config(wd = wd),
-    stdOutErr, stdOutErr, stdIn, None, Map.empty
+    stdOutErr, stdOutErr, stdIn, None, Map.empty,
+    b => ()
   )
   def eval(s: String*) = {
     if (!fork) runner.runScript(workspacePath / buildPath , s.toList)
diff --git a/scratch/build.sc b/scratch/build.sc
new file mode 100644
index 00000000..9ee78d63
--- /dev/null
+++ b/scratch/build.sc
@@ -0,0 +1,3 @@
+def thingy = T {
+  1234567
+}
\ No newline at end of file
diff --git a/scratch/readme.md b/scratch/readme.md
new file mode 100644
index 00000000..6f6ffef8
--- /dev/null
+++ b/scratch/readme.md
@@ -0,0 +1,15 @@
+Scratch folder for testing out the local checkout of Mill on ad-hoc builds.
+
+To run the current checkout of Mill on the build in the scratch folder, use:
+
+```
+mill -i dev.run scratch -w show thingy
+```
+
+If you want to avoid having the bootstrap Mill process running while your
+locally compiled Mill process is running on the build within `scratch` (e.g. to
+simplify debugging) you can use:
+
+```
+mill -i dev.launcher && (cd scratch && ../out/dev/launcher/dest/run -w show thingy)
+```
\ No newline at end of file