[ISSUE-186][Feature] Use I/O cost time to select storage paths

coordinator/pom.xml

@@ -73,6 +73,10 @@
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-common</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.hadoop</groupId>
+      <artifactId>hadoop-minicluster</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-inline</artifactId>

coordinator/src/main/java/org/apache/uniffle/coordinator/AppBalanceSelectStorageStrategy.java

@@ -0,0 +1,144 @@
+/*
+ * 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.uniffle.coordinator;
+
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.coordinator.LowestIOSampleCostSelectStorageStrategy.RankValue;
+
+/**
+ * AppBalanceSelectStorageStrategy will consider the number of apps allocated on each remote path is balanced.
+ */
+public class AppBalanceSelectStorageStrategy implements SelectStorageStrategy {
+
+  private static final Logger LOG = LoggerFactory.getLogger(AppBalanceSelectStorageStrategy.class);
+  /**
+   * store appId -> remote path to make sure all shuffle data of the same application
+   * will be written to the same remote storage
+   */
+  private final Map<String, RemoteStorageInfo> appIdToRemoteStorageInfo;
+  /**
+   * store remote path -> application count for assignment strategy
+   */
+  private final Map<String, RankValue> remoteStoragePathCounter;
+  private final Map<String, RemoteStorageInfo> availableRemoteStorageInfo;
+
+  public AppBalanceSelectStorageStrategy() {
+    this.appIdToRemoteStorageInfo = Maps.newConcurrentMap();
+    this.remoteStoragePathCounter = Maps.newConcurrentMap();
+    this.availableRemoteStorageInfo = Maps.newHashMap();
+  }
+
+  /**
+   * the strategy of pick remote storage is according to assignment count
+   */
+  @Override
+  public RemoteStorageInfo pickRemoteStorage(String appId) {
+    if (appIdToRemoteStorageInfo.containsKey(appId)) {
+      return appIdToRemoteStorageInfo.get(appId);
+    }
+
+    // create list for sort
+    List<Map.Entry<String, RankValue>> sizeList =
+        Lists.newArrayList(remoteStoragePathCounter.entrySet()).stream().filter(Objects::nonNull)
+            .sorted(Comparator.comparingInt(entry -> entry.getValue().getAppNum().get())).collect(Collectors.toList());
+
+    for (Map.Entry<String, RankValue> entry : sizeList) {
+      String storagePath = entry.getKey();
+      if (availableRemoteStorageInfo.containsKey(storagePath)) {
+        appIdToRemoteStorageInfo.putIfAbsent(appId, availableRemoteStorageInfo.get(storagePath));
+        incRemoteStorageCounter(storagePath);
+        break;
+      }
+    }
+    return appIdToRemoteStorageInfo.get(appId);
+  }
+
+  @Override
+  @VisibleForTesting
+  public synchronized void incRemoteStorageCounter(String remoteStoragePath) {
+    RankValue counter = remoteStoragePathCounter.get(remoteStoragePath);
+    if (counter != null) {
+      counter.getAppNum().incrementAndGet();
+    } else {
+      // it may be happened when assignment remote storage
+      // and refresh remote storage at the same time
+      LOG.warn("Remote storage path lost during assignment: %s doesn't exist, reset it to 1",
+          remoteStoragePath);
+      remoteStoragePathCounter.put(remoteStoragePath, new RankValue(1));
+    }
+  }
+
+  @Override
+  @VisibleForTesting
+  public synchronized void decRemoteStorageCounter(String storagePath) {
+    if (!StringUtils.isEmpty(storagePath)) {
+      RankValue atomic = remoteStoragePathCounter.get(storagePath);
+      if (atomic != null) {
+        double count = atomic.getAppNum().decrementAndGet();
+        if (count < 0) {
+          LOG.warn("Unexpected counter for remote storage: %s, which is %i, reset to 0",
+              storagePath, count);
+          atomic.getAppNum().set(0);
+        }
+      } else {
+        LOG.warn("Can't find counter for remote storage: {}", storagePath);
+        remoteStoragePathCounter.putIfAbsent(storagePath, new RankValue(0));
+      }
+      if (remoteStoragePathCounter.get(storagePath).getAppNum().get() == 0
+          && !availableRemoteStorageInfo.containsKey(storagePath)) {
+        remoteStoragePathCounter.remove(storagePath);
+      }
+    }
+  }
+
+  @Override
+  public synchronized void removePathFromCounter(String storagePath) {
+    RankValue atomic = remoteStoragePathCounter.get(storagePath);
+    if (atomic != null && atomic.getAppNum().get() == 0) {
+      remoteStoragePathCounter.remove(storagePath);
+    }
+  }
+
+  @Override
+  public Map<String, RemoteStorageInfo> getAppIdToRemoteStorageInfo() {
+    return appIdToRemoteStorageInfo;
+  }
+
+  @Override
+  public Map<String, RankValue> getRemoteStoragePathRankValue() {
+    return remoteStoragePathCounter;
+  }
+
+  @Override
+  public Map<String, RemoteStorageInfo> getAvailableRemoteStorageInfo() {
+    return availableRemoteStorageInfo;
+  }
+}

coordinator/src/main/java/org/apache/uniffle/coordinator/ApplicationManager.java

@@ -19,16 +19,12 @@
 
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.stream.Collectors;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Lists;
@@ -41,24 +37,38 @@
 import org.apache.uniffle.common.RemoteStorageInfo;
 import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.ThreadUtils;
+import org.apache.uniffle.coordinator.LowestIOSampleCostSelectStorageStrategy.RankValue;
 
 public class ApplicationManager {
 
   private static final Logger LOG = LoggerFactory.getLogger(ApplicationManager.class);
   private long expired;
+  private StrategyName storageStrategy;
   private Map<String, Long> appIds = Maps.newConcurrentMap();
+  private SelectStorageStrategy selectStorageStrategy;
   // store appId -> remote path to make sure all shuffle data of the same application
   // will be written to the same remote storage
-  private Map<String, RemoteStorageInfo> appIdToRemoteStorageInfo = Maps.newConcurrentMap();
+  private Map<String, RemoteStorageInfo> appIdToRemoteStorageInfo;
   // store remote path -> application count for assignment strategy
-  private Map<String, AtomicInteger> remoteStoragePathCounter = Maps.newConcurrentMap();
+  private Map<String, RankValue> remoteStoragePathRankValue;
   private Map<String, String> remoteStorageToHost = Maps.newConcurrentMap();
-  private Map<String, RemoteStorageInfo> availableRemoteStorageInfo = Maps.newHashMap();
+  private Map<String, RemoteStorageInfo> availableRemoteStorageInfo;
   private ScheduledExecutorService scheduledExecutorService;
   // it's only for test case to check if status check has problem
   private boolean hasErrorInStatusCheck = false;
 
   public ApplicationManager(CoordinatorConf conf) {
+    storageStrategy = conf.get(CoordinatorConf.COORDINATOR_REMOTE_STORAGE_SELECT_STRATEGY);
+    if (StrategyName.IO_SAMPLE == storageStrategy) {
+      selectStorageStrategy = new LowestIOSampleCostSelectStorageStrategy(conf);
+    } else if (StrategyName.APP_BALANCE == storageStrategy) {
+      selectStorageStrategy = new AppBalanceSelectStorageStrategy();
+    } else {
+      throw new UnsupportedOperationException("Unsupported selected storage strategy.");
+    }
+    appIdToRemoteStorageInfo = selectStorageStrategy.getAppIdToRemoteStorageInfo();
+    remoteStoragePathRankValue = selectStorageStrategy.getRemoteStoragePathRankValue();
+    availableRemoteStorageInfo = selectStorageStrategy.getAvailableRemoteStorageInfo();
     expired = conf.getLong(CoordinatorConf.COORDINATOR_APP_EXPIRED);
     // the thread for checking application status
     scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(
@@ -83,7 +93,7 @@ public void refreshRemoteStorage(String remoteStoragePath, String remoteStorageC
       // add remote path if not exist
       for (String path : paths) {
         if (!availableRemoteStorageInfo.containsKey(path)) {
-          remoteStoragePathCounter.putIfAbsent(path, new AtomicInteger(0));
+          remoteStoragePathRankValue.putIfAbsent(path, new RankValue(0));
           // refreshRemoteStorage is designed without multiple thread problem
           // metrics shouldn't be added duplicated
           addRemoteStorageMetrics(path);
@@ -117,81 +127,31 @@ public void refreshRemoteStorage(String remoteStoragePath, String remoteStorageC
   // the strategy of pick remote storage is according to assignment count
   // todo: better strategy with workload balance
   public RemoteStorageInfo pickRemoteStorage(String appId) {
-    if (appIdToRemoteStorageInfo.containsKey(appId)) {
-      return appIdToRemoteStorageInfo.get(appId);
-    }
-
-    // create list for sort
-    List<Map.Entry<String, AtomicInteger>> sizeList =
-        Lists.newArrayList(remoteStoragePathCounter.entrySet()).stream().filter(Objects::nonNull)
-            .sorted(Comparator.comparingInt(entry -> entry.getValue().get())).collect(Collectors.toList());
-
-    for (Map.Entry<String, AtomicInteger> entry : sizeList) {
-      String storagePath = entry.getKey();
-      if (availableRemoteStorageInfo.containsKey(storagePath)) {
-        appIdToRemoteStorageInfo.putIfAbsent(appId, availableRemoteStorageInfo.get(storagePath));
-        incRemoteStorageCounter(storagePath);
-        break;
-      }
-    }
+    selectStorageStrategy.pickRemoteStorage(appId);
     return appIdToRemoteStorageInfo.get(appId);
   }
 
-  @VisibleForTesting
-  protected synchronized void incRemoteStorageCounter(String remoteStoragePath) {
-    AtomicInteger counter = remoteStoragePathCounter.get(remoteStoragePath);
-    if (counter != null) {
-      counter.incrementAndGet();
-    } else {
-      // it may be happened when assignment remote storage
-      // and refresh remote storage at the same time
-      LOG.warn("Remote storage path lost during assignment: %s doesn't exist, reset it to 1",
-          remoteStoragePath);
-      remoteStoragePathCounter.put(remoteStoragePath, new AtomicInteger(1));
-    }
-  }
-
   @VisibleForTesting
   protected synchronized void decRemoteStorageCounter(String storagePath) {
-    if (!StringUtils.isEmpty(storagePath)) {
-      AtomicInteger atomic = remoteStoragePathCounter.get(storagePath);
-      if (atomic != null) {
-        int count = atomic.decrementAndGet();
-        if (count < 0) {
-          LOG.warn("Unexpected counter for remote storage: %s, which is %i, reset to 0",
-              storagePath, count);
-          atomic.set(0);
-        }
-      } else {
-        LOG.warn("Can't find counter for remote storage: {}", storagePath);
-        remoteStoragePathCounter.putIfAbsent(storagePath, new AtomicInteger(0));
-      }
-      if (remoteStoragePathCounter.get(storagePath).get() == 0
-          && !availableRemoteStorageInfo.containsKey(storagePath)) {
-        remoteStoragePathCounter.remove(storagePath);
-      }
-    }
+    selectStorageStrategy.decRemoteStorageCounter(storagePath);
   }
 
   private synchronized void removePathFromCounter(String storagePath) {
-    AtomicInteger atomic = remoteStoragePathCounter.get(storagePath);
-    if (atomic != null && atomic.get() == 0) {
-      remoteStoragePathCounter.remove(storagePath);
-    }
+    selectStorageStrategy.removePathFromCounter(storagePath);
   }
 
   public Set<String> getAppIds() {
     return appIds.keySet();
   }
 
   @VisibleForTesting
-  protected Map<String, RemoteStorageInfo> getAppIdToRemoteStorageInfo() {
-    return appIdToRemoteStorageInfo;
+  protected Map<String, RankValue> getRemoteStoragePathRankValue() {
+    return remoteStoragePathRankValue;
   }
 
   @VisibleForTesting
-  protected Map<String, AtomicInteger> getRemoteStoragePathCounter() {
-    return remoteStoragePathCounter;
+  public SelectStorageStrategy getSelectStorageStrategy() {
+    return selectStorageStrategy;
   }
 
   @VisibleForTesting
@@ -237,7 +197,7 @@ private void updateRemoteStorageMetrics() {
       try {
         String storageHost = getStorageHost(remoteStoragePath);
         CoordinatorMetrics.updateDynamicGaugeForRemoteStorage(storageHost,
-            remoteStoragePathCounter.get(remoteStoragePath).get());
+            remoteStoragePathRankValue.get(remoteStoragePath).getAppNum().get());
       } catch (Exception e) {
         LOG.warn("Update remote storage metrics for {} failed ", remoteStoragePath);
       }
@@ -266,4 +226,9 @@ private String getStorageHost(String remoteStoragePath) {
     }
     return storageHost;
   }
+
+  public enum StrategyName {
+    APP_BALANCE,
+    IO_SAMPLE
+  }
 }

coordinator/src/main/java/org/apache/uniffle/coordinator/CoordinatorConf.java

@@ -26,6 +26,7 @@
 import org.apache.uniffle.common.config.RssBaseConf;
 import org.apache.uniffle.common.util.RssUtils;
 
+import static org.apache.uniffle.coordinator.ApplicationManager.StrategyName.APP_BALANCE;
 import static org.apache.uniffle.coordinator.AssignmentStrategyFactory.StrategyName.PARTITION_BALANCE;
 
 /**
@@ -128,7 +129,26 @@ public class CoordinatorConf extends RssBaseConf {
       .stringType()
       .noDefaultValue()
       .withDescription("Remote Storage Cluster related conf with format $clusterId,$key=$value, separated by ';'");
-
+  public static final ConfigOption<ApplicationManager.StrategyName> COORDINATOR_REMOTE_STORAGE_SELECT_STRATEGY =
+      ConfigOptions.key("rss.coordinator.remote.storage.select.strategy")
+      .enumType(ApplicationManager.StrategyName.class)
+      .defaultValue(APP_BALANCE)
+      .withDescription("Strategy for selecting the remote path");
+  public static final ConfigOption<Long> COORDINATOR_REMOTE_STORAGE_IO_SAMPLE_SCHEDULE_TIME = ConfigOptions
+      .key("rss.coordinator.remote.storage.io.sample.schedule.time")
+      .longType()
+      .defaultValue(60 * 1000L)
+      .withDescription("The time of scheduling the read and write time of the paths to obtain different HDFS");
+  public static final ConfigOption<Integer> COORDINATOR_REMOTE_STORAGE_IO_SAMPLE_FILE_SIZE = ConfigOptions
+      .key("rss.coordinator.remote.storage.io.sample.file.size")
+      .intType()
+      .defaultValue(204800 * 1000)
+      .withDescription("The size of the file that the scheduled thread reads and writes");
+  public static final ConfigOption<Integer> COORDINATOR_REMOTE_STORAGE_IO_SAMPLE_ACCESS_TIMES = ConfigOptions
+      .key("rss.coordinator.remote.storage.io.sample.access.times")
+      .intType()
+      .defaultValue(3)
+      .withDescription("The number of times to read and write HDFS files");
 
   public CoordinatorConf() {
   }