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

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

@@ -25,20 +25,24 @@
 import java.util.function.BooleanSupplier;
 
 /**
- * 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
- * {@code Process}/{@code Thread} exits before the port becomes ready.
+ * Test helpers that block until a federated worker is accepting TCP connections on its port. 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 {@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. This budget has to cover a full TCP handshake, i.e.,
+	 * two traversals of the network, so it must not be set close to the round trip time of the link. Sizing this
+	 * generously is free while the worker is still starting up, as the kernel refuses a closed port immediately
+	 * (ECONNREFUSED), so the budget only applies once a handshake is actually in flight. 2s also covers one lost SYN,
+	 * which Linux retransmits after ~1s.
+	 */
+	private static final int CONNECT_TIMEOUT_MS = 2000;
 
 	/**
 	 * Minimum value applied to the caller-supplied {@code timeoutMs}. The wait returns as soon as the
@@ -76,7 +80,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();
@@ -145,15 +149,17 @@ 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++) {
+			// the deadline is rechecked per port, since a sweep over many ports can now 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,16 +180,33 @@ 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 int timeout = attemptTimeout(deadline - System.currentTimeMillis());
+		if(timeout == 0) // out of time, do not start another attempt
+			return false;
 		try(Socket s = new Socket()) {
-			s.connect(new InetSocketAddress("localhost", port), CONNECT_TIMEOUT_MS);
+			s.connect(new InetSocketAddress("localhost", port), timeout);
 			return true;
 		}
-		catch(IOException e) {
+		catch(IOException e) { // closed port, or a handshake that outlasted the budget
 			return false;
 		}
 	}
 
+	/**
+	 * Budget for a single connect attempt, capped by the time left until the overall deadline so that one slow attempt
+	 * cannot substantially exceed the limit.
+	 *
+	 * @param remainingMs time left until the deadline, in ms
+	 * @return the timeout to pass to {@link Socket#connect}, or 0 if no attempt should be made. Never returns 0 while
+	 *         time is left, because {@code connect} reads a timeout of 0 as 'infinite'.
+	 */
+	public static int attemptTimeout(long remainingMs) {
+		if(remainingMs <= 0)
+			return 0;
+		return (int) Math.min(CONNECT_TIMEOUT_MS, remainingMs);
+	}
+
 	private static void sleepQuietly() {
 		try {
 			Thread.sleep(POLL_INTERVAL_MS);

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

@@ -0,0 +1,93 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.sysds.test.component.federated;
+
+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.apache.sysds.test.FederatedWorkerUtils;
+import org.junit.Test;
+
+/**
+ * Tests for the readiness probe that blocks until a federated worker accepts connections.
+ */
+public class FederatedWorkerUtilsTest {
+
+	/** Round-trip time the probe should tolerate without giving up, in ms. */
+	private static final int TOLERATED_HANDSHAKE_MS = 1000;
+
+	@Test
+	public void attemptBudgetCoversADelayedHandshake() {
+		// with the sufficient time budget, a single attempt must be allowed to outlast a slow handshake
+		assertTrue("per-attempt connect budget is too small to complete a delayed TCP handshake",
+			FederatedWorkerUtils.attemptTimeout(Long.MAX_VALUE) >= TOLERATED_HANDSHAKE_MS);
+	}
+
+	@Test
+	public void attemptBudgetIsCappedByRemainingTime() {
+		assertEquals(5, FederatedWorkerUtils.attemptTimeout(5));
+		assertEquals(1, FederatedWorkerUtils.attemptTimeout(1));
+	}
+
+	@Test
+	public void attemptBudgetIsZeroWhenOutOfTime() {
+		// 0 must only mean 'do not attempt', since `Socket.connect` reads a timeout of 0 as infinite
+		assertEquals(0, FederatedWorkerUtils.attemptTimeout(0));
+		assertEquals(0, FederatedWorkerUtils.attemptTimeout(-1));
+		assertEquals(0, FederatedWorkerUtils.attemptTimeout(Long.MIN_VALUE));
+	}
+
+	@Test
+	public void attemptBudgetIsNeverZeroWhileTimeIsLeft() {
+		for(long remaining = 1; remaining < 10000; remaining += 7)
+			assertTrue("a positive remaining time must not produce an infinite connect timeout",
+				FederatedWorkerUtils.attemptTimeout(remaining) > 0);
+	}
+
+	@Test
+	public void waitReturnsForAListeningPort() throws IOException {
+		try(ServerSocket listening = new ServerSocket(0)) {
+			// returns 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 reduced to a minute
+			assertTrue("the dead worker was not reported promptly", System.currentTimeMillis() - t0 < 10000);
+		}
+	}
+}