Accept sysprops via cmdline and forward to mill server

See https://github.com/lihaoyi/mill/issues/498
See https://github.com/lihaoyi/mill/issues/499

7 files changed, 128 insertions, 78 deletions

diff --git a/main/src/MillMain.scala b/main/src/MillMain.scala
index f1a7a9e7..e6403e7f 100644
--- a/main/src/MillMain.scala
+++ b/main/src/MillMain.scala
@@ -23,54 +23,57 @@ object MillMain {
       System.out,
       System.err,
       System.getenv().asScala.toMap,
-      b => ()
+      b => (),
+      System.getProperties().asScala.toMap
     )
-    System.exit(if(result) 0 else 1)
+    System.exit(if (result) 0 else 1)
   }
 
-  def main0(args: Array[String],
-            stateCache: Option[Evaluator.State],
-            mainInteractive: Boolean,
-            stdin: InputStream,
-            stdout: PrintStream,
-            stderr: PrintStream,
-            env: Map[String, String],
-            setIdle: Boolean => Unit): (Boolean, Option[Evaluator.State]) = {
+  def main0(
+    args: Array[String],
+    stateCache: Option[Evaluator.State],
+    mainInteractive: Boolean,
+    stdin: InputStream,
+    stdout: PrintStream,
+    stderr: PrintStream,
+    env: Map[String, String],
+    setIdle: Boolean => Unit,
+    initialSystemProperties: Map[String, String]
+  ): (Boolean, Option[Evaluator.State]) = {
     import ammonite.main.Cli
-    
+
     val millHome = mill.api.Ctx.defaultHome
 
     val removed = Set("predef-code", "no-home-predef")
+
     var interactive = false
     val interactiveSignature = Arg[Config, Unit](
       "interactive", Some('i'),
       "Run Mill in interactive mode, suitable for opening REPLs and taking user input. In this mode, no mill server will be used.",
-      (c, v) =>{
+      (c, v) => {
         interactive = true
         c
       }
     )
 
-
-
     var disableTicker = false
     val disableTickerSignature = Arg[Config, Unit](
-      "disable-ticker", None,
-      "Disable ticker log (e.g. short-lived prints of stages and progress bars)",
-      (c, v) =>{
-        disableTicker = true
-        c
-      }
+      name = "disable-ticker", shortName = None,
+      doc = "Disable ticker log (e.g. short-lived prints of stages and progress bars)",
+      action = (c, v) => {
+      disableTicker = true
+      c
+    }
     )
 
     var debugLog = false
     val debugLogSignature = Arg[Config, Unit](
       name = "debug", shortName = Some('d'),
       doc = "Show debug output on STDOUT",
-      (c, v) => {
-        debugLog = true
-        c
-      }
+      action = (c, v) => {
+      debugLog = true
+      c
+    }
     )
 
     var keepGoing = false
@@ -82,27 +85,46 @@ object MillMain {
       }
     )
 
+    var extraSystemProperties = Map[String, String]()
+    val extraSystemPropertiesSignature = Arg[Config, String](
+      name = "define", shortName = Some('D'),
+      doc = "Define (or overwrite) a system property",
+      action = { (c, v) =>
+      extraSystemProperties += (v.split("[=]", 2) match {
+        case Array(k, v) => k -> v
+        case Array(k) => k -> ""
+      })
+      c
+    }
+    )
+
     val millArgSignature =
       Cli.genericSignature.filter(a => !removed(a.name)) ++
-        Seq(interactiveSignature, disableTickerSignature, debugLogSignature, keepGoingSignature)
+        Seq(
+          interactiveSignature,
+          disableTickerSignature,
+          debugLogSignature,
+          keepGoingSignature,
+          extraSystemPropertiesSignature
+        )
 
     Cli.groupArgs(
       args.toList,
       millArgSignature,
       Cli.Config(home = millHome, remoteLogging = false)
-    ) match{
-      case _ if interactive =>
-        stderr.println("-i/--interactive must be passed in as the first argument")
-        (false, None)
-      case Left(msg) =>
-        stderr.println(msg)
-        (false, None)
-      case Right((cliConfig, _)) if cliConfig.help =>
-        val leftMargin = millArgSignature.map(ammonite.main.Cli.showArg(_).length).max + 2
-        stdout.println(
-        s"""Mill Build Tool
-           |usage: mill [mill-options] [target [target-options]]
-           |
+    ) match {
+        case _ if interactive =>
+          stderr.println("-i/--interactive must be passed in as the first argument")
+          (false, None)
+        case Left(msg) =>
+          stderr.println(msg)
+          (false, None)
+        case Right((cliConfig, _)) if cliConfig.help =>
+          val leftMargin = millArgSignature.map(ammonite.main.Cli.showArg(_).length).max + 2
+          stdout.println(
+            s"""Mill Build Tool
+             |usage: mill [mill-options] [target [target-options]]
+             |
            |${formatBlock(millArgSignature, leftMargin).mkString(ammonite.util.Util.newLine)}""".stripMargin
         )
         (true, None)
@@ -126,42 +148,44 @@ object MillMain {
                   |  build.millSelf.get,
                   |  build.millDiscover,
                   |  $debugLog,
-                  |  keepGoing = $keepGoing
+                  |  keepGoing = $keepGoing,
+                  |  ${initialSystemProperties ++ extraSystemProperties}
                   |)
                   |repl.pprinter() = replApplyHandler.pprinter
                   |import replApplyHandler.generatedEval._
                   |
                 """.stripMargin,
-              welcomeBanner = None
+                welcomeBanner = None
+              )
+
+            val runner = new mill.main.MainRunner(
+              config.copy(colored = config.colored orElse Option(mainInteractive)),
+              disableTicker,
+              stdout, stderr, stdin,
+              stateCache,
+              env,
+              setIdle,
+              debugLog,
+              keepGoing = keepGoing,
+              initialSystemProperties ++ extraSystemProperties
             )
 
-          val runner = new mill.main.MainRunner(
-            config.copy(colored = config.colored orElse Option(mainInteractive)),
-            disableTicker,
-            stdout, stderr, stdin,
-            stateCache,
-            env,
-            setIdle,
-            debugLog,
-            keepGoing = keepGoing
-          )
-
-          if (mill.main.client.Util.isJava9OrAbove) {
-            val rt = cliConfig.home / Export.rtJarName
-            if (!os.exists(rt)) {
-              runner.printInfo(s"Preparing Java ${System.getProperty("java.version")} runtime; this may take a minute or two ...")
-              Export.rtTo(rt.toIO, false)
+            if (mill.main.client.Util.isJava9OrAbove) {
+              val rt = cliConfig.home / Export.rtJarName
+              if (!os.exists(rt)) {
+                runner.printInfo(s"Preparing Java ${System.getProperty("java.version")} runtime; this may take a minute or two ...")
+                Export.rtTo(rt.toIO, false)
+              }
             }
-          }
 
-          if (repl){
-            runner.printInfo("Loading...")
-            (runner.watchLoop(isRepl = true, printing = false, _.run()), runner.stateCache)
-          } else {
-            (runner.runScript(os.pwd / "build.sc", leftoverArgs), runner.stateCache)
+            if (repl) {
+              runner.printInfo("Loading...")
+              (runner.watchLoop(isRepl = true, printing = false, _.run()), runner.stateCache)
+            } else {
+              (runner.runScript(os.pwd / "build.sc", leftoverArgs), runner.stateCache)
+            }
           }
-      }
 
-    }
+      }
   }
 }
