From 2ddb2689f3082c3d5d40b475a77170bde7ab950f Mon Sep 17 00:00:00 2001 From: Grigorii Turchenko Date: Tue, 11 Aug 2026 14:53:43 +0200 Subject: [PATCH 1/3] fix(federated): increase the worker readiness probe for delayed connections --- .../sysds/test/FederatedWorkerUtils.java | 51 +++++++--- .../federated/FederatedWorkerUtilsTest.java | 93 +++++++++++++++++++ 2 files changed, 130 insertions(+), 14 deletions(-) create mode 100644 src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java diff --git a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java b/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java index d604d7dcab4..445ec509a45 100644 --- a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java +++ b/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. - * - *

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,7 +149,9 @@ 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; } @@ -153,7 +159,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,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); diff --git a/src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java b/src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java new file mode 100644 index 00000000000..147c57b9df1 --- /dev/null +++ b/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); + } + } +} From a4b2adc2c09f20bd732cec975c9846707cff5324 Mon Sep 17 00:00:00 2001 From: Grigorii Turchenko Date: Fri, 21 Aug 2026 11:28:56 +0200 Subject: [PATCH 2/3] fix(FederatedWorkerUtils): remove `attemptTimeout()`, change "equals" to "less or equal" in `tryConnect()`, shorten comments --- .../sysds/test/FederatedWorkerUtils.java | 56 ++++++------------- 1 file changed, 17 insertions(+), 39 deletions(-) diff --git a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java b/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java index 445ec509a45..dfd68cef841 100644 --- a/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java +++ b/src/test/java/org/apache/sysds/test/FederatedWorkerUtils.java @@ -25,10 +25,11 @@ import java.util.function.BooleanSupplier; /** - * 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. + * 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. 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 { @@ -36,20 +37,14 @@ public final class FederatedWorkerUtils { private static final int POLL_INTERVAL_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. + * 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; @@ -100,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 @@ -138,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) { @@ -149,8 +142,7 @@ 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) { - // the deadline is rechecked per port, since a sweep over many ports can now spend up to - // CONNECT_TIMEOUT_MS on each of them + // 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; @@ -181,11 +173,11 @@ public static void waitForWorkers(int[] ports, int timeoutMs, java.util.function } 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 + 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), timeout); + s.connect(new InetSocketAddress("localhost", port), (int) Math.min(CONNECT_TIMEOUT_MS, remaining)); return true; } catch(IOException e) { // closed port, or a handshake that outlasted the budget @@ -193,20 +185,6 @@ private static boolean tryConnect(int port, long deadline) { } } - /** - * 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); From c58a3b125deea6a003b5fc0a6e6093dfe36eb3d2 Mon Sep 17 00:00:00 2001 From: Grigorii Turchenko Date: Fri, 21 Aug 2026 11:29:37 +0200 Subject: [PATCH 3/3] fix(tests): remove federated worker utils tests, move the rest to `FederatedUrlParserTest` --- .../federated/FederatedUrlParserTest.java | 31 +++++++ .../federated/FederatedWorkerUtilsTest.java | 93 ------------------- 2 files changed, 31 insertions(+), 93 deletions(-) delete mode 100644 src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java 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); + } + } } diff --git a/src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java b/src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java deleted file mode 100644 index 147c57b9df1..00000000000 --- a/src/test/java/org/apache/sysds/test/component/federated/FederatedWorkerUtilsTest.java +++ /dev/null @@ -1,93 +0,0 @@ -/* - * 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); - } - } -}