From d027a24f0349b60efa5125c330058f123771748f Mon Sep 17 00:00:00 2001
From: Weiwei Yang <wwei@apache.org>
Date: Sun, 18 Nov 2018 23:18:26 +0800
Subject: [PATCH] YARN-8833. Avoid potential integer overflow when computing
 fair shares. Contributed by liyakun.

---
 .../fair/policies/ComputeFairShares.java        | 15 +++++++++------
 .../scheduler/fair/TestComputeFairShares.java   | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/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 b/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
index 0a21b026714..3fe0c68986b 100644
--- a/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
+++ b/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 @@ private static void computeSharesInternal(
     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 @@ private static void computeSharesInternal(
    * 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 @@ private static int resourceUsedWithWeightToResourceRatio(double w2rRatio,
    * 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;
   }
 
   /**
diff --git a/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 b/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
index c3bcb3b2179..b1666d61a59 100644
--- a/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
+++ b/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 @@ private void verifyCPUShares(int... shares) {
       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);
+  }
 }