瀏覽代碼

YARN-8833. Avoid potential integer overflow when computing fair shares. Contributed by liyakun.

(cherry picked from commit d027a24f0349b60efa5125c330058f123771748f)
Weiwei Yang 6 年之前
父節點
當前提交
067e39f781

+ 9 - 6
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/fair/policies/ComputeFairShares.java

@@ -145,7 +145,7 @@ public class ComputeFairShares {
     double right = rMax;
     for (int i = 0; i < COMPUTE_FAIR_SHARES_ITERATIONS; i++) {
       double mid = (left + right) / 2.0;
-      int plannedResourceUsed = resourceUsedWithWeightToResourceRatio(
+      long plannedResourceUsed = resourceUsedWithWeightToResourceRatio(
           mid, schedulables, type);
       if (plannedResourceUsed == totalResource) {
         right = mid;
@@ -174,11 +174,14 @@ public class ComputeFairShares {
    * Compute the resources that would be used given a weight-to-resource ratio
    * w2rRatio, for use in the computeFairShares algorithm as described in #
    */
-  private static int resourceUsedWithWeightToResourceRatio(double w2rRatio,
+  private static long resourceUsedWithWeightToResourceRatio(double w2rRatio,
       Collection<? extends Schedulable> schedulables, String type) {
-    int resourcesTaken = 0;
+    long resourcesTaken = 0;
     for (Schedulable sched : schedulables) {
-      int share = computeShare(sched, w2rRatio, type);
+      long share = computeShare(sched, w2rRatio, type);
+      if (Long.MAX_VALUE - resourcesTaken < share) {
+        return Long.MAX_VALUE;
+      }
       resourcesTaken += share;
     }
     return resourcesTaken;
@@ -188,12 +191,12 @@ public class ComputeFairShares {
    * Compute the resources assigned to a Schedulable given a particular
    * weight-to-resource ratio w2rRatio.
    */
-  private static int computeShare(Schedulable sched, double w2rRatio,
+  private static long computeShare(Schedulable sched, double w2rRatio,
       String type) {
     double share = sched.getWeight() * w2rRatio;
     share = Math.max(share, sched.getMinShare().getResourceValue(type));
     share = Math.min(share, sched.getMaxShare().getResourceValue(type));
-    return (int) share;
+    return (long) share;
   }
 
   /**

+ 17 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/fair/TestComputeFairShares.java

@@ -19,7 +19,10 @@
 package org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
+
+import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceInformation;
 
 import org.junit.Assert;
@@ -205,4 +208,18 @@ public class TestComputeFairShares {
       Assert.assertEquals(shares[i], scheds.get(i).getFairShare().getVirtualCores());
     }
   }
+
+  /**
+   * Test computeShares will not enter into infinite loop.
+   */
+  @Test(timeout = 10000)
+  public void testResourceUsedWithWeightToResourceRatio() {
+    Collection<Schedulable> schedulables = new ArrayList<>();
+    schedulables.add(new FakeSchedulable(Integer.MAX_VALUE));
+    schedulables.add(new FakeSchedulable(Integer.MAX_VALUE));
+
+    Resource totalResource = Resource.newInstance(Integer.MAX_VALUE, 0);
+    ComputeFairShares.computeShares(
+        schedulables, totalResource, ResourceInformation.MEMORY_URI);
+  }
 }