diff --git a/main/src/main/MainRunner.scala b/main/src/main/MainRunner.scala
index e08905a6..c773087c 100644
--- a/main/src/main/MainRunner.scala
+++ b/main/src/main/MainRunner.scala
@@ -25,7 +25,8 @@ class MainRunner(val config: ammonite.main.Cli.Config,
                  env : Map[String, String],
                  setIdle: Boolean => Unit,
                  debugLog: Boolean,
-                 keepGoing: Boolean)
+                 keepGoing: Boolean,
+                 systemProperties: Map[String, String])
   extends ammonite.MainRunner(
     config, outprintStream, errPrintStream,
     stdIn, outprintStream, errPrintStream
@@ -85,7 +86,8 @@ class MainRunner(val config: ammonite.main.Cli.Config,
             debugEnabled = debugLog
           ),
           env,
-          keepGoing = keepGoing
+          keepGoing = keepGoing,
+          systemProperties
         )
 
         result match{
diff --git a/main/src/main/MillServerMain.scala b/main/src/main/MillServerMain.scala
index 862daaf7..500c3e8f 100644
--- a/main/src/main/MillServerMain.scala
+++ b/main/src/main/MillServerMain.scala
@@ -21,7 +21,8 @@ trait MillServerMain[T]{
             stdout: PrintStream,
             stderr: PrintStream,
             env : Map[String, String],
-            setIdle: Boolean => Unit): (Boolean, Option[T])
+            setIdle: Boolean => Unit,
+            systemProperties: Map[String, String]): (Boolean, Option[T])
 }
 
 object MillServerMain extends mill.main.MillServerMain[Evaluator.State]{
@@ -44,6 +45,7 @@ object MillServerMain extends mill.main.MillServerMain[Evaluator.State]{
       mill.main.client.Locks.files(args0(0))
     ).run()
   }
+
   def main0(args: Array[String],
             stateCache: Option[Evaluator.State],
             mainInteractive: Boolean,
@@ -51,7 +53,8 @@ object MillServerMain extends mill.main.MillServerMain[Evaluator.State]{
             stdout: PrintStream,
             stderr: PrintStream,
             env : Map[String, String],
-            setIdle: Boolean => Unit) = {
+            setIdle: Boolean => Unit,
+            systemProperties: Map[String, String]) = {
     MillMain.main0(
       args,
       stateCache,
@@ -60,7 +63,8 @@ object MillServerMain extends mill.main.MillServerMain[Evaluator.State]{
       stdout,
       stderr,
       env,
-      setIdle = setIdle
+      setIdle = setIdle,
+      systemProperties
     )
   }
 }
@@ -132,6 +136,7 @@ class Server[T](lockBase: String,
     }
     val args = Util.parseArgs(argStream)
     val env = Util.parseMap(argStream)
+    val systemProperties = Util.parseMap(argStream)
     argStream.close()
 
     @volatile var done = false
@@ -147,6 +152,7 @@ class Server[T](lockBase: String,
           stderr,
           env.asScala.toMap,
           idle = _,
+          systemProperties.asScala.toMap
         )
 
         sm.stateCache = newStateCache
diff --git a/main/src/main/ReplApplyHandler.scala b/main/src/main/ReplApplyHandler.scala
index 6f1e060d..7f959929 100644
--- a/main/src/main/ReplApplyHandler.scala
+++ b/main/src/main/ReplApplyHandler.scala
@@ -1,14 +1,13 @@
 package mill.main
 
+import scala.collection.mutable
 
+import mill.api.Strict.Agg
 import mill.define.Applicative.ApplyHandler
 import mill.define.Segment.Label
 import mill.define._
 import mill.eval.{Evaluator, Result}
 
-import mill.api.Strict.Agg
-
-import scala.collection.mutable
 object ReplApplyHandler{
   def apply[T](home: os.Path,
                disableTicker: Boolean,
@@ -17,7 +16,8 @@ object ReplApplyHandler{
                rootModule: mill.define.BaseModule,
                discover: Discover[_],
                debugLog: Boolean,
-               keepGoing: Boolean) = {
+               keepGoing: Boolean,
+               systemProperties: Map[String, String]): ReplApplyHandler = {
     new ReplApplyHandler(
       pprinter0,
       new Evaluator(
@@ -36,7 +36,8 @@ object ReplApplyHandler{
           debugEnabled = debugLog
         ),
         failFast = !keepGoing
-      )
+      ),
+      systemProperties
     )
   }
   def pprintCross(c: mill.define.Cross[_], evaluator: Evaluator) = {
@@ -113,8 +114,15 @@ object ReplApplyHandler{
   }
 
 }
+
 class ReplApplyHandler(pprinter0: pprint.PPrinter,
-                       val evaluator: Evaluator) extends ApplyHandler[Task] {
+                       val evaluator: Evaluator,
+                       systemProperties: Map[String, String]) extends ApplyHandler[Task] {
+
+  systemProperties.foreach {case (k,v) =>
+    System.setProperty(k,v)
+  }
+
   // Evaluate classLoaderSig only once in the REPL to avoid busting caches
   // as the user enters more REPL commands and changes the classpath
   val classLoaderSig = Evaluator.classLoaderSig
diff --git a/main/src/main/RunScript.scala b/main/src/main/RunScript.scala
index ea8e554f..ab53aa1a 100644
--- a/main/src/main/RunScript.scala
+++ b/main/src/main/RunScript.scala
@@ -30,9 +30,14 @@ object RunScript{
                 stateCache: Option[Evaluator.State],
                 log: Logger,
                 env : Map[String, String],
-                keepGoing: Boolean)
+                keepGoing: Boolean,
+                systemProperties: Map[String, String])
   : (Res[(Evaluator, Seq[PathRef], Either[String, Seq[ujson.Value]])], Seq[(os.Path, Long)]) = {
 
+    systemProperties.foreach {case (k,v) =>
+      System.setProperty(k, v)
+    }
+
     val (evalState, interpWatched) = stateCache match{
       case Some(s) if watchedSigUnchanged(s.watched) => Res.Success(s) -> s.watched
       case _ =>
diff --git a/main/test/src/main/ClientServerTests.scala b/main/test/src/main/ClientServerTests.scala
index 05238a5f..6d918b30 100644
--- a/main/test/src/main/ClientServerTests.scala
+++ b/main/test/src/main/ClientServerTests.scala
@@ -13,7 +13,8 @@ class EchoServer extends MillServerMain[Int]{
             stdout: PrintStream,
             stderr: PrintStream,
             env: Map[String, String],
-            setIdle: Boolean => Unit) = {
+            setIdle: Boolean => Unit,
+            systemProperties: Map[String, String]) = {
 
     val reader = new BufferedReader(new InputStreamReader(stdin))
     val str = reader.readLine()
@@ -23,6 +24,9 @@ class EchoServer extends MillServerMain[Int]{
     env.toSeq.sortBy(_._1).foreach{
       case (key, value) => stdout.println(s"$key=$value")
     }
+    systemProperties.toSeq.sortBy(_._1).foreach{
+      case (key, value) => stdout.println(s"$key=$value")
+    }
     stdout.flush()
     if (args.nonEmpty){
       stderr.println(str.toUpperCase + args(0))
diff --git a/main/test/src/util/ScriptTestSuite.scala b/main/test/src/util/ScriptTestSuite.scala
index 92f57c4f..f3850dda 100644
--- a/main/test/src/util/ScriptTestSuite.scala
+++ b/main/test/src/util/ScriptTestSuite.scala
@@ -16,10 +16,11 @@ abstract class ScriptTestSuite(fork: Boolean) extends TestSuite{
   val disableTicker = false
   val debugLog = false
   val keepGoing = false
+  val systemProperties = Map[String, String]()
   lazy val runner = new mill.main.MainRunner(
     ammonite.main.Cli.Config(wd = wd), disableTicker,
     stdOutErr, stdOutErr, stdIn, None, Map.empty,
-    b => (), debugLog, keepGoing = keepGoing
+    b => (), debugLog, keepGoing = keepGoing, systemProperties
   )
   def eval(s: String*) = {
     if (!fork) runner.runScript(workspacePath / buildPath , s.toList)