Add fractional and subtraction for jobs flag (#3554)

Resolve #3482 Allow jobs param to have following formats: 1. An integer indicates the parallel level. 2. A float N followed by character "C" indicates uses (N * all available cores). e.g. "0.5C" uses half of the available cores. 3. "C-" followed by an integer N indicates uses (all available cores - N). e.g. "C-1" leaves 1 core and uses all the other cores. --------- Co-authored-by: Li Haoyi <[email protected]>
com-lihaoyi · Sep 16, 2024 · 1ede694 · 1ede694
1 parent 37de5b0
commit 1ede694
Show file tree

Hide file tree

Showing 3 changed files with 116 additions and 10 deletions.
diff --git a/runner/src/mill/runner/MillCliConfig.scala b/runner/src/mill/runner/MillCliConfig.scala
@@ -62,10 +62,19 @@ class MillCliConfig private (
       name = "jobs",
       short = 'j',
       doc =
-        """Allow processing N targets in parallel.
+        """Allow processing N targets in parallel. It can has following formats:
+
+           1. An integer indicates the parallel level.
+           2. A float N followed by character "C" indicates uses (N * all available cores).
+              e.g. "0.5C" uses half of the available cores.
+           3. "C-" followed by an integer N indicates uses (all available cores - N).
+              e.g. "C-1" leaves 1 core and uses all the other cores.
+
+           If calculated result is less than 1, will round up to 1. This doesn't include pass in 0 directly.
+
            Use 1 to disable parallel and 0 to use as much threads as available processors."""
     )
-    val threadCountRaw: Option[Int],
+    val threadCountRaw: Option[String],
     @arg(
       name = "import",
       doc = """Additional ivy dependencies to load into mill, e.g. plugins."""
@@ -175,7 +184,7 @@ object MillCliConfig {
       debugLog: Flag = Flag(),
       keepGoing: Flag = Flag(),
       extraSystemProperties: Map[String, String] = Map(),
-      threadCountRaw: Option[Int] = None,
+      threadCountRaw: Option[String] = None,
       imports: Seq[String] = Seq(),
       interactive: Flag = Flag(),
       help: Flag = Flag(),
@@ -224,7 +233,7 @@ object MillCliConfig {
       debugLog: Flag,
       keepGoing: Flag,
       extraSystemProperties: Map[String, String],
-      threadCountRaw: Option[Int],
+      threadCountRaw: Option[String],
       imports: Seq[String],
       interactive: Flag,
       help: Flag,
@@ -273,7 +282,7 @@ object MillCliConfig {
       debugLog: Flag,
       keepGoing: Flag,
       extraSystemProperties: Map[String, String],
-      threadCountRaw: Option[Int],
+      threadCountRaw: Option[String],
       imports: Seq[String],
       interactive: Flag,
       help: Flag,

diff --git a/runner/src/mill/runner/MillMain.scala b/runner/src/mill/runner/MillMain.scala
@@ -169,6 +169,8 @@ object MillMain {
 
             // special BSP mode, in which we spawn a server and register the current evaluator when-ever we start to eval a dedicated command
             val bspMode = config.bsp.value && config.leftoverArgs.value.isEmpty
+            val maybeThreadCount =
+              parseThreadCount(config.threadCountRaw, Runtime.getRuntime.availableProcessors())
 
             val (success, nextStateCache) = {
               if (config.repl.value) {
@@ -179,15 +181,15 @@ object MillMain {
                 logger.error("A target must be provided.")
                 (false, stateCache)
 
+              } else if (maybeThreadCount.isLeft) {
+                logger.error(maybeThreadCount.swap.toOption.get)
+                (false, stateCache)
+
               } else {
                 val userSpecifiedProperties =
                   userSpecifiedProperties0 ++ config.extraSystemProperties
 
-                val threadCount = config.threadCountRaw match {
-                  case None => None
-                  case Some(0) => None
-                  case Some(n) => Some(n)
-                }
+                val threadCount = Some(maybeThreadCount.toOption.get)
 
                 if (mill.main.client.Util.isJava9OrAbove) {
                   val rt = config.home / Export.rtJarName
@@ -275,6 +277,26 @@ object MillMain {
     }
   }
 
+  private[runner] def parseThreadCount(
+      threadCountRaw: Option[String],
+      availableCores: Int
+  ): Either[String, Int] = {
+    def err(detail: String) =
+      s"Invalid value \"${threadCountRaw.getOrElse("")}\" for flag -j/--jobs: $detail"
+    (threadCountRaw match {
+      case None => Right(availableCores)
+      case Some("0") => Right(availableCores)
+      case Some(s"${n}C") => n.toDoubleOption
+          .toRight(err("Failed to find a float number before \"C\"."))
+          .map(m => (m * availableCores).toInt)
+      case Some(s"C-${n}") => n.toIntOption
+          .toRight(err("Failed to find a int number after \"C-\"."))
+          .map(availableCores - _)
+      case Some(n) => n.toIntOption
+          .toRight(err("Failed to find a int number"))
+    }).map { x => if (x < 1) 1 else x }
+  }
+
   def getLogger(
       streams: SystemStreams,
       config: MillCliConfig,

diff --git a/runner/test/src/mill/runner/MillMainTests.scala b/runner/test/src/mill/runner/MillMainTests.scala
@@ -0,0 +1,75 @@
+package mill.runner
+
+import utest._
+
+object MillMainTests extends TestSuite {
+
+  private def assertParseErr(result: Either[String, Int], msg: String): Unit = {
+    assert(result.isLeft)
+    assert(result.swap.toOption.get.contains(msg))
+  }
+
+  def tests: Tests = Tests {
+
+    test("Parsing --jobs/-j flag") {
+
+      test("parse none") {
+        assert(MillMain.parseThreadCount(None, 10) == Right(10))
+      }
+
+      test("parse int number") {
+        assert(MillMain.parseThreadCount(Some("1"), 10) == Right(1))
+        assert(MillMain.parseThreadCount(Some("11"), 10) == Right(11))
+
+        assertParseErr(MillMain.parseThreadCount(Some("1.0"), 10), "Failed to find a int number")
+        assertParseErr(MillMain.parseThreadCount(Some("1.1"), 10), "Failed to find a int number")
+        assertParseErr(MillMain.parseThreadCount(Some("0.1"), 10), "Failed to find a int number")
+        assert(MillMain.parseThreadCount(Some("0"), 10) == Right(10))
+        assert(MillMain.parseThreadCount(Some("-1"), 10) == Right(1))
+      }
+
+      test("parse fraction number") {
+        assert(MillMain.parseThreadCount(Some("0.5C"), 10) == Right(5))
+        assert(MillMain.parseThreadCount(Some("0.54C"), 10) == Right(5))
+        assert(MillMain.parseThreadCount(Some("0.59C"), 10) == Right(5))
+        assert(MillMain.parseThreadCount(Some(".5C"), 10) == Right(5))
+        assert(MillMain.parseThreadCount(Some("1.0C"), 10) == Right(10))
+        assert(MillMain.parseThreadCount(Some("1.5C"), 10) == Right(15))
+        assert(MillMain.parseThreadCount(Some("0.09C"), 10) == Right(1))
+        assert(MillMain.parseThreadCount(Some("-0.5C"), 10) == Right(1))
+        assertParseErr(
+          MillMain.parseThreadCount(Some("0.5.4C"), 10),
+          "Failed to find a float number before \"C\""
+        )
+      }
+
+      test("parse subtraction") {
+        assert(MillMain.parseThreadCount(Some("C-1"), 10) == Right(9))
+        assert(MillMain.parseThreadCount(Some("C-10"), 10) == Right(1))
+        assert(MillMain.parseThreadCount(Some("C-11"), 10) == Right(1))
+
+        assertParseErr(
+          MillMain.parseThreadCount(Some("C-1.1"), 10),
+          "Failed to find a int number after \"C-\""
+        )
+        assertParseErr(
+          MillMain.parseThreadCount(Some("11-C"), 10),
+          "Failed to find a float number before \"C\""
+        )
+      }
+
+      test("parse invalid input") {
+        assertParseErr(
+          MillMain.parseThreadCount(Some("CCCC"), 10),
+          "Failed to find a float number before \"C\""
+        )
+        assertParseErr(
+          MillMain.parseThreadCount(Some("abcdefg"), 10),
+          "Failed to find a int number"
+        )
+      }
+
+    }
+
+  }
+}