[#1608] feat: Introduce ExpiringClosableSupplier and refactor ShuffleManagerClient creation

client-spark/common/src/main/java/org/apache/spark/shuffle/reader/RssFetchFailedIterator.java

@@ -17,7 +17,6 @@
 
 package org.apache.spark.shuffle.reader;
 
-import java.io.IOException;
 import java.util.Objects;
 
 import scala.Product2;
@@ -30,10 +29,8 @@
 import org.slf4j.LoggerFactory;
 
 import org.apache.uniffle.client.api.ShuffleManagerClient;
-import org.apache.uniffle.client.factory.ShuffleManagerClientFactory;
 import org.apache.uniffle.client.request.RssReportShuffleFetchFailureRequest;
 import org.apache.uniffle.client.response.RssReportShuffleFetchFailureResponse;
-import org.apache.uniffle.common.ClientType;
 import org.apache.uniffle.common.exception.RssException;
 import org.apache.uniffle.common.exception.RssFetchFailedException;
 
@@ -52,8 +49,7 @@ public static class Builder {
     private int shuffleId;
     private int partitionId;
     private int stageAttemptId;
-    private String reportServerHost;
-    private int reportServerPort;
+    private ShuffleManagerClient shuffleManagerClient;
 
     private Builder() {}
 
@@ -77,19 +73,13 @@ Builder stageAttemptId(int stageAttemptId) {
       return this;
     }
 
-    Builder reportServerHost(String host) {
-      this.reportServerHost = host;
-      return this;
-    }
-
-    Builder port(int port) {
-      this.reportServerPort = port;
+    Builder shuffleManagerClient(ShuffleManagerClient shuffleManagerClient) {
+      this.shuffleManagerClient = shuffleManagerClient;
       return this;
     }
 
     <K, C> RssFetchFailedIterator<K, C> build(Iterator<Product2<K, C>> iter) {
       Objects.requireNonNull(this.appId);
-      Objects.requireNonNull(this.reportServerHost);
       return new RssFetchFailedIterator<>(this, iter);
     }
   }
@@ -98,37 +88,23 @@ static Builder newBuilder() {
     return new Builder();
   }
 
-  private static ShuffleManagerClient createShuffleManagerClient(String host, int port)
-      throws IOException {
-    ClientType grpc = ClientType.GRPC;
-    // host is passed from spark.driver.bindAddress, which would be set when SparkContext is
-    // constructed.
-    return ShuffleManagerClientFactory.getInstance().createShuffleManagerClient(grpc, host, port);
-  }
-
   private RssException generateFetchFailedIfNecessary(RssFetchFailedException e) {
-    String driver = builder.reportServerHost;
-    int port = builder.reportServerPort;
-    // todo: reuse this manager client if this is a bottleneck.
-    try (ShuffleManagerClient client = createShuffleManagerClient(driver, port)) {
-      RssReportShuffleFetchFailureRequest req =
-          new RssReportShuffleFetchFailureRequest(
-              builder.appId,
-              builder.shuffleId,
-              builder.stageAttemptId,
-              builder.partitionId,
-              e.getMessage());
-      RssReportShuffleFetchFailureResponse response = client.reportShuffleFetchFailure(req);
-      if (response.getReSubmitWholeStage()) {
-        // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is
-        // provided.
-        FetchFailedException ffe =
-            RssSparkShuffleUtils.createFetchFailedException(
-                builder.shuffleId, -1, builder.partitionId, e);
-        return new RssException(ffe);
-      }
-    } catch (IOException ioe) {
-      LOG.info("Error closing shuffle manager client with error:", ioe);
+    RssReportShuffleFetchFailureRequest req =
+        new RssReportShuffleFetchFailureRequest(
+            builder.appId,
+            builder.shuffleId,
+            builder.stageAttemptId,
+            builder.partitionId,
+            e.getMessage());
+    RssReportShuffleFetchFailureResponse response =
+        builder.shuffleManagerClient.reportShuffleFetchFailure(req);
+    if (response.getReSubmitWholeStage()) {
+      // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is
+      // provided.
+      FetchFailedException ffe =
+          RssSparkShuffleUtils.createFetchFailedException(
+              builder.shuffleId, -1, builder.partitionId, e);
+      return new RssException(ffe);
     }
     return e;
   }

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -446,6 +446,7 @@ public <K, V> ShuffleWriter<K, V> getWriter(
           this,
           sparkConf,
           shuffleWriteClient,
+          shuffleManagerClient,
           rssHandle,
           this::markFailedTask,
           context,
@@ -550,7 +551,8 @@ public <K, C> ShuffleReader<K, C> getReader(
           blockIdBitmap,
           taskIdBitmap,
           RssSparkConfig.toRssConf(sparkConf),
-          partitionToServers);
+          partitionToServers,
+          shuffleManagerClient);
     } else {
       throw new RssException("Unexpected ShuffleHandle:" + handle.getClass().getName());
     }

client-spark/spark2/src/main/java/org/apache/spark/shuffle/reader/RssShuffleReader.java

@@ -47,6 +47,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.uniffle.client.api.ShuffleManagerClient;
 import org.apache.uniffle.client.api.ShuffleReadClient;
 import org.apache.uniffle.client.factory.ShuffleClientFactory;
 import org.apache.uniffle.client.util.RssClientConfig;
@@ -77,6 +78,7 @@ public class RssShuffleReader<K, C> implements ShuffleReader<K, C> {
   private List<ShuffleServerInfo> shuffleServerInfoList;
   private Configuration hadoopConf;
   private RssConf rssConf;
+  private ShuffleManagerClient shuffleManagerClient;
 
   public RssShuffleReader(
       int startPartition,
@@ -90,7 +92,8 @@ public RssShuffleReader(
       Roaring64NavigableMap blockIdBitmap,
       Roaring64NavigableMap taskIdBitmap,
       RssConf rssConf,
-      Map<Integer, List<ShuffleServerInfo>> partitionToServers) {
+      Map<Integer, List<ShuffleServerInfo>> partitionToServers,
+      ShuffleManagerClient shuffleManagerClient) {
     this.appId = rssShuffleHandle.getAppId();
     this.startPartition = startPartition;
     this.endPartition = endPartition;
@@ -107,6 +110,7 @@ public RssShuffleReader(
     this.hadoopConf = hadoopConf;
     this.shuffleServerInfoList = (List<ShuffleServerInfo>) (partitionToServers.get(startPartition));
     this.rssConf = rssConf;
+    this.shuffleManagerClient = shuffleManagerClient;
     expectedTaskIdsBitmapFilterEnable = shuffleServerInfoList.size() > 1;
   }
 
@@ -235,16 +239,12 @@ public BoxedUnit apply() {
     // stage re-compute and shuffle manager server port are both set
     if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED)
         && rssConf.getInteger(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT, 0) > 0) {
-      String driver = rssConf.getString("driver.host", "");
-      int port = rssConf.get(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT);
       resultIter =
           RssFetchFailedIterator.newBuilder()
               .appId(appId)
               .shuffleId(shuffleId)
               .partitionId(startPartition)
               .stageAttemptId(context.stageAttemptNumber())
-              .reportServerHost(driver)
-              .port(port)
               .build(resultIter);
     }
     return resultIter;

client-spark/spark2/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -17,7 +17,6 @@
 
 package org.apache.spark.shuffle.writer;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -64,17 +63,13 @@
 
 import org.apache.uniffle.client.api.ShuffleManagerClient;
 import org.apache.uniffle.client.api.ShuffleWriteClient;
-import org.apache.uniffle.client.factory.ShuffleManagerClientFactory;
 import org.apache.uniffle.client.impl.FailedBlockSendTracker;
 import org.apache.uniffle.client.request.RssReassignServersRequest;
 import org.apache.uniffle.client.request.RssReportShuffleWriteFailureRequest;
 import org.apache.uniffle.client.response.RssReassignServersResponse;
 import org.apache.uniffle.client.response.RssReportShuffleWriteFailureResponse;
-import org.apache.uniffle.common.ClientType;
 import org.apache.uniffle.common.ShuffleBlockInfo;
 import org.apache.uniffle.common.ShuffleServerInfo;
-import org.apache.uniffle.common.config.RssClientConf;
-import org.apache.uniffle.common.config.RssConf;
 import org.apache.uniffle.common.exception.RssException;
 import org.apache.uniffle.common.exception.RssSendFailedException;
 import org.apache.uniffle.common.exception.RssWaitFailedException;
@@ -113,6 +108,7 @@ public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
   private final Set<Long> blockIds = Sets.newConcurrentHashSet();
   private TaskContext taskContext;
   private SparkConf sparkConf;
+  private ShuffleManagerClient shuffleManagerClient;
 
   public RssShuffleWriter(
       String appId,
@@ -124,6 +120,7 @@ public RssShuffleWriter(
       RssShuffleManager shuffleManager,
       SparkConf sparkConf,
       ShuffleWriteClient shuffleWriteClient,
+      ShuffleManagerClient shuffleManagerClient,
       RssShuffleHandle<K, V, C> rssHandle,
       SimpleShuffleHandleInfo shuffleHandleInfo,
       TaskContext context) {
@@ -136,6 +133,7 @@ public RssShuffleWriter(
         shuffleManager,
         sparkConf,
         shuffleWriteClient,
+        shuffleManagerClient,
         rssHandle,
         (tid) -> true,
         shuffleHandleInfo,
@@ -152,6 +150,7 @@ private RssShuffleWriter(
       RssShuffleManager shuffleManager,
       SparkConf sparkConf,
       ShuffleWriteClient shuffleWriteClient,
+      ShuffleManagerClient shuffleManagerClient,
       RssShuffleHandle<K, V, C> rssHandle,
       Function<String, Boolean> taskFailureCallback,
       ShuffleHandleInfo shuffleHandleInfo,
@@ -171,6 +170,7 @@ private RssShuffleWriter(
     this.bitmapSplitNum = sparkConf.get(RssSparkConfig.RSS_CLIENT_BITMAP_SPLIT_NUM);
     this.serverToPartitionToBlockIds = Maps.newHashMap();
     this.shuffleWriteClient = shuffleWriteClient;
+    this.shuffleManagerClient = shuffleManagerClient;
     this.shuffleServersForData = shuffleHandleInfo.getServers();
     this.partitionToServers = shuffleHandleInfo.getAvailablePartitionServersForWriter();
     this.isMemoryShuffleEnabled =
@@ -189,6 +189,7 @@ public RssShuffleWriter(
       RssShuffleManager shuffleManager,
       SparkConf sparkConf,
       ShuffleWriteClient shuffleWriteClient,
+      ShuffleManagerClient shuffleManagerClient,
       RssShuffleHandle<K, V, C> rssHandle,
       Function<String, Boolean> taskFailureCallback,
       TaskContext context,
@@ -202,6 +203,7 @@ public RssShuffleWriter(
         shuffleManager,
         sparkConf,
         shuffleWriteClient,
+        shuffleManagerClient,
         rssHandle,
         taskFailureCallback,
         shuffleHandleInfo,
@@ -527,14 +529,6 @@ protected ShuffleWriteMetrics getShuffleWriteMetrics() {
     return shuffleWriteMetrics;
   }
 
-  private static ShuffleManagerClient createShuffleManagerClient(String host, int port)
-      throws IOException {
-    ClientType grpc = ClientType.GRPC;
-    // Host can be inferred from `spark.driver.bindAddress`, which would be set when SparkContext is
-    // constructed.
-    return ShuffleManagerClientFactory.getInstance().createShuffleManagerClient(grpc, host, port);
-  }
-
   private void throwFetchFailedIfNecessary(Exception e) {
     // The shuffleServer is registered only when a Block fails to be sent
     if (e instanceof RssSendFailedException) {
@@ -549,34 +543,27 @@ private void throwFetchFailedIfNecessary(Exception e) {
               taskContext.stageAttemptNumber(),
               shuffleServerInfos,
               e.getMessage());
-      RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
-      String driver = rssConf.getString("driver.host", "");
-      int port = rssConf.get(RssClientConf.SHUFFLE_MANAGER_GRPC_PORT);
-      try (ShuffleManagerClient shuffleManagerClient = createShuffleManagerClient(driver, port)) {
-        RssReportShuffleWriteFailureResponse response =
-            shuffleManagerClient.reportShuffleWriteFailure(req);
-        if (response.getReSubmitWholeStage()) {
-          // The shuffle server is reassigned.
-          RssReassignServersRequest rssReassignServersRequest =
-              new RssReassignServersRequest(
-                  taskContext.stageId(),
-                  taskContext.stageAttemptNumber(),
-                  shuffleId,
-                  partitioner.numPartitions());
-          RssReassignServersResponse rssReassignServersResponse =
-              shuffleManagerClient.reassignOnStageResubmit(rssReassignServersRequest);
-          LOG.info(
-              "Whether the reassignment is successful: {}",
-              rssReassignServersResponse.isNeedReassign());
-          // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is
-          // provided.
-          FetchFailedException ffe =
-              RssSparkShuffleUtils.createFetchFailedException(
-                  shuffleId, -1, taskContext.stageAttemptNumber(), e);
-          throw new RssException(ffe);
-        }
-      } catch (IOException ioe) {
-        LOG.info("Error closing shuffle manager client with error:", ioe);
+      RssReportShuffleWriteFailureResponse response =
+          shuffleManagerClient.reportShuffleWriteFailure(req);
+      if (response.getReSubmitWholeStage()) {
+        // The shuffle server is reassigned.
+        RssReassignServersRequest rssReassignServersRequest =
+            new RssReassignServersRequest(
+                taskContext.stageId(),
+                taskContext.stageAttemptNumber(),
+                shuffleId,
+                partitioner.numPartitions());
+        RssReassignServersResponse rssReassignServersResponse =
+            shuffleManagerClient.reassignOnStageResubmit(rssReassignServersRequest);
+        LOG.info(
+            "Whether the reassignment is successful: {}",
+            rssReassignServersResponse.isNeedReassign());
+        // since we are going to roll out the whole stage, mapIndex shouldn't matter, hence -1 is
+        // provided.
+        FetchFailedException ffe =
+            RssSparkShuffleUtils.createFetchFailedException(
+                shuffleId, -1, taskContext.stageAttemptNumber(), e);
+        throw new RssException(ffe);
       }
     }
     throw new RssException(e);

client-spark/spark2/src/test/java/org/apache/spark/shuffle/reader/RssShuffleReaderTest.java

@@ -35,6 +35,7 @@
 import org.junit.jupiter.api.Test;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 
+import org.apache.uniffle.client.api.ShuffleManagerClient;
 import org.apache.uniffle.common.ShuffleServerInfo;
 import org.apache.uniffle.common.config.RssClientConf;
 import org.apache.uniffle.common.config.RssConf;
@@ -85,6 +86,7 @@ public void readTest() throws Exception {
     rssConf.set(RssClientConf.RSS_STORAGE_TYPE, StorageType.HDFS.name());
     rssConf.set(RssClientConf.RSS_INDEX_READ_LIMIT, 1000);
     rssConf.set(RssClientConf.RSS_CLIENT_READ_BUFFER_SIZE, "1000");
+    final ShuffleManagerClient mockShuffleManagerClient = mock(ShuffleManagerClient.class);
     RssShuffleReader<String, String> rssShuffleReaderSpy =
         spy(
             new RssShuffleReader<>(
@@ -99,7 +101,8 @@ public void readTest() throws Exception {
                 blockIdBitmap,
                 taskIdBitmap,
                 rssConf,
-                partitionToServers));
+                partitionToServers,
+                mockShuffleManagerClient));
 
     validateResult(rssShuffleReaderSpy.read(), expectedData, 10);
   }