`--watch` now works in client-server mode, with the server probing the lockfile to see if the client is still alive

3 files changed, 42 insertions, 9 deletions

diff --git a/main/src/mill/Main.scala b/main/src/mill/Main.scala
index b281f573..a0ef9120 100644
--- a/main/src/mill/Main.scala
+++ b/main/src/mill/Main.scala
@@ -16,10 +16,12 @@ object Main {
                     watch: Boolean = false)
 
   def main(args: Array[String]): Unit = {
-    val (result, _) = main0(args, None)
+    val (result, _) = main0(args, None, () => false)
     System.exit(if(result) 0 else 1)
   }
-  def main0(args: Array[String], mainRunner: Option[(Cli.Config, MainRunner)]): (Boolean, Option[(Cli.Config, MainRunner)]) = {
+  def main0(args: Array[String],
+            mainRunner: Option[(Cli.Config, MainRunner)],
+            watchInterrupted: () => Boolean): (Boolean, Option[(Cli.Config, MainRunner)]) = {
     import ammonite.main.Cli
 
     val removed = Set("predef-code", "home", "no-home-predef")
@@ -65,6 +67,7 @@ object Main {
         val runner = new mill.main.MainRunner(
           config.copy(home = pwd / "out" / ".ammonite"),
           System.out, System.err, System.in,
+          watchInterrupted,
           mainRunner match{
             case Some((c, mr)) if c.copy(storageBackend = null) == cliConfig.copy(storageBackend = null) =>
               mr.lastEvaluator
diff --git a/main/src/mill/ServerClient.scala b/main/src/mill/ServerClient.scala
index 9fa8f70a..4a80eef5 100644
--- a/main/src/mill/ServerClient.scala
+++ b/main/src/mill/ServerClient.scala
@@ -124,6 +124,9 @@ object Server{
     var currentIn = System.in
     var currentOut: OutputStream = System.out
     var currentErr: OutputStream = System.err
+    val raf = new RandomAccessFile(lockBase + "/lock", "rw")
+    val channel = raf.getChannel
+
     System.setOut(new PrintStream(new ProxyOutputStream(currentOut), true))
     System.setErr(new PrintStream(new ProxyOutputStream(currentErr), true))
     System.setIn(new ProxyInputStream(currentIn))
@@ -140,14 +143,27 @@ object Server{
           val args = new String(Files.readAllBytes(runFile)).split('\n')
           try {
 
-            val (_, mr) = mill.Main.main0(args, mainRunner)
+            val (_, mr) = mill.Main.main0(
+              args,
+              mainRunner,
+              () => {
+                channel.tryLock()  match{
+                  case null =>
+                    false
+                  case lock =>
+                    lock.release()
+                    true
+                }
+              }
+            )
             val end = System.currentTimeMillis()
 
             mainRunner = mr
-            System.out.flush()
-            System.err.flush()
-            pprint.log(end - start)
-          } finally {
+          } catch{case MainRunner.WatchInterrupted(mr) =>
+            mainRunner = Some((mr.config, mr))
+          } finally{
+            currentOut.flush()
+            currentErr.flush()
             Files.delete(runFile)
             lastRun = System.currentTimeMillis()
           }
diff --git a/main/src/mill/main/MainRunner.scala b/main/src/mill/main/MainRunner.scala
index 7087dbc7..f6aec9c8 100644
--- a/main/src/mill/main/MainRunner.scala
+++ b/main/src/mill/main/MainRunner.scala
@@ -7,26 +7,40 @@ import ammonite.ops.Path
 import ammonite.util._
 import mill.define.Discover
 import mill.eval.{Evaluator, PathRef}
+import mill.main.MainRunner.WatchInterrupted
 import mill.util.PrintLogger
 import mill.main.RunScript
 import upickle.Js
-
+object MainRunner{
+  case class WatchInterrupted(mr: MainRunner) extends Exception
+}
 /**
   * Customized version of [[ammonite.MainRunner]], allowing us to run Mill
   * `build.sc` scripts with mill-specific tweaks such as a custom
   * `scriptCodeWrapper` or with a persistent evaluator between runs.
   */
-class MainRunner(config: ammonite.main.Cli.Config,
+class MainRunner(val config: ammonite.main.Cli.Config,
                  outprintStream: PrintStream,
                  errPrintStream: PrintStream,
                  stdIn: InputStream,
+                 interruptWatch: () => Boolean,
                  var lastEvaluator: Option[(Seq[(Path, Long)], Evaluator[Any])] = None)
   extends ammonite.MainRunner(
     config, outprintStream, errPrintStream,
     stdIn, outprintStream, errPrintStream
   ){
 
+  override def watchAndWait(watched: Seq[(Path, Long)]) = {
+    printInfo(s"Watching for changes to ${watched.length} files... (Ctrl-C to exit)")
+    def statAll() = watched.forall{ case (file, lastMTime) =>
+      Interpreter.pathSignature(file) == lastMTime
+    }
 
+    while(statAll()) {
+      if (interruptWatch()) throw WatchInterrupted(this)
+      Thread.sleep(100)
+    }
+  }
 
   override def runScript(scriptPath: Path, scriptArgs: List[String]) =
     watchLoop(