[SYSTEMDS-3960] Fix federated worker readiness probe on delayed connections

src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java

@@ -27,25 +27,24 @@
 /**
  * Test helpers that block until a federated worker is accepting TCP connections on its port.
  *
- * <p>The federated worker opens its TCP port after Netty's {@code bind().sync()} returns; a successful
- * TCP connect to that port therefore indicates that the worker is ready to accept requests. The methods
- * here poll for that signal and throw {@link RuntimeException} on timeout or if the underlying
+ * The federated worker opens its TCP port after Netty's {@code bind().sync()} returns. The methods here poll for a
+ * successful TCP connection and throw {@link RuntimeException} on timeout or if the underlying
  * {@code Process}/{@code Thread} exits before the port becomes ready.
  */
 public final class FederatedWorkerUtils {
 
 	/** Sleep between successive poll rounds, in milliseconds. */
 	private static final int POLL_INTERVAL_MS = 25;
 
-	/** Per-attempt {@link Socket#connect} timeout, in milliseconds. */
-	private static final int CONNECT_TIMEOUT_MS = 25;
+	/**
+	 * Per-attempt {@link Socket#connect} timeout, in milliseconds. Covers a full TCP handshake, so it must not be close
+	 * to the round trip time of the link.
+	 */
+	private static final int CONNECT_TIMEOUT_MS = 2000;
 
 	/**
-	 * Minimum value applied to the caller-supplied {@code timeoutMs}. The wait returns as soon as the
-	 * worker accepts a connection, so this only affects the upper bound used when a worker never becomes
-	 * ready. Set to 60s to accommodate cold JVM startup on heavily contended CI runners: tests starting
-	 * four workers in parallel can have all four still pending after 30s when the runner is CPU-starved,
-	 * and burning a surefire retry costs more wall time than padding this clamp.
+	 * Minimum value applied to the caller-supplied {@code timeoutMs}, returns as soon as the worker accepts a
+	 * connection.
 	 */
 	private static final int MIN_TIMEOUT_MS = 60_000;
 
@@ -76,7 +75,7 @@ public static void waitForWorker(int port, int timeoutMs, BooleanSupplier aliveC
 				throw new RuntimeException(
 					"Federated " + workerKind + " on port " + port + " died before becoming ready.");
 			}
-			if(tryConnect(port)) {
+			if(tryConnect(port, deadline)) {
 				return;
 			}
 			sleepQuietly();
@@ -96,9 +95,8 @@ public static void waitForWorker(Thread thread, int port, int timeoutMs) {
 	}
 
 	/**
-	 * Block until every listed federated worker is accepting TCP connections. All ports are polled in
-	 * one shared loop, so the wall-clock wait is bounded by the slowest worker rather than the sum of
-	 * individual waits.
+	 * Block until every listed federated worker is accepting TCP connections. All ports are polled in one shared loop,
+	 * so the wall-clock wait is bounded by the slowest worker.
 	 *
 	 * @param ports     ports the workers are expected to bind
 	 * @param timeoutMs upper bound on the wait, in ms; raised to {@link #MIN_TIMEOUT_MS} if smaller
@@ -134,9 +132,8 @@ public static void waitForWorkers(Thread[] threads, int[] ports, int timeoutMs)
 	}
 
 	/**
-	 * Bulk variant taking a per-index liveness predicate so callers can plug in either {@code Process}
-	 * or {@code Thread} liveness. Each port flips to ready as soon as it accepts a connection; the loop
-	 * yields between sweeps so a still-pending worker is not starved by repeated probes on the same CPU.
+	 * Bulk variant taking a per-index liveness predicate so callers can plug in either {@code Process} or
+	 * {@code Thread} liveness. Each port flips to ready as soon as it accepts a connection.
 	 */
 	public static void waitForWorkers(int[] ports, int timeoutMs, java.util.function.IntPredicate aliveCheck,
 		String workerKind) {
@@ -145,15 +142,16 @@ public static void waitForWorkers(int[] ports, int timeoutMs, java.util.function
 		final boolean[] ready = new boolean[ports.length];
 		int remaining = ports.length;
 		while(remaining > 0 && System.currentTimeMillis() < deadline) {
-			for(int i = 0; i < ports.length; i++) {
+			// recheck the deadline per port, a sweep can spend up to CONNECT_TIMEOUT_MS on each of them
+			for(int i = 0; i < ports.length && System.currentTimeMillis() < deadline; i++) {
 				if(ready[i]) {
 					continue;
 				}
 				if(!aliveCheck.test(i)) {
 					throw new RuntimeException("Federated " + workerKind + " on port " + ports[i]
 						+ " died before becoming ready.");
 				}
-				if(tryConnect(ports[i])) {
+				if(tryConnect(ports[i], deadline)) {
 					ready[i] = true;
 					remaining--;
 				}
@@ -174,12 +172,15 @@ public static void waitForWorkers(int[] ports, int timeoutMs, java.util.function
 		}
 	}
 
-	private static boolean tryConnect(int port) {
+	private static boolean tryConnect(int port, long deadline) {
+		final long remaining = deadline - System.currentTimeMillis();
+		if(remaining <= 0) // out of time => connect reads a timeout of 0 as "infinite"
+			return false;
 		try(Socket s = new Socket()) {
-			s.connect(new InetSocketAddress("localhost", port), CONNECT_TIMEOUT_MS);
+			s.connect(new InetSocketAddress("localhost", port), (int) Math.min(CONNECT_TIMEOUT_MS, remaining));
 			return true;
 		}
-		catch(IOException e) {
+		catch(IOException e) { // closed port, or a handshake that outlasted the budget
 			return false;
 		}
 	}

src/test/java/org/apache/sysds/test/component/federated/FederatedUrlParserTest.java

@@ -21,9 +21,14 @@
 
 import org.apache.sysds.runtime.instructions.fed.InitFEDInstruction;
 import org.apache.sysds.conf.DMLConfig;
+import org.apache.sysds.test.FederatedWorkerUtils;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.io.IOException;
+import java.net.ServerSocket;
 
 import org.junit.Test;
 
@@ -167,4 +172,30 @@ public void checkDefaultPortIsValid() {
 		assertTrue(defaultPort <= IANA_limit);
 		assertTrue(defaultPort > 0);
 	}
+
+	@Test
+	public void waitReturnsForAListeningPort() throws IOException {
+		try(ServerSocket listening = new ServerSocket(0)) {
+			// Return as soon as the port accepts, the timeout is only the upper bound.
+			FederatedWorkerUtils.waitForWorker(listening.getLocalPort(), 1000);
+		}
+	}
+
+	@Test
+	public void waitFailsFastWhenTheWorkerDied() throws IOException {
+		final int port;
+		try(ServerSocket closed = new ServerSocket(0)) {
+			port = closed.getLocalPort();
+		}
+		final long t0 = System.currentTimeMillis();
+		try {
+			FederatedWorkerUtils.waitForWorker(port, 1000, () -> false, "worker");
+			fail("expected the wait to report the dead worker");
+		}
+		catch(RuntimeException e) {
+			assertTrue(e.getMessage(), e.getMessage().contains("died before becoming ready"));
+			// Must not sit out the timeout, which is clamped up to a minute.
+			assertTrue("the dead worker was not reported promptly", System.currentTimeMillis() - t0 < 10000);
+		}
+	}
 }