From 905ee1d01c69b11a81fb986cb589e2ce7ef1e317 Mon Sep 17 00:00:00 2001
From: Lorenzo Gabriele <lorenzolespaul@gmail.com>
Date: Thu, 16 Nov 2023 12:29:44 +0100
Subject: [PATCH] Recover from client connection errors

Attempt to fix https://github.com/com-lihaoyi/mill/issues/2805
We now catch the exception, delete the content of the mill-worker-xyx
directory and try again.
Seems to be working correcly from manual tests
---
 .../src/mill/main/client/MillClientMain.java  | 99 +++++++++++--------
 1 file changed, 56 insertions(+), 43 deletions(-)

diff --git a/main/client/src/mill/main/client/MillClientMain.java b/main/client/src/mill/main/client/MillClientMain.java
index 16cb257329b..798ce58fcdc 100644
--- a/main/client/src/mill/main/client/MillClientMain.java
+++ b/main/client/src/mill/main/client/MillClientMain.java
@@ -101,45 +101,58 @@ public static int main0(String[] args) throws Exception {
 
         int index = 0;
         while (index < serverProcessesLimit) {
-            index += 1;
+            index++;
             final String lockBase = "out/mill-worker-" + versionAndJvmHomeEncoding + "-" + index;
-            new java.io.File(lockBase).mkdirs();
-
-            final File stdout = new java.io.File(lockBase + "/stdout");
-            final File stderr = new java.io.File(lockBase + "/stderr");
-            final int refeshIntervalMillis = 2;
-
-            try (
-                Locks locks = Locks.files(lockBase);
-                final FileToStreamTailer stdoutTailer = new FileToStreamTailer(stdout, System.out, refeshIntervalMillis);
-                final FileToStreamTailer stderrTailer = new FileToStreamTailer(stderr, System.err, refeshIntervalMillis);
-            ) {
-                Locked clientLock = locks.clientLock.tryLock();
-                if (clientLock != null) {
-                    stdoutTailer.start();
-                    stderrTailer.start();
-                    final int exitCode = run(
-                        lockBase,
-                        () -> {
-                            try {
-                                initServer(lockBase, setJnaNoSys);
-                            } catch (Exception e) {
-                                throw new RuntimeException(e);
-                            }
-                        },
-                        locks,
-                        System.in,
-                        System.out,
-                        System.err,
-                        args,
-                        System.getenv()
-                    );
-
-                    // Here, we ensure we process the tails of the output files before interrupting the threads
-                    stdoutTailer.flush();
-                    stderrTailer.flush();
-                    clientLock.release();
-                    return exitCode;
+            java.io.File lockBaseFile = new java.io.File(lockBase);
+            final File stdout = new java.io.File(lockBaseFile, "stdout");
+            final File stderr = new java.io.File(lockBaseFile, "stderr");
+
+            int attempts = 0;
+            while (attempts < 3) {
+                try {
+                    lockBaseFile.mkdirs();
+
+                    final int refeshIntervalMillis = 2;
+
+                    try (
+                        Locks locks = Locks.files(lockBase);
+                        final FileToStreamTailer stdoutTailer =new FileToStreamTailer(stdout, System.out, refeshIntervalMillis);
+                        final FileToStreamTailer stderrTailer = new FileToStreamTailer(stderr, System.err, refeshIntervalMillis);
+                    ) {
+                        Locked clientLock = locks.clientLock.tryLock();
+                        if (clientLock != null) {
+                            stdoutTailer.start();
+                            stderrTailer.start();
+                            final int exitCode = run(
+                                    lockBase,
+                                    () -> {
+                                        try {
+                                            initServer(lockBase, setJnaNoSys);
+                                        } catch (Exception e) {
+                                            throw new RuntimeException(e);
+                                        }
+                                    },
+                                    locks,
+                                    System.in,
+                                    System.out,
+                                    System.err,
+                                    args,
+                                    System.getenv());
+
+                            // Here, we ensure we process the tails of the output files before interrupting
+                            // the threads
+                            stdoutTailer.flush();
+                            stderrTailer.flush();
+                            clientLock.release();
+                            return exitCode;
+                        }
+                    }
+                } catch (Exception e) {
+                    for (File file : lockBaseFile.listFiles()) {
+                        file.delete();
+                    }
+                } finally {
+                    attempts++;
                 }
             }
         }
@@ -176,14 +189,14 @@ public static int run(
         }
         while (locks.processLock.probe()) Thread.sleep(3);
 
+        String socketName = lockBase + "/mill-" + Util.md5hex(new File(lockBase).getCanonicalPath()) + "-io";
+        AFUNIXSocketAddress addr = AFUNIXSocketAddress.of(new File(socketName));
+
+        long retryStart = System.currentTimeMillis();
         Socket ioSocket = null;
         Throwable socketThrowable = null;
-        long retryStart = System.currentTimeMillis();
-
-        while (ioSocket == null && System.currentTimeMillis() - retryStart < 5000) {
+        while (ioSocket == null && System.currentTimeMillis() - retryStart < 1000) {
             try {
-                String socketName = lockBase + "/mill-" + Util.md5hex(new File(lockBase).getCanonicalPath()) + "-io";
-                AFUNIXSocketAddress addr = AFUNIXSocketAddress.of(new File(socketName));
                 ioSocket = AFUNIXSocket.connectTo(addr);
             } catch (Throwable e) {
                 socketThrowable = e;