From 2ddb2689f3082c3d5d40b475a77170bde7ab950f Mon Sep 17 00:00:00 2001 From: Grigorii Turchenko Date: Tue, 11 Aug 2026 14:53:43 +0200 Subject: [PATCH] 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); + } + } +}