Simplify launchOrConnectToServer to eliminate nested retries (#5872)

Previously we needed a `retryWithTimeout` around the reading of the port
because the initial process launch could take >10s, which would make the
outer `retryWithTimeout` not take effect. This is because we were
performing the coursier runner classpath resolution as part of that
launch, which can take arbitrarily long, overshooting the retry timeout.

This PR fixes it by moving the coursier runner classpath resolution out
of the retry loop, so the retry loop only contains the logic around
launching the process. This should help keep the time of each attempt to
a more reasonable amount, allowing the outer retries to work and
avoiding the need for nested retries around the individual steps of the
process launch workflow

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: lihaoyi

integration/failure/invalid-build-header-no-space/src/InvalidBuildHeaderNoSpaceTests.scala

@@ -5,6 +5,8 @@ import mill.testkit.UtestIntegrationTestSuite
 import utest._
 
 object InvalidBuildHeaderNoSpaceTests extends UtestIntegrationTestSuite {
+  override def cleanupProcessIdFile =
+    false // process never launches due to yaml header syntax error
   val tests: Tests = Tests {
     test - integrationTest { tester =>
       import tester._

integration/failure/invalid-build-header/src/InvalidBuildHeaderTests.scala

@@ -5,6 +5,8 @@ import mill.testkit.UtestIntegrationTestSuite
 import utest._
 
 object InvalidBuildHeaderTests extends UtestIntegrationTestSuite {
+  override def cleanupProcessIdFile =
+    false // process never launches due to yaml header syntax error
   val tests: Tests = Tests {
     test - integrationTest { tester =>
       import tester._

libs/daemon/client/src/mill/client/ServerLauncher.java

@@ -142,23 +142,21 @@ public static Launched launchOrConnectToServer(
     try (var ignored = locks.launcherLock.lock()) {
       return retryWithTimeout(serverInitWaitMillis, "server launch failed", () -> {
         try {
+          log.accept("launchOrConnectToServer attempt");
+
           var result =
               ensureServerIsRunning(locks, daemonDir, initServer, serverInitWaitMillis / 3, log);
           if (result instanceof ServerLaunchResult.Success
               || result instanceof ServerLaunchResult.AlreadyRunning) {
             log.accept("Reading server port: " + daemonDir.toAbsolutePath());
-            var port = retryWithTimeout(serverInitWaitMillis / 10, "Reading server port", () -> {
-              try {
-                return Optional.of(
-                    Integer.parseInt(Files.readString(daemonDir.resolve(DaemonFiles.socketPort))));
-              } catch (IOException e) {
-                throw new RuntimeException(e);
-              }
-            });
+            var port =
+                Integer.parseInt(Files.readString(daemonDir.resolve(DaemonFiles.socketPort)));
+
             var launched = new Launched();
             launched.port = port;
-            log.accept("Read server port, connecting: " + port);
+            log.accept("Read server port");
             if (openSocket) {
+              log.accept("Connecting: " + port);
               var connected = new Socket(InetAddress.getLoopbackAddress(), port);
               launched.socket = connected;
             }
@@ -276,6 +274,7 @@ public void kill() {
                     };
             return Optional.of(new ServerLaunchResult.AlreadyRunning(launchedServer));
           } catch (IOException | NumberFormatException e) {
+            log.accept("Read PID exception: " + e);
             return Optional.empty();
           }
         }

runner/launcher/src/mill/launcher/MillLauncherMain.java

@@ -1,13 +1,18 @@
 package mill.launcher;
 
+import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Optional;
 import java.util.function.Consumer;
 import mill.client.*;
 import mill.client.lock.Locks;
+import mill.constants.BuildInfo;
 import mill.constants.EnvVars;
 import mill.constants.OutFiles;
 import mill.constants.OutFolderMode;
@@ -70,16 +75,30 @@ public static void main(String[] args) throws Exception {
                   + " make it less responsive.");
     }
 
+    String[] runnerClasspath = MillProcessLauncher.cachedComputedValue0(
+        outMode,
+        "resolve-runner",
+        BuildInfo.millVersion,
+        () -> CoursierClient.resolveMillDaemon(),
+        arr -> {
+          for (String s : arr) {
+            if (!Files.exists(Paths.get(s))) return false;
+          }
+          return true;
+        });
+
     if (runNoDaemon) {
       // start in no-server mode
-      System.exit(MillProcessLauncher.launchMillNoDaemon(args, outMode));
+      System.exit(MillProcessLauncher.launchMillNoDaemon(args, outMode, runnerClasspath));
     } else {
       var logs = new java.util.ArrayList<String>();
       try {
         // start in client-server mode
         var optsArgs = new java.util.ArrayList<>(MillProcessLauncher.millOpts(outMode));
         Collections.addAll(optsArgs, args);
-
+        var formatter =
+            DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss'Z'").withZone(ZoneId.of("UTC"));
+        Consumer<String> log = (s) -> logs.add(formatter.format(Instant.now()) + " " + s);
         MillServerLauncher launcher =
             new MillServerLauncher(
                 new MillServerLauncher.Streams(System.in, System.out, System.err),
@@ -89,13 +108,14 @@ public static void main(String[] args) throws Exception {
                 -1) {
               public LaunchedServer initServer(Path daemonDir, Locks locks) throws Exception {
                 return new LaunchedServer.OsProcess(
-                    MillProcessLauncher.launchMillDaemon(daemonDir, outMode).toHandle());
+                    MillProcessLauncher.launchMillDaemon(daemonDir, outMode, runnerClasspath)
+                        .toHandle());
               }
             };
 
         var daemonDir0 = Paths.get(outDir, OutFiles.millDaemon);
         String javaHome = MillProcessLauncher.javaHome(outMode);
-        Consumer<String> log = logs::add;
+
         var exitCode = launcher.run(daemonDir0, javaHome, log).exitCode;
         if (exitCode == ClientUtil.ExitServerCodeWhenVersionMismatch()) {
           exitCode = launcher.run(daemonDir0, javaHome, log).exitCode;

runner/launcher/src/mill/launcher/MillProcessLauncher.java

@@ -19,13 +19,14 @@
 
 public class MillProcessLauncher {
 
-  static int launchMillNoDaemon(String[] args, OutFolderMode outMode) throws Exception {
+  static int launchMillNoDaemon(String[] args, OutFolderMode outMode, String[] runnerClasspath)
+      throws Exception {
     final String sig = String.format("%08x", UUID.randomUUID().hashCode());
     final Path processDir =
         Paths.get(".").resolve(outFor(outMode)).resolve(millNoDaemon).resolve(sig);
 
     final List<String> l = new ArrayList<>();
-    l.addAll(millLaunchJvmCommand(outMode));
+    l.addAll(millLaunchJvmCommand(outMode, runnerClasspath));
     Map<String, String> propsMap = ClientUtil.getUserSetProperties();
     for (String key : propsMap.keySet()) l.add("-D" + key + "=" + propsMap.get(key));
     l.add("mill.daemon.MillNoDaemonMain");
@@ -56,8 +57,9 @@ static int launchMillNoDaemon(String[] args, OutFolderMode outMode) throws Excep
     }
   }
 
-  static Process launchMillDaemon(Path daemonDir, OutFolderMode outMode) throws Exception {
-    List<String> l = new ArrayList<>(millLaunchJvmCommand(outMode));
+  static Process launchMillDaemon(Path daemonDir, OutFolderMode outMode, String[] runnerClasspath)
+      throws Exception {
+    List<String> l = new ArrayList<>(millLaunchJvmCommand(outMode, runnerClasspath));
     l.add("mill.daemon.MillDaemonMain");
     l.add(daemonDir.toFile().getCanonicalPath());
     l.add(outMode.asString());
@@ -212,7 +214,8 @@ static String javaExe(OutFolderMode outMode) throws Exception {
     }
   }
 
-  static List<String> millLaunchJvmCommand(OutFolderMode outMode) throws Exception {
+  static List<String> millLaunchJvmCommand(OutFolderMode outMode, String[] runnerClasspath)
+      throws Exception {
     final List<String> vmOptions = new ArrayList<>();
 
     // Java executable
@@ -234,17 +237,7 @@ static List<String> millLaunchJvmCommand(OutFolderMode outMode) throws Exception
 
     vmOptions.add("-XX:+HeapDumpOnOutOfMemoryError");
     vmOptions.add("-cp");
-    String[] runnerClasspath = cachedComputedValue0(
-        outMode,
-        "resolve-runner",
-        BuildInfo.millVersion,
-        () -> CoursierClient.resolveMillDaemon(),
-        arr -> {
-          for (String s : arr) {
-            if (!Files.exists(Paths.get(s))) return false;
-          }
-          return true;
-        });
+
     vmOptions.add(String.join(File.pathSeparator, runnerClasspath));
 
     return vmOptions;

testkit/src/mill/testkit/ExampleTester.scala

@@ -76,7 +76,8 @@ class ExampleTester(
     millExecutable: os.Path,
     bashExecutable: String = ExampleTester.defaultBashExecutable(),
     val baseWorkspacePath: os.Path,
-    val propagateJavaHome: Boolean = true
+    val propagateJavaHome: Boolean = true,
+    val cleanupProcessIdFile: Boolean = true
 ) extends IntegrationTesterBase {
 
   def processCommandBlock(commandBlock: String): Unit = {

testkit/src/mill/testkit/IntegrationTestSuite.scala

@@ -35,6 +35,7 @@ trait IntegrationTestSuite {
    */
   protected def propagateJavaHome: Boolean = true
 
+  protected def cleanupProcessIdFile: Boolean = true
   def debugLog: Boolean = false
 
   /**
@@ -53,7 +54,8 @@ trait IntegrationTestSuite {
         millExecutable,
         debugLog,
         baseWorkspacePath = os.pwd,
-        propagateJavaHome = propagateJavaHome
+        propagateJavaHome = propagateJavaHome,
+        cleanupProcessIdFile = cleanupProcessIdFile
       )
       try block(tester)
       finally tester.close()

testkit/src/mill/testkit/IntegrationTester.scala

@@ -26,7 +26,8 @@ class IntegrationTester(
     val millExecutable: os.Path,
     override val debugLog: Boolean = false,
     val baseWorkspacePath: os.Path = os.pwd,
-    val propagateJavaHome: Boolean = true
+    val propagateJavaHome: Boolean = true,
+    val cleanupProcessIdFile: Boolean = true
 ) extends IntegrationTester.Impl {
   initWorkspace()
 }

testkit/src/mill/testkit/IntegrationTesterBase.scala

@@ -6,6 +6,7 @@ import mill.util.Retry
 trait IntegrationTesterBase {
   def workspaceSourcePath: os.Path
   def daemonMode: Boolean
+  def cleanupProcessIdFile: Boolean
 
   def propagateJavaHome: Boolean
 
@@ -75,7 +76,7 @@ trait IntegrationTesterBase {
   def removeProcessIdFile(): Unit = {
     if (!sys.env.contains("MILL_TEST_SHARED_OUTPUT_DIR")) {
       val outDir = os.Path(out, workspacePath)
-      if (os.exists(outDir)) {
+      if (os.exists(outDir) && cleanupProcessIdFile) {
         if (daemonMode) {
           val serverPath = outDir / millDaemon
           os.remove(serverPath / processId)