(uniffle) branch master updated: [#2749] fix(test): stabilize shuffle server assignment retry test for `ShuffleWithRssClientTest` (#2795)

[zu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zuston pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new 99a8a10af [#2749] fix(test): stabilize shuffle server assignment retry test for `ShuffleWithRssClientTest` (#2795)
99a8a10af is described below

commit 99a8a10afa2c29889f2b14e26cc6659545ec92e6
Author: Junfan Zhang <zu...@apache.org>
AuthorDate: Tue Sep 22 14:12:48 2026 +0800

    [#2749] fix(test): stabilize shuffle server assignment retry test for `ShuffleWithRssClientTest` (#2795)
    
    ### What changes were proposed in this pull request?
    
    fix the flaky test of `ShuffleWithRssClientTest`
    
    ### Why are the changes needed?
    
    Fix: #2749
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    unit tests
---
 .../uniffle/test/ShuffleWithRssClientTest.java     | 28 +++++++++++++++++-----
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleWithRssClientTest.java b/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleWithRssClientTest.java
index c5e754d4b..4b9f4797c 100644
--- a/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleWithRssClientTest.java
+++ b/integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleWithRssClientTest.java
@@ -21,11 +21,13 @@ import java.io.File;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
+import org.awaitility.Awaitility;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -50,6 +52,7 @@ import org.apache.uniffle.common.util.BlockIdLayout;
 import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.RetryUtils;
 import org.apache.uniffle.coordinator.CoordinatorConf;
+import org.apache.uniffle.server.ShuffleServer;
 import org.apache.uniffle.server.ShuffleServerConf;
 import org.apache.uniffle.storage.util.StorageType;
 
@@ -454,18 +457,30 @@ public class ShuffleWithRssClientTest extends ShuffleReadWriteBase {
     String appId = "app-1";
     RemoteStorageInfo remoteStorage = new RemoteStorageInfo("");
     ShuffleAssignmentsInfo response = null;
-    ShuffleServerConf shuffleServerConf =
-        shuffleServerConfWithoutPort(0, null, ServerType.GRPC_NETTY);
-    int heartbeatInterval = shuffleServerConf.getInteger("rss.server.heartbeat.interval", 1000);
-    Thread.sleep(heartbeatInterval * 2);
+    long heartbeatInterval =
+        nettyShuffleServers
+            .get(0)
+            .getShuffleServerConf()
+            .get(ShuffleServerConf.SERVER_HEARTBEAT_INTERVAL);
+    Set<String> requiredTags = Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION);
+    // Both servers must be assignable before stopping the first assigned server.
+    Awaitility.await()
+        .atMost(10, TimeUnit.SECONDS)
+        .untilAsserted(
+            () -> {
+              nettyShuffleServers.forEach(ShuffleServer::sendHeartbeat);
+              coordinators.forEach(
+                  c ->
+                      assertEquals(
+                          maxTryTime, c.getClusterManager().getServerList(requiredTags).size()));
+            });
     shuffleWriteClientImpl.registerCoordinators(getQuorum());
     response =
         RetryUtils.retry(
             () -> {
               int currentTryTime = tryTime.incrementAndGet();
               ShuffleAssignmentsInfo shuffleAssignments =
-                  shuffleWriteClientImpl.getShuffleAssignments(
-                      appId, 1, 1, 1, Sets.newHashSet(Constants.SHUFFLE_SERVER_VERSION), 1, -1);
+                  shuffleWriteClientImpl.getShuffleAssignments(appId, 1, 1, 1, requiredTags, 1, -1);
 
               Map<ShuffleServerInfo, List<PartitionRange>> serverToPartitionRanges =
                   shuffleAssignments.getServerToPartitionRanges();
@@ -505,5 +520,6 @@ public class ShuffleWithRssClientTest extends ShuffleReadWriteBase {
             maxTryTime);
 
     assertNotNull(response);
+    assertEquals(maxTryTime, tryTime.get());
   }
 }