Added MILL_TEST_FREE_PORT

core/constants/src/mill/constants/EnvVars.java

@@ -46,4 +46,11 @@ public class EnvVars {
    * e.g. to turn on additional testing/debug/log-related code
    */
   public static final String MILL_TEST_SUITE = "MILL_TEST_SUITE";
+
+  /**
+   * Used to store free ports for Mill's tests to use.
+   * Ports are stored in comma-separated values, with no prefix or suffix.
+   * e.g 1024,9999,10953
+   */
+  public static final String MILL_TEST_FREE_PORT = "MILL_TEST_FREE_PORT";
 }

integration/feature/ports-system-prop/resources/build.mill

@@ -0,0 +1 @@
+

integration/feature/ports-system-prop/src/ShutdownExitCodeTests.scala

@@ -0,0 +1,22 @@
+package mill.integration
+
+import mill.testkit.UtestIntegrationTestSuite
+
+import utest._
+
+object ShutdownExitCodeTests extends UtestIntegrationTestSuite {
+  // Ensure that `shutdown` succeeds even if the prior command failed
+  val tests: Tests = Tests {
+    test("test") - integrationTest { tester =>
+      val result1 = tester.eval(("resolve", "_"))
+      assert(result1.isSuccess == true)
+      val result2 = tester.eval("shutdown")
+      assert(result2.isSuccess == true)
+
+      val result3 = tester.eval("doesnt-exit")
+      assert(result3.isSuccess == false)
+      val result4 = tester.eval("shutdown")
+      assert(result4.isSuccess == true)
+    }
+  }
+}

runner/src/mill/runner/MillCliConfig.scala

@@ -155,7 +155,19 @@ case class MillCliConfig(
           |There are currently no guarantees that modules don't attempt to fetch remote sources."""
           .stripMargin
     )
-    offline: Flag = Flag()
+    offline: Flag = Flag(),
+    @arg(
+      doc =
+        """Provides specified ports for Mill to propagate into System Properties. 
+          |The ports must be between 1024 to 65535.
+          |The ports must be free at the time of Mill launching, and should remain free.
+          |If no ports are specified, free ports will be determined on startup.
+          |Ports are specified in comma-seperated values, e.g 8080,4859,2000.
+          """
+          .stripMargin,
+      name = "ports"
+    )
+    ports: Option[String]
 )
 
 import mainargs.ParserForClass

runner/src/mill/runner/MillMain.scala

@@ -20,6 +20,8 @@ import java.lang.reflect.InvocationTargetException
 import scala.util.control.NonFatal
 import scala.util.Using
 import mill.server.Server
+import java.lang.NumberFormatException
+import java.net.{ServerSocket, BindException}
 
 @internal
 object MillMain {
@@ -149,6 +151,7 @@ object MillMain {
             (false, RunnerState.empty)
 
           case Result.Success(config) =>
+
             val noColorViaEnv = env.get("NO_COLOR").exists(_.nonEmpty)
             val colored = config.color.getOrElse(mainInteractive && !noColorViaEnv)
             val colors =
@@ -178,8 +181,19 @@ object MillMain {
                 (false, stateCache)
 
               } else {
+                val ports = validatePorts(config, streams)
+                if (ports == None) {
+                  return (true, stateCache)
+                }
+                val portStr = ports.get.mkString(",")
+
+                val portJvmProperties = (
+                  "MILL_TEST_FREE_PORT",
+                  portStr
+                )
+
                 val userSpecifiedProperties =
-                  userSpecifiedProperties0 ++ config.extraSystemProperties
+                  userSpecifiedProperties0 ++ config.extraSystemProperties + portJvmProperties
 
                 val threadCount = Some(maybeThreadCount.toOption.get)
 
@@ -300,6 +314,71 @@ object MillMain {
     }
   }
 
+  val tartgetNumberOfPorts = 6
+
+  def validatePorts(config: MillCliConfig, streams: SystemStreams): Option[Set[Int]] = {
+    val result =
+      if (config.ports.isDefined) {
+        var ports = Set.empty[Int]
+        try {
+          ports = config.ports.get.split(",").map(_.toInt).toSet
+        } catch {
+          case e: NumberFormatException =>
+            streams.err.println(
+              "Failed to parse  --ports, please provide comma sepperated values. E.g 4090,5000,2910"
+            )
+            None
+        }
+        // Port validation
+        ports.foreach { case port =>
+          // 1024 is the min port on windows
+          if (port < 1024 || port > 65535) {
+            streams.err.println(
+              s"Failed to parse $port. Ports provided must be within 1024 to 65535."
+            )
+            None
+          }
+          try {
+            new ServerSocket(port).close()
+          } catch {
+            case e: BindException => {
+              streams.err.println(
+                s"Failed to bind to $port! Please only provide free ports."
+              )
+              None
+            }
+          }
+        }
+        Some(ports)
+      } else {
+        var i = 0
+        var ports = Set.empty[Int]
+        for (_ <- 1 to 100) {
+          if (i == tartgetNumberOfPorts) {
+            return Some(ports)
+          }
+          val socket = new ServerSocket(0)
+          try {
+            val port = socket.getLocalPort
+            if (!ports.contains(port)) {
+              ports = ports + port
+              i += 1
+            }
+
+          } finally {
+            socket.close()
+          }
+        }
+
+        streams.err.println(
+          s"Failed to generate the target number of ports used by tests ($tartgetNumberOfPorts). Either there isn't 6 free ports on your computer, or there is a firewall issue. "
+        )
+
+        None
+      }
+    result
+  }
+
   def runBspSession(
       streams0: SystemStreams,
       runMillBootstrap: (Option[RunnerState], SystemStreams) => RunnerState