Stream exit code over ProxyStreams rather than writing to a file on disk (#5131)

Fixes #5130

We already send a `ProxyStream.END` byte for any normal termination, so
it's straightforward to make the server send one byte after representing
the exit code, and have the client read the exit code byte before
closing the connection

This also explicitly closes the `serverSocket` and releases the
`serverLock` when shutting down, as leaving those open sometimes keeps
the server running 500-1000ms longer than expected. This should fix the
occasional `exitCode file not found` errors we have been seeing since
time immemorial

Author: lihaoyi

core/constants/src/mill/constants/ProxyStream.java

@@ -37,9 +37,10 @@ public class ProxyStream {
   public static final int END = 0;
   public static final int HEARTBEAT = 127;
 
-  public static void sendEnd(OutputStream out) throws IOException {
+  public static void sendEnd(OutputStream out, int exitCode) throws IOException {
     synchronized (out) {
       out.write(ProxyStream.END);
+      out.write(exitCode);
       out.flush();
     }
   }
@@ -113,6 +114,7 @@ public static class Pumper implements Runnable {
     private final OutputStream destOut;
     private final OutputStream destErr;
     private final Object synchronizer;
+    public volatile int exitCode = 255;
 
     public Pumper(
         InputStream src, OutputStream destOut, OutputStream destErr, Object synchronizer) {
@@ -144,8 +146,11 @@ public void run() {
           // that only header values > 0 represent actual data to read:
           // - sign((byte)header) represents which stream the data should be sent to
           // - abs((byte)header) represents the length of the data to read and send
-          if (header == -1 || header == END) break;
-          else if (header == HEARTBEAT) continue;
+          if (header == -1) break;
+          else if (header == END) {
+            exitCode = src.read();
+            break;
+          } else if (header == HEARTBEAT) continue;
           else {
             int stream = (byte) header > 0 ? 1 : -1;
             int quantity0 = (byte) header;

core/constants/src/mill/constants/ServerFiles.java

@@ -48,12 +48,6 @@ public static String pipe(String base) {
    */
   public static final String serverLog = "server.log";
 
-  /**
-   * File the server writes to pass the exit code of a completed run back to the
-   * client
-   */
-  public static final String exitCode = "exitCode";
-
   /**
    * Where the server's stdout is piped to
    */

core/constants/test/src/mill/client/ProxyStreamTests.java

@@ -77,7 +77,7 @@ public void test0(byte[] outData, byte[] errData, int repeats, boolean gracefulE
               srcErr.write(errData);
             }
 
-            if (gracefulEnd) ProxyStream.sendEnd(pipedOutputStream);
+            if (gracefulEnd) ProxyStream.sendEnd(pipedOutputStream, 0);
             else {
               pipedOutputStream.close();
             }

core/internal/src/mill/internal/PromptLogger.scala

@@ -337,7 +337,10 @@ private[mill] object PromptLogger {
       // Close the write side of the pipe first but do not close the read side, so
       // the `pumperThread` can continue reading remaining text in the pipe buffer
       // before terminating on its own
-      ProxyStream.sendEnd(pipe.output)
+      ProxyStream.sendEnd(
+        pipe.output,
+        0 // exit code value is not used since this ProxyStream doesn't wrap a subprocess
+      )
       pipe.output.close()
       pumperThread.join()
     }

integration/feature/shutdown-exit-code/src/ShutdownExitCodeTests.scala

@@ -10,11 +10,13 @@ object ShutdownExitCodeTests extends UtestIntegrationTestSuite {
     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)
+      assert(result3.exitCode == 1)
+
       val result4 = tester.eval("shutdown")
       assert(result4.isSuccess == true)
     }

integration/ide/bsp-modules/src/BspModulesTests.scala

@@ -18,6 +18,8 @@ object BspModulesTests extends UtestIntegrationTestSuite {
           os.read(workspacePath / Constants.bspDir / s"${Constants.serverName}.json")
         )
 
+        eval("shutdown")
+        Thread.sleep(1000)
         val executable = json("argv").arr(0).str
         val checkRes = os.call((executable, "checkExecutable"), cwd = workspacePath)
         assert(checkRes.exitCode == 0)

integration/invalidation/version-change/src/VersionChangeTests.scala

@@ -14,7 +14,7 @@ object VersionChangeTests extends UtestIntegrationTestSuite {
       os.write.over(workspacePath / ".mill-jvm-version", "temurin:19.0.2")
 
       val javaVersion2 = eval(("show", "javaVersion"))
-      assert(javaVersion2.out == "\"19.0.2\"")
+      assert(javaVersion2.out.contains("\"19.0.2\""))
 
     }
   }

runner/client/src/mill/client/ServerLauncher.java

@@ -1,6 +1,5 @@
 package mill.client;
 
-import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.PrintStream;
@@ -90,17 +89,17 @@ public Result run(Path serverDir, String javaHome) throws Exception {
 
     Socket ioSocket = launchConnectToServer(serverDir);
 
+    Result result = new Result();
     try {
-      Thread outPumperThread = startStreamPumpers(ioSocket, javaHome);
+      PumperThread outPumperThread = startStreamPumpers(ioSocket, javaHome);
       forceTestFailure(serverDir);
       outPumperThread.join();
+      result.exitCode = outPumperThread.exitCode();
+      result.serverDir = serverDir;
     } finally {
       ioSocket.close();
     }
 
-    Result result = new Result();
-    result.exitCode = readExitCode(serverDir);
-    result.serverDir = serverDir;
     return result;
   }
 
@@ -137,7 +136,20 @@ private void forceTestFailure(Path serverDir) throws Exception {
     }
   }
 
-  Thread startStreamPumpers(Socket ioSocket, String javaHome) throws Exception {
+  class PumperThread extends Thread {
+    ProxyStream.Pumper runnable;
+
+    public PumperThread(ProxyStream.Pumper runnable, String name) {
+      super(runnable, name);
+      this.runnable = runnable;
+    }
+
+    public int exitCode() {
+      return runnable.exitCode;
+    }
+  }
+
+  PumperThread startStreamPumpers(Socket ioSocket, String javaHome) throws Exception {
     InputStream outErr = ioSocket.getInputStream();
     OutputStream in = ioSocket.getOutputStream();
     in.write(Util.hasConsole() ? 1 : 0);
@@ -147,22 +159,12 @@ Thread startStreamPumpers(Socket ioSocket, String javaHome) throws Exception {
     ClientUtil.writeMap(env, in);
     ProxyStream.Pumper outPumper = new ProxyStream.Pumper(outErr, stdout, stderr);
     InputPumper inPump = new InputPumper(() -> stdin, () -> in, true);
-    Thread outPumperThread = new Thread(outPumper, "outPump");
+    PumperThread outPumperThread = new PumperThread(outPumper, "outPump");
     outPumperThread.setDaemon(true);
     Thread inThread = new Thread(inPump, "inPump");
     inThread.setDaemon(true);
     outPumperThread.start();
     inThread.start();
     return outPumperThread;
   }
-
-  int readExitCode(Path serverDir) throws IOException {
-    Path exitCodeFile = serverDir.resolve(ServerFiles.exitCode);
-    if (Files.exists(exitCodeFile)) {
-      return Integer.parseInt(Files.readAllLines(exitCodeFile).get(0));
-    } else {
-      System.err.println("mill-server/ exitCode file not found");
-      return 1;
-    }
-  }
 }

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

@@ -364,7 +364,6 @@ public static boolean checkTputExists() {
   public static void prepareMillRunFolder(Path serverDir) throws Exception {
     // Clear out run-related files from the server folder to make sure we
     // never hit issues where we are reading the files from a previous run
-    Files.deleteIfExists(serverDir.resolve(ServerFiles.exitCode));
     Files.deleteIfExists(serverDir.resolve(ServerFiles.terminfo));
 
     Path sandbox = serverDir.resolve(ServerFiles.sandbox);

runner/server/src/mill/server/Server.scala

@@ -9,12 +9,14 @@ import mill.constants.InputPumper
 import mill.constants.ProxyStream
 
 import java.io.*
-import java.net.{InetAddress, Socket}
+import java.net.{InetAddress, Socket, ServerSocket}
 import scala.jdk.CollectionConverters.*
 import scala.util.Try
 import scala.util.Using
 import mill.constants.OutFiles
 
+import java.util.concurrent.atomic.AtomicBoolean
+
 /**
  * Models a long-lived server that receives requests from a client and calls a [[main0]]
  * method to run the commands in-process. Provides the command args, env variables,
@@ -52,7 +54,7 @@ abstract class Server[T](
     val initialSystemProperties = sys.props.toMap
 
     try {
-      Server.tryLockBlock(locks.serverLock) {
+      Server.tryLockBlock(locks.serverLock) { locked =>
         serverLog("server file locked")
         Server.watchProcessIdFile(
           serverDir / ServerFiles.processId,
@@ -64,27 +66,42 @@ abstract class Server[T](
           }
         )
         val serverSocket = new java.net.ServerSocket(0, 0, InetAddress.getByName(null))
-        os.write.over(serverDir / ServerFiles.socketPort, serverSocket.getLocalPort.toString)
-        serverLog("listening on port " + serverSocket.getLocalPort)
-        while (
-          running && {
-            interruptWithTimeout(() => serverSocket.close(), () => serverSocket.accept()) match {
-              case None => false
-              case Some(sock) =>
-                serverLog("handling run")
-                new Thread(
-                  () =>
-                    try handleRun(sock, initialSystemProperties)
-                    catch {
-                      case e: Throwable =>
-                        serverLog(e.toString + "\n" + e.getStackTrace.mkString("\n"))
-                    } finally sock.close();,
-                  "HandleRunThread"
-                ).start()
-                true
-            }
+        try {
+          os.write.over(serverDir / ServerFiles.socketPort, serverSocket.getLocalPort.toString)
+          serverLog("listening on port " + serverSocket.getLocalPort)
+
+          def systemExit(exitCode: Int) = {
+            // Explicitly close serverSocket before exiting otherwise it can keep the
+            // server alive 500-1000ms before letting it exit properly
+            serverSocket.close()
+            // Explicitly release process lock to indicate this serverwill not be
+            // taking any more requests, and a new server should be spawned if necessary.
+            // Otherwise launchers may continue trying to connect to the server and
+            // failing since the socket is closed.
+            locked.release()
+            sys.exit(exitCode)
           }
-        ) ()
+
+          while (
+            running && {
+              interruptWithTimeout(() => serverSocket.close(), () => serverSocket.accept()) match {
+                case None => false
+                case Some(sock) =>
+                  serverLog("handling run")
+                  new Thread(
+                    () =>
+                      try handleRun(systemExit, sock, initialSystemProperties)
+                      catch {
+                        case e: Throwable =>
+                          serverLog(e.toString + "\n" + e.getStackTrace.mkString("\n"))
+                      } finally sock.close();,
+                    "HandleRunThread"
+                  ).start()
+                  true
+              }
+            }
+          ) ()
+        } finally serverSocket.close()
         serverLog("server loop ended")
       }.getOrElse(throw new Exception("Mill server process already present"))
     } catch {
@@ -142,9 +159,19 @@ abstract class Server[T](
     }
   }
 
-  def handleRun(clientSocket: Socket, initialSystemProperties: Map[String, String]): Unit = {
-
+  def handleRun(
+      systemExit: Int => Nothing,
+      clientSocket: Socket,
+      initialSystemProperties: Map[String, String]
+  ): Unit = {
     val currentOutErr = clientSocket.getOutputStream
+    val writtenExitCode = AtomicBoolean()
+    def writeExitCode(code: Int) = {
+      if (!writtenExitCode.getAndSet(true)) {
+        ProxyStream.sendEnd(currentOutErr, code)
+      }
+    }
+
     var clientDisappeared = false
     // We cannot use Socket#{isConnected, isClosed, isBound} because none of these
     // detect client-side connection closing, so instead we send a no-op heartbeat
@@ -181,6 +208,7 @@ abstract class Server[T](
 
       val millVersionChanged = lastMillVersion.exists(_ != clientMillVersion)
       val javaVersionChanged = lastJavaVersion.exists(_ != clientJavaVersion)
+
       if (millVersionChanged || javaVersionChanged) {
         Server.withOutLock(
           noBuildLock = false,
@@ -204,11 +232,9 @@ abstract class Server[T](
               s"Java version changed ($lastJavaVersion -> $clientJavaVersion), re-starting server"
             )
           }
-          os.write(
-            serverDir / ServerFiles.exitCode,
-            ClientUtil.ExitServerCodeWhenVersionMismatch().toString.getBytes()
-          )
-          System.exit(ClientUtil.ExitServerCodeWhenVersionMismatch())
+
+          writeExitCode(ClientUtil.ExitServerCodeWhenVersionMismatch())
+          systemExit(ClientUtil.ExitServerCodeWhenVersionMismatch())
         }
       }
       lastMillVersion = Some(clientMillVersion)
@@ -228,15 +254,15 @@ abstract class Server[T](
               Map(),
               initialSystemProperties,
               systemExit = exitCode => {
-                os.write.over(serverDir / ServerFiles.exitCode, exitCode.toString)
-                sys.exit(exitCode)
+                writeExitCode(exitCode)
+                systemExit(exitCode)
               }
             )
 
             stateCache = newStateCache
-            val exitCode = if (result) "0" else "1"
+            val exitCode = if (result) 0 else 1
             serverLog("exitCode " + exitCode)
-            os.write.over(serverDir / ServerFiles.exitCode, exitCode)
+            writeExitCode(exitCode)
           } finally {
             done = true
             idle = true
@@ -271,7 +297,10 @@ abstract class Server[T](
       System.err.flush()
 
     } finally {
-      if (!clientDisappeared) ProxyStream.sendEnd(currentOutErr) // Send a termination
+      try writeExitCode(1) // Send a termination if it has not already happened
+      catch {
+        case e: Throwable => /*donothing*/
+      }
     }
   }
 
@@ -329,12 +358,12 @@ object Server {
     processIdThread.start()
   }
 
-  def tryLockBlock[T](lock: Lock)(t: => T): Option[T] = {
+  def tryLockBlock[T](lock: Lock)(block: mill.client.lock.TryLocked => T): Option[T] = {
     lock.tryLock() match {
       case null => None
       case l =>
         if (l.isLocked) {
-          try Some(t)
+          try Some(block(l))
           finally l.release()
         } else {
           None