diff --git a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java b/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java index d604d7dcab4..dfd68cef841 100644 --- a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java +++ b/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java @@ -27,9 +27,8 @@ /** * 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 + * 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 { @@ -37,15 +36,15 @@ 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,7 +142,8 @@ 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; } @@ -153,7 +151,7 @@ public static void waitForWorkers(int[] ports, int timeoutMs, java.util.function 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; } } diff --git a/src/test/java/org/apache/sysds/test/component/federated/FederatedUrlParserTest.java b/src/test/java/org/apache/sysds/test/component/federated/FederatedUrlParserTest.java index 10e1e6b549d..79971f965f1 100644 --- a/src/test/java/org/apache/sysds/test/component/federated/FederatedUrlParserTest.java +++ b/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); + } + } }