From 3d586f811c1afed24726ccf15f564289777a28c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Wed, 13 Aug 2025 16:57:57 +0300 Subject: [PATCH 01/25] Extra logging to files for JvmWorker/ZincWorker --- .../mill/javalib/worker/JvmWorkerImpl.scala | 122 ++++++++++-- .../javalib/zinc/ZincWorkerRpcServer.scala | 182 ++++++++++-------- libs/util/src/mill/util/Timed.scala | 18 ++ 3 files changed, 217 insertions(+), 105 deletions(-) create mode 100644 libs/util/src/mill/util/Timed.scala diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index f0264bdfaf08..868b70b5cd9b 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -11,13 +11,14 @@ import mill.javalib.internal.{JvmWorkerArgs, RpcCompileProblemReporterMessage} import mill.javalib.zinc.ZincWorkerRpcServer.ReporterMode import mill.javalib.zinc.{ZincApi, ZincWorker, ZincWorkerRpcServer} import mill.rpc.{MillRpcChannel, MillRpcClient, MillRpcWireTransport} -import mill.util.{CachedFactoryWithInitData, HexFormat, Jvm} +import mill.util.{CachedFactoryWithInitData, HexFormat, Jvm, Timed} import org.apache.logging.log4j.core.util.NullOutputStream import sbt.internal.util.ConsoleOut import java.io.* import java.nio.file.FileSystemException import java.security.MessageDigest +import java.time.LocalDateTime import scala.concurrent.duration.* import scala.util.Using @@ -25,6 +26,18 @@ import scala.util.Using class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoCloseable { import args.* + private def fileLog(s: String): Unit = + os.write.append( + compilerBridge.workspace / "jvm-worker.log", + s"[${LocalDateTime.now()}] $s\n", + createFolders = true + ) + + private def fileAndDebugLog(log: Logger.Actions, s: String): Unit = { + fileLog(s) + log.debug(s) + } + /** The local Zinc instance which is used when we do not want to override Java home or runtime options. */ private val zincLocalWorker = ZincWorker( @@ -41,8 +54,12 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo reporter: Option[CompileProblemReporter], reportCachedProblems: Boolean )(using ctx: JvmWorkerApi.Ctx): Result[CompilationResult] = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, javaRuntimeOptions) - zinc.compileJava(op, reporter = reporter, reportCachedProblems = reportCachedProblems) + val result = + Timed(zinc.compileJava(op, reporter = reporter, reportCachedProblems = reportCachedProblems)) + fileLog(s"Compilation took ${result.durationPretty}") + result.result } override def compileMixed( @@ -52,16 +69,23 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo reporter: Option[CompileProblemReporter], reportCachedProblems: Boolean )(using ctx: JvmWorkerApi.Ctx): Result[CompilationResult] = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, javaRuntimeOptions) - zinc.compileMixed(op, reporter = reporter, reportCachedProblems = reportCachedProblems) + val result = + Timed(zinc.compileMixed(op, reporter = reporter, reportCachedProblems = reportCachedProblems)) + fileLog(s"Compilation took ${result.durationPretty}") + result.result } def scaladocJar( op: ZincScaladocJar, javaHome: Option[os.Path] )(using ctx: JvmWorkerApi.Ctx): Boolean = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, JavaRuntimeOptions(Seq.empty)) - zinc.scaladocJar(op) + val result = Timed(zinc.scaladocJar(op)) + fileLog(s"Scaladoc took ${result.durationPretty}") + result.result } override def close(): Unit = { @@ -88,8 +112,20 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo logPromptColored = log.prompt.colored ) - if (javaRuntimeOptions.options.isEmpty && javaHome.isEmpty) localZincApi(zincCtx, log) - else subprocessZincApi(javaHome, javaRuntimeOptions, zincCtx, log) + if (javaRuntimeOptions.options.isEmpty && javaHome.isEmpty) { + fileLog("Using local Zinc instance") + localZincApi(zincCtx, log) + } else { + fileLog( + s"""Using remote Zinc instance: + | javaHome: $javaHome + | javaRuntimeOptions: $javaRuntimeOptions + |""".stripMargin + ) + val result = Timed(subprocessZincApi(javaHome, javaRuntimeOptions, zincCtx, log)) + fileLog(s"Remote Zinc instance acquired in ${result.durationPretty}") + result.result + } } private case class SubprocessCacheKey( @@ -186,34 +222,47 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ) } - log.debug(s"Checking if $mainClass is already running for $key") - val result = ServerLauncher.ensureServerIsRunning( + fileAndDebugLog(log, s"Checking if $mainClass is already running for $key") + val result = Timed(ServerLauncher.ensureServerIsRunning( locks, daemonDir.toNIO, () => { - log.debug(s"Starting JVM subprocess for $mainClass for $key") - val process = Jvm.spawnProcess( + fileAndDebugLog(log, s"Starting JVM subprocess for $mainClass for $key") + val process = Timed(Jvm.spawnProcess( mainClass = mainClass, mainArgs = Seq(daemonDir.toString), javaHome = key.javaHome, jvmArgs = key.runtimeOptions.options, classPath = classPath + )) + fileAndDebugLog( + log, + s"Starting JVM subprocess for $mainClass for $key took ${process.durationPretty}" ) - LaunchedServer.OsProcess(process.wrapped.toHandle) + LaunchedServer.OsProcess(process.result.wrapped.toHandle) }, - log.debug + fileAndDebugLog(log, _) + )) + fileAndDebugLog( + log, + s"Ensuring that server is running for $key took ${result.durationPretty}" ) def onSuccess(launched: LaunchedServer) = { val serverInitWaitMillis = 5.seconds.toMillis val startTime = System.nanoTime() - log.debug(s"Reading server port: $daemonDir") - val port = ServerLauncher.readServerPort(daemonDir.toNIO, startTime, serverInitWaitMillis) - log.debug(s"Started $mainClass for $key on port $port") - SubprocessCacheValue(port, daemonDir, launched) + fileAndDebugLog(log, s"Reading server port: $daemonDir") + val port = + Timed(ServerLauncher.readServerPort(daemonDir.toNIO, startTime, serverInitWaitMillis)) + fileAndDebugLog( + log, + s"Reading server port for $daemonDir took ${port.durationPretty}." + ) + fileAndDebugLog(log, s"Started $mainClass for $key on port ${port.result}") + SubprocessCacheValue(port.result, daemonDir, launched) } - result.fold( + result.result.fold( success => onSuccess(success.server), alreadyRunning => onSuccess(alreadyRunning.server), processDied => @@ -270,7 +319,7 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ) { case SubprocessCacheValue(port, daemonDir, _) => Using.Manager { use => val startTimeNanos = System.nanoTime() - log.debug(s"Connecting to $daemonDir on port $port") + fileAndDebugLog(log, s"Connecting to $daemonDir on port $port") val socket = use(ServerLauncher.connectToServer( startTimeNanos, 5.seconds.toMillis, @@ -304,13 +353,46 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo compileToJar = compileToJar, zincLogDebug = zincLogDebug ) + fileAndDebugLog( + log, + s"Connected to $daemonDir on port $port, sending init: ${pprint(init)}" + ) + val clientLogger = new Logger.Actions { + override def info(s: String): Unit = { + fileLog(s"[LOGGER:INFO] $s") + log.info(s) + } + + override def debug(s: String): Unit = { + fileLog(s"[LOGGER:DEBUG] $s") + log.debug(s) + } + + override def warn(s: String): Unit = { + fileLog(s"[LOGGER:WARN] $s") + log.warn(s) + } + + override def error(s: String): Unit = { + fileLog(s"[LOGGER:ERROR] $s") + log.error(s) + } + + override def ticker(s: String): Unit = { + fileLog(s"[LOGGER:TICKER] $s") + log.ticker(s) + } + } val client = MillRpcClient.create[ ZincWorkerRpcServer.Initialize, ZincWorkerRpcServer.ClientToServer, ZincWorkerRpcServer.ServerToClient - ](init, wireTransport, log)(handler) + ](init, wireTransport, clientLogger)(handler) - f(client) + fileAndDebugLog(log, "Running command.") + val result = Timed(f(client)) + fileAndDebugLog(log, s"Command finished in ${result.durationPretty}") + result.result } ).result }.get diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala index 48cf8bf510ed..08014a2f0a4d 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala @@ -8,6 +8,7 @@ import mill.javalib.api.internal.{ZincCompileJava, ZincCompileMixed, ZincScalado import mill.javalib.internal.{RpcCompileProblemReporterMessage, ZincCompilerBridgeProvider} import mill.rpc.* import mill.server.Server +import mill.util.Timed import org.apache.logging.log4j.core.util.NullOutputStream import sbt.internal.util.ConsoleOut import upickle.default.ReadWriter @@ -33,100 +34,111 @@ class ZincWorkerRpcServer( clientStderr: RpcConsole, serverToClient: MillRpcChannel[ServerToClient] ): MillRpcChannel[ClientToServer] = { - val zincCompilerBridge = ZincCompilerBridgeProvider[MillRpcRequestId]( - workspace = initialize.compilerBridgeWorkspace, - logInfo = log.info, - acquire = (scalaVersion, scalaOrganization, clientRequestId) => - serverToClient( - clientRequestId, - ServerToClient.AcquireZincCompilerBridge( - scalaVersion = scalaVersion, - scalaOrganization = scalaOrganization + val result = Timed { + val zincCompilerBridge = ZincCompilerBridgeProvider[MillRpcRequestId]( + workspace = initialize.compilerBridgeWorkspace, + logInfo = log.info, + acquire = (scalaVersion, scalaOrganization, clientRequestId) => + serverToClient( + clientRequestId, + ServerToClient.AcquireZincCompilerBridge( + scalaVersion = scalaVersion, + scalaOrganization = scalaOrganization + ) ) - ) - ) - val worker = ZincWorker( - zincCompilerBridge, - jobs = initialize.jobs, - compileToJar = initialize.compileToJar, - zincLogDebug = initialize.zincLogDebug - ) - - val deps = { - // This is an ugly hack. `ConsoleOut` is sealed but we need to provide a way to send these logs to the Mill server - // over RPC, so we hijack `PrintStream` by overriding the methods that `ConsoleOut` uses. - // - // This is obviously extra fragile, but I couldn't find a better way to do it. - val consoleOut = ConsoleOut.printStreamOut(new PrintStream(NullOutputStream.getInstance()) { - override def print(s: String): Unit = clientStderr.print(s) - override def println(s: String): Unit = print(s + "\n") - override def println(): Unit = print("\n") - override def flush(): Unit = clientStderr.flush() - }) - - ZincWorker.InvocationDependencies(log, consoleOut) - } + ) + val worker = ZincWorker( + zincCompilerBridge, + jobs = initialize.jobs, + compileToJar = initialize.compileToJar, + zincLogDebug = initialize.zincLogDebug + ) + + val deps = { + // This is an ugly hack. `ConsoleOut` is sealed but we need to provide a way to send these logs to the Mill server + // over RPC, so we hijack `PrintStream` by overriding the methods that `ConsoleOut` uses. + // + // This is obviously extra fragile, but I couldn't find a better way to do it. + val consoleOut = ConsoleOut.printStreamOut(new PrintStream(NullOutputStream.getInstance()) { + override def print(s: String): Unit = clientStderr.print(s) + + override def println(s: String): Unit = print(s + "\n") + + override def println(): Unit = print("\n") + + override def flush(): Unit = clientStderr.flush() + }) + + ZincWorker.InvocationDependencies(log, consoleOut) + } - def reporter(clientRequestId: MillRpcRequestId, maxErrors: Int) = RpcCompileProblemReporter( - maxErrors = maxErrors, - send = msg => - serverToClient( - clientRequestId, - ServerToClient.ReportCompilationProblem(clientRequestId, msg) - ) - ) - - def reporterAsOption( - clientRequestId: MillRpcRequestId, - mode: ReporterMode - ): Option[CompileProblemReporter] = mode match { - case ReporterMode.NoReporter => None - case r: ReporterMode.Reporter => - Some(reporter(clientRequestId = clientRequestId, maxErrors = r.maxErrors)) - } + def reporter(clientRequestId: MillRpcRequestId, maxErrors: Int) = RpcCompileProblemReporter( + maxErrors = maxErrors, + send = msg => + serverToClient( + clientRequestId, + ServerToClient.ReportCompilationProblem(clientRequestId, msg) + ) + ) + + def reporterAsOption( + clientRequestId: MillRpcRequestId, + mode: ReporterMode + ): Option[CompileProblemReporter] = mode match { + case ReporterMode.NoReporter => None + case r: ReporterMode.Reporter => + Some(reporter(clientRequestId = clientRequestId, maxErrors = r.maxErrors)) + } - new MillRpcChannel[ClientToServer] { - override def apply(requestId: MillRpcRequestId, input: ClientToServer): input.Response = - setIdle.doWork { - input match { - case msg: ClientToServer.CompileJava => - compileJava(requestId, msg).asInstanceOf[input.Response] - case msg: ClientToServer.CompileMixed => - compileMixed(requestId, msg).asInstanceOf[input.Response] - case msg: ClientToServer.ScaladocJar => - docJar(requestId, msg).asInstanceOf[input.Response] + new MillRpcChannel[ClientToServer] { + override def apply(requestId: MillRpcRequestId, input: ClientToServer): input.Response = + setIdle.doWork { + val result = Timed { + input match { + case msg: ClientToServer.CompileJava => + compileJava(requestId, msg).asInstanceOf[input.Response] + case msg: ClientToServer.CompileMixed => + compileMixed(requestId, msg).asInstanceOf[input.Response] + case msg: ClientToServer.ScaladocJar => + docJar(requestId, msg).asInstanceOf[input.Response] + } + } + writeToLocalLog(s"$requestId with data $input processed in ${result.durationPretty}") + result.result } + + private def compileJava( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.CompileJava + ): msg.Response = { + worker.compileJava( + op = msg.op, + reporter = reporterAsOption(clientRequestId, msg.reporterMode), + reportCachedProblems = msg.reporterMode.reportCachedProblems + )(using msg.ctx, deps) } - private def compileJava( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.CompileJava - ): msg.Response = { - worker.compileJava( - op = msg.op, - reporter = reporterAsOption(clientRequestId, msg.reporterMode), - reportCachedProblems = msg.reporterMode.reportCachedProblems - )(using msg.ctx, deps) - } + private def compileMixed( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.CompileMixed + ): msg.Response = { + worker.compileMixed( + msg.op, + reporter = reporterAsOption(clientRequestId, msg.reporterMode), + reportCachedProblems = msg.reporterMode.reportCachedProblems, + compilerBridgeData = clientRequestId + )(using msg.ctx, deps) + } - private def compileMixed( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.CompileMixed - ): msg.Response = { - worker.compileMixed( - msg.op, - reporter = reporterAsOption(clientRequestId, msg.reporterMode), - reportCachedProblems = msg.reporterMode.reportCachedProblems, - compilerBridgeData = clientRequestId - )(using msg.ctx, deps) + private def docJar( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.ScaladocJar + ): msg.Response = + worker.scaladocJar(msg.op, compilerBridgeData = clientRequestId) } - - private def docJar( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.ScaladocJar - ): msg.Response = - worker.scaladocJar(msg.op, compilerBridgeData = clientRequestId) } + writeToLocalLog(s"Initialized in ${result.durationPretty}.") + result.result } } object ZincWorkerRpcServer { diff --git a/libs/util/src/mill/util/Timed.scala b/libs/util/src/mill/util/Timed.scala new file mode 100644 index 000000000000..32ea13e117d4 --- /dev/null +++ b/libs/util/src/mill/util/Timed.scala @@ -0,0 +1,18 @@ +package mill.util + +import java.util.concurrent.TimeUnit +import scala.concurrent.duration.FiniteDuration + +object Timed { + case class Result[+A](result: A, duration: FiniteDuration) { + def durationPretty: String = duration.toMillis + "ms" + } + + /** Computes how long it took to compute `f`. */ + def apply[A](f: => A): Result[A] = { + val start = System.nanoTime() + val result = f + val duration = FiniteDuration(System.nanoTime() - start, TimeUnit.NANOSECONDS) + Result(result, duration) + } +} From 352fa72e52c9021c8f28cfc75b46793e4b298810 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Wed, 13 Aug 2025 20:46:42 +0300 Subject: [PATCH 02/25] WIP: cleaning up stream pumpers. --- .../src/mill/constants/InputPumper.java | 14 +- .../src/mill/constants/ProxyStream.java | 168 ++++++++++-------- .../src/mill/client/ProxyStreamTests.java | 6 +- .../src/mill/internal/PromptLogger.scala | 8 +- .../client/src/mill/client/ClientUtil.java | 6 +- .../test/src/mill/client/ClientTests.java | 4 +- .../server/src/mill/server/Server.scala | 23 ++- .../mill/javalib/worker/JvmWorkerImpl.scala | 57 +++--- .../mill/javalib/zinc/ZincWorkerMain.scala | 2 +- .../src/mill/rpc/MillRpcWireTransport.scala | 1 + .../src/mill/server/MillDaemonServer.scala | 2 +- 11 files changed, 161 insertions(+), 130 deletions(-) diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index bc016b89302a..8baf1083cd29 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -31,26 +31,26 @@ public void run() { InputStream src = src0.get(); OutputStream dest = dest0.get(); - byte[] buffer = new byte[1024]; + var buffer = new byte[256 * 1024 /* 256kb, otherwise it's too slow. */]; try { while (running) { if (checkAvailable && src.available() == 0) //noinspection BusyWait Thread.sleep(1); else { - int n; + int bytesRead; try { - n = src.read(buffer); + bytesRead = src.read(buffer); } catch (Exception e) { - n = -1; + bytesRead = -1; } - if (n == -1) running = false; - else if (n == 0) + if (bytesRead == -1) running = false; + else if (bytesRead == 0) //noinspection BusyWait Thread.sleep(1); else { try { - dest.write(buffer, 0, n); + dest.write(buffer, 0, bytesRead); dest.flush(); } catch (java.io.IOException e) { running = false; diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index 64f4b2a8a6f1..82407930cc71 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -1,9 +1,9 @@ package mill.constants; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; +import java.io.*; import java.net.SocketException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; /// Logic to capture a pair of streams (typically stdout and stderr), combining /// them into a single stream, and splitting it back into two streams later while @@ -32,11 +32,32 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with `header == 0` public class ProxyStream { + private static final int MAX_CHUNK_SIZE = 256 * 1024; // 256kb - public static final int OUT = 1; - public static final int ERR = -1; - public static final int END = 0; - public static final int HEARTBEAT = 127; + /** Indicates the end of the connection. */ + private static final byte HEADER_END = 0; + + /** The key for the output stream */ + private static final byte HEADER_STREAM_OUT = 1; + + /** The key for the error stream */ + private static final byte HEADER_STREAM_ERR = 2; + + /** A heartbeat packet to keep the connection alive. */ + private static final byte HEADER_HEARTBEAT = 127; + + public enum StreamType { + /** The output stream */ + OUT(ProxyStream.HEADER_STREAM_OUT), + /** The error stream */ + ERR(ProxyStream.HEADER_STREAM_ERR); + + public final byte header; + + StreamType(byte key) { + this.header = key; + } + } private static boolean clientHasClosedConnection(SocketException e) { var message = e.getMessage(); @@ -46,8 +67,9 @@ private static boolean clientHasClosedConnection(SocketException e) { public static void sendEnd(OutputStream out, int exitCode) throws IOException { synchronized (out) { try { - out.write(ProxyStream.END); - out.write(exitCode); + var buffer = new byte[5]; + ByteBuffer.wrap(buffer).order(ByteOrder.BIG_ENDIAN).put(ProxyStream.HEADER_END).putInt(exitCode); + out.write(buffer); out.flush(); } catch (SocketException e) { // If the client has already closed the connection, we don't really care about sending the @@ -59,25 +81,26 @@ public static void sendEnd(OutputStream out, int exitCode) throws IOException { public static void sendHeartbeat(OutputStream out) throws IOException { synchronized (out) { - out.write(ProxyStream.HEARTBEAT); + out.write(ProxyStream.HEADER_HEARTBEAT); out.flush(); } } public static class Output extends java.io.OutputStream { - private final java.io.OutputStream destination; - private final int key; + private final DataOutputStream destination; + private final StreamType streamType; - public Output(java.io.OutputStream out, int key) { - this.destination = out; - this.key = key; + public Output(java.io.OutputStream out, StreamType streamType) { + this.destination = new DataOutputStream(out); + this.streamType = streamType; } @Override public void write(int b) throws IOException { synchronized (destination) { - destination.write(key); - destination.write(b); + destination.writeByte(streamType.header); + destination.writeInt(1); // 1 byte + destination.writeByte(b); } } @@ -92,15 +115,26 @@ public void write(byte[] b) throws IOException { @Override public void write(byte[] b, int off, int len) throws IOException { + // Validate arguments once at the beginning, which is cleaner + // and standard practice for public methods. + if (b == null) throw new NullPointerException("byte array is null"); + else if (off < 0 || off > b.length || len < 0 || off + len > b.length || off + len < 0) { + throw new IndexOutOfBoundsException( + "Write indexes out of range: off=" + off + ", len=" + len + ", b.length=" + b.length + ); + } synchronized (destination) { - int i = 0; - while (i < len && i + off < b.length) { - int chunkLength = Math.min(len - i, 126); + var bytesWritten = 0; + + while (bytesWritten < len) { + var chunkLength = Math.min(len - bytesWritten, MAX_CHUNK_SIZE); + if (chunkLength > 0) { - destination.write(chunkLength * key); - destination.write(b, off + i, Math.min(b.length - off - i, chunkLength)); - i += chunkLength; + destination.writeByte(streamType.header); + destination.writeInt(chunkLength); + destination.write(b, off + bytesWritten, chunkLength); + bytesWritten += chunkLength; } } } @@ -122,7 +156,7 @@ public void close() throws IOException { } public static class Pumper implements Runnable { - private final InputStream src; + private final DataInputStream src; private final OutputStream destOut; private final OutputStream destErr; private final Object synchronizer; @@ -130,7 +164,7 @@ public static class Pumper implements Runnable { public Pumper( InputStream src, OutputStream destOut, OutputStream destErr, Object synchronizer) { - this.src = src; + this.src = new DataInputStream(src); this.destOut = destOut; this.destErr = destErr; this.synchronizer = synchronizer; @@ -140,63 +174,33 @@ public Pumper(InputStream src, OutputStream destOut, OutputStream destErr) { this(src, destOut, destErr, new Object()); } - public void preRead(InputStream src) {} - - public void write(OutputStream dest, byte[] buffer, int length) throws IOException { - dest.write(buffer, 0, length); - } + public void preRead(DataInputStream src) {} @Override public void run() { - - byte[] buffer = new byte[1024]; - while (true) { - try { + var buffer = new byte[MAX_CHUNK_SIZE]; + try { + while (true) { this.preRead(src); - int header = src.read(); - // -1 means socket was closed, 0 means a ProxyStream.END was sent. Note - // 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) break; - else if (header == END) { - exitCode = src.read(); + var header = src.readByte(); + + if (header == HEADER_END) { + exitCode = src.readInt(); break; - } else if (header == HEARTBEAT) continue; - else { - int stream = (byte) header > 0 ? 1 : -1; - int quantity0 = (byte) header; - int quantity = Math.abs(quantity0); - int offset = 0; - int delta = -1; - while (offset < quantity) { - this.preRead(src); - delta = src.read(buffer, offset, quantity - offset); - if (delta == -1) { - break; - } else { - offset += delta; - } - } - - if (delta != -1) { - synchronized (synchronizer) { - switch (stream) { - case ProxyStream.OUT: - this.write(destOut, buffer, offset); - break; - case ProxyStream.ERR: - this.write(destErr, buffer, offset); - break; - } - } - } } - } catch (IOException e) { - // This happens when the upstream pipe was closed - break; + else if (header == HEADER_HEARTBEAT) continue; + else if (header == HEADER_STREAM_OUT) pumpData(buffer, destOut); + else if (header == HEADER_STREAM_ERR) pumpData(buffer, destErr); + else throw new IllegalStateException("Unexpected header: " + header); } } + catch (EOFException ignored) { + // This is a normal and expected way for the loop to terminate + // when the other side closes the connection. + } + catch (IOException ignored) { + // This happens when the upstream pipe was closed + } try { synchronized (synchronizer) { @@ -207,6 +211,22 @@ else if (header == END) { } } + private void pumpData(byte[] buffer, OutputStream stream) throws IOException { + var quantity = src.readInt(); + + if (quantity > buffer.length) { + // Handle error: received chunk is larger than buffer + throw new IOException("Received chunk of size " + quantity + + " is larger than buffer of size " + buffer.length); + } + + src.readFully(buffer, 0, quantity); + + synchronized (synchronizer) { + stream.write(buffer, 0, quantity); + } + } + public void flush() throws IOException { synchronized (synchronizer) { destOut.flush(); diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index 9db033e521bd..2b4b1767786b 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/core/constants/test/src/mill/client/ProxyStreamTests.java @@ -56,8 +56,8 @@ public void test0(byte[] outData, byte[] errData, int repeats, boolean gracefulE pipedInputStream.connect(pipedOutputStream); - ProxyStream.Output srcOut = new ProxyStream.Output(pipedOutputStream, ProxyStream.OUT); - ProxyStream.Output srcErr = new ProxyStream.Output(pipedOutputStream, ProxyStream.ERR); + var srcOut = new ProxyStream.Output(pipedOutputStream, ProxyStream.StreamType.OUT); + var srcErr = new ProxyStream.Output(pipedOutputStream, ProxyStream.StreamType.ERR); // Capture both the destOut/destErr from the pumper, as well as the destCombined // to ensure the individual streams contain the right data and combined stream @@ -77,7 +77,7 @@ public void test0(byte[] outData, byte[] errData, int repeats, boolean gracefulE srcErr.write(errData); } - if (gracefulEnd) ProxyStream.sendEnd(pipedOutputStream, 0); + if (gracefulEnd) ProxyStream.sendEnd(pipedOutputStream, (byte) 0); else { pipedOutputStream.close(); } diff --git a/core/internal/src/mill/internal/PromptLogger.scala b/core/internal/src/mill/internal/PromptLogger.scala index 8040837460e3..3fd2be4cb91a 100644 --- a/core/internal/src/mill/internal/PromptLogger.scala +++ b/core/internal/src/mill/internal/PromptLogger.scala @@ -286,9 +286,9 @@ private[mill] object PromptLogger { // `ProxyStream`, as we need to preserve the ordering of writes to each individual // stream, and also need to know when *both* streams are quiescent so that we can // print the prompt at the bottom - val pipe = new PipeStreams() - val proxyOut = new ProxyStream.Output(pipe.output, ProxyStream.OUT) - val proxyErr: ProxyStream.Output = new ProxyStream.Output(pipe.output, ProxyStream.ERR) + val pipe = PipeStreams() + val proxyOut = ProxyStream.Output(pipe.output, ProxyStream.StreamType.OUT) + val proxyErr = ProxyStream.Output(pipe.output, ProxyStream.StreamType.ERR) val proxySystemStreams = new SystemStreams( new PrintStream(proxyOut), new PrintStream(proxyErr), @@ -330,7 +330,7 @@ private[mill] object PromptLogger { private var lastCharWritten = 0.toChar // Make sure we synchronize everywhere - override def preRead(src: InputStream): Unit = synchronizer.synchronized { + override def preRead(src: DataInputStream): Unit = synchronizer.synchronized { if ( enableTicker && diff --git a/libs/daemon/client/src/mill/client/ClientUtil.java b/libs/daemon/client/src/mill/client/ClientUtil.java index 98e733cd1443..e66768bbb552 100644 --- a/libs/daemon/client/src/mill/client/ClientUtil.java +++ b/libs/daemon/client/src/mill/client/ClientUtil.java @@ -8,15 +8,15 @@ public class ClientUtil { // use methods instead of constants to avoid inlining by compiler - public static int ExitClientCodeCannotReadFromExitCodeFile() { + public static byte ExitClientCodeCannotReadFromExitCodeFile() { return 1; } - public static int ExitServerCodeWhenIdle() { + public static byte ExitServerCodeWhenIdle() { return 0; } - public static int ExitServerCodeWhenVersionMismatch() { + public static byte ExitServerCodeWhenVersionMismatch() { return 101; } diff --git a/libs/daemon/client/test/src/mill/client/ClientTests.java b/libs/daemon/client/test/src/mill/client/ClientTests.java index db8fbd6f2276..193666177f27 100644 --- a/libs/daemon/client/test/src/mill/client/ClientTests.java +++ b/libs/daemon/client/test/src/mill/client/ClientTests.java @@ -127,8 +127,8 @@ public void proxyInputOutputStreams(byte[] samples1, byte[] samples2, int chunkM throws Exception { ByteArrayOutputStream pipe = new ByteArrayOutputStream(); - OutputStream src1 = new ProxyStream.Output(pipe, ProxyStream.OUT); - OutputStream src2 = new ProxyStream.Output(pipe, ProxyStream.ERR); + var src1 = new ProxyStream.Output(pipe, ProxyStream.StreamType.OUT); + var src2 = new ProxyStream.Output(pipe, ProxyStream.StreamType.ERR); Random random = new Random(31337); diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index ffb008ae9136..57307af804be 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -4,7 +4,14 @@ import mill.client.lock.{Lock, Locks} import mill.constants.{DaemonFiles, InputPumper, ProxyStream} import sun.misc.{Signal, SignalHandler} -import java.io.{InputStream, PipedInputStream, PipedOutputStream, PrintStream} +import java.io.{ + BufferedInputStream, + BufferedOutputStream, + InputStream, + PipedInputStream, + PipedOutputStream, + PrintStream +} import java.net.{InetAddress, Socket, SocketAddress} import java.time.LocalDateTime import java.time.format.DateTimeFormatter @@ -248,12 +255,12 @@ abstract class Server( ): Unit = { /** stdout and stderr combined into one stream */ - val currentOutErr = clientSocket.getOutputStream + val currentOutErr = BufferedOutputStream(clientSocket.getOutputStream) val writtenExitCode = AtomicBoolean() - def writeExitCode(code: Int): Unit = { + def writeExitCode(exitCode: Int): Unit = { if (!writtenExitCode.getAndSet(true)) { - ProxyStream.sendEnd(currentOutErr, code) + ProxyStream.sendEnd(currentOutErr, exitCode) } } @@ -280,10 +287,12 @@ abstract class Server( } try { - val socketIn = clientSocket.getInputStream + val socketIn = BufferedInputStream(clientSocket.getInputStream) - val stdout = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.OUT), true) - val stderr = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.ERR), true) + val stdout = + new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.OUT), true) + val stderr = + new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.ERR), true) val data = preHandleConnection( socketInfo = socketInfo, diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index 868b70b5cd9b..b3373e753c79 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -309,6 +309,33 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ): ZincApi = { val cacheKey = SubprocessCacheKey(javaHome, runtimeOptions) + def makeClientLogger() = new Logger.Actions { + override def info(s: String): Unit = { + fileLog(s"[LOGGER:INFO] $s") + log.info(s) + } + + override def debug(s: String): Unit = { + fileLog(s"[LOGGER:DEBUG] $s") + log.debug(s) + } + + override def warn(s: String): Unit = { + fileLog(s"[LOGGER:WARN] $s") + log.warn(s) + } + + override def error(s: String): Unit = { + fileLog(s"[LOGGER:ERROR] $s") + log.error(s) + } + + override def ticker(s: String): Unit = { + fileLog(s"[LOGGER:TICKER] $s") + log.ticker(s) + } + } + def withRpcClient[R]( handler: MillRpcChannel[ZincWorkerRpcServer.ServerToClient] )(f: MillRpcClient[ZincWorkerRpcServer.ClientToServer, ZincWorkerRpcServer.ServerToClient] => R) @@ -340,7 +367,7 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo /* sendInitData */ _ => {}, () => { val serverToClient = use(BufferedReader(InputStreamReader(PipedInputStream(stdout)))) - val clientToServer = use(PrintStream(PipedOutputStream(stdin))) + val clientToServer = use(PrintStream(BufferedOutputStream(PipedOutputStream(stdin)))) val wireTransport = MillRpcWireTransport.ViaStreams( s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}", serverToClient, @@ -357,37 +384,11 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo log, s"Connected to $daemonDir on port $port, sending init: ${pprint(init)}" ) - val clientLogger = new Logger.Actions { - override def info(s: String): Unit = { - fileLog(s"[LOGGER:INFO] $s") - log.info(s) - } - - override def debug(s: String): Unit = { - fileLog(s"[LOGGER:DEBUG] $s") - log.debug(s) - } - - override def warn(s: String): Unit = { - fileLog(s"[LOGGER:WARN] $s") - log.warn(s) - } - - override def error(s: String): Unit = { - fileLog(s"[LOGGER:ERROR] $s") - log.error(s) - } - - override def ticker(s: String): Unit = { - fileLog(s"[LOGGER:TICKER] $s") - log.ticker(s) - } - } val client = MillRpcClient.create[ ZincWorkerRpcServer.Initialize, ZincWorkerRpcServer.ClientToServer, ZincWorkerRpcServer.ServerToClient - ](init, wireTransport, clientLogger)(handler) + ](init, wireTransport, makeClientLogger())(handler) fileAndDebugLog(log, "Running command.") val result = Timed(f(client)) diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index 30e896c21954..ac00a7b9b2ed 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -58,7 +58,7 @@ object ZincWorkerMain { setIdle: Server.SetIdle, initialSystemProperties: Map[String, String], data: PreHandleConnectionData - ): Int = { + ) = { val serverName = s"$className{${socketInfo.remote} -> ${socketInfo.local}}" Using.resource(BufferedReader(InputStreamReader(stdin))) { stdin => val transport = MillRpcWireTransport.ViaStreams(serverName, stdin, stdout) diff --git a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala index ea165da82772..a65a49b4015f 100644 --- a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala +++ b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala @@ -43,6 +43,7 @@ trait MillRpcWireTransport extends AutoCloseable { val serialized = upickle.default.write(message) log(s"Sending: $serialized") write(serialized) + log(s"Sent: $serialized") } } diff --git a/runner/server/src/mill/server/MillDaemonServer.scala b/runner/server/src/mill/server/MillDaemonServer.scala index 1462b2d0982b..93c1a572acd9 100644 --- a/runner/server/src/mill/server/MillDaemonServer.scala +++ b/runner/server/src/mill/server/MillDaemonServer.scala @@ -111,7 +111,7 @@ abstract class MillDaemonServer[State]( setIdle: Server.SetIdle, initialSystemProperties: Map[String, String], data: ClientInitData - ): Int = { + ) = { val (result, newStateCache) = main0( data.args, stateCache, From 9758861cdc339b33ea001afddfc74d54b2a65399 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Wed, 13 Aug 2025 23:14:01 +0300 Subject: [PATCH 03/25] Things seem to work. --- .../src/mill/constants/InputPumper.java | 13 ++- .../src/mill/constants/ProxyStream.java | 102 +++++++++++------- .../src/mill/client/ProxyStreamTests.java | 4 +- .../src/mill/client/ServerLauncher.java | 11 +- .../test/src/mill/client/ClientTests.java | 10 +- .../server/src/mill/server/Server.scala | 49 +++------ .../mill/javalib/worker/JvmWorkerImpl.scala | 22 ++-- .../mill/javalib/zinc/ZincWorkerMain.scala | 6 +- .../src/mill/client/MillServerLauncher.java | 2 + .../src/mill/server/MillDaemonServer.scala | 4 +- 10 files changed, 117 insertions(+), 106 deletions(-) diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index 8baf1083cd29..7654ff42f6ba 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -4,7 +4,7 @@ import java.io.OutputStream; import java.util.function.Supplier; -/// A `Runnable` that reads from `src` and writes to `dest`. +/** A `Runnable` that reads from `src` and writes to `dest`. */ public class InputPumper implements Runnable { private final Supplier src0; private final Supplier dest0; @@ -13,12 +13,11 @@ public class InputPumper implements Runnable { /// /// We need to do that because if we call `.read` /// and there is nothing to read, [it can unnecessarily delay the JVM exit by 350ms]( - /// - // https://stackoverflow.com/questions/48951611/blocking-on-stdin-makes-java-process-take-350ms-more-to-exit) - private final Boolean checkAvailable; + /// https://stackoverflow.com/questions/48951611/blocking-on-stdin-makes-java-process-take-350ms-more-to-exit) + private final boolean checkAvailable; public InputPumper( - Supplier src, Supplier dest, Boolean checkAvailable) { + Supplier src, Supplier dest, boolean checkAvailable) { this.src0 = src; this.dest0 = dest; this.checkAvailable = checkAvailable; @@ -28,8 +27,8 @@ public InputPumper( @Override public void run() { - InputStream src = src0.get(); - OutputStream dest = dest0.get(); + var src = src0.get(); + var dest = dest0.get(); var buffer = new byte[256 * 1024 /* 256kb, otherwise it's too slow. */]; try { diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index 82407930cc71..a43dcd1d4ee4 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -15,47 +15,56 @@ /// the form: /// ``` /// 1 byte n bytes -/// | header | body | +/// | header | frame | /// ``` /// /// Where header is a single byte of the form: /// -/// - header more than 0 indicating that this packet is for the `OUT` stream -/// - header less than 0 indicating that this packet is for the `ERR` stream -/// - abs(header) indicating the length of the packet body, in bytes -/// - header == 0 indicating the end of the stream +/// - [#HEADER_STREAM_OUT]/[#HEADER_STREAM_ERR] respectively indicating that this packet is for the `OUT`/`ERR` +/// stream, and it will be followed by 4 bytes for the length of the body and then the body. +/// - [#HEADER_STREAM_OUT_SINGLE_BYTE]/[#HEADER_STREAM_ERR_SINGLE_BYTE] respectively indicating that this packet is +/// for the `OUT`/`ERR` stream, and it will be followed by a single byte for the body +/// - [#HEADER_HEARTBEAT] indicating that this packet is a heartbeat and will be ignored +/// - [#HEADER_END] indicating the end of the stream /// /// /// Writes to either of the two `Output`s are synchronized on the shared /// `destination` stream, ensuring that they always arrive complete and without /// interleaving. On the other side, a `Pumper` reads from the combined /// stream, forwards each packet to its respective destination stream, or terminates -/// when it hits a packet with `header == 0` +/// when it hits a packet with [#HEADER_END]. public class ProxyStream { - private static final int MAX_CHUNK_SIZE = 256 * 1024; // 256kb + public static final int MAX_CHUNK_SIZE = 256 * 1024; // 256kb /** Indicates the end of the connection. */ private static final byte HEADER_END = 0; - /** The key for the output stream */ + /** The header for the output stream */ private static final byte HEADER_STREAM_OUT = 1; - /** The key for the error stream */ - private static final byte HEADER_STREAM_ERR = 2; + /** The header for the output stream when a single byte is sent. */ + private static final byte HEADER_STREAM_OUT_SINGLE_BYTE = 2; + + /** The header for the error stream */ + private static final byte HEADER_STREAM_ERR = 3; + + /** The header for the error stream when a single byte is sent. */ + private static final byte HEADER_STREAM_ERR_SINGLE_BYTE = 4; /** A heartbeat packet to keep the connection alive. */ private static final byte HEADER_HEARTBEAT = 127; public enum StreamType { /** The output stream */ - OUT(ProxyStream.HEADER_STREAM_OUT), + OUT(ProxyStream.HEADER_STREAM_OUT, ProxyStream.HEADER_STREAM_OUT_SINGLE_BYTE), /** The error stream */ - ERR(ProxyStream.HEADER_STREAM_ERR); + ERR(ProxyStream.HEADER_STREAM_ERR, ProxyStream.HEADER_STREAM_ERR_SINGLE_BYTE); - public final byte header; + public final byte header, headerSingleByte; - StreamType(byte key) { - this.header = key; + StreamType(byte header, byte headerSingleByte) { + this.header = header; + this.headerSingleByte = headerSingleByte; } } @@ -68,7 +77,9 @@ public static void sendEnd(OutputStream out, int exitCode) throws IOException { synchronized (out) { try { var buffer = new byte[5]; - ByteBuffer.wrap(buffer).order(ByteOrder.BIG_ENDIAN).put(ProxyStream.HEADER_END).putInt(exitCode); + ByteBuffer.wrap(buffer).order(ByteOrder.BIG_ENDIAN) + .put(ProxyStream.HEADER_END) + .putInt(exitCode); out.write(buffer); out.flush(); } catch (SocketException e) { @@ -98,19 +109,17 @@ public Output(java.io.OutputStream out, StreamType streamType) { @Override public void write(int b) throws IOException { synchronized (destination) { - destination.writeByte(streamType.header); - destination.writeInt(1); // 1 byte - destination.writeByte(b); + destination.write(streamType.headerSingleByte); + destination.write(b); } } @Override public void write(byte[] b) throws IOException { - if (b.length > 0) { - synchronized (destination) { - write(b, 0, b.length); - } - } + if (b.length <= 0) return; + + if (b.length == 1) write(b[0]); + else write(b, 0, b.length); } @Override @@ -127,13 +136,14 @@ else if (off < 0 || off > b.length || len < 0 || off + len > b.length || off + l synchronized (destination) { var bytesWritten = 0; - while (bytesWritten < len) { + while (bytesWritten < len && bytesWritten + off < b.length) { var chunkLength = Math.min(len - bytesWritten, MAX_CHUNK_SIZE); if (chunkLength > 0) { - destination.writeByte(streamType.header); + destination.write(streamType.header); destination.writeInt(chunkLength); - destination.write(b, off + bytesWritten, chunkLength); + var bytesToWrite = Math.min(b.length - off - bytesWritten, chunkLength); + destination.write(b, off + bytesWritten, bytesToWrite); bytesWritten += chunkLength; } } @@ -176,22 +186,40 @@ public Pumper(InputStream src, OutputStream destOut, OutputStream destErr) { public void preRead(DataInputStream src) {} + public void write(OutputStream dest, byte[] buffer, int length) throws IOException { + dest.write(buffer, 0, length); + } + @Override public void run() { var buffer = new byte[MAX_CHUNK_SIZE]; try { + readLoop: while (true) { this.preRead(src); var header = src.readByte(); - if (header == HEADER_END) { - exitCode = src.readInt(); - break; + switch (header) { + case HEADER_END: + exitCode = src.readInt(); + break readLoop; + case HEADER_HEARTBEAT: + continue; + case HEADER_STREAM_OUT: + pumpData(buffer, false, destOut); + break; + case HEADER_STREAM_OUT_SINGLE_BYTE: + pumpData(buffer, true, destOut); + break; + case HEADER_STREAM_ERR: + pumpData(buffer, false, destErr); + break; + case HEADER_STREAM_ERR_SINGLE_BYTE: + pumpData(buffer, true, destErr); + break; + default: + throw new IllegalStateException("Unexpected header: " + header); } - else if (header == HEADER_HEARTBEAT) continue; - else if (header == HEADER_STREAM_OUT) pumpData(buffer, destOut); - else if (header == HEADER_STREAM_ERR) pumpData(buffer, destErr); - else throw new IllegalStateException("Unexpected header: " + header); } } catch (EOFException ignored) { @@ -211,8 +239,8 @@ public void run() { } } - private void pumpData(byte[] buffer, OutputStream stream) throws IOException { - var quantity = src.readInt(); + private void pumpData(byte[] buffer, boolean singleByte, OutputStream stream) throws IOException { + var quantity = singleByte ? 1 : src.readInt(); if (quantity > buffer.length) { // Handle error: received chunk is larger than buffer @@ -223,7 +251,7 @@ private void pumpData(byte[] buffer, OutputStream stream) throws IOException { src.readFully(buffer, 0, quantity); synchronized (synchronizer) { - stream.write(buffer, 0, quantity); + write(stream, buffer, quantity); } } diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index 2b4b1767786b..218c127b8ca1 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/core/constants/test/src/mill/client/ProxyStreamTests.java @@ -18,7 +18,9 @@ public void test() throws Exception { // are likely sizes to have bugs since we write data in chunks of size 127 int[] interestingLengths = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 100, 126, 127, 128, 129, 130, 253, 254, 255, - 256, 257, 1000, 2000, 4000, 8000 + 256, 257, 1000, 2000, 4000, 8000, + ProxyStream.MAX_CHUNK_SIZE / 16, ProxyStream.MAX_CHUNK_SIZE / 8, ProxyStream.MAX_CHUNK_SIZE / 4, + ProxyStream.MAX_CHUNK_SIZE / 2, ProxyStream.MAX_CHUNK_SIZE }; byte[] interestingBytes = { -1, -127, -126, -120, -100, -80, -60, -40, -20, -10, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 10, diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 2629bed613d2..38bdfca867bf 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -59,6 +59,7 @@ public RunWithConnectionResult(A result, int exitCode) { /// @param runClientLogic the client logic to run /// @return the exit code that the server sent back public static RunWithConnectionResult runWithConnection( + String debugName, Socket connection, Streams streams, boolean closeConnectionAfterClientLogic, @@ -69,7 +70,7 @@ public static RunWithConnectionResult runWithConnection( var socketOutputStream = connection.getOutputStream(); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); - var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams); + var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); var result = runClientLogic.run(); if (closeConnectionAfterClientLogic) socketInputStream.close(); pumperThread.join(); @@ -323,12 +324,14 @@ private static ServerLaunchOutputs readOutputs(Path daemonDir) throws IOExceptio * @return a PumperThread that processes the output/error streams from the server */ static PumperThread startStreamPumpers( - InputStream socketInputStream, OutputStream socketOutputStream, Streams streams) { + InputStream socketInputStream, OutputStream socketOutputStream, Streams streams, + String name + ) { var outPumper = new ProxyStream.Pumper(socketInputStream, streams.stdout, streams.stderr); var inPump = new InputPumper(() -> streams.stdin, () -> socketOutputStream, true); - var outPumperThread = new PumperThread(outPumper, "outPump"); + var outPumperThread = new PumperThread(outPumper, "outPump-" + name); outPumperThread.setDaemon(true); - var inThread = new Thread(inPump, "inPump"); + var inThread = new Thread(inPump, "inPump-" + name); inThread.setDaemon(true); outPumperThread.start(); inThread.start(); diff --git a/libs/daemon/client/test/src/mill/client/ClientTests.java b/libs/daemon/client/test/src/mill/client/ClientTests.java index 193666177f27..96166f6ddad5 100644 --- a/libs/daemon/client/test/src/mill/client/ClientTests.java +++ b/libs/daemon/client/test/src/mill/client/ClientTests.java @@ -1,15 +1,13 @@ package mill.client; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertTrue; - import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; -import java.io.OutputStream; import java.util.*; import mill.constants.ProxyStream; import org.junit.Test; +import static org.junit.Assert.*; + public class ClientTests { @org.junit.Rule @@ -154,7 +152,7 @@ public void proxyInputOutputStreams(byte[] samples1, byte[] samples2, int chunkM ProxyStream.Pumper pumper = new ProxyStream.Pumper(new ByteArrayInputStream(bytes), dest1, dest2); pumper.run(); - assertTrue(Arrays.equals(samples1, dest1.toByteArray())); - assertTrue(Arrays.equals(samples2, dest2.toByteArray())); + assertArrayEquals(samples1, dest1.toByteArray()); + assertArrayEquals(samples2, dest2.toByteArray()); } } diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 57307af804be..71d1ccf5421b 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -1,17 +1,10 @@ package mill.server import mill.client.lock.{Lock, Locks} -import mill.constants.{DaemonFiles, InputPumper, ProxyStream} +import mill.constants.{DaemonFiles, ProxyStream} import sun.misc.{Signal, SignalHandler} -import java.io.{ - BufferedInputStream, - BufferedOutputStream, - InputStream, - PipedInputStream, - PipedOutputStream, - PrintStream -} +import java.io.{BufferedInputStream, PrintStream} import java.net.{InetAddress, Socket, SocketAddress} import java.time.LocalDateTime import java.time.format.DateTimeFormatter @@ -55,7 +48,7 @@ abstract class Server( */ protected def preHandleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, @@ -69,7 +62,7 @@ abstract class Server( */ protected def handleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, @@ -190,10 +183,7 @@ abstract class Server( socketOpt match { case Some(sock) => - val socketInfo = Server.SocketInfo( - remote = sock.getRemoteSocketAddress, - local = sock.getLocalSocketAddress - ) + val socketInfo = Server.SocketInfo(sock) serverLog(s"handling run for $socketInfo") new Thread( () => @@ -255,7 +245,7 @@ abstract class Server( ): Unit = { /** stdout and stderr combined into one stream */ - val currentOutErr = BufferedOutputStream(clientSocket.getOutputStream) + val currentOutErr = clientSocket.getOutputStream val writtenExitCode = AtomicBoolean() def writeExitCode(exitCode: Int): Unit = { @@ -287,7 +277,7 @@ abstract class Server( } try { - val socketIn = BufferedInputStream(clientSocket.getInputStream) + val socketIn = BufferedInputStream(clientSocket.getInputStream, ProxyStream.MAX_CHUNK_SIZE) val stdout = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.OUT), true) @@ -309,18 +299,9 @@ abstract class Server( val t = new Thread( () => try { - // Proxy the input stream through a pair of Piped**putStream via a pumper, - // as the `UnixDomainSocketInputStream` we get directly from the socket does - // not properly implement `available(): Int` and thus messes up polling logic - // that relies on that method - // - // TODO: this seems to be a leftover from the times when unix sockets were used, we should try removing it - // in a separate PR. - val proxiedSocketInput = Server.proxyInputStreamThroughPumper(socketIn) - val exitCode = handleConnection( socketInfo = socketInfo, - stdin = proxiedSocketInput, + stdin = socketIn, stdout = stdout, stderr = stderr, stopServer = @@ -393,6 +374,10 @@ object Server { case class SocketInfo(remote: SocketAddress, local: SocketAddress) { override def toString: String = s"SocketInfo(remote=$remote, local=$local)" } + object SocketInfo { + def apply(socket: Socket): SocketInfo = + apply(socket.getRemoteSocketAddress, socket.getLocalSocketAddress) + } /** Immediately stops the server reporting the provided exit code to all clients. */ @FunctionalInterface trait StopServer { @@ -477,14 +462,4 @@ object Server { } } } - - private def proxyInputStreamThroughPumper(in: InputStream): PipedInputStream = { - val pipedInput = new PipedInputStream() - val pipedOutput = new PipedOutputStream(pipedInput) - val pumper = new InputPumper(() => in, () => pipedOutput, /* checkAvailable */ false) - val pumperThread = new Thread(pumper, "proxyInputStreamThroughPumper") - pumperThread.setDaemon(true) - pumperThread.start() - pipedInput - } } diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index b3373e753c79..bc956e247695 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -4,7 +4,7 @@ import mill.api.* import mill.api.daemon.internal.{CompileProblemReporter, internal} import mill.client.{LaunchedServer, ServerLauncher} import mill.client.lock.{DoubleLock, Locks, MemoryLock} -import mill.constants.DaemonFiles +import mill.constants.{DaemonFiles, ProxyStream} import mill.javalib.api.CompilationResult import mill.javalib.api.internal.* import mill.javalib.internal.{JvmWorkerArgs, RpcCompileProblemReporterMessage} @@ -353,7 +353,7 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo port, s"From '${getClass.getName}'. Daemon directory: $daemonDir" )) - val stdin = use(PipedInputStream()) + val stdin = use(PipedInputStream(ProxyStream.MAX_CHUNK_SIZE)) val stdout = use(PipedOutputStream()) val streams = ServerLauncher.Streams( stdin, @@ -361,18 +361,22 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo // stderr stream is not used in this case NullOutputStream.getInstance() ) + val debugName = + s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}" ServerLauncher.runWithConnection( + debugName, socket, - streams, /* closeConnectionAfterCommand */ true, + streams, + /* closeConnectionAfterCommand */ true, /* sendInitData */ _ => {}, () => { - val serverToClient = use(BufferedReader(InputStreamReader(PipedInputStream(stdout)))) + val serverToClient = use(BufferedReader(InputStreamReader(PipedInputStream( + stdout, + ProxyStream.MAX_CHUNK_SIZE + )))) val clientToServer = use(PrintStream(BufferedOutputStream(PipedOutputStream(stdin)))) - val wireTransport = MillRpcWireTransport.ViaStreams( - s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}", - serverToClient, - clientToServer - ) + val wireTransport = + MillRpcWireTransport.ViaStreams(debugName, serverToClient, clientToServer) val init = ZincWorkerRpcServer.Initialize( compilerBridgeWorkspace = compilerBridge.workspace, diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index ac00a7b9b2ed..a4dc8e253817 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -7,7 +7,7 @@ import mill.rpc.MillRpcWireTransport import mill.server.Server import pprint.{TPrint, TPrintColors} -import java.io.{BufferedReader, InputStream, InputStreamReader, PrintStream} +import java.io.{BufferedInputStream, BufferedReader, InputStreamReader, PrintStream} import scala.util.Using import scala.util.control.NonFatal @@ -42,7 +42,7 @@ object ZincWorkerMain { override protected def preHandleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, @@ -51,7 +51,7 @@ object ZincWorkerMain { override protected def handleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, diff --git a/runner/client/src/mill/client/MillServerLauncher.java b/runner/client/src/mill/client/MillServerLauncher.java index 7fb70282c34b..8809a16e752a 100644 --- a/runner/client/src/mill/client/MillServerLauncher.java +++ b/runner/client/src/mill/client/MillServerLauncher.java @@ -80,6 +80,8 @@ public Result run(Path daemonDir, String javaHome, Consumer log) throws log)) { log.accept("runWithConnection: " + connection); var result = runWithConnection( + "MillServerLauncher[" + + connection.getLocalSocketAddress() + " -> " + connection.getRemoteSocketAddress() + "]", connection, streams, false, diff --git a/runner/server/src/mill/server/MillDaemonServer.scala b/runner/server/src/mill/server/MillDaemonServer.scala index 93c1a572acd9..8f1ce9eb7b54 100644 --- a/runner/server/src/mill/server/MillDaemonServer.scala +++ b/runner/server/src/mill/server/MillDaemonServer.scala @@ -47,7 +47,7 @@ abstract class MillDaemonServer[State]( override protected def preHandleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, @@ -104,7 +104,7 @@ abstract class MillDaemonServer[State]( override protected def handleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, From 9dbb65b29252a5deae3a7331298e37018cd97121 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 11:09:55 +0300 Subject: [PATCH 04/25] Things seem to work for real now. --- .../src/mill/constants/InputPumper.java | 2 +- .../src/mill/constants/ProxyStream.java | 109 ++++++++++++------ .../src/mill/client/ServerLauncher.java | 5 +- .../server/src/mill/server/Server.scala | 42 ++++--- 4 files changed, 103 insertions(+), 55 deletions(-) diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index 7654ff42f6ba..18d1f8ef44a2 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -30,7 +30,7 @@ public void run() { var src = src0.get(); var dest = dest0.get(); - var buffer = new byte[256 * 1024 /* 256kb, otherwise it's too slow. */]; + var buffer = new byte[16 * 1024 /* 16kb, otherwise it's too slow. */]; try { while (running) { if (checkAvailable && src.available() == 0) diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index a43dcd1d4ee4..40fe1bf9d22e 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -34,25 +34,27 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with [#HEADER_END]. public class ProxyStream { - public static final int MAX_CHUNK_SIZE = 256 * 1024; // 256kb + public static final int MAX_CHUNK_SIZE = 32 * 1024; // 32kb - /** Indicates the end of the connection. */ - private static final byte HEADER_END = 0; + // The values are picked to make it a bit easier to spot when debugging the hex dump. /** The header for the output stream */ - private static final byte HEADER_STREAM_OUT = 1; + private static final byte HEADER_STREAM_OUT = 26; // 0x1A /** The header for the output stream when a single byte is sent. */ - private static final byte HEADER_STREAM_OUT_SINGLE_BYTE = 2; + private static final byte HEADER_STREAM_OUT_SINGLE_BYTE = 27; // 0x1B, B as in BYTE /** The header for the error stream */ - private static final byte HEADER_STREAM_ERR = 3; + private static final byte HEADER_STREAM_ERR = 42; // 0x2A /** The header for the error stream when a single byte is sent. */ - private static final byte HEADER_STREAM_ERR_SINGLE_BYTE = 4; + private static final byte HEADER_STREAM_ERR_SINGLE_BYTE = 43; // 0x2B, B as in BYTE /** A heartbeat packet to keep the connection alive. */ - private static final byte HEADER_HEARTBEAT = 127; + private static final byte HEADER_HEARTBEAT = 123; // 0x7B, B as in BEAT + + /** Indicates the end of the connection. */ + private static final byte HEADER_END = 126; // 0x7E, E as in END public enum StreamType { /** The output stream */ @@ -68,9 +70,13 @@ public enum StreamType { } } - private static boolean clientHasClosedConnection(SocketException e) { + public static boolean clientHasClosedConnection(SocketException e) { var message = e.getMessage(); - return message != null && message.contains("Broken pipe"); + return message != null && ( + message.contains("Broken pipe") + || message.contains("Socket closed") + || message.contains("Connection reset by peer") + ); } public static void sendEnd(OutputStream out, int exitCode) throws IOException { @@ -98,17 +104,26 @@ public static void sendHeartbeat(OutputStream out) throws IOException { } public static class Output extends java.io.OutputStream { + /** + * Object used for synchronization so that our writes wouldn't interleave. + *

+ * We can't use {@link #destination} because it's a private object that we create here and {@link #sendEnd} + * and {@link #sendHeartbeat} use a different object. + **/ + private final java.io.OutputStream synchronizer; + private final DataOutputStream destination; private final StreamType streamType; public Output(java.io.OutputStream out, StreamType streamType) { + this.synchronizer = out; this.destination = new DataOutputStream(out); this.streamType = streamType; } @Override public void write(int b) throws IOException { - synchronized (destination) { + synchronized (synchronizer) { destination.write(streamType.headerSingleByte); destination.write(b); } @@ -116,50 +131,56 @@ public void write(int b) throws IOException { @Override public void write(byte[] b) throws IOException { - if (b.length <= 0) return; - - if (b.length == 1) write(b[0]); - else write(b, 0, b.length); + switch (b.length) { + case 0: + return; + case 1: + write(b[0]); + break; + default: + write(b, 0, b.length); + break; + } } @Override - public void write(byte[] b, int off, int len) throws IOException { + public void write(byte[] sourceBuffer, int offset, int len) throws IOException { // Validate arguments once at the beginning, which is cleaner // and standard practice for public methods. - if (b == null) throw new NullPointerException("byte array is null"); - else if (off < 0 || off > b.length || len < 0 || off + len > b.length || off + len < 0) { - throw new IndexOutOfBoundsException( - "Write indexes out of range: off=" + off + ", len=" + len + ", b.length=" + b.length - ); - } + if (sourceBuffer == null) throw new NullPointerException("byte array is null"); + if (offset < 0 || offset > sourceBuffer.length) throw new IndexOutOfBoundsException("Write offset out of range: " + offset); + if (len < 0) throw new IndexOutOfBoundsException("Write length is negative: " + len); + if (offset + len > sourceBuffer.length) throw new IndexOutOfBoundsException( + "Write goes beyond end of buffer: offset=" + offset + ", len=" + len + ", end=" + (offset + len) + " > " + sourceBuffer.length + ); + + synchronized (synchronizer) { + var bytesRemaining = len; + var currentOffset = offset; - synchronized (destination) { - var bytesWritten = 0; + while (bytesRemaining > 0) { + var chunkSize = Math.min(bytesRemaining, MAX_CHUNK_SIZE); - while (bytesWritten < len && bytesWritten + off < b.length) { - var chunkLength = Math.min(len - bytesWritten, MAX_CHUNK_SIZE); + destination.writeByte(streamType.header); + destination.writeInt(chunkSize); + destination.write(sourceBuffer, currentOffset, chunkSize); - if (chunkLength > 0) { - destination.write(streamType.header); - destination.writeInt(chunkLength); - var bytesToWrite = Math.min(b.length - off - bytesWritten, chunkLength); - destination.write(b, off + bytesWritten, bytesToWrite); - bytesWritten += chunkLength; - } + bytesRemaining -= chunkSize; + currentOffset += chunkSize; } } } @Override public void flush() throws IOException { - synchronized (destination) { + synchronized (synchronizer) { destination.flush(); } } @Override public void close() throws IOException { - synchronized (destination) { + synchronized (synchronizer) { destination.close(); } } @@ -248,10 +269,22 @@ private void pumpData(byte[] buffer, boolean singleByte, OutputStream stream) th " is larger than buffer of size " + buffer.length); } - src.readFully(buffer, 0, quantity); + var totalBytesRead = 0; + var bytesReadThisIteration = -1; + while (totalBytesRead < quantity) { + this.preRead(src); + bytesReadThisIteration = src.read(buffer, totalBytesRead, quantity - totalBytesRead); + if (bytesReadThisIteration == -1) { + break; + } else { + totalBytesRead += bytesReadThisIteration; + } + } - synchronized (synchronizer) { - write(stream, buffer, quantity); + if (bytesReadThisIteration != -1) { + synchronized (synchronizer) { + this.write(stream, buffer, totalBytesRead); + } } } diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 38bdfca867bf..2771e2c10dde 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -1,12 +1,11 @@ package mill.client; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; +import java.io.*; import java.net.InetAddress; import java.net.Socket; import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.Paths; import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 71d1ccf5421b..97f4f5975e51 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -5,7 +5,7 @@ import mill.constants.{DaemonFiles, ProxyStream} import sun.misc.{Signal, SignalHandler} import java.io.{BufferedInputStream, PrintStream} -import java.net.{InetAddress, Socket, SocketAddress} +import java.net.{InetAddress, Socket, SocketAddress, SocketException} import java.time.LocalDateTime import java.time.format.DateTimeFormatter import java.util.concurrent.atomic.AtomicBoolean @@ -257,17 +257,21 @@ abstract class Server( // We cannot use Socket#{isConnected, isClosed, isBound} because none of these // detect client-side connection closing, so instead we send a no-op heartbeat // message to see if the socket can receive data. + @volatile var lastClientAlive = true def checkClientAlive() = { - try { - ProxyStream.sendHeartbeat(currentOutErr) - true - } catch { - case NonFatal(err) => - serverLog( - s"error sending heartbeat, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" - ) - false - } + val result = + try { + ProxyStream.sendHeartbeat(currentOutErr) + true + } catch { + case NonFatal(err) => + serverLog( + s"error sending heartbeat, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" + ) + false + } + lastClientAlive = result + result } def stopServer(from: String, reason: String, exitCode: Int) = { @@ -296,6 +300,7 @@ abstract class Server( @volatile var done = false @volatile var idle = false + @volatile var writingExceptionLog = false val t = new Thread( () => try { @@ -314,9 +319,16 @@ abstract class Server( serverLog(s"connection handler finished, sending exitCode $exitCode to client") writeExitCode(exitCode) } catch { + case e: SocketException if ProxyStream.clientHasClosedConnection(e) => // do nothing case e: Throwable => - serverLog("connection handler error: " + e) - serverLog("connection handler stack trace: " + e.getStackTrace.mkString("\n")) + writingExceptionLog = true + try { + serverLog( + s"""connection handler error: $e + |connection handler stack trace: ${e.getStackTrace.mkString("\n")} + |""".stripMargin + ) + } finally writingExceptionLog = false writeExitCode(1) } finally { done = true @@ -336,6 +348,10 @@ abstract class Server( serverSocketClose() } + serverLog(s"done=$done, idle=$idle, lastClientAlive=$lastClientAlive") + + // Wait until exception log writer finishes. + while (writingExceptionLog) Thread.sleep(1) t.interrupt() // Try to give thread a moment to stop before we kill it for real // From f9ae7e4324859410ef3cc4f43b2cc3ed3966eb54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 15:14:34 +0300 Subject: [PATCH 05/25] WIP --- .../src/mill/constants/InputPumper.java | 2 +- .../src/mill/constants/ProxyStream.java | 2 +- .../src/mill/internal/PromptLoggerTests.scala | 4 +- .../src/mill/client/ServerLauncher.java | 8 ++- .../client/debug/DebuggingInputStream.java | 51 +++++++++++++++++++ .../client/debug/DebuggingOutputStream.java | 48 +++++++++++++++++ 6 files changed, 109 insertions(+), 6 deletions(-) create mode 100644 libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java create mode 100644 libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index 18d1f8ef44a2..cd7b9377f5b6 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -30,7 +30,7 @@ public void run() { var src = src0.get(); var dest = dest0.get(); - var buffer = new byte[16 * 1024 /* 16kb, otherwise it's too slow. */]; + var buffer = new byte[1024 /* 1kb */]; try { while (running) { if (checkAvailable && src.available() == 0) diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index 40fe1bf9d22e..90be86e5c1aa 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -34,7 +34,7 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with [#HEADER_END]. public class ProxyStream { - public static final int MAX_CHUNK_SIZE = 32 * 1024; // 32kb + public static final int MAX_CHUNK_SIZE = 126; // 32 * 1024; // 32kb // The values are picked to make it a bit easier to spot when debugging the hex dump. diff --git a/core/internal/test/src/mill/internal/PromptLoggerTests.scala b/core/internal/test/src/mill/internal/PromptLoggerTests.scala index 0843dcb55601..536652f7da3d 100644 --- a/core/internal/test/src/mill/internal/PromptLoggerTests.scala +++ b/core/internal/test/src/mill/internal/PromptLoggerTests.scala @@ -9,8 +9,8 @@ object PromptLoggerTests extends TestSuite { def setup(now: () => Long, terminfoPath: os.Path) = { val baos = new ByteArrayOutputStream() - val baosOut = new PrintStream(new ProxyStream.Output(baos, ProxyStream.OUT)) - val baosErr = new PrintStream(new ProxyStream.Output(baos, ProxyStream.ERR)) + val baosOut = new PrintStream(new ProxyStream.Output(baos, ProxyStream.StreamType.OUT)) + val baosErr = new PrintStream(new ProxyStream.Output(baos, ProxyStream.StreamType.ERR)) val promptLogger = new PromptLogger( colored = false, enableTicker = true, diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 2771e2c10dde..0cfe02cadef1 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -9,6 +9,9 @@ import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; + +import mill.client.debug.DebuggingInputStream; +import mill.client.debug.DebuggingOutputStream; import mill.client.lock.Lock; import mill.client.lock.Locks; import mill.constants.DaemonFiles; @@ -65,8 +68,9 @@ public static RunWithConnectionResult runWithConnection( Consumer sendInitData, RunClientLogic runClientLogic) throws Exception { - var socketInputStream = connection.getInputStream(); - var socketOutputStream = connection.getOutputStream(); + var wd = Paths.get("."); + var socketInputStream = new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true); + var socketOutputStream = new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java new file mode 100644 index 000000000000..c40c3885bf1c --- /dev/null +++ b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java @@ -0,0 +1,51 @@ +package mill.client.debug; +import java.io.*; +import java.nio.file.Path; +import java.time.LocalDateTime; + +public class DebuggingInputStream extends FilterInputStream { + + private final OutputStream debugOutput; + private final boolean writeSeparateOps; + + public DebuggingInputStream(InputStream in, Path workingDir, String name, boolean writeSeparateOps) { + super(in); + this.writeSeparateOps = writeSeparateOps; + try { + this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + } catch (FileNotFoundException e) { + throw new RuntimeException(e); + } + } + + @Override + public int read() throws IOException { + int b = super.read(); + if (b != -1) { + if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " read(): " + b + "\n").getBytes()); + else debugOutput.write(b); + } + return b; + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + int bytesRead = super.read(b, off, len); + if (bytesRead != -1) { + if (writeSeparateOps) debugOutput.write( + (LocalDateTime.now() + "readArray(off=" + off + ", len=" + len + ", bytesRead=" + bytesRead + "): " + new String(b, off, bytesRead) + "\n").getBytes() + ); + else debugOutput.write(b, off, bytesRead); + } + return bytesRead; + } + + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + debugOutput.close(); + } + } +} diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java new file mode 100644 index 000000000000..bbf83adc88c6 --- /dev/null +++ b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java @@ -0,0 +1,48 @@ +package mill.client.debug; + +import java.io.*; +import java.nio.file.Path; +import java.time.LocalDateTime; +import java.util.Arrays; +import java.util.stream.Collectors; + +public class DebuggingOutputStream extends FilterOutputStream { + private final OutputStream debugOutput; + private final boolean writeSeparateOps; + + public DebuggingOutputStream(OutputStream out, Path workingDir, String name, boolean writeSeparateOps) { + super(out); + + this.writeSeparateOps = writeSeparateOps; + try { + this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + } catch (FileNotFoundException e) { + throw new RuntimeException(e); + } + } + + @Override + public void write(int b) throws IOException { + if (writeSeparateOps) debugOutput.write( + ( + LocalDateTime.now() + " write(): " + b + "\n" /*+ + Arrays.stream(new Exception().getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n"))*/ + ).getBytes() + ); + else debugOutput.write(b); + + super.write(b); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + if (writeSeparateOps) debugOutput.write( + ( + LocalDateTime.now() + " write(off=" + off + ", len=" + len + "): " + new String(b, off, len) + "\n" + ).getBytes() + ); + else debugOutput.write(b, off, len); + + super.write(b, off, len); + } +} From 2db6ef710d0a0a8a0f665fec0cdc2a8bb72d2420 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 17:47:20 +0300 Subject: [PATCH 06/25] WIP: debugging --- .../src/mill/constants/ProxyStream.java | 2 +- .../src/mill/client/ServerLauncher.java | 41 ++++++++++++++++++- .../client/debug/DebuggingInputStream.java | 10 +++-- .../client/debug/DebuggingOutputStream.java | 14 ++++--- .../server/src/mill/server/Server.scala | 9 +++- 5 files changed, 62 insertions(+), 14 deletions(-) diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index 90be86e5c1aa..ec1137e90594 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -34,7 +34,7 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with [#HEADER_END]. public class ProxyStream { - public static final int MAX_CHUNK_SIZE = 126; // 32 * 1024; // 32kb + public static final int MAX_CHUNK_SIZE = 1024; // 32 * 1024; // 32kb // The values are picked to make it a bit easier to spot when debugging the hex dump. diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 0cfe02cadef1..bbb1ad951837 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -69,8 +69,12 @@ public static RunWithConnectionResult runWithConnection( RunClientLogic runClientLogic) throws Exception { var wd = Paths.get("."); - var socketInputStream = new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true); - var socketOutputStream = new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true); + var socketInputStream = new BufferedInputStream( + new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true) + ); + var socketOutputStream = new BufferedOutputStream( + new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true) + ); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); @@ -80,6 +84,34 @@ public static RunWithConnectionResult runWithConnection( return new RunWithConnectionResult<>(result, pumperThread.exitCode()); } + /// Run a client logic with a connection established to a Mill server (via [#connectToServer]). + /// + /// @param connection the socket connected to the server + /// @param closeConnectionAfterClientLogic whether to close the connection after running the + // client logic + /// @param runClientLogic the client logic to run + /// @return the exit code that the server sent back + public static RunWithConnectionResult runWithConnection( + String debugName, + Socket connection, + boolean closeConnectionAfterClientLogic, + Consumer sendInitData, + RunClientLogicWithStreams runClientLogic) + throws Exception { + var wd = Paths.get("."); + var socketInputStream = new BufferedInputStream( + new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true) + ); + var socketOutputStream = new BufferedOutputStream( + new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true) + ); + sendInitData.accept(socketOutputStream); + socketOutputStream.flush(); + var result = runClientLogic.run(socketInputStream, socketOutputStream); + if (closeConnectionAfterClientLogic) socketInputStream.close(); + return new RunWithConnectionResult<>(result, pumperThread.exitCode()); + } + /** * Establishes a connection to the Mill server by acquiring necessary locks and potentially * starting a new server process if one is not already running. @@ -350,6 +382,11 @@ public interface RunClientLogic { /// Runs the client logic. A run() throws Exception; } + + public interface RunClientLogicWithStreams { + /// Runs the client logic. + A run(InputStream inStream, OutputStream outStream) throws Exception; + } public static class Result { public final int exitCode; diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java index c40c3885bf1c..e284f8473cab 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java @@ -1,15 +1,17 @@ package mill.client.debug; + import java.io.*; import java.nio.file.Path; import java.time.LocalDateTime; -public class DebuggingInputStream extends FilterInputStream { +public class DebuggingInputStream extends InputStream { + private final InputStream in; private final OutputStream debugOutput; private final boolean writeSeparateOps; public DebuggingInputStream(InputStream in, Path workingDir, String name, boolean writeSeparateOps) { - super(in); + this.in = in; this.writeSeparateOps = writeSeparateOps; try { this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); @@ -20,7 +22,7 @@ public DebuggingInputStream(InputStream in, Path workingDir, String name, boolea @Override public int read() throws IOException { - int b = super.read(); + int b = in.read(); if (b != -1) { if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " read(): " + b + "\n").getBytes()); else debugOutput.write(b); @@ -30,7 +32,7 @@ public int read() throws IOException { @Override public int read(byte[] b, int off, int len) throws IOException { - int bytesRead = super.read(b, off, len); + int bytesRead = in.read(b, off, len); if (bytesRead != -1) { if (writeSeparateOps) debugOutput.write( (LocalDateTime.now() + "readArray(off=" + off + ", len=" + len + ", bytesRead=" + bytesRead + "): " + new String(b, off, bytesRead) + "\n").getBytes() diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java index bbf83adc88c6..2ff4f3ac3212 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java @@ -6,13 +6,13 @@ import java.util.Arrays; import java.util.stream.Collectors; -public class DebuggingOutputStream extends FilterOutputStream { +public class DebuggingOutputStream extends OutputStream { + private final OutputStream out; private final OutputStream debugOutput; private final boolean writeSeparateOps; public DebuggingOutputStream(OutputStream out, Path workingDir, String name, boolean writeSeparateOps) { - super(out); - + this.out = out; this.writeSeparateOps = writeSeparateOps; try { this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); @@ -31,7 +31,9 @@ public void write(int b) throws IOException { ); else debugOutput.write(b); - super.write(b); + out.write(b); + + if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " done\n").getBytes()); } @Override @@ -43,6 +45,8 @@ public void write(byte[] b, int off, int len) throws IOException { ); else debugOutput.write(b, off, len); - super.write(b, off, len); + out.write(b, off, len); + + if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " done\n").getBytes()); } } diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 97f4f5975e51..1a8ceaf685aa 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -1,5 +1,6 @@ package mill.server +import mill.client.debug.{DebuggingInputStream, DebuggingOutputStream} import mill.client.lock.{Lock, Locks} import mill.constants.{DaemonFiles, ProxyStream} import sun.misc.{Signal, SignalHandler} @@ -245,7 +246,8 @@ abstract class Server( ): Unit = { /** stdout and stderr combined into one stream */ - val currentOutErr = clientSocket.getOutputStream + val currentOutErr = + DebuggingOutputStream(clientSocket.getOutputStream, daemonDir.toNIO, "out_server", true) val writtenExitCode = AtomicBoolean() def writeExitCode(exitCode: Int): Unit = { @@ -281,7 +283,10 @@ abstract class Server( } try { - val socketIn = BufferedInputStream(clientSocket.getInputStream, ProxyStream.MAX_CHUNK_SIZE) + val socketIn = BufferedInputStream( + DebuggingInputStream(clientSocket.getInputStream, daemonDir.toNIO, "in_server", true), + ProxyStream.MAX_CHUNK_SIZE + ) val stdout = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.OUT), true) From f19bad6a4d5d27d451af147424a6d8bc7f821fd4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 19:16:10 +0300 Subject: [PATCH 07/25] JvmWorkerImpl: a version without pumpers --- .../src/mill/constants/ProxyStream.java | 13 +- .../src/mill/constants/SocketUtil.java | 14 ++ .../src/mill/client/ServerLauncher.java | 8 +- .../src/mill/server/ProxyStreamServer.scala | 174 +++++++++++++++ .../server/src/mill/server/Server.scala | 203 ++++++++++-------- .../mill/javalib/worker/JvmWorkerImpl.scala | 27 +-- .../mill/javalib/zinc/ZincWorkerMain.scala | 73 +++++-- .../src/mill/rpc/MillRpcWireTransport.scala | 20 +- .../src/mill/server/MillDaemonServer.scala | 9 +- 9 files changed, 392 insertions(+), 149 deletions(-) create mode 100644 core/constants/src/mill/constants/SocketUtil.java create mode 100644 libs/daemon/server/src/mill/server/ProxyStreamServer.scala diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index ec1137e90594..0deed745f488 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -34,7 +34,7 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with [#HEADER_END]. public class ProxyStream { - public static final int MAX_CHUNK_SIZE = 1024; // 32 * 1024; // 32kb + public static final int MAX_CHUNK_SIZE = 4 * 1024; // 4kb // The values are picked to make it a bit easier to spot when debugging the hex dump. @@ -70,15 +70,6 @@ public enum StreamType { } } - public static boolean clientHasClosedConnection(SocketException e) { - var message = e.getMessage(); - return message != null && ( - message.contains("Broken pipe") - || message.contains("Socket closed") - || message.contains("Connection reset by peer") - ); - } - public static void sendEnd(OutputStream out, int exitCode) throws IOException { synchronized (out) { try { @@ -91,7 +82,7 @@ public static void sendEnd(OutputStream out, int exitCode) throws IOException { } catch (SocketException e) { // If the client has already closed the connection, we don't really care about sending the // exit code to it. - if (!clientHasClosedConnection(e)) throw e; + if (!SocketUtil.clientHasClosedConnection(e)) throw e; } } } diff --git a/core/constants/src/mill/constants/SocketUtil.java b/core/constants/src/mill/constants/SocketUtil.java new file mode 100644 index 000000000000..20db26ff1637 --- /dev/null +++ b/core/constants/src/mill/constants/SocketUtil.java @@ -0,0 +1,14 @@ +package mill.constants; + +import java.net.SocketException; + +public class SocketUtil { + public static boolean clientHasClosedConnection(SocketException e) { + var message = e.getMessage(); + return message != null && ( + message.contains("Broken pipe") + || message.contains("Socket closed") + || message.contains("Connection reset by peer") + ); + } +} diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index bbb1ad951837..215d30e04eec 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -91,7 +91,7 @@ public static RunWithConnectionResult runWithConnection( // client logic /// @param runClientLogic the client logic to run /// @return the exit code that the server sent back - public static RunWithConnectionResult runWithConnection( + public static A runWithConnection( String debugName, Socket connection, boolean closeConnectionAfterClientLogic, @@ -109,7 +109,7 @@ public static RunWithConnectionResult runWithConnection( socketOutputStream.flush(); var result = runClientLogic.run(socketInputStream, socketOutputStream); if (closeConnectionAfterClientLogic) socketInputStream.close(); - return new RunWithConnectionResult<>(result, pumperThread.exitCode()); + return result; } /** @@ -382,10 +382,10 @@ public interface RunClientLogic { /// Runs the client logic. A run() throws Exception; } - + public interface RunClientLogicWithStreams { /// Runs the client logic. - A run(InputStream inStream, OutputStream outStream) throws Exception; + A run(BufferedInputStream inStream, BufferedOutputStream outStream) throws Exception; } public static class Result { diff --git a/libs/daemon/server/src/mill/server/ProxyStreamServer.scala b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala new file mode 100644 index 000000000000..8b23be6ccc90 --- /dev/null +++ b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala @@ -0,0 +1,174 @@ +package mill.server + +import mill.constants.ProxyStream + +import java.io.{BufferedInputStream, BufferedOutputStream, PrintStream} +import java.util.concurrent.atomic.AtomicBoolean +import scala.util.control.NonFatal + +/** [[Server]] that incorporates [[ProxyStream]] functionality. */ +abstract class ProxyStreamServer(args: Server.Args) extends Server(args) { self => + + /** Replaces [[PreHandleConnectionData]] in this class. */ + protected type PreHandleConnectionCustomData + + override protected type PreHandleConnectionData = ProxyStreamServerData + + case class ProxyStreamServerData( + stdout: PrintStream, + stderr: PrintStream, + writtenExitCode: AtomicBoolean, + customData: PreHandleConnectionCustomData + ) { + def writeExitCode(serverToClient: BufferedOutputStream, exitCode: Int): Unit = { + self.writeExitCode(serverToClient, exitCode, writtenExitCode) + } + } + + private def writeExitCode( + serverToClient: BufferedOutputStream, + exitCode: Int, + guard: AtomicBoolean + ): Unit = { + if (!guard.getAndSet(true)) { + ProxyStream.sendEnd(serverToClient, exitCode) + } + } + + override protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Boolean = { + ProxyStream.sendHeartbeat(connectionData.serverToClient) + true + } + + override protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit = { + // Notify the client that the server is shutting down with the given exit code + data match { + case Some(data) => data.writeExitCode(connectionData.serverToClient, exitCode) + case None => writeExitCode(connectionData.serverToClient, exitCode, new AtomicBoolean(false)) + } + } + + /** + * Invoked before a thread that runs [[handleConnection]] is spawned. + */ + override final protected def preHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer + ): PreHandleConnectionData = { + val stdout = + new PrintStream( + new ProxyStream.Output(connectionData.serverToClient, ProxyStream.StreamType.OUT), + true + ) + val stderr = + new PrintStream( + new ProxyStream.Output(connectionData.serverToClient, ProxyStream.StreamType.ERR), + true + ) + + val customData = preHandleConnection( + connectionData.socketInfo, + connectionData.clientToServer, + stdout, + stderr, + stopServer, + connectionData.initialSystemProperties + ) + + ProxyStreamServerData( + stdout = stdout, + stderr = stderr, + writtenExitCode = AtomicBoolean(false), + customData = customData + ) + } + + /** + * Invoked before a thread that runs [[handleConnection]] is spawned. + */ + protected def preHandleConnection( + socketInfo: Server.SocketInfo, + stdin: BufferedInputStream, + stdout: PrintStream, + stderr: PrintStream, + stopServer: Server.StopServer, + initialSystemProperties: Map[String, String] + ): PreHandleConnectionCustomData + + /** + * Handle a single client connection in a separate thread. + * + * @return the exit code to return to the client + */ + protected def handleConnection( + socketInfo: Server.SocketInfo, + stdin: BufferedInputStream, + stdout: PrintStream, + stderr: PrintStream, + stopServer: Server.StopServer, + setIdle: Server.SetIdle, + initialSystemProperties: Map[String, String], + data: PreHandleConnectionCustomData + ): Int + + override final protected def handleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + setIdle: Server.SetIdle, + data: PreHandleConnectionData + ): Unit = { + val exitCode = handleConnection( + connectionData.socketInfo, + connectionData.clientToServer, + data.stdout, + data.stderr, + stopServer, + setIdle, + connectionData.initialSystemProperties, + data.customData + ) + + serverLog(s"connection handler finished, sending exitCode $exitCode to client") + data.writeExitCode(connectionData.serverToClient, exitCode) + } + + override protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData, + exception: Throwable + ): Unit = { + data.writeExitCode(connectionData.serverToClient, 1) + } + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData + ): Unit = { + // TODO review: move to mill daemon + // flush before closing the socket + System.out.flush() + System.err.flush() + + try + // Send a termination if it has not already happened + data.writeExitCode(connectionData.serverToClient, 1) + catch { + case NonFatal(err) => + serverLog( + s"error sending exit code 1, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" + ) + } + } +} diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 1a8ceaf685aa..470964889590 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -2,10 +2,10 @@ package mill.server import mill.client.debug.{DebuggingInputStream, DebuggingOutputStream} import mill.client.lock.{Lock, Locks} -import mill.constants.{DaemonFiles, ProxyStream} +import mill.constants.{DaemonFiles, SocketUtil} import sun.misc.{Signal, SignalHandler} -import java.io.{BufferedInputStream, PrintStream} +import java.io.{BufferedInputStream, BufferedOutputStream, PrintStream} import java.net.{InetAddress, Socket, SocketAddress, SocketException} import java.time.LocalDateTime import java.time.format.DateTimeFormatter @@ -17,16 +17,10 @@ import scala.util.control.NonFatal /** * Implementation of a server that binds to a random port, informs a client of the port, and accepts a client * connections. - * - * @param daemonDir directory used for exchanging pre-TCP data with a client - * @param acceptTimeout shuts down after this timeout if no clients are connected */ -abstract class Server( - daemonDir: os.Path, - acceptTimeout: Option[FiniteDuration], - locks: Locks, - testLogEvenWhenServerIdWrong: Boolean = false -) { +abstract class Server(args: Server.Args) { + import args.* + val processId: Long = Server.computeProcessId() private val acceptTimeoutMillis = acceptTimeout.map(_.toMillis) private val handlerName = getClass.getName @@ -44,37 +38,68 @@ abstract class Server( protected type PreHandleConnectionData + protected case class ConnectionData( + socketInfo: Server.SocketInfo, + clientToServer: BufferedInputStream, + serverToClient: BufferedOutputStream, + initialSystemProperties: Map[String, String] + ) + /** * Invoked before a thread that runs [[handleConnection]] is spawned. */ protected def preHandleConnection( - socketInfo: Server.SocketInfo, - stdin: BufferedInputStream, - stdout: PrintStream, - stderr: PrintStream, - stopServer: Server.StopServer, - initialSystemProperties: Map[String, String] + connectionData: ConnectionData, + stopServer: Server.StopServer ): PreHandleConnectionData /** * Handle a single client connection in a separate thread. - * - * @return the exit code to return to the client */ protected def handleConnection( - socketInfo: Server.SocketInfo, - stdin: BufferedInputStream, - stdout: PrintStream, - stderr: PrintStream, + connectionData: ConnectionData, stopServer: Server.StopServer, setIdle: Server.SetIdle, - initialSystemProperties: Map[String, String], data: PreHandleConnectionData - ): Int + ): Unit + + /** Invoked in [[handleConnection]] results in an exception. */ + protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData, + exception: Throwable + ): Unit + + protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Unit protected def connectionHandlerThreadName(socket: Socket): String = s"ConnectionHandler($handlerName, ${socket.getInetAddress}:${socket.getPort})" + /** Returns true if the client is still alive. Invoked from another thread. */ + protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Boolean + + /** + * Invoked when the server is stopped. + * + * @param data will be [[None]] if [[preHandleConnection]] haven't been invoked yet. + */ + protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit + def run(): Unit = { serverLog(s"running server in $daemonDir") serverLog(s"acceptTimeout=$acceptTimeout") @@ -244,31 +269,49 @@ abstract class Server( initialSystemProperties: Map[String, String], serverSocketClose: () => Unit ): Unit = { + val clientToServer = BufferedInputStream( + DebuggingInputStream(clientSocket.getInputStream, daemonDir.toNIO, "in_server", true), + bufferSize + ) - /** stdout and stderr combined into one stream */ - val currentOutErr = - DebuggingOutputStream(clientSocket.getOutputStream, daemonDir.toNIO, "out_server", true) - val writtenExitCode = AtomicBoolean() + val serverToClient = BufferedOutputStream( + DebuggingOutputStream(clientSocket.getOutputStream, daemonDir.toNIO, "out_server", true), + bufferSize + ) - def writeExitCode(exitCode: Int): Unit = { - if (!writtenExitCode.getAndSet(true)) { - ProxyStream.sendEnd(currentOutErr, exitCode) - } + val connectionData = + ConnectionData(socketInfo, clientToServer, serverToClient, initialSystemProperties) + + def stopServer( + from: String, + reason: String, + exitCode: Int, + data: Option[PreHandleConnectionData] + ): Nothing = { + serverLog(s"$from invoked `stopServer` (reason: $reason), exitCode $exitCode") + onStopServer(from, reason, exitCode, connectionData, data) + systemExit0(reason, exitCode) } + val data = preHandleConnection( + connectionData, + stopServer = + (reason, exitCode) => stopServer("pre-connection handler", reason, exitCode, data = None) + ) + // We cannot use Socket#{isConnected, isClosed, isBound} because none of these // detect client-side connection closing, so instead we send a no-op heartbeat // message to see if the socket can receive data. @volatile var lastClientAlive = true + val stopServerFromCheckClientAlive: Server.StopServer = (reason, exitCode) => + stopServer("checkClientAlive", reason, exitCode, Some(data)) def checkClientAlive() = { val result = - try { - ProxyStream.sendHeartbeat(currentOutErr) - true - } catch { + try checkIfClientAlive(connectionData, stopServerFromCheckClientAlive, data) + catch { case NonFatal(err) => serverLog( - s"error sending heartbeat, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" + s"error checking for client liveness, assuming client to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" ) false } @@ -276,55 +319,23 @@ abstract class Server( result } - def stopServer(from: String, reason: String, exitCode: Int) = { - serverLog(s"$from invoked `stopServer` (reason: $reason), exitCode $exitCode") - writeExitCode(exitCode) - systemExit0(reason, exitCode) - } - try { - val socketIn = BufferedInputStream( - DebuggingInputStream(clientSocket.getInputStream, daemonDir.toNIO, "in_server", true), - ProxyStream.MAX_CHUNK_SIZE - ) - - val stdout = - new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.OUT), true) - val stderr = - new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.StreamType.ERR), true) - - val data = preHandleConnection( - socketInfo = socketInfo, - stdin = socketIn, - stdout = stdout, - stderr = stderr, - stopServer = - (reason, exitCode) => stopServer("pre-connection handler", reason, exitCode), - initialSystemProperties = initialSystemProperties - ) - @volatile var done = false @volatile var idle = false @volatile var writingExceptionLog = false val t = new Thread( () => try { - val exitCode = handleConnection( - socketInfo = socketInfo, - stdin = socketIn, - stdout = stdout, - stderr = stderr, + handleConnection( + connectionData, stopServer = - (reason, exitCode) => stopServer("connection handler", reason, exitCode), + (reason, exitCode) => + stopServer("connection handler", reason, exitCode, Some(data)), setIdle = idle = _, - initialSystemProperties = initialSystemProperties, data = data ) - - serverLog(s"connection handler finished, sending exitCode $exitCode to client") - writeExitCode(exitCode) } catch { - case e: SocketException if ProxyStream.clientHasClosedConnection(e) => // do nothing + case e: SocketException if SocketUtil.clientHasClosedConnection(e) => // do nothing case e: Throwable => writingExceptionLog = true try { @@ -334,7 +345,14 @@ abstract class Server( |""".stripMargin ) } finally writingExceptionLog = false - writeExitCode(1) + onExceptionInHandleConnection( + connectionData, + stopServer = + (reason, exitCode) => + stopServer("onExceptionInHandleConnection", reason, exitCode, Some(data)), + data = data, + exception = e + ) } finally { done = true idle = true @@ -357,6 +375,7 @@ abstract class Server( // Wait until exception log writer finishes. while (writingExceptionLog) Thread.sleep(1) + t.interrupt() // Try to give thread a moment to stop before we kill it for real // @@ -371,23 +390,31 @@ abstract class Server( case e: java.lang.Error if e.getMessage.contains("Cleaner terminated abnormally") => // ignore this error and do nothing; seems benign } - - // flush before closing the socket - System.out.flush() - System.err.flush() } finally { - try writeExitCode(1) // Send a termination if it has not already happened - catch { - case NonFatal(err) => - serverLog( - s"error sending exit code 1, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" - ) - } + beforeSocketClose( + connectionData, + stopServer = + (reason, exitCode) => stopServer("beforeSocketClose", reason, exitCode, Some(data)), + data + ) } } } object Server { + /** + * @param daemonDir directory used for exchanging pre-TCP data with a client + * @param acceptTimeout shuts down after this timeout if no clients are connected + * @param bufferSize size of the buffer used to read/write from/to the client + */ + case class Args( + daemonDir: os.Path, + acceptTimeout: Option[FiniteDuration], + locks: Locks, + bufferSize: Int, + testLogEvenWhenServerIdWrong: Boolean = false + ) + /** * @param remote the address of the client * @param local the address of the server @@ -400,7 +427,7 @@ object Server { apply(socket.getRemoteSocketAddress, socket.getLocalSocketAddress) } - /** Immediately stops the server reporting the provided exit code to all clients. */ + /** Immediately stops the server with the given exit code. */ @FunctionalInterface trait StopServer { def apply(reason: String, exitCode: Int): Nothing } diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index bc956e247695..6b4de7e092f0 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -353,30 +353,23 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo port, s"From '${getClass.getName}'. Daemon directory: $daemonDir" )) - val stdin = use(PipedInputStream(ProxyStream.MAX_CHUNK_SIZE)) - val stdout = use(PipedOutputStream()) - val streams = ServerLauncher.Streams( - stdin, - stdout, - // stderr stream is not used in this case - NullOutputStream.getInstance() - ) val debugName = s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}" ServerLauncher.runWithConnection( debugName, socket, - streams, /* closeConnectionAfterCommand */ true, /* sendInitData */ _ => {}, - () => { - val serverToClient = use(BufferedReader(InputStreamReader(PipedInputStream( - stdout, - ProxyStream.MAX_CHUNK_SIZE - )))) - val clientToServer = use(PrintStream(BufferedOutputStream(PipedOutputStream(stdin)))) + (in, out) => { + val serverToClient = use(BufferedReader(InputStreamReader(in))) + val clientToServer = use(PrintStream(out)) val wireTransport = - MillRpcWireTransport.ViaStreams(debugName, serverToClient, clientToServer) + MillRpcWireTransport.ViaStreams( + debugName, + serverToClient, + clientToServer, + writeSynchronizer = clientToServer + ) val init = ZincWorkerRpcServer.Initialize( compilerBridgeWorkspace = compilerBridge.workspace, @@ -399,7 +392,7 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo fileAndDebugLog(log, s"Command finished in ${result.durationPretty}") result.result } - ).result + ) }.get } } diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index a4dc8e253817..680c3f725cc0 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -30,38 +30,38 @@ object ZincWorkerMain { } } - private class ZincWorkerTcpServer(daemonDir: os.Path) extends Server( + private class ZincWorkerTcpServer(daemonDir: os.Path) extends Server(Server.Args( daemonDir, // The worker kills the process when it needs to. acceptTimeout = None, - Locks.files(daemonDir.toString) - ) { + Locks.files(daemonDir.toString), + bufferSize = 4 * 1024 + )) { private val className = summon[TPrint[ZincWorkerTcpServer]].render(using TPrintColors.Colors) - override protected type PreHandleConnectionData = Unit + protected class WriteSynchronizer + + override protected type PreHandleConnectionData = WriteSynchronizer override protected def preHandleConnection( - socketInfo: Server.SocketInfo, - stdin: BufferedInputStream, - stdout: PrintStream, - stderr: PrintStream, - stopServer: Server.StopServer, - initialSystemProperties: Map[String, String] - ): PreHandleConnectionData = () + connectionData: ConnectionData, + stopServer: Server.StopServer + ): WriteSynchronizer = new WriteSynchronizer override protected def handleConnection( - socketInfo: Server.SocketInfo, - stdin: BufferedInputStream, - stdout: PrintStream, - stderr: PrintStream, + connectionData: ConnectionData, stopServer: Server.StopServer, setIdle: Server.SetIdle, - initialSystemProperties: Map[String, String], - data: PreHandleConnectionData + writeSynchronizer: WriteSynchronizer ) = { + import connectionData.socketInfo + val serverName = s"$className{${socketInfo.remote} -> ${socketInfo.local}}" - Using.resource(BufferedReader(InputStreamReader(stdin))) { stdin => - val transport = MillRpcWireTransport.ViaStreams(serverName, stdin, stdout) + Using.Manager { use => + val stdin = use(BufferedReader(InputStreamReader(connectionData.clientToServer))) + val stdout = use(PrintStream(connectionData.serverToClient)) + val transport = + MillRpcWireTransport.ViaStreams(serverName, stdin, stdout, writeSynchronizer) try { val server = ZincWorkerRpcServer(serverName, transport, setIdle, serverLog) @@ -74,14 +74,45 @@ object ZincWorkerMain { serverLog("server.run() starting") server.run() serverLog("server.run() finished") - 0 } } catch { case NonFatal(err) => serverLog(s"$socketInfo failed: $err") - 1 } + }.get + } + + override protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer, + exception: Throwable + ): Unit = {} + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer + ): Unit = {} + + override protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer + ): Boolean = { + writeSynchronizer.synchronized { + connectionData.serverToClient.write('\n'.toInt) + connectionData.serverToClient.flush() + true } } + + override protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit = {} } } diff --git a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala index a65a49b4015f..fd9d788a4330 100644 --- a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala +++ b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala @@ -5,6 +5,7 @@ import upickle.default.{Reader, Writer} import java.io.{BufferedReader, PrintStream} import java.util.concurrent.BlockingQueue +import scala.annotation.tailrec trait MillRpcWireTransport extends AutoCloseable { @@ -19,13 +20,20 @@ trait MillRpcWireTransport extends AutoCloseable { readAndTryToParse[A](typeName.render(using TPrintColors.Colors).render, log) /** Helper that reads a message from the wire and tries to parse it, logging along the way. */ - def readAndTryToParse[A: Reader](typeName: String, log: String => Unit): Option[A] = { + @tailrec final def readAndTryToParse[A: Reader]( + typeName: String, + log: String => Unit + ): Option[A] = { log(s"Trying to read $typeName") read() match { case None => log("Transport wire broken.") None + // Empty line means a heartbeat message + case Some("") => + readAndTryToParse[A](typeName, log) + case Some(line) => log(s"Received, will try to parse as $typeName: $line") val parsed = upickle.default.read(line) @@ -101,17 +109,21 @@ object MillRpcWireTransport { /** * @param serverToClient server to client stream * @param clientToServer client to server stream + * @param writeSynchronizer synchronizer for writes */ class ViaStreams( override val name: String, serverToClient: BufferedReader, - clientToServer: PrintStream + clientToServer: PrintStream, + writeSynchronizer: AnyRef ) extends MillRpcWireTransport { override def read(): Option[String] = Option(serverToClient.readLine()) override def write(message: String): Unit = { - clientToServer.println(message) - clientToServer.flush() + writeSynchronizer.synchronized { + clientToServer.println(message) + clientToServer.flush() + } } override def close(): Unit = { diff --git a/runner/server/src/mill/server/MillDaemonServer.scala b/runner/server/src/mill/server/MillDaemonServer.scala index 8f1ce9eb7b54..01765a90448f 100644 --- a/runner/server/src/mill/server/MillDaemonServer.scala +++ b/runner/server/src/mill/server/MillDaemonServer.scala @@ -23,12 +23,13 @@ abstract class MillDaemonServer[State]( acceptTimeout: FiniteDuration, locks: Locks, testLogEvenWhenServerIdWrong: Boolean = false -) extends Server( +) extends ProxyStreamServer(Server.Args( daemonDir = daemonDir, acceptTimeout = Some(acceptTimeout), locks = locks, - testLogEvenWhenServerIdWrong = testLogEvenWhenServerIdWrong - ) { + testLogEvenWhenServerIdWrong = testLogEvenWhenServerIdWrong, + bufferSize = 4 * 1024 + )) { def outLock: mill.client.lock.Lock def out: os.Path @@ -43,7 +44,7 @@ abstract class MillDaemonServer[State]( override protected def connectionHandlerThreadName(socket: Socket): String = s"MillServerActionRunner(${socket.getInetAddress}:${socket.getPort})" - protected override type PreHandleConnectionData = ClientInitData + override protected type PreHandleConnectionCustomData = ClientInitData override protected def preHandleConnection( socketInfo: Server.SocketInfo, From aafd0b6dd27bab1747d6126de769f2cbce60a844 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 19:24:40 +0300 Subject: [PATCH 08/25] move code --- .../server/src/mill/server/ProxyStreamServer.scala | 5 ----- runner/server/src/mill/server/MillDaemonServer.scala | 12 ++++++++++++ 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/libs/daemon/server/src/mill/server/ProxyStreamServer.scala b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala index 8b23be6ccc90..ec369825e480 100644 --- a/libs/daemon/server/src/mill/server/ProxyStreamServer.scala +++ b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala @@ -156,11 +156,6 @@ abstract class ProxyStreamServer(args: Server.Args) extends Server(args) { self stopServer: Server.StopServer, data: ProxyStreamServerData ): Unit = { - // TODO review: move to mill daemon - // flush before closing the socket - System.out.flush() - System.err.flush() - try // Send a termination if it has not already happened data.writeExitCode(connectionData.serverToClient, 1) diff --git a/runner/server/src/mill/server/MillDaemonServer.scala b/runner/server/src/mill/server/MillDaemonServer.scala index 01765a90448f..fbd43f958d5d 100644 --- a/runner/server/src/mill/server/MillDaemonServer.scala +++ b/runner/server/src/mill/server/MillDaemonServer.scala @@ -141,6 +141,18 @@ abstract class MillDaemonServer[State]( initialSystemProperties: Map[String, String], stopServer: Server.StopServer ): (Boolean, State) + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData + ): Unit = { + // flush before closing the socket + System.out.flush() + System.err.flush() + + super.beforeSocketClose(connectionData, stopServer, data) + } } object MillDaemonServer { From ca69b6c2588e0f142b974ab3de35b06ec30c6df0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 19:27:36 +0300 Subject: [PATCH 09/25] Remove stream debugging. --- .../client/src/mill/client/ServerLauncher.java | 18 ++++-------------- .../client/debug/DebuggingInputStream.java | 1 + .../client/debug/DebuggingOutputStream.java | 1 + .../daemon/server/src/mill/server/Server.scala | 11 ++--------- 4 files changed, 8 insertions(+), 23 deletions(-) diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 215d30e04eec..11bf984804ce 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -68,13 +68,8 @@ public static RunWithConnectionResult runWithConnection( Consumer sendInitData, RunClientLogic runClientLogic) throws Exception { - var wd = Paths.get("."); - var socketInputStream = new BufferedInputStream( - new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true) - ); - var socketOutputStream = new BufferedOutputStream( - new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true) - ); + var socketInputStream = new BufferedInputStream(connection.getInputStream()); + var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); @@ -98,13 +93,8 @@ public static A runWithConnection( Consumer sendInitData, RunClientLogicWithStreams runClientLogic) throws Exception { - var wd = Paths.get("."); - var socketInputStream = new BufferedInputStream( - new DebuggingInputStream(connection.getInputStream(), wd, "in_" + debugName, true) - ); - var socketOutputStream = new BufferedOutputStream( - new DebuggingOutputStream(connection.getOutputStream(), wd, "out_" + debugName, true) - ); + var socketInputStream = new BufferedInputStream(connection.getInputStream()); + var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); var result = runClientLogic.run(socketInputStream, socketOutputStream); diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java index e284f8473cab..4a3666b35967 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java @@ -4,6 +4,7 @@ import java.nio.file.Path; import java.time.LocalDateTime; +/** Writes everything you read from the input stream to a file as well. */ public class DebuggingInputStream extends InputStream { private final InputStream in; diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java index 2ff4f3ac3212..a3f97f13ddcd 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java @@ -6,6 +6,7 @@ import java.util.Arrays; import java.util.stream.Collectors; +/* Writes everything you write to the output stream to a file as well. */ public class DebuggingOutputStream extends OutputStream { private final OutputStream out; private final OutputStream debugOutput; diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 470964889590..c18a2cec2ba3 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -269,15 +269,8 @@ abstract class Server(args: Server.Args) { initialSystemProperties: Map[String, String], serverSocketClose: () => Unit ): Unit = { - val clientToServer = BufferedInputStream( - DebuggingInputStream(clientSocket.getInputStream, daemonDir.toNIO, "in_server", true), - bufferSize - ) - - val serverToClient = BufferedOutputStream( - DebuggingOutputStream(clientSocket.getOutputStream, daemonDir.toNIO, "out_server", true), - bufferSize - ) + val clientToServer = BufferedInputStream(clientSocket.getInputStream, bufferSize) + val serverToClient = BufferedOutputStream(clientSocket.getOutputStream, bufferSize) val connectionData = ConnectionData(socketInfo, clientToServer, serverToClient, initialSystemProperties) From 63a79ec9d83dda4faf775ee3d2ad12d364d55f27 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Thu, 14 Aug 2025 19:32:41 +0300 Subject: [PATCH 10/25] Optimize imports. --- libs/daemon/client/src/mill/client/ServerLauncher.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 11bf984804ce..4f64138eabe3 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -5,13 +5,10 @@ import java.net.Socket; import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.Paths; import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; -import mill.client.debug.DebuggingInputStream; -import mill.client.debug.DebuggingOutputStream; import mill.client.lock.Lock; import mill.client.lock.Locks; import mill.constants.DaemonFiles; From 1257817053eb5c64c4413e2e49cf7e1bb4c35e0f Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 14 Aug 2025 16:41:14 +0000 Subject: [PATCH 11/25] [autofix.ci] apply automated fixes --- .../src/mill/constants/InputPumper.java | 3 +- .../src/mill/constants/ProxyStream.java | 36 +++++++++--------- .../src/mill/constants/SocketUtil.java | 9 ++--- .../src/mill/client/ProxyStreamTests.java | 38 +++++++++++++++++-- .../src/mill/client/ServerLauncher.java | 11 +++--- .../client/debug/DebuggingInputStream.java | 16 +++++--- .../client/debug/DebuggingOutputStream.java | 26 ++++++------- .../test/src/mill/client/ClientTests.java | 4 +- .../server/src/mill/server/Server.scala | 4 +- .../mill/javalib/worker/JvmWorkerImpl.scala | 3 +- .../mill/javalib/zinc/ZincWorkerMain.scala | 2 +- .../src/mill/client/MillServerLauncher.java | 4 +- 12 files changed, 93 insertions(+), 63 deletions(-) diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index cd7b9377f5b6..9e06355e952f 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -13,7 +13,8 @@ public class InputPumper implements Runnable { /// /// We need to do that because if we call `.read` /// and there is nothing to read, [it can unnecessarily delay the JVM exit by 350ms]( - /// https://stackoverflow.com/questions/48951611/blocking-on-stdin-makes-java-process-take-350ms-more-to-exit) + /// + // https://stackoverflow.com/questions/48951611/blocking-on-stdin-makes-java-process-take-350ms-more-to-exit) private final boolean checkAvailable; public InputPumper( diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index 0deed745f488..bf7e8f1de33c 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -20,9 +20,11 @@ /// /// Where header is a single byte of the form: /// -/// - [#HEADER_STREAM_OUT]/[#HEADER_STREAM_ERR] respectively indicating that this packet is for the `OUT`/`ERR` +/// - [#HEADER_STREAM_OUT]/[#HEADER_STREAM_ERR] respectively indicating that this packet is for +// the `OUT`/`ERR` /// stream, and it will be followed by 4 bytes for the length of the body and then the body. -/// - [#HEADER_STREAM_OUT_SINGLE_BYTE]/[#HEADER_STREAM_ERR_SINGLE_BYTE] respectively indicating that this packet is +/// - [#HEADER_STREAM_OUT_SINGLE_BYTE]/[#HEADER_STREAM_ERR_SINGLE_BYTE] respectively indicating +// that this packet is /// for the `OUT`/`ERR` stream, and it will be followed by a single byte for the body /// - [#HEADER_HEARTBEAT] indicating that this packet is a heartbeat and will be ignored /// - [#HEADER_END] indicating the end of the stream @@ -61,7 +63,6 @@ public enum StreamType { OUT(ProxyStream.HEADER_STREAM_OUT, ProxyStream.HEADER_STREAM_OUT_SINGLE_BYTE), /** The error stream */ ERR(ProxyStream.HEADER_STREAM_ERR, ProxyStream.HEADER_STREAM_ERR_SINGLE_BYTE); - public final byte header, headerSingleByte; StreamType(byte header, byte headerSingleByte) { @@ -74,9 +75,10 @@ public static void sendEnd(OutputStream out, int exitCode) throws IOException { synchronized (out) { try { var buffer = new byte[5]; - ByteBuffer.wrap(buffer).order(ByteOrder.BIG_ENDIAN) - .put(ProxyStream.HEADER_END) - .putInt(exitCode); + ByteBuffer.wrap(buffer) + .order(ByteOrder.BIG_ENDIAN) + .put(ProxyStream.HEADER_END) + .putInt(exitCode); out.write(buffer); out.flush(); } catch (SocketException e) { @@ -139,11 +141,12 @@ public void write(byte[] sourceBuffer, int offset, int len) throws IOException { // Validate arguments once at the beginning, which is cleaner // and standard practice for public methods. if (sourceBuffer == null) throw new NullPointerException("byte array is null"); - if (offset < 0 || offset > sourceBuffer.length) throw new IndexOutOfBoundsException("Write offset out of range: " + offset); + if (offset < 0 || offset > sourceBuffer.length) + throw new IndexOutOfBoundsException("Write offset out of range: " + offset); if (len < 0) throw new IndexOutOfBoundsException("Write length is negative: " + len); - if (offset + len > sourceBuffer.length) throw new IndexOutOfBoundsException( - "Write goes beyond end of buffer: offset=" + offset + ", len=" + len + ", end=" + (offset + len) + " > " + sourceBuffer.length - ); + if (offset + len > sourceBuffer.length) + throw new IndexOutOfBoundsException("Write goes beyond end of buffer: offset=" + offset + + ", len=" + len + ", end=" + (offset + len) + " > " + sourceBuffer.length); synchronized (synchronizer) { var bytesRemaining = len; @@ -233,12 +236,10 @@ public void run() { throw new IllegalStateException("Unexpected header: " + header); } } - } - catch (EOFException ignored) { + } catch (EOFException ignored) { // This is a normal and expected way for the loop to terminate // when the other side closes the connection. - } - catch (IOException ignored) { + } catch (IOException ignored) { // This happens when the upstream pipe was closed } @@ -251,13 +252,14 @@ public void run() { } } - private void pumpData(byte[] buffer, boolean singleByte, OutputStream stream) throws IOException { + private void pumpData(byte[] buffer, boolean singleByte, OutputStream stream) + throws IOException { var quantity = singleByte ? 1 : src.readInt(); if (quantity > buffer.length) { // Handle error: received chunk is larger than buffer - throw new IOException("Received chunk of size " + quantity + - " is larger than buffer of size " + buffer.length); + throw new IOException("Received chunk of size " + quantity + + " is larger than buffer of size " + buffer.length); } var totalBytesRead = 0; diff --git a/core/constants/src/mill/constants/SocketUtil.java b/core/constants/src/mill/constants/SocketUtil.java index 20db26ff1637..8bbf758c073a 100644 --- a/core/constants/src/mill/constants/SocketUtil.java +++ b/core/constants/src/mill/constants/SocketUtil.java @@ -5,10 +5,9 @@ public class SocketUtil { public static boolean clientHasClosedConnection(SocketException e) { var message = e.getMessage(); - return message != null && ( - message.contains("Broken pipe") - || message.contains("Socket closed") - || message.contains("Connection reset by peer") - ); + return message != null + && (message.contains("Broken pipe") + || message.contains("Socket closed") + || message.contains("Connection reset by peer")); } } diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index 218c127b8ca1..ca926e08089f 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/core/constants/test/src/mill/client/ProxyStreamTests.java @@ -17,10 +17,40 @@ public void test() throws Exception { // Test writes of sizes around 1, around 127, around 255, and much larger. These // are likely sizes to have bugs since we write data in chunks of size 127 int[] interestingLengths = { - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 100, 126, 127, 128, 129, 130, 253, 254, 255, - 256, 257, 1000, 2000, 4000, 8000, - ProxyStream.MAX_CHUNK_SIZE / 16, ProxyStream.MAX_CHUNK_SIZE / 8, ProxyStream.MAX_CHUNK_SIZE / 4, - ProxyStream.MAX_CHUNK_SIZE / 2, ProxyStream.MAX_CHUNK_SIZE + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 20, + 30, + 40, + 50, + 100, + 126, + 127, + 128, + 129, + 130, + 253, + 254, + 255, + 256, + 257, + 1000, + 2000, + 4000, + 8000, + ProxyStream.MAX_CHUNK_SIZE / 16, + ProxyStream.MAX_CHUNK_SIZE / 8, + ProxyStream.MAX_CHUNK_SIZE / 4, + ProxyStream.MAX_CHUNK_SIZE / 2, + ProxyStream.MAX_CHUNK_SIZE }; byte[] interestingBytes = { -1, -127, -126, -120, -100, -80, -60, -40, -20, -10, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 10, diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 4f64138eabe3..1499862676e0 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -8,7 +8,6 @@ import java.util.Optional; import java.util.function.Consumer; import java.util.function.Supplier; - import mill.client.lock.Lock; import mill.client.lock.Locks; import mill.constants.DaemonFiles; @@ -69,7 +68,8 @@ public static RunWithConnectionResult runWithConnection( var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); - var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); + var pumperThread = + startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); var result = runClientLogic.run(); if (closeConnectionAfterClientLogic) socketInputStream.close(); pumperThread.join(); @@ -346,9 +346,10 @@ private static ServerLaunchOutputs readOutputs(Path daemonDir) throws IOExceptio * @return a PumperThread that processes the output/error streams from the server */ static PumperThread startStreamPumpers( - InputStream socketInputStream, OutputStream socketOutputStream, Streams streams, - String name - ) { + InputStream socketInputStream, + OutputStream socketOutputStream, + Streams streams, + String name) { var outPumper = new ProxyStream.Pumper(socketInputStream, streams.stdout, streams.stderr); var inPump = new InputPumper(() -> streams.stdin, () -> socketOutputStream, true); var outPumperThread = new PumperThread(outPumper, "outPump-" + name); diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java index 4a3666b35967..5bca9eb1ef51 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java @@ -11,11 +11,13 @@ public class DebuggingInputStream extends InputStream { private final OutputStream debugOutput; private final boolean writeSeparateOps; - public DebuggingInputStream(InputStream in, Path workingDir, String name, boolean writeSeparateOps) { + public DebuggingInputStream( + InputStream in, Path workingDir, String name, boolean writeSeparateOps) { this.in = in; this.writeSeparateOps = writeSeparateOps; try { - this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + this.debugOutput = + new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); } catch (FileNotFoundException e) { throw new RuntimeException(e); } @@ -25,7 +27,8 @@ public DebuggingInputStream(InputStream in, Path workingDir, String name, boolea public int read() throws IOException { int b = in.read(); if (b != -1) { - if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " read(): " + b + "\n").getBytes()); + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " read(): " + b + "\n").getBytes()); else debugOutput.write(b); } return b; @@ -35,9 +38,10 @@ public int read() throws IOException { public int read(byte[] b, int off, int len) throws IOException { int bytesRead = in.read(b, off, len); if (bytesRead != -1) { - if (writeSeparateOps) debugOutput.write( - (LocalDateTime.now() + "readArray(off=" + off + ", len=" + len + ", bytesRead=" + bytesRead + "): " + new String(b, off, bytesRead) + "\n").getBytes() - ); + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + "readArray(off=" + off + ", len=" + len + + ", bytesRead=" + bytesRead + "): " + new String(b, off, bytesRead) + "\n") + .getBytes()); else debugOutput.write(b, off, bytesRead); } return bytesRead; diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java index a3f97f13ddcd..000580161427 100644 --- a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java +++ b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java @@ -3,8 +3,6 @@ import java.io.*; import java.nio.file.Path; import java.time.LocalDateTime; -import java.util.Arrays; -import java.util.stream.Collectors; /* Writes everything you write to the output stream to a file as well. */ public class DebuggingOutputStream extends OutputStream { @@ -12,11 +10,13 @@ public class DebuggingOutputStream extends OutputStream { private final OutputStream debugOutput; private final boolean writeSeparateOps; - public DebuggingOutputStream(OutputStream out, Path workingDir, String name, boolean writeSeparateOps) { + public DebuggingOutputStream( + OutputStream out, Path workingDir, String name, boolean writeSeparateOps) { this.out = out; this.writeSeparateOps = writeSeparateOps; try { - this.debugOutput = new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + this.debugOutput = + new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); } catch (FileNotFoundException e) { throw new RuntimeException(e); } @@ -24,12 +24,9 @@ public DebuggingOutputStream(OutputStream out, Path workingDir, String name, boo @Override public void write(int b) throws IOException { - if (writeSeparateOps) debugOutput.write( - ( - LocalDateTime.now() + " write(): " + b + "\n" /*+ - Arrays.stream(new Exception().getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n"))*/ - ).getBytes() - ); + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " write(): " + b + "\n" /*+ + Arrays.stream(new Exception().getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n"))*/).getBytes()); else debugOutput.write(b); out.write(b); @@ -39,11 +36,10 @@ public void write(int b) throws IOException { @Override public void write(byte[] b, int off, int len) throws IOException { - if (writeSeparateOps) debugOutput.write( - ( - LocalDateTime.now() + " write(off=" + off + ", len=" + len + "): " + new String(b, off, len) + "\n" - ).getBytes() - ); + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " write(off=" + off + ", len=" + len + "): " + + new String(b, off, len) + "\n") + .getBytes()); else debugOutput.write(b, off, len); out.write(b, off, len); diff --git a/libs/daemon/client/test/src/mill/client/ClientTests.java b/libs/daemon/client/test/src/mill/client/ClientTests.java index 96166f6ddad5..2e46188c1f4a 100644 --- a/libs/daemon/client/test/src/mill/client/ClientTests.java +++ b/libs/daemon/client/test/src/mill/client/ClientTests.java @@ -1,13 +1,13 @@ package mill.client; +import static org.junit.Assert.*; + import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; import java.util.*; import mill.constants.ProxyStream; import org.junit.Test; -import static org.junit.Assert.*; - public class ClientTests { @org.junit.Rule diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index c18a2cec2ba3..63cc93fc5b4f 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -1,15 +1,13 @@ package mill.server -import mill.client.debug.{DebuggingInputStream, DebuggingOutputStream} import mill.client.lock.{Lock, Locks} import mill.constants.{DaemonFiles, SocketUtil} import sun.misc.{Signal, SignalHandler} -import java.io.{BufferedInputStream, BufferedOutputStream, PrintStream} +import java.io.{BufferedInputStream, BufferedOutputStream} import java.net.{InetAddress, Socket, SocketAddress, SocketException} import java.time.LocalDateTime import java.time.format.DateTimeFormatter -import java.util.concurrent.atomic.AtomicBoolean import scala.concurrent.duration.FiniteDuration import scala.util.Try import scala.util.control.NonFatal diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index 6b4de7e092f0..9dc40430d63f 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -4,7 +4,7 @@ import mill.api.* import mill.api.daemon.internal.{CompileProblemReporter, internal} import mill.client.{LaunchedServer, ServerLauncher} import mill.client.lock.{DoubleLock, Locks, MemoryLock} -import mill.constants.{DaemonFiles, ProxyStream} +import mill.constants.DaemonFiles import mill.javalib.api.CompilationResult import mill.javalib.api.internal.* import mill.javalib.internal.{JvmWorkerArgs, RpcCompileProblemReporterMessage} @@ -12,7 +12,6 @@ import mill.javalib.zinc.ZincWorkerRpcServer.ReporterMode import mill.javalib.zinc.{ZincApi, ZincWorker, ZincWorkerRpcServer} import mill.rpc.{MillRpcChannel, MillRpcClient, MillRpcWireTransport} import mill.util.{CachedFactoryWithInitData, HexFormat, Jvm, Timed} -import org.apache.logging.log4j.core.util.NullOutputStream import sbt.internal.util.ConsoleOut import java.io.* diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index 680c3f725cc0..860cf3c60cf8 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -7,7 +7,7 @@ import mill.rpc.MillRpcWireTransport import mill.server.Server import pprint.{TPrint, TPrintColors} -import java.io.{BufferedInputStream, BufferedReader, InputStreamReader, PrintStream} +import java.io.{BufferedReader, InputStreamReader, PrintStream} import scala.util.Using import scala.util.control.NonFatal diff --git a/runner/client/src/mill/client/MillServerLauncher.java b/runner/client/src/mill/client/MillServerLauncher.java index 8809a16e752a..055e23b70ba1 100644 --- a/runner/client/src/mill/client/MillServerLauncher.java +++ b/runner/client/src/mill/client/MillServerLauncher.java @@ -80,8 +80,8 @@ public Result run(Path daemonDir, String javaHome, Consumer log) throws log)) { log.accept("runWithConnection: " + connection); var result = runWithConnection( - "MillServerLauncher[" + - connection.getLocalSocketAddress() + " -> " + connection.getRemoteSocketAddress() + "]", + "MillServerLauncher[" + connection.getLocalSocketAddress() + " -> " + + connection.getRemoteSocketAddress() + "]", connection, streams, false, From 1b6c4f60499716affce4beee3e5f968c1b0c19f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Sun, 17 Aug 2025 11:49:43 +0300 Subject: [PATCH 12/25] Update golden tests. --- .../leak-hygiene/src/LeakHygieneTests.scala | 24 +++++++------------ 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala index 3ddf541f8e2d..9434866a260a 100644 --- a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala +++ b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala @@ -80,8 +80,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "MillServerTimeoutThread", "Process ID Checker Thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -109,8 +108,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -139,8 +137,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } @@ -165,8 +162,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "MillServerTimeoutThread", "Process ID Checker Thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -194,8 +190,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } @@ -225,8 +220,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -258,8 +252,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -289,8 +282,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Timer", "leaked thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } From fdd3543b86f2f3b3beeae5b39fd0ab3e362c155a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Sun, 17 Aug 2025 11:59:29 +0300 Subject: [PATCH 13/25] binary compatibility --- .../src/mill/constants/InputPumper.java | 2 +- .../src/mill/constants/ProxyStream.java | 35 ++++++++++++++++--- .../src/mill/client/ProxyStreamTests.java | 10 +++--- .../src/mill/internal/PromptLogger.scala | 2 +- 4 files changed, 38 insertions(+), 11 deletions(-) diff --git a/core/constants/src/mill/constants/InputPumper.java b/core/constants/src/mill/constants/InputPumper.java index 9e06355e952f..2eb006747739 100644 --- a/core/constants/src/mill/constants/InputPumper.java +++ b/core/constants/src/mill/constants/InputPumper.java @@ -18,7 +18,7 @@ public class InputPumper implements Runnable { private final boolean checkAvailable; public InputPumper( - Supplier src, Supplier dest, boolean checkAvailable) { + Supplier src, Supplier dest, Boolean checkAvailable) { this.src0 = src; this.dest0 = dest; this.checkAvailable = checkAvailable; diff --git a/core/constants/src/mill/constants/ProxyStream.java b/core/constants/src/mill/constants/ProxyStream.java index bf7e8f1de33c..d5970e9668c1 100644 --- a/core/constants/src/mill/constants/ProxyStream.java +++ b/core/constants/src/mill/constants/ProxyStream.java @@ -21,10 +21,10 @@ /// Where header is a single byte of the form: /// /// - [#HEADER_STREAM_OUT]/[#HEADER_STREAM_ERR] respectively indicating that this packet is for -// the `OUT`/`ERR` +/// the `OUT`/`ERR` /// stream, and it will be followed by 4 bytes for the length of the body and then the body. /// - [#HEADER_STREAM_OUT_SINGLE_BYTE]/[#HEADER_STREAM_ERR_SINGLE_BYTE] respectively indicating -// that this packet is +/// that this packet is /// for the `OUT`/`ERR` stream, and it will be followed by a single byte for the body /// - [#HEADER_HEARTBEAT] indicating that this packet is a heartbeat and will be ignored /// - [#HEADER_END] indicating the end of the stream @@ -36,28 +36,44 @@ /// stream, forwards each packet to its respective destination stream, or terminates /// when it hits a packet with [#HEADER_END]. public class ProxyStream { - public static final int MAX_CHUNK_SIZE = 4 * 1024; // 4kb + private static final int MAX_CHUNK_SIZE = 4 * 1024; // 4kb // The values are picked to make it a bit easier to spot when debugging the hex dump. /** The header for the output stream */ private static final byte HEADER_STREAM_OUT = 26; // 0x1A + // bincompat forwarder + @SuppressWarnings("unused") + public static final int OUT = HEADER_STREAM_OUT; + /** The header for the output stream when a single byte is sent. */ private static final byte HEADER_STREAM_OUT_SINGLE_BYTE = 27; // 0x1B, B as in BYTE /** The header for the error stream */ private static final byte HEADER_STREAM_ERR = 42; // 0x2A + // bincompat forwarder + @SuppressWarnings("unused") + public static final int ERR = HEADER_STREAM_ERR; + /** The header for the error stream when a single byte is sent. */ private static final byte HEADER_STREAM_ERR_SINGLE_BYTE = 43; // 0x2B, B as in BYTE /** A heartbeat packet to keep the connection alive. */ private static final byte HEADER_HEARTBEAT = 123; // 0x7B, B as in BEAT + // bincompat forwarder + @SuppressWarnings("unused") + public static final int HEARTBEAT = HEADER_HEARTBEAT; + /** Indicates the end of the connection. */ private static final byte HEADER_END = 126; // 0x7E, E as in END + // bincompat forwarder + @SuppressWarnings("unused") + public static final int END = HEADER_END; + public enum StreamType { /** The output stream */ OUT(ProxyStream.HEADER_STREAM_OUT, ProxyStream.HEADER_STREAM_OUT_SINGLE_BYTE), @@ -114,6 +130,11 @@ public Output(java.io.OutputStream out, StreamType streamType) { this.streamType = streamType; } + // bincompat forwarder + public Output(java.io.OutputStream out, int key) { + this(out, key == OUT ? StreamType.OUT : StreamType.ERR); + } + @Override public void write(int b) throws IOException { synchronized (synchronizer) { @@ -199,7 +220,13 @@ public Pumper(InputStream src, OutputStream destOut, OutputStream destErr) { this(src, destOut, destErr, new Object()); } - public void preRead(DataInputStream src) {} + protected void preRead(DataInputStream src) {} + + @Deprecated(forRemoval = true, since = "1.0.4") + public void preRead(InputStream src) { + if (src instanceof DataInputStream) preRead((DataInputStream) src); + else throw new UnsupportedOperationException("preRead(InputStream) is deprecated"); + } public void write(OutputStream dest, byte[] buffer, int length) throws IOException { dest.write(buffer, 0, length); diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index ca926e08089f..fbd5ced61136 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/core/constants/test/src/mill/client/ProxyStreamTests.java @@ -46,11 +46,11 @@ public void test() throws Exception { 2000, 4000, 8000, - ProxyStream.MAX_CHUNK_SIZE / 16, - ProxyStream.MAX_CHUNK_SIZE / 8, - ProxyStream.MAX_CHUNK_SIZE / 4, - ProxyStream.MAX_CHUNK_SIZE / 2, - ProxyStream.MAX_CHUNK_SIZE + 16000, + 32000, + 64000, + 128000, + 256000, }; byte[] interestingBytes = { -1, -127, -126, -120, -100, -80, -60, -40, -20, -10, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 10, diff --git a/core/internal/src/mill/internal/PromptLogger.scala b/core/internal/src/mill/internal/PromptLogger.scala index 3fd2be4cb91a..29658235d8ae 100644 --- a/core/internal/src/mill/internal/PromptLogger.scala +++ b/core/internal/src/mill/internal/PromptLogger.scala @@ -330,7 +330,7 @@ private[mill] object PromptLogger { private var lastCharWritten = 0.toChar // Make sure we synchronize everywhere - override def preRead(src: DataInputStream): Unit = synchronizer.synchronized { + override protected def preRead(src: DataInputStream): Unit = synchronizer.synchronized { if ( enableTicker && From b860b240a7708c680ff88d9062ff0f3e4ab95780 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Sun, 17 Aug 2025 12:04:55 +0300 Subject: [PATCH 14/25] Make `Timed` `private[mill]` --- libs/util/src/mill/util/Timed.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/util/src/mill/util/Timed.scala b/libs/util/src/mill/util/Timed.scala index 32ea13e117d4..e3dfab0699af 100644 --- a/libs/util/src/mill/util/Timed.scala +++ b/libs/util/src/mill/util/Timed.scala @@ -3,7 +3,7 @@ package mill.util import java.util.concurrent.TimeUnit import scala.concurrent.duration.FiniteDuration -object Timed { +private[mill] object Timed { case class Result[+A](result: A, duration: FiniteDuration) { def durationPretty: String = duration.toMillis + "ms" } From 75ba99014119e11492d3842948678a6705af3309 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Sun, 17 Aug 2025 09:16:55 +0000 Subject: [PATCH 15/25] [autofix.ci] apply automated fixes --- .../src/mill/client/ProxyStreamTests.java | 36 ++----------------- 1 file changed, 2 insertions(+), 34 deletions(-) diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index fbd5ced61136..2f3a09ae858c 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/core/constants/test/src/mill/client/ProxyStreamTests.java @@ -17,40 +17,8 @@ public void test() throws Exception { // Test writes of sizes around 1, around 127, around 255, and much larger. These // are likely sizes to have bugs since we write data in chunks of size 127 int[] interestingLengths = { - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - 10, - 20, - 30, - 40, - 50, - 100, - 126, - 127, - 128, - 129, - 130, - 253, - 254, - 255, - 256, - 257, - 1000, - 2000, - 4000, - 8000, - 16000, - 32000, - 64000, - 128000, - 256000, + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 20, 30, 40, 50, 100, 126, 127, 128, 129, 130, 253, 254, 255, + 256, 257, 1000, 2000, 4000, 8000, 16000, 32000, 64000, 128000, 256000, }; byte[] interestingBytes = { -1, -127, -126, -120, -100, -80, -60, -40, -20, -10, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 10, From 56534b7b1af8538f2d0ac3ee24887771f6f9cda8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Sun, 17 Aug 2025 12:25:19 +0300 Subject: [PATCH 16/25] Minimized version. --- .../src/mill/constants/SocketUtil.java | 13 + .../src/mill/client/ServerLauncher.java | 49 +++- .../client/debug/DebuggingInputStream.java | 58 ++++ .../client/debug/DebuggingOutputStream.java | 49 ++++ .../src/mill/server/ProxyStreamServer.scala | 169 ++++++++++++ .../server/src/mill/server/Server.scala | 253 ++++++++++-------- .../mill/javalib/worker/JvmWorkerImpl.scala | 157 ++++++++--- .../mill/javalib/zinc/ZincWorkerMain.scala | 77 ++++-- .../javalib/zinc/ZincWorkerRpcServer.scala | 182 +++++++------ .../src/mill/rpc/MillRpcWireTransport.scala | 21 +- libs/util/src/mill/util/Timed.scala | 18 ++ .../src/mill/client/MillServerLauncher.java | 2 + .../src/mill/server/MillDaemonServer.scala | 27 +- 13 files changed, 793 insertions(+), 282 deletions(-) create mode 100644 core/constants/src/mill/constants/SocketUtil.java create mode 100644 libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java create mode 100644 libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java create mode 100644 libs/daemon/server/src/mill/server/ProxyStreamServer.scala create mode 100644 libs/util/src/mill/util/Timed.scala diff --git a/core/constants/src/mill/constants/SocketUtil.java b/core/constants/src/mill/constants/SocketUtil.java new file mode 100644 index 000000000000..083b56ca066d --- /dev/null +++ b/core/constants/src/mill/constants/SocketUtil.java @@ -0,0 +1,13 @@ +package mill.constants; + +import java.net.SocketException; + +public class SocketUtil { + public static boolean clientHasClosedConnection(SocketException e) { + var message = e.getMessage(); + return message != null + && (message.contains("Broken pipe") + || message.contains("Socket closed") + || message.contains("Connection reset by peer")); + } +} diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 2629bed613d2..1499862676e0 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -1,8 +1,6 @@ package mill.client; -import java.io.IOException; -import java.io.InputStream; -import java.io.OutputStream; +import java.io.*; import java.net.InetAddress; import java.net.Socket; import java.nio.file.Files; @@ -59,23 +57,48 @@ public RunWithConnectionResult(A result, int exitCode) { /// @param runClientLogic the client logic to run /// @return the exit code that the server sent back public static RunWithConnectionResult runWithConnection( + String debugName, Socket connection, Streams streams, boolean closeConnectionAfterClientLogic, Consumer sendInitData, RunClientLogic runClientLogic) throws Exception { - var socketInputStream = connection.getInputStream(); - var socketOutputStream = connection.getOutputStream(); + var socketInputStream = new BufferedInputStream(connection.getInputStream()); + var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); sendInitData.accept(socketOutputStream); socketOutputStream.flush(); - var pumperThread = startStreamPumpers(socketInputStream, socketOutputStream, streams); + var pumperThread = + startStreamPumpers(socketInputStream, socketOutputStream, streams, debugName); var result = runClientLogic.run(); if (closeConnectionAfterClientLogic) socketInputStream.close(); pumperThread.join(); return new RunWithConnectionResult<>(result, pumperThread.exitCode()); } + /// Run a client logic with a connection established to a Mill server (via [#connectToServer]). + /// + /// @param connection the socket connected to the server + /// @param closeConnectionAfterClientLogic whether to close the connection after running the + // client logic + /// @param runClientLogic the client logic to run + /// @return the exit code that the server sent back + public static A runWithConnection( + String debugName, + Socket connection, + boolean closeConnectionAfterClientLogic, + Consumer sendInitData, + RunClientLogicWithStreams runClientLogic) + throws Exception { + var socketInputStream = new BufferedInputStream(connection.getInputStream()); + var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); + sendInitData.accept(socketOutputStream); + socketOutputStream.flush(); + var result = runClientLogic.run(socketInputStream, socketOutputStream); + if (closeConnectionAfterClientLogic) socketInputStream.close(); + return result; + } + /** * Establishes a connection to the Mill server by acquiring necessary locks and potentially * starting a new server process if one is not already running. @@ -323,12 +346,15 @@ private static ServerLaunchOutputs readOutputs(Path daemonDir) throws IOExceptio * @return a PumperThread that processes the output/error streams from the server */ static PumperThread startStreamPumpers( - InputStream socketInputStream, OutputStream socketOutputStream, Streams streams) { + InputStream socketInputStream, + OutputStream socketOutputStream, + Streams streams, + String name) { var outPumper = new ProxyStream.Pumper(socketInputStream, streams.stdout, streams.stderr); var inPump = new InputPumper(() -> streams.stdin, () -> socketOutputStream, true); - var outPumperThread = new PumperThread(outPumper, "outPump"); + var outPumperThread = new PumperThread(outPumper, "outPump-" + name); outPumperThread.setDaemon(true); - var inThread = new Thread(inPump, "inPump"); + var inThread = new Thread(inPump, "inPump-" + name); inThread.setDaemon(true); outPumperThread.start(); inThread.start(); @@ -345,6 +371,11 @@ public interface RunClientLogic { A run() throws Exception; } + public interface RunClientLogicWithStreams { + /// Runs the client logic. + A run(BufferedInputStream inStream, BufferedOutputStream outStream) throws Exception; + } + public static class Result { public final int exitCode; public final Path daemonDir; diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java new file mode 100644 index 000000000000..5bca9eb1ef51 --- /dev/null +++ b/libs/daemon/client/src/mill/client/debug/DebuggingInputStream.java @@ -0,0 +1,58 @@ +package mill.client.debug; + +import java.io.*; +import java.nio.file.Path; +import java.time.LocalDateTime; + +/** Writes everything you read from the input stream to a file as well. */ +public class DebuggingInputStream extends InputStream { + + private final InputStream in; + private final OutputStream debugOutput; + private final boolean writeSeparateOps; + + public DebuggingInputStream( + InputStream in, Path workingDir, String name, boolean writeSeparateOps) { + this.in = in; + this.writeSeparateOps = writeSeparateOps; + try { + this.debugOutput = + new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + } catch (FileNotFoundException e) { + throw new RuntimeException(e); + } + } + + @Override + public int read() throws IOException { + int b = in.read(); + if (b != -1) { + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " read(): " + b + "\n").getBytes()); + else debugOutput.write(b); + } + return b; + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + int bytesRead = in.read(b, off, len); + if (bytesRead != -1) { + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + "readArray(off=" + off + ", len=" + len + + ", bytesRead=" + bytesRead + "): " + new String(b, off, bytesRead) + "\n") + .getBytes()); + else debugOutput.write(b, off, bytesRead); + } + return bytesRead; + } + + @Override + public void close() throws IOException { + try { + super.close(); + } finally { + debugOutput.close(); + } + } +} diff --git a/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java new file mode 100644 index 000000000000..000580161427 --- /dev/null +++ b/libs/daemon/client/src/mill/client/debug/DebuggingOutputStream.java @@ -0,0 +1,49 @@ +package mill.client.debug; + +import java.io.*; +import java.nio.file.Path; +import java.time.LocalDateTime; + +/* Writes everything you write to the output stream to a file as well. */ +public class DebuggingOutputStream extends OutputStream { + private final OutputStream out; + private final OutputStream debugOutput; + private final boolean writeSeparateOps; + + public DebuggingOutputStream( + OutputStream out, Path workingDir, String name, boolean writeSeparateOps) { + this.out = out; + this.writeSeparateOps = writeSeparateOps; + try { + this.debugOutput = + new FileOutputStream(workingDir.resolve(name.replaceAll("\\W", "_")).toFile()); + } catch (FileNotFoundException e) { + throw new RuntimeException(e); + } + } + + @Override + public void write(int b) throws IOException { + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " write(): " + b + "\n" /*+ + Arrays.stream(new Exception().getStackTrace()).map(StackTraceElement::toString).collect(Collectors.joining("\n"))*/).getBytes()); + else debugOutput.write(b); + + out.write(b); + + if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " done\n").getBytes()); + } + + @Override + public void write(byte[] b, int off, int len) throws IOException { + if (writeSeparateOps) + debugOutput.write((LocalDateTime.now() + " write(off=" + off + ", len=" + len + "): " + + new String(b, off, len) + "\n") + .getBytes()); + else debugOutput.write(b, off, len); + + out.write(b, off, len); + + if (writeSeparateOps) debugOutput.write((LocalDateTime.now() + " done\n").getBytes()); + } +} diff --git a/libs/daemon/server/src/mill/server/ProxyStreamServer.scala b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala new file mode 100644 index 000000000000..0f5bd36a18c2 --- /dev/null +++ b/libs/daemon/server/src/mill/server/ProxyStreamServer.scala @@ -0,0 +1,169 @@ +package mill.server + +import mill.constants.ProxyStream + +import java.io.{BufferedInputStream, BufferedOutputStream, PrintStream} +import java.util.concurrent.atomic.AtomicBoolean +import scala.util.control.NonFatal + +/** [[Server]] that incorporates [[ProxyStream]] functionality. */ +abstract class ProxyStreamServer(args: Server.Args) extends Server(args) { self => + + /** Replaces [[PreHandleConnectionData]] in this class. */ + protected type PreHandleConnectionCustomData + + override protected type PreHandleConnectionData = ProxyStreamServerData + + case class ProxyStreamServerData( + stdout: PrintStream, + stderr: PrintStream, + writtenExitCode: AtomicBoolean, + customData: PreHandleConnectionCustomData + ) { + def writeExitCode(serverToClient: BufferedOutputStream, exitCode: Int): Unit = { + self.writeExitCode(serverToClient, exitCode, writtenExitCode) + } + } + + private def writeExitCode( + serverToClient: BufferedOutputStream, + exitCode: Int, + guard: AtomicBoolean + ): Unit = { + if (!guard.getAndSet(true)) { + ProxyStream.sendEnd(serverToClient, exitCode) + } + } + + override protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Boolean = { + ProxyStream.sendHeartbeat(connectionData.serverToClient) + true + } + + override protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit = { + // Notify the client that the server is shutting down with the given exit code + data match { + case Some(data) => data.writeExitCode(connectionData.serverToClient, exitCode) + case None => writeExitCode(connectionData.serverToClient, exitCode, new AtomicBoolean(false)) + } + } + + /** + * Invoked before a thread that runs [[handleConnection]] is spawned. + */ + override final protected def preHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer + ): PreHandleConnectionData = { + val stdout = + new PrintStream( + new ProxyStream.Output(connectionData.serverToClient, ProxyStream.OUT), + true + ) + val stderr = + new PrintStream( + new ProxyStream.Output(connectionData.serverToClient, ProxyStream.ERR), + true + ) + + val customData = preHandleConnection( + connectionData.socketInfo, + connectionData.clientToServer, + stdout, + stderr, + stopServer, + connectionData.initialSystemProperties + ) + + ProxyStreamServerData( + stdout = stdout, + stderr = stderr, + writtenExitCode = AtomicBoolean(false), + customData = customData + ) + } + + /** + * Invoked before a thread that runs [[handleConnection]] is spawned. + */ + protected def preHandleConnection( + socketInfo: Server.SocketInfo, + stdin: BufferedInputStream, + stdout: PrintStream, + stderr: PrintStream, + stopServer: Server.StopServer, + initialSystemProperties: Map[String, String] + ): PreHandleConnectionCustomData + + /** + * Handle a single client connection in a separate thread. + * + * @return the exit code to return to the client + */ + protected def handleConnection( + socketInfo: Server.SocketInfo, + stdin: BufferedInputStream, + stdout: PrintStream, + stderr: PrintStream, + stopServer: Server.StopServer, + setIdle: Server.SetIdle, + initialSystemProperties: Map[String, String], + data: PreHandleConnectionCustomData + ): Int + + override final protected def handleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + setIdle: Server.SetIdle, + data: PreHandleConnectionData + ): Unit = { + val exitCode = handleConnection( + connectionData.socketInfo, + connectionData.clientToServer, + data.stdout, + data.stderr, + stopServer, + setIdle, + connectionData.initialSystemProperties, + data.customData + ) + + serverLog(s"connection handler finished, sending exitCode $exitCode to client") + data.writeExitCode(connectionData.serverToClient, exitCode) + } + + override protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData, + exception: Throwable + ): Unit = { + data.writeExitCode(connectionData.serverToClient, 1) + } + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData + ): Unit = { + try + // Send a termination if it has not already happened + data.writeExitCode(connectionData.serverToClient, 1) + catch { + case NonFatal(err) => + serverLog( + s"error sending exit code 1, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" + ) + } + } +} diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index ffb008ae9136..63cc93fc5b4f 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -1,14 +1,13 @@ package mill.server import mill.client.lock.{Lock, Locks} -import mill.constants.{DaemonFiles, InputPumper, ProxyStream} +import mill.constants.{DaemonFiles, SocketUtil} import sun.misc.{Signal, SignalHandler} -import java.io.{InputStream, PipedInputStream, PipedOutputStream, PrintStream} -import java.net.{InetAddress, Socket, SocketAddress} +import java.io.{BufferedInputStream, BufferedOutputStream} +import java.net.{InetAddress, Socket, SocketAddress, SocketException} import java.time.LocalDateTime import java.time.format.DateTimeFormatter -import java.util.concurrent.atomic.AtomicBoolean import scala.concurrent.duration.FiniteDuration import scala.util.Try import scala.util.control.NonFatal @@ -16,16 +15,10 @@ import scala.util.control.NonFatal /** * Implementation of a server that binds to a random port, informs a client of the port, and accepts a client * connections. - * - * @param daemonDir directory used for exchanging pre-TCP data with a client - * @param acceptTimeout shuts down after this timeout if no clients are connected */ -abstract class Server( - daemonDir: os.Path, - acceptTimeout: Option[FiniteDuration], - locks: Locks, - testLogEvenWhenServerIdWrong: Boolean = false -) { +abstract class Server(args: Server.Args) { + import args.* + val processId: Long = Server.computeProcessId() private val acceptTimeoutMillis = acceptTimeout.map(_.toMillis) private val handlerName = getClass.getName @@ -43,37 +36,68 @@ abstract class Server( protected type PreHandleConnectionData + protected case class ConnectionData( + socketInfo: Server.SocketInfo, + clientToServer: BufferedInputStream, + serverToClient: BufferedOutputStream, + initialSystemProperties: Map[String, String] + ) + /** * Invoked before a thread that runs [[handleConnection]] is spawned. */ protected def preHandleConnection( - socketInfo: Server.SocketInfo, - stdin: InputStream, - stdout: PrintStream, - stderr: PrintStream, - stopServer: Server.StopServer, - initialSystemProperties: Map[String, String] + connectionData: ConnectionData, + stopServer: Server.StopServer ): PreHandleConnectionData /** * Handle a single client connection in a separate thread. - * - * @return the exit code to return to the client */ protected def handleConnection( - socketInfo: Server.SocketInfo, - stdin: InputStream, - stdout: PrintStream, - stderr: PrintStream, + connectionData: ConnectionData, stopServer: Server.StopServer, setIdle: Server.SetIdle, - initialSystemProperties: Map[String, String], data: PreHandleConnectionData - ): Int + ): Unit + + /** Invoked in [[handleConnection]] results in an exception. */ + protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData, + exception: Throwable + ): Unit + + protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Unit protected def connectionHandlerThreadName(socket: Socket): String = s"ConnectionHandler($handlerName, ${socket.getInetAddress}:${socket.getPort})" + /** Returns true if the client is still alive. Invoked from another thread. */ + protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: PreHandleConnectionData + ): Boolean + + /** + * Invoked when the server is stopped. + * + * @param data will be [[None]] if [[preHandleConnection]] haven't been invoked yet. + */ + protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit + def run(): Unit = { serverLog(s"running server in $daemonDir") serverLog(s"acceptTimeout=$acceptTimeout") @@ -183,10 +207,7 @@ abstract class Server( socketOpt match { case Some(sock) => - val socketInfo = Server.SocketInfo( - remote = sock.getRemoteSocketAddress, - local = sock.getLocalSocketAddress - ) + val socketInfo = Server.SocketInfo(sock) serverLog(s"handling run for $socketInfo") new Thread( () => @@ -246,88 +267,83 @@ abstract class Server( initialSystemProperties: Map[String, String], serverSocketClose: () => Unit ): Unit = { - - /** stdout and stderr combined into one stream */ - val currentOutErr = clientSocket.getOutputStream - val writtenExitCode = AtomicBoolean() - - def writeExitCode(code: Int): Unit = { - if (!writtenExitCode.getAndSet(true)) { - ProxyStream.sendEnd(currentOutErr, code) - } + val clientToServer = BufferedInputStream(clientSocket.getInputStream, bufferSize) + val serverToClient = BufferedOutputStream(clientSocket.getOutputStream, bufferSize) + + val connectionData = + ConnectionData(socketInfo, clientToServer, serverToClient, initialSystemProperties) + + def stopServer( + from: String, + reason: String, + exitCode: Int, + data: Option[PreHandleConnectionData] + ): Nothing = { + serverLog(s"$from invoked `stopServer` (reason: $reason), exitCode $exitCode") + onStopServer(from, reason, exitCode, connectionData, data) + systemExit0(reason, exitCode) } + val data = preHandleConnection( + connectionData, + stopServer = + (reason, exitCode) => stopServer("pre-connection handler", reason, exitCode, data = None) + ) + // We cannot use Socket#{isConnected, isClosed, isBound} because none of these // detect client-side connection closing, so instead we send a no-op heartbeat // message to see if the socket can receive data. + @volatile var lastClientAlive = true + val stopServerFromCheckClientAlive: Server.StopServer = (reason, exitCode) => + stopServer("checkClientAlive", reason, exitCode, Some(data)) def checkClientAlive() = { - try { - ProxyStream.sendHeartbeat(currentOutErr) - true - } catch { - case NonFatal(err) => - serverLog( - s"error sending heartbeat, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" - ) - false - } - } - - def stopServer(from: String, reason: String, exitCode: Int) = { - serverLog(s"$from invoked `stopServer` (reason: $reason), exitCode $exitCode") - writeExitCode(exitCode) - systemExit0(reason, exitCode) + val result = + try checkIfClientAlive(connectionData, stopServerFromCheckClientAlive, data) + catch { + case NonFatal(err) => + serverLog( + s"error checking for client liveness, assuming client to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" + ) + false + } + lastClientAlive = result + result } try { - val socketIn = clientSocket.getInputStream - - val stdout = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.OUT), true) - val stderr = new PrintStream(new ProxyStream.Output(currentOutErr, ProxyStream.ERR), true) - - val data = preHandleConnection( - socketInfo = socketInfo, - stdin = socketIn, - stdout = stdout, - stderr = stderr, - stopServer = - (reason, exitCode) => stopServer("pre-connection handler", reason, exitCode), - initialSystemProperties = initialSystemProperties - ) - @volatile var done = false @volatile var idle = false + @volatile var writingExceptionLog = false val t = new Thread( () => try { - // Proxy the input stream through a pair of Piped**putStream via a pumper, - // as the `UnixDomainSocketInputStream` we get directly from the socket does - // not properly implement `available(): Int` and thus messes up polling logic - // that relies on that method - // - // TODO: this seems to be a leftover from the times when unix sockets were used, we should try removing it - // in a separate PR. - val proxiedSocketInput = Server.proxyInputStreamThroughPumper(socketIn) - - val exitCode = handleConnection( - socketInfo = socketInfo, - stdin = proxiedSocketInput, - stdout = stdout, - stderr = stderr, + handleConnection( + connectionData, stopServer = - (reason, exitCode) => stopServer("connection handler", reason, exitCode), + (reason, exitCode) => + stopServer("connection handler", reason, exitCode, Some(data)), setIdle = idle = _, - initialSystemProperties = initialSystemProperties, data = data ) - - serverLog(s"connection handler finished, sending exitCode $exitCode to client") - writeExitCode(exitCode) } catch { + case e: SocketException if SocketUtil.clientHasClosedConnection(e) => // do nothing case e: Throwable => - serverLog("connection handler error: " + e) - serverLog("connection handler stack trace: " + e.getStackTrace.mkString("\n")) - writeExitCode(1) + writingExceptionLog = true + try { + serverLog( + s"""connection handler error: $e + |connection handler stack trace: ${e.getStackTrace.mkString("\n")} + |""".stripMargin + ) + } finally writingExceptionLog = false + onExceptionInHandleConnection( + connectionData, + stopServer = + (reason, exitCode) => + stopServer("onExceptionInHandleConnection", reason, exitCode, Some(data)), + data = data, + exception = e + ) } finally { done = true idle = true @@ -346,6 +362,11 @@ abstract class Server( serverSocketClose() } + serverLog(s"done=$done, idle=$idle, lastClientAlive=$lastClientAlive") + + // Wait until exception log writer finishes. + while (writingExceptionLog) Thread.sleep(1) + t.interrupt() // Try to give thread a moment to stop before we kill it for real // @@ -360,23 +381,31 @@ abstract class Server( case e: java.lang.Error if e.getMessage.contains("Cleaner terminated abnormally") => // ignore this error and do nothing; seems benign } - - // flush before closing the socket - System.out.flush() - System.err.flush() } finally { - try writeExitCode(1) // Send a termination if it has not already happened - catch { - case NonFatal(err) => - serverLog( - s"error sending exit code 1, client seems to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" - ) - } + beforeSocketClose( + connectionData, + stopServer = + (reason, exitCode) => stopServer("beforeSocketClose", reason, exitCode, Some(data)), + data + ) } } } object Server { + /** + * @param daemonDir directory used for exchanging pre-TCP data with a client + * @param acceptTimeout shuts down after this timeout if no clients are connected + * @param bufferSize size of the buffer used to read/write from/to the client + */ + case class Args( + daemonDir: os.Path, + acceptTimeout: Option[FiniteDuration], + locks: Locks, + bufferSize: Int, + testLogEvenWhenServerIdWrong: Boolean = false + ) + /** * @param remote the address of the client * @param local the address of the server @@ -384,8 +413,12 @@ object Server { case class SocketInfo(remote: SocketAddress, local: SocketAddress) { override def toString: String = s"SocketInfo(remote=$remote, local=$local)" } + object SocketInfo { + def apply(socket: Socket): SocketInfo = + apply(socket.getRemoteSocketAddress, socket.getLocalSocketAddress) + } - /** Immediately stops the server reporting the provided exit code to all clients. */ + /** Immediately stops the server with the given exit code. */ @FunctionalInterface trait StopServer { def apply(reason: String, exitCode: Int): Nothing } @@ -468,14 +501,4 @@ object Server { } } } - - private def proxyInputStreamThroughPumper(in: InputStream): PipedInputStream = { - val pipedInput = new PipedInputStream() - val pipedOutput = new PipedOutputStream(pipedInput) - val pumper = new InputPumper(() => in, () => pipedOutput, /* checkAvailable */ false) - val pumperThread = new Thread(pumper, "proxyInputStreamThroughPumper") - pumperThread.setDaemon(true) - pumperThread.start() - pipedInput - } } diff --git a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala index f0264bdfaf08..9dc40430d63f 100644 --- a/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala +++ b/libs/javalib/worker/src/mill/javalib/worker/JvmWorkerImpl.scala @@ -11,13 +11,13 @@ import mill.javalib.internal.{JvmWorkerArgs, RpcCompileProblemReporterMessage} import mill.javalib.zinc.ZincWorkerRpcServer.ReporterMode import mill.javalib.zinc.{ZincApi, ZincWorker, ZincWorkerRpcServer} import mill.rpc.{MillRpcChannel, MillRpcClient, MillRpcWireTransport} -import mill.util.{CachedFactoryWithInitData, HexFormat, Jvm} -import org.apache.logging.log4j.core.util.NullOutputStream +import mill.util.{CachedFactoryWithInitData, HexFormat, Jvm, Timed} import sbt.internal.util.ConsoleOut import java.io.* import java.nio.file.FileSystemException import java.security.MessageDigest +import java.time.LocalDateTime import scala.concurrent.duration.* import scala.util.Using @@ -25,6 +25,18 @@ import scala.util.Using class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoCloseable { import args.* + private def fileLog(s: String): Unit = + os.write.append( + compilerBridge.workspace / "jvm-worker.log", + s"[${LocalDateTime.now()}] $s\n", + createFolders = true + ) + + private def fileAndDebugLog(log: Logger.Actions, s: String): Unit = { + fileLog(s) + log.debug(s) + } + /** The local Zinc instance which is used when we do not want to override Java home or runtime options. */ private val zincLocalWorker = ZincWorker( @@ -41,8 +53,12 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo reporter: Option[CompileProblemReporter], reportCachedProblems: Boolean )(using ctx: JvmWorkerApi.Ctx): Result[CompilationResult] = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, javaRuntimeOptions) - zinc.compileJava(op, reporter = reporter, reportCachedProblems = reportCachedProblems) + val result = + Timed(zinc.compileJava(op, reporter = reporter, reportCachedProblems = reportCachedProblems)) + fileLog(s"Compilation took ${result.durationPretty}") + result.result } override def compileMixed( @@ -52,16 +68,23 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo reporter: Option[CompileProblemReporter], reportCachedProblems: Boolean )(using ctx: JvmWorkerApi.Ctx): Result[CompilationResult] = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, javaRuntimeOptions) - zinc.compileMixed(op, reporter = reporter, reportCachedProblems = reportCachedProblems) + val result = + Timed(zinc.compileMixed(op, reporter = reporter, reportCachedProblems = reportCachedProblems)) + fileLog(s"Compilation took ${result.durationPretty}") + result.result } def scaladocJar( op: ZincScaladocJar, javaHome: Option[os.Path] )(using ctx: JvmWorkerApi.Ctx): Boolean = { + fileLog(pprint.apply(op).render) val zinc = zincApi(javaHome, JavaRuntimeOptions(Seq.empty)) - zinc.scaladocJar(op) + val result = Timed(zinc.scaladocJar(op)) + fileLog(s"Scaladoc took ${result.durationPretty}") + result.result } override def close(): Unit = { @@ -88,8 +111,20 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo logPromptColored = log.prompt.colored ) - if (javaRuntimeOptions.options.isEmpty && javaHome.isEmpty) localZincApi(zincCtx, log) - else subprocessZincApi(javaHome, javaRuntimeOptions, zincCtx, log) + if (javaRuntimeOptions.options.isEmpty && javaHome.isEmpty) { + fileLog("Using local Zinc instance") + localZincApi(zincCtx, log) + } else { + fileLog( + s"""Using remote Zinc instance: + | javaHome: $javaHome + | javaRuntimeOptions: $javaRuntimeOptions + |""".stripMargin + ) + val result = Timed(subprocessZincApi(javaHome, javaRuntimeOptions, zincCtx, log)) + fileLog(s"Remote Zinc instance acquired in ${result.durationPretty}") + result.result + } } private case class SubprocessCacheKey( @@ -186,34 +221,47 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ) } - log.debug(s"Checking if $mainClass is already running for $key") - val result = ServerLauncher.ensureServerIsRunning( + fileAndDebugLog(log, s"Checking if $mainClass is already running for $key") + val result = Timed(ServerLauncher.ensureServerIsRunning( locks, daemonDir.toNIO, () => { - log.debug(s"Starting JVM subprocess for $mainClass for $key") - val process = Jvm.spawnProcess( + fileAndDebugLog(log, s"Starting JVM subprocess for $mainClass for $key") + val process = Timed(Jvm.spawnProcess( mainClass = mainClass, mainArgs = Seq(daemonDir.toString), javaHome = key.javaHome, jvmArgs = key.runtimeOptions.options, classPath = classPath + )) + fileAndDebugLog( + log, + s"Starting JVM subprocess for $mainClass for $key took ${process.durationPretty}" ) - LaunchedServer.OsProcess(process.wrapped.toHandle) + LaunchedServer.OsProcess(process.result.wrapped.toHandle) }, - log.debug + fileAndDebugLog(log, _) + )) + fileAndDebugLog( + log, + s"Ensuring that server is running for $key took ${result.durationPretty}" ) def onSuccess(launched: LaunchedServer) = { val serverInitWaitMillis = 5.seconds.toMillis val startTime = System.nanoTime() - log.debug(s"Reading server port: $daemonDir") - val port = ServerLauncher.readServerPort(daemonDir.toNIO, startTime, serverInitWaitMillis) - log.debug(s"Started $mainClass for $key on port $port") - SubprocessCacheValue(port, daemonDir, launched) + fileAndDebugLog(log, s"Reading server port: $daemonDir") + val port = + Timed(ServerLauncher.readServerPort(daemonDir.toNIO, startTime, serverInitWaitMillis)) + fileAndDebugLog( + log, + s"Reading server port for $daemonDir took ${port.durationPretty}." + ) + fileAndDebugLog(log, s"Started $mainClass for $key on port ${port.result}") + SubprocessCacheValue(port.result, daemonDir, launched) } - result.fold( + result.result.fold( success => onSuccess(success.server), alreadyRunning => onSuccess(alreadyRunning.server), processDied => @@ -260,6 +308,33 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ): ZincApi = { val cacheKey = SubprocessCacheKey(javaHome, runtimeOptions) + def makeClientLogger() = new Logger.Actions { + override def info(s: String): Unit = { + fileLog(s"[LOGGER:INFO] $s") + log.info(s) + } + + override def debug(s: String): Unit = { + fileLog(s"[LOGGER:DEBUG] $s") + log.debug(s) + } + + override def warn(s: String): Unit = { + fileLog(s"[LOGGER:WARN] $s") + log.warn(s) + } + + override def error(s: String): Unit = { + fileLog(s"[LOGGER:ERROR] $s") + log.error(s) + } + + override def ticker(s: String): Unit = { + fileLog(s"[LOGGER:TICKER] $s") + log.ticker(s) + } + } + def withRpcClient[R]( handler: MillRpcChannel[ZincWorkerRpcServer.ServerToClient] )(f: MillRpcClient[ZincWorkerRpcServer.ClientToServer, ZincWorkerRpcServer.ServerToClient] => R) @@ -270,33 +345,30 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo ) { case SubprocessCacheValue(port, daemonDir, _) => Using.Manager { use => val startTimeNanos = System.nanoTime() - log.debug(s"Connecting to $daemonDir on port $port") + fileAndDebugLog(log, s"Connecting to $daemonDir on port $port") val socket = use(ServerLauncher.connectToServer( startTimeNanos, 5.seconds.toMillis, port, s"From '${getClass.getName}'. Daemon directory: $daemonDir" )) - val stdin = use(PipedInputStream()) - val stdout = use(PipedOutputStream()) - val streams = ServerLauncher.Streams( - stdin, - stdout, - // stderr stream is not used in this case - NullOutputStream.getInstance() - ) + val debugName = + s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}" ServerLauncher.runWithConnection( + debugName, socket, - streams, /* closeConnectionAfterCommand */ true, + /* closeConnectionAfterCommand */ true, /* sendInitData */ _ => {}, - () => { - val serverToClient = use(BufferedReader(InputStreamReader(PipedInputStream(stdout)))) - val clientToServer = use(PrintStream(PipedOutputStream(stdin))) - val wireTransport = MillRpcWireTransport.ViaStreams( - s"ZincWorker,TCP ${socket.getRemoteSocketAddress} -> ${socket.getLocalSocketAddress}", - serverToClient, - clientToServer - ) + (in, out) => { + val serverToClient = use(BufferedReader(InputStreamReader(in))) + val clientToServer = use(PrintStream(out)) + val wireTransport = + MillRpcWireTransport.ViaStreams( + debugName, + serverToClient, + clientToServer, + writeSynchronizer = clientToServer + ) val init = ZincWorkerRpcServer.Initialize( compilerBridgeWorkspace = compilerBridge.workspace, @@ -304,15 +376,22 @@ class JvmWorkerImpl(args: JvmWorkerArgs[Unit]) extends JvmWorkerApi with AutoClo compileToJar = compileToJar, zincLogDebug = zincLogDebug ) + fileAndDebugLog( + log, + s"Connected to $daemonDir on port $port, sending init: ${pprint(init)}" + ) val client = MillRpcClient.create[ ZincWorkerRpcServer.Initialize, ZincWorkerRpcServer.ClientToServer, ZincWorkerRpcServer.ServerToClient - ](init, wireTransport, log)(handler) + ](init, wireTransport, makeClientLogger())(handler) - f(client) + fileAndDebugLog(log, "Running command.") + val result = Timed(f(client)) + fileAndDebugLog(log, s"Command finished in ${result.durationPretty}") + result.result } - ).result + ) }.get } } diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index 30e896c21954..860cf3c60cf8 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -7,7 +7,7 @@ import mill.rpc.MillRpcWireTransport import mill.server.Server import pprint.{TPrint, TPrintColors} -import java.io.{BufferedReader, InputStream, InputStreamReader, PrintStream} +import java.io.{BufferedReader, InputStreamReader, PrintStream} import scala.util.Using import scala.util.control.NonFatal @@ -30,38 +30,38 @@ object ZincWorkerMain { } } - private class ZincWorkerTcpServer(daemonDir: os.Path) extends Server( + private class ZincWorkerTcpServer(daemonDir: os.Path) extends Server(Server.Args( daemonDir, // The worker kills the process when it needs to. acceptTimeout = None, - Locks.files(daemonDir.toString) - ) { + Locks.files(daemonDir.toString), + bufferSize = 4 * 1024 + )) { private val className = summon[TPrint[ZincWorkerTcpServer]].render(using TPrintColors.Colors) - override protected type PreHandleConnectionData = Unit + protected class WriteSynchronizer + + override protected type PreHandleConnectionData = WriteSynchronizer override protected def preHandleConnection( - socketInfo: Server.SocketInfo, - stdin: InputStream, - stdout: PrintStream, - stderr: PrintStream, - stopServer: Server.StopServer, - initialSystemProperties: Map[String, String] - ): PreHandleConnectionData = () + connectionData: ConnectionData, + stopServer: Server.StopServer + ): WriteSynchronizer = new WriteSynchronizer override protected def handleConnection( - socketInfo: Server.SocketInfo, - stdin: InputStream, - stdout: PrintStream, - stderr: PrintStream, + connectionData: ConnectionData, stopServer: Server.StopServer, setIdle: Server.SetIdle, - initialSystemProperties: Map[String, String], - data: PreHandleConnectionData - ): Int = { + writeSynchronizer: WriteSynchronizer + ) = { + import connectionData.socketInfo + val serverName = s"$className{${socketInfo.remote} -> ${socketInfo.local}}" - Using.resource(BufferedReader(InputStreamReader(stdin))) { stdin => - val transport = MillRpcWireTransport.ViaStreams(serverName, stdin, stdout) + Using.Manager { use => + val stdin = use(BufferedReader(InputStreamReader(connectionData.clientToServer))) + val stdout = use(PrintStream(connectionData.serverToClient)) + val transport = + MillRpcWireTransport.ViaStreams(serverName, stdin, stdout, writeSynchronizer) try { val server = ZincWorkerRpcServer(serverName, transport, setIdle, serverLog) @@ -74,14 +74,45 @@ object ZincWorkerMain { serverLog("server.run() starting") server.run() serverLog("server.run() finished") - 0 } } catch { case NonFatal(err) => serverLog(s"$socketInfo failed: $err") - 1 } + }.get + } + + override protected def onExceptionInHandleConnection( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer, + exception: Throwable + ): Unit = {} + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer + ): Unit = {} + + override protected def checkIfClientAlive( + connectionData: ConnectionData, + stopServer: Server.StopServer, + writeSynchronizer: WriteSynchronizer + ): Boolean = { + writeSynchronizer.synchronized { + connectionData.serverToClient.write('\n'.toInt) + connectionData.serverToClient.flush() + true } } + + override protected def onStopServer( + from: String, + reason: String, + exitCode: Int, + connectionData: ConnectionData, + data: Option[PreHandleConnectionData] + ): Unit = {} } } diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala index 48cf8bf510ed..08014a2f0a4d 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala @@ -8,6 +8,7 @@ import mill.javalib.api.internal.{ZincCompileJava, ZincCompileMixed, ZincScalado import mill.javalib.internal.{RpcCompileProblemReporterMessage, ZincCompilerBridgeProvider} import mill.rpc.* import mill.server.Server +import mill.util.Timed import org.apache.logging.log4j.core.util.NullOutputStream import sbt.internal.util.ConsoleOut import upickle.default.ReadWriter @@ -33,100 +34,111 @@ class ZincWorkerRpcServer( clientStderr: RpcConsole, serverToClient: MillRpcChannel[ServerToClient] ): MillRpcChannel[ClientToServer] = { - val zincCompilerBridge = ZincCompilerBridgeProvider[MillRpcRequestId]( - workspace = initialize.compilerBridgeWorkspace, - logInfo = log.info, - acquire = (scalaVersion, scalaOrganization, clientRequestId) => - serverToClient( - clientRequestId, - ServerToClient.AcquireZincCompilerBridge( - scalaVersion = scalaVersion, - scalaOrganization = scalaOrganization + val result = Timed { + val zincCompilerBridge = ZincCompilerBridgeProvider[MillRpcRequestId]( + workspace = initialize.compilerBridgeWorkspace, + logInfo = log.info, + acquire = (scalaVersion, scalaOrganization, clientRequestId) => + serverToClient( + clientRequestId, + ServerToClient.AcquireZincCompilerBridge( + scalaVersion = scalaVersion, + scalaOrganization = scalaOrganization + ) ) - ) - ) - val worker = ZincWorker( - zincCompilerBridge, - jobs = initialize.jobs, - compileToJar = initialize.compileToJar, - zincLogDebug = initialize.zincLogDebug - ) - - val deps = { - // This is an ugly hack. `ConsoleOut` is sealed but we need to provide a way to send these logs to the Mill server - // over RPC, so we hijack `PrintStream` by overriding the methods that `ConsoleOut` uses. - // - // This is obviously extra fragile, but I couldn't find a better way to do it. - val consoleOut = ConsoleOut.printStreamOut(new PrintStream(NullOutputStream.getInstance()) { - override def print(s: String): Unit = clientStderr.print(s) - override def println(s: String): Unit = print(s + "\n") - override def println(): Unit = print("\n") - override def flush(): Unit = clientStderr.flush() - }) - - ZincWorker.InvocationDependencies(log, consoleOut) - } + ) + val worker = ZincWorker( + zincCompilerBridge, + jobs = initialize.jobs, + compileToJar = initialize.compileToJar, + zincLogDebug = initialize.zincLogDebug + ) + + val deps = { + // This is an ugly hack. `ConsoleOut` is sealed but we need to provide a way to send these logs to the Mill server + // over RPC, so we hijack `PrintStream` by overriding the methods that `ConsoleOut` uses. + // + // This is obviously extra fragile, but I couldn't find a better way to do it. + val consoleOut = ConsoleOut.printStreamOut(new PrintStream(NullOutputStream.getInstance()) { + override def print(s: String): Unit = clientStderr.print(s) + + override def println(s: String): Unit = print(s + "\n") + + override def println(): Unit = print("\n") + + override def flush(): Unit = clientStderr.flush() + }) + + ZincWorker.InvocationDependencies(log, consoleOut) + } - def reporter(clientRequestId: MillRpcRequestId, maxErrors: Int) = RpcCompileProblemReporter( - maxErrors = maxErrors, - send = msg => - serverToClient( - clientRequestId, - ServerToClient.ReportCompilationProblem(clientRequestId, msg) - ) - ) - - def reporterAsOption( - clientRequestId: MillRpcRequestId, - mode: ReporterMode - ): Option[CompileProblemReporter] = mode match { - case ReporterMode.NoReporter => None - case r: ReporterMode.Reporter => - Some(reporter(clientRequestId = clientRequestId, maxErrors = r.maxErrors)) - } + def reporter(clientRequestId: MillRpcRequestId, maxErrors: Int) = RpcCompileProblemReporter( + maxErrors = maxErrors, + send = msg => + serverToClient( + clientRequestId, + ServerToClient.ReportCompilationProblem(clientRequestId, msg) + ) + ) + + def reporterAsOption( + clientRequestId: MillRpcRequestId, + mode: ReporterMode + ): Option[CompileProblemReporter] = mode match { + case ReporterMode.NoReporter => None + case r: ReporterMode.Reporter => + Some(reporter(clientRequestId = clientRequestId, maxErrors = r.maxErrors)) + } - new MillRpcChannel[ClientToServer] { - override def apply(requestId: MillRpcRequestId, input: ClientToServer): input.Response = - setIdle.doWork { - input match { - case msg: ClientToServer.CompileJava => - compileJava(requestId, msg).asInstanceOf[input.Response] - case msg: ClientToServer.CompileMixed => - compileMixed(requestId, msg).asInstanceOf[input.Response] - case msg: ClientToServer.ScaladocJar => - docJar(requestId, msg).asInstanceOf[input.Response] + new MillRpcChannel[ClientToServer] { + override def apply(requestId: MillRpcRequestId, input: ClientToServer): input.Response = + setIdle.doWork { + val result = Timed { + input match { + case msg: ClientToServer.CompileJava => + compileJava(requestId, msg).asInstanceOf[input.Response] + case msg: ClientToServer.CompileMixed => + compileMixed(requestId, msg).asInstanceOf[input.Response] + case msg: ClientToServer.ScaladocJar => + docJar(requestId, msg).asInstanceOf[input.Response] + } + } + writeToLocalLog(s"$requestId with data $input processed in ${result.durationPretty}") + result.result } + + private def compileJava( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.CompileJava + ): msg.Response = { + worker.compileJava( + op = msg.op, + reporter = reporterAsOption(clientRequestId, msg.reporterMode), + reportCachedProblems = msg.reporterMode.reportCachedProblems + )(using msg.ctx, deps) } - private def compileJava( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.CompileJava - ): msg.Response = { - worker.compileJava( - op = msg.op, - reporter = reporterAsOption(clientRequestId, msg.reporterMode), - reportCachedProblems = msg.reporterMode.reportCachedProblems - )(using msg.ctx, deps) - } + private def compileMixed( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.CompileMixed + ): msg.Response = { + worker.compileMixed( + msg.op, + reporter = reporterAsOption(clientRequestId, msg.reporterMode), + reportCachedProblems = msg.reporterMode.reportCachedProblems, + compilerBridgeData = clientRequestId + )(using msg.ctx, deps) + } - private def compileMixed( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.CompileMixed - ): msg.Response = { - worker.compileMixed( - msg.op, - reporter = reporterAsOption(clientRequestId, msg.reporterMode), - reportCachedProblems = msg.reporterMode.reportCachedProblems, - compilerBridgeData = clientRequestId - )(using msg.ctx, deps) + private def docJar( + clientRequestId: MillRpcRequestId, + msg: ClientToServer.ScaladocJar + ): msg.Response = + worker.scaladocJar(msg.op, compilerBridgeData = clientRequestId) } - - private def docJar( - clientRequestId: MillRpcRequestId, - msg: ClientToServer.ScaladocJar - ): msg.Response = - worker.scaladocJar(msg.op, compilerBridgeData = clientRequestId) } + writeToLocalLog(s"Initialized in ${result.durationPretty}.") + result.result } } object ZincWorkerRpcServer { diff --git a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala index ea165da82772..fd9d788a4330 100644 --- a/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala +++ b/libs/rpc/src/mill/rpc/MillRpcWireTransport.scala @@ -5,6 +5,7 @@ import upickle.default.{Reader, Writer} import java.io.{BufferedReader, PrintStream} import java.util.concurrent.BlockingQueue +import scala.annotation.tailrec trait MillRpcWireTransport extends AutoCloseable { @@ -19,13 +20,20 @@ trait MillRpcWireTransport extends AutoCloseable { readAndTryToParse[A](typeName.render(using TPrintColors.Colors).render, log) /** Helper that reads a message from the wire and tries to parse it, logging along the way. */ - def readAndTryToParse[A: Reader](typeName: String, log: String => Unit): Option[A] = { + @tailrec final def readAndTryToParse[A: Reader]( + typeName: String, + log: String => Unit + ): Option[A] = { log(s"Trying to read $typeName") read() match { case None => log("Transport wire broken.") None + // Empty line means a heartbeat message + case Some("") => + readAndTryToParse[A](typeName, log) + case Some(line) => log(s"Received, will try to parse as $typeName: $line") val parsed = upickle.default.read(line) @@ -43,6 +51,7 @@ trait MillRpcWireTransport extends AutoCloseable { val serialized = upickle.default.write(message) log(s"Sending: $serialized") write(serialized) + log(s"Sent: $serialized") } } @@ -100,17 +109,21 @@ object MillRpcWireTransport { /** * @param serverToClient server to client stream * @param clientToServer client to server stream + * @param writeSynchronizer synchronizer for writes */ class ViaStreams( override val name: String, serverToClient: BufferedReader, - clientToServer: PrintStream + clientToServer: PrintStream, + writeSynchronizer: AnyRef ) extends MillRpcWireTransport { override def read(): Option[String] = Option(serverToClient.readLine()) override def write(message: String): Unit = { - clientToServer.println(message) - clientToServer.flush() + writeSynchronizer.synchronized { + clientToServer.println(message) + clientToServer.flush() + } } override def close(): Unit = { diff --git a/libs/util/src/mill/util/Timed.scala b/libs/util/src/mill/util/Timed.scala new file mode 100644 index 000000000000..e3dfab0699af --- /dev/null +++ b/libs/util/src/mill/util/Timed.scala @@ -0,0 +1,18 @@ +package mill.util + +import java.util.concurrent.TimeUnit +import scala.concurrent.duration.FiniteDuration + +private[mill] object Timed { + case class Result[+A](result: A, duration: FiniteDuration) { + def durationPretty: String = duration.toMillis + "ms" + } + + /** Computes how long it took to compute `f`. */ + def apply[A](f: => A): Result[A] = { + val start = System.nanoTime() + val result = f + val duration = FiniteDuration(System.nanoTime() - start, TimeUnit.NANOSECONDS) + Result(result, duration) + } +} diff --git a/runner/client/src/mill/client/MillServerLauncher.java b/runner/client/src/mill/client/MillServerLauncher.java index 7fb70282c34b..055e23b70ba1 100644 --- a/runner/client/src/mill/client/MillServerLauncher.java +++ b/runner/client/src/mill/client/MillServerLauncher.java @@ -80,6 +80,8 @@ public Result run(Path daemonDir, String javaHome, Consumer log) throws log)) { log.accept("runWithConnection: " + connection); var result = runWithConnection( + "MillServerLauncher[" + connection.getLocalSocketAddress() + " -> " + + connection.getRemoteSocketAddress() + "]", connection, streams, false, diff --git a/runner/server/src/mill/server/MillDaemonServer.scala b/runner/server/src/mill/server/MillDaemonServer.scala index 1462b2d0982b..fbd43f958d5d 100644 --- a/runner/server/src/mill/server/MillDaemonServer.scala +++ b/runner/server/src/mill/server/MillDaemonServer.scala @@ -23,12 +23,13 @@ abstract class MillDaemonServer[State]( acceptTimeout: FiniteDuration, locks: Locks, testLogEvenWhenServerIdWrong: Boolean = false -) extends Server( +) extends ProxyStreamServer(Server.Args( daemonDir = daemonDir, acceptTimeout = Some(acceptTimeout), locks = locks, - testLogEvenWhenServerIdWrong = testLogEvenWhenServerIdWrong - ) { + testLogEvenWhenServerIdWrong = testLogEvenWhenServerIdWrong, + bufferSize = 4 * 1024 + )) { def outLock: mill.client.lock.Lock def out: os.Path @@ -43,11 +44,11 @@ abstract class MillDaemonServer[State]( override protected def connectionHandlerThreadName(socket: Socket): String = s"MillServerActionRunner(${socket.getInetAddress}:${socket.getPort})" - protected override type PreHandleConnectionData = ClientInitData + override protected type PreHandleConnectionCustomData = ClientInitData override protected def preHandleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, @@ -104,14 +105,14 @@ abstract class MillDaemonServer[State]( override protected def handleConnection( socketInfo: Server.SocketInfo, - stdin: InputStream, + stdin: BufferedInputStream, stdout: PrintStream, stderr: PrintStream, stopServer: Server.StopServer, setIdle: Server.SetIdle, initialSystemProperties: Map[String, String], data: ClientInitData - ): Int = { + ) = { val (result, newStateCache) = main0( data.args, stateCache, @@ -140,6 +141,18 @@ abstract class MillDaemonServer[State]( initialSystemProperties: Map[String, String], stopServer: Server.StopServer ): (Boolean, State) + + override protected def beforeSocketClose( + connectionData: ConnectionData, + stopServer: Server.StopServer, + data: ProxyStreamServerData + ): Unit = { + // flush before closing the socket + System.out.flush() + System.err.flush() + + super.beforeSocketClose(connectionData, stopServer, data) + } } object MillDaemonServer { From 82cad20e582e818b50ee61629511b7e6f3a13aba Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Sun, 17 Aug 2025 09:34:54 +0000 Subject: [PATCH 17/25] [autofix.ci] apply automated fixes --- core/constants/src/mill/constants/SocketUtil.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/constants/src/mill/constants/SocketUtil.java b/core/constants/src/mill/constants/SocketUtil.java index 083b56ca066d..8bbf758c073a 100644 --- a/core/constants/src/mill/constants/SocketUtil.java +++ b/core/constants/src/mill/constants/SocketUtil.java @@ -6,8 +6,8 @@ public class SocketUtil { public static boolean clientHasClosedConnection(SocketException e) { var message = e.getMessage(); return message != null - && (message.contains("Broken pipe") - || message.contains("Socket closed") - || message.contains("Connection reset by peer")); + && (message.contains("Broken pipe") + || message.contains("Socket closed") + || message.contains("Connection reset by peer")); } } From 559a39f05c4b950b754a0362913a241a0c8b259e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 09:12:39 +0300 Subject: [PATCH 18/25] Updated tests. --- .../feature/leak-hygiene/src/LeakHygieneTests.scala | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala index 3ddf541f8e2d..92cfb157150c 100644 --- a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala +++ b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala @@ -80,8 +80,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "MillServerTimeoutThread", "Process ID Checker Thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -109,8 +108,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -139,8 +137,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } From a3e34e5a35e228c9a231aed6a57570524d473edf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 09:40:26 +0300 Subject: [PATCH 19/25] Fix. --- .../startup-shutdown/resources/build.mill | 1 + .../src/StartupShutdownTests.scala | 29 +++++++++++++ .../server/src/mill/server/Server.scala | 43 ++++++++++++++----- 3 files changed, 63 insertions(+), 10 deletions(-) create mode 100644 integration/feature/startup-shutdown/resources/build.mill create mode 100644 integration/feature/startup-shutdown/src/StartupShutdownTests.scala diff --git a/integration/feature/startup-shutdown/resources/build.mill b/integration/feature/startup-shutdown/resources/build.mill new file mode 100644 index 000000000000..8b137891791f --- /dev/null +++ b/integration/feature/startup-shutdown/resources/build.mill @@ -0,0 +1 @@ + diff --git a/integration/feature/startup-shutdown/src/StartupShutdownTests.scala b/integration/feature/startup-shutdown/src/StartupShutdownTests.scala new file mode 100644 index 000000000000..79c2d51b5ccd --- /dev/null +++ b/integration/feature/startup-shutdown/src/StartupShutdownTests.scala @@ -0,0 +1,29 @@ +package mill.integration + +import mill.testkit.UtestIntegrationTestSuite + +import utest._ + +object StartupShutdownTests extends UtestIntegrationTestSuite { + // Ensure that you can quickly cycle through startup and shutdown without failures. + val tests: Tests = Tests { + test("rapid_cycle") - integrationTest { tester => + if (daemonMode) { + val NoCycles = 7 + + for (index <- 1 to NoCycles) { + println(s"Cycle $index") + + val result1 = tester.eval(("resolve", "_")) + assert(result1.isSuccess) + + // No point in running shutdown on the last cycle, as the test suite will shut it down. + if (index != NoCycles) { + val result2 = tester.eval("shutdown") + assert(result2.isSuccess) + } + } + } + } + } +} diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index ffb008ae9136..5f005ed7787c 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -80,8 +80,24 @@ abstract class Server( val initialSystemProperties = sys.props.toMap val socketPortFile = daemonDir / DaemonFiles.socketPort + def removeSocketPortFile(): Unit = { + // Make sure no one would read the old socket port file + serverLog(s"removing $socketPortFile") + os.remove(socketPortFile) + serverLog(s"removed $socketPortFile") + } + try { - Server.tryLockBlock(locks.daemonLock) { locked => + Server.tryLockBlock( + locks.daemonLock, + beforeClose = () => { + removeSocketPortFile() + serverLog("releasing daemonLock") + }, + afterClose = () => { + serverLog("daemonLock released") + } + ) { locked => serverLog("server file locked") val serverSocket = new java.net.ServerSocket(0, 0, InetAddress.getByName(null)) Server.watchProcessIdFile( @@ -157,8 +173,7 @@ abstract class Server( // 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() - serverLog("daemonLock released") + locked.close() sys.exit(exitCode) } @@ -225,10 +240,6 @@ abstract class Server( serverLog("server loop stack trace: " + e.getStackTrace.mkString("\n")) throw e } finally { - // Make sure no one would read the old socket port file - serverLog(s"removing $socketPortFile") - os.remove(socketPortFile) - serverLog("exiting server") } } @@ -456,13 +467,25 @@ object Server { processIdThread.start() } - def tryLockBlock[T](lock: Lock)(block: mill.client.lock.TryLocked => T): Option[T] = { + def tryLockBlock[T]( + lock: Lock, + beforeClose: () => Unit, + afterClose: () => Unit + )(block: AutoCloseable => T): Option[T] = { lock.tryLock() match { case null => None case l => if (l.isLocked) { - try Some(block(l)) - finally l.release() + val autoCloseable = new AutoCloseable { + override def close(): Unit = { + beforeClose() + l.release() + afterClose() + } + } + + try Some(block(autoCloseable)) + finally autoCloseable.close() } else { None } From bfa08c3a1dd505376bffe85393605d41c33b33b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 10:00:51 +0300 Subject: [PATCH 20/25] Update tests. --- .../leak-hygiene/src/LeakHygieneTests.scala | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala index 92cfb157150c..9434866a260a 100644 --- a/integration/feature/leak-hygiene/src/LeakHygieneTests.scala +++ b/integration/feature/leak-hygiene/src/LeakHygieneTests.scala @@ -162,8 +162,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "MillServerTimeoutThread", "Process ID Checker Thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -191,8 +190,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } @@ -222,8 +220,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -255,8 +252,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Process ID Checker Thread", "Timer", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) @@ -286,8 +282,7 @@ object LeakHygieneTests extends UtestIntegrationTestSuite { "Timer", "leaked thread", "main", - "prompt-logger-stream-pumper-thread", - "proxyInputStreamThroughPumper" + "prompt-logger-stream-pumper-thread" ) ) } From 14a2986c7776f9cfd8d6e38dbd36b39f1ea5cc49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 12:00:30 +0300 Subject: [PATCH 21/25] Improvements. --- .../server/src/mill/server/Server.scala | 27 +++++++++++++----- .../mill/javalib/zinc/ZincWorkerMain.scala | 28 ++++++++----------- .../javalib/zinc/ZincWorkerRpcServer.scala | 2 +- 3 files changed, 32 insertions(+), 25 deletions(-) diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index e7454e675cbb..cfa04a6a8f8f 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -6,6 +6,7 @@ import sun.misc.{Signal, SignalHandler} import java.io.{BufferedInputStream, BufferedOutputStream} import java.net.{InetAddress, Socket, SocketAddress, SocketException} +import java.nio.channels.ClosedByInterruptException import java.time.LocalDateTime import java.time.format.DateTimeFormatter import scala.concurrent.duration.FiniteDuration @@ -28,7 +29,10 @@ abstract class Server(args: Server.Args) { def serverLog0(s: String): Unit = { if (os.exists(daemonDir) || testLogEvenWhenServerIdWrong) { - os.write.append(daemonDir / DaemonFiles.serverLog, s"$s\n", createFolders = true) + try os.write.append(daemonDir / DaemonFiles.serverLog, s"$s\n", createFolders = true) + catch { + case _: ClosedByInterruptException => // write was interrupted with a thread interrupt + } } } @@ -311,6 +315,9 @@ abstract class Server(args: Server.Args) { val result = try checkIfClientAlive(connectionData, stopServerFromCheckClientAlive, data) catch { + case e: SocketException if SocketUtil.clientHasClosedConnection(e) => + serverLog(s"client has closed connection") + false case NonFatal(err) => serverLog( s"error checking for client liveness, assuming client to be dead: $err\n\n${err.getStackTrace.mkString("\n")}" @@ -339,13 +346,19 @@ abstract class Server(args: Server.Args) { } catch { case e: SocketException if SocketUtil.clientHasClosedConnection(e) => // do nothing case e: Throwable => + val msg = + s"""connection handler for $socketInfo error: $e + |connection handler stack trace: ${e.getStackTrace.mkString("\n")} + |""".stripMargin writingExceptionLog = true - try { - serverLog( - s"""connection handler error: $e - |connection handler stack trace: ${e.getStackTrace.mkString("\n")} - |""".stripMargin - ) + try serverLog(msg) + catch { + // If we get interrupted while writing the exception log, there is not much we can do, so just print + // it to stderr + case _: ClosedByInterruptException => + Console.err.println( + s"Interrupted while writing to server log, writing to stderr instead:\n${msg.indent(4)}" + ) } finally writingExceptionLog = false onExceptionInHandleConnection( connectionData, diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala index 860cf3c60cf8..0663d9ba338d 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerMain.scala @@ -1,7 +1,7 @@ package mill.javalib.zinc -import mill.api.daemon.{DummyInputStream, SystemStreams} import mill.api.SystemStreamsUtils +import mill.api.daemon.{DummyInputStream, SystemStreams} import mill.client.lock.Locks import mill.rpc.MillRpcWireTransport import mill.server.Server @@ -9,7 +9,6 @@ import pprint.{TPrint, TPrintColors} import java.io.{BufferedReader, InputStreamReader, PrintStream} import scala.util.Using -import scala.util.control.NonFatal /** Entry point for the Zinc worker subprocess. */ object ZincWorkerMain { @@ -62,22 +61,17 @@ object ZincWorkerMain { val stdout = use(PrintStream(connectionData.serverToClient)) val transport = MillRpcWireTransport.ViaStreams(serverName, stdin, stdout, writeSynchronizer) - try { - val server = ZincWorkerRpcServer(serverName, transport, setIdle, serverLog) + val server = ZincWorkerRpcServer(serverName, transport, setIdle, serverLog) - // Make sure stdout and stderr is sent to the client - SystemStreamsUtils.withStreams(SystemStreams( - out = PrintStream(server.clientStdout.asStream), - err = PrintStream(server.clientStderr.asStream), - in = DummyInputStream - )) { - serverLog("server.run() starting") - server.run() - serverLog("server.run() finished") - } - } catch { - case NonFatal(err) => - serverLog(s"$socketInfo failed: $err") + // Make sure stdout and stderr is sent to the client + SystemStreamsUtils.withStreams(SystemStreams( + out = PrintStream(server.clientStdout.asStream), + err = PrintStream(server.clientStderr.asStream), + in = DummyInputStream + )) { + serverLog("server.run() starting") + server.run() + serverLog("server.run() finished") } }.get } diff --git a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala index 08014a2f0a4d..76e5a7bfa000 100644 --- a/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala +++ b/libs/javalib/worker/src/mill/javalib/zinc/ZincWorkerRpcServer.scala @@ -33,7 +33,7 @@ class ZincWorkerRpcServer( clientStdout: RpcConsole, clientStderr: RpcConsole, serverToClient: MillRpcChannel[ServerToClient] - ): MillRpcChannel[ClientToServer] = { + ): MillRpcChannel[ClientToServer] = setIdle.doWork { val result = Timed { val zincCompilerBridge = ZincCompilerBridgeProvider[MillRpcRequestId]( workspace = initialize.compilerBridgeWorkspace, From b3bb0683e3399e1f30e056b670bec44e5e19a73b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 12:05:18 +0300 Subject: [PATCH 22/25] Comments. --- libs/daemon/client/src/mill/client/ServerLauncher.java | 3 +++ libs/daemon/server/src/mill/server/Server.scala | 3 +++ 2 files changed, 6 insertions(+) diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 1499862676e0..387a64185167 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -64,6 +64,9 @@ public static RunWithConnectionResult runWithConnection( Consumer sendInitData, RunClientLogic runClientLogic) throws Exception { + // According to https://pzemtsov.github.io/2015/01/19/on-the-benefits-of-stream-buffering-in-Java.html it seems that + // buffering on the application level is still beneficial due to syscall overhead, even if kernel has its own + // socket buffers. var socketInputStream = new BufferedInputStream(connection.getInputStream()); var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); sendInitData.accept(socketOutputStream); diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index cfa04a6a8f8f..5193767fe36e 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -282,6 +282,9 @@ abstract class Server(args: Server.Args) { initialSystemProperties: Map[String, String], serverSocketClose: () => Unit ): Unit = { + // According to https://pzemtsov.github.io/2015/01/19/on-the-benefits-of-stream-buffering-in-Java.html it seems that + // buffering on the application level is still beneficial due to syscall overhead, even if kernel has its own + // socket buffers. val clientToServer = BufferedInputStream(clientSocket.getInputStream, bufferSize) val serverToClient = BufferedOutputStream(clientSocket.getOutputStream, bufferSize) From 65086d2174e467f9c6b470971e8530a72e28f45b Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Mon, 18 Aug 2025 09:13:35 +0000 Subject: [PATCH 23/25] [autofix.ci] apply automated fixes --- libs/daemon/client/src/mill/client/ServerLauncher.java | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/libs/daemon/client/src/mill/client/ServerLauncher.java b/libs/daemon/client/src/mill/client/ServerLauncher.java index 387a64185167..6d16a02c548a 100644 --- a/libs/daemon/client/src/mill/client/ServerLauncher.java +++ b/libs/daemon/client/src/mill/client/ServerLauncher.java @@ -64,8 +64,11 @@ public static RunWithConnectionResult runWithConnection( Consumer sendInitData, RunClientLogic runClientLogic) throws Exception { - // According to https://pzemtsov.github.io/2015/01/19/on-the-benefits-of-stream-buffering-in-Java.html it seems that - // buffering on the application level is still beneficial due to syscall overhead, even if kernel has its own + // According to + // https://pzemtsov.github.io/2015/01/19/on-the-benefits-of-stream-buffering-in-Java.html it + // seems that + // buffering on the application level is still beneficial due to syscall overhead, even if + // kernel has its own // socket buffers. var socketInputStream = new BufferedInputStream(connection.getInputStream()); var socketOutputStream = new BufferedOutputStream(connection.getOutputStream()); From 98a598a1dd8f75dc5869c6350133bdeff9bfe2b6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 12:14:28 +0300 Subject: [PATCH 24/25] Fix compilation on JDK 11. --- libs/daemon/server/src/mill/server/Server.scala | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/daemon/server/src/mill/server/Server.scala b/libs/daemon/server/src/mill/server/Server.scala index 5193767fe36e..6fe7f0483edf 100644 --- a/libs/daemon/server/src/mill/server/Server.scala +++ b/libs/daemon/server/src/mill/server/Server.scala @@ -360,7 +360,9 @@ abstract class Server(args: Server.Args) { // it to stderr case _: ClosedByInterruptException => Console.err.println( - s"Interrupted while writing to server log, writing to stderr instead:\n${msg.indent(4)}" + s"Interrupted while writing to server log, writing to stderr instead:\n${ + msg.linesIterator.map(s => s" $s").mkString("\n") + }" ) } finally writingExceptionLog = false onExceptionInHandleConnection( From fbf3c3e3d614a5b673bf39843a719340afd0e6e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Art=C5=ABras=20=C5=A0lajus?= Date: Mon, 18 Aug 2025 13:55:24 +0300 Subject: [PATCH 25/25] PR updates. --- core/constants/test/src/mill/client/ProxyStreamTests.java | 2 +- libs/daemon/client/src/mill/client/ClientUtil.java | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/core/constants/test/src/mill/client/ProxyStreamTests.java b/core/constants/test/src/mill/client/ProxyStreamTests.java index 2f3a09ae858c..cdaa732d0491 100644 --- a/core/constants/test/src/mill/client/ProxyStreamTests.java +++ b/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, (byte) 0); + if (gracefulEnd) ProxyStream.sendEnd(pipedOutputStream, 0); else { pipedOutputStream.close(); } diff --git a/libs/daemon/client/src/mill/client/ClientUtil.java b/libs/daemon/client/src/mill/client/ClientUtil.java index e66768bbb552..98e733cd1443 100644 --- a/libs/daemon/client/src/mill/client/ClientUtil.java +++ b/libs/daemon/client/src/mill/client/ClientUtil.java @@ -8,15 +8,15 @@ public class ClientUtil { // use methods instead of constants to avoid inlining by compiler - public static byte ExitClientCodeCannotReadFromExitCodeFile() { + public static int ExitClientCodeCannotReadFromExitCodeFile() { return 1; } - public static byte ExitServerCodeWhenIdle() { + public static int ExitServerCodeWhenIdle() { return 0; } - public static byte ExitServerCodeWhenVersionMismatch() { + public static int ExitServerCodeWhenVersionMismatch() { return 101; }