YARN-12007. Avoid intermediate overflow in YARN unit conversion

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/UnitsConversionUtil.java

@@ -137,29 +137,27 @@ public static long convert(String fromUnit, String toUnit, long fromValue) {
     }
     Converter fc = getConverter(fromUnit);
     Converter tc = getConverter(toUnit);
-    long numerator = fc.numerator * tc.denominator;
-    long denominator = fc.denominator * tc.numerator;
-    long numeratorMultiplierLimit = Long.MAX_VALUE / numerator;
-    if (numerator < denominator) {
-      if (numeratorMultiplierLimit < fromValue) {
-        String overflowMsg =
+    if (fromValue == 0) {
+      return 0;
+    }
+    try {
+      long numerator = Math.multiplyExact(fc.numerator, tc.denominator);
+      long denominator = Math.multiplyExact(fc.denominator, tc.numerator);
+      return Math.multiplyExact(fromValue, numerator) / denominator;
+    } catch (ArithmeticException e) {
+      BigInteger numerator = BigInteger.valueOf(fromValue)
+          .multiply(BigInteger.valueOf(fc.numerator))
+          .multiply(BigInteger.valueOf(tc.denominator));
+      BigInteger denominator = BigInteger.valueOf(fc.denominator)
+          .multiply(BigInteger.valueOf(tc.numerator));
+      try {
+        return numerator.divide(denominator).longValueExact();
+      } catch (ArithmeticException overflow) {
+        throw new IllegalArgumentException(
             "Converting " + fromValue + " from '" + fromUnit + "' to '" + toUnit
-                + "' will result in an overflow of Long";
-        throw new IllegalArgumentException(overflowMsg);
+                + "' will result in an overflow of Long", overflow);
       }
-      return (fromValue * numerator) / denominator;
-    }
-    if (numeratorMultiplierLimit > fromValue) {
-      return (numerator * fromValue) / denominator;
-    }
-    long tmp = numerator / denominator;
-    if ((Long.MAX_VALUE / tmp) < fromValue) {
-      String overflowMsg =
-          "Converting " + fromValue + " from '" + fromUnit + "' to '" + toUnit
-              + "' will result in an overflow of Long";
-      throw new IllegalArgumentException(overflowMsg);
     }
-    return fromValue * tmp;
   }
 
   /**

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/util/TestUnitsConversionUtil.java

@@ -21,6 +21,7 @@
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
 
 /**
@@ -112,6 +113,38 @@ void testOverflow() {
     }
   }
 
+  @Test
+  void testIntermediateOverflow() {
+    assertEquals(0, UnitsConversionUtil.convert("p", "G", Long.MAX_VALUE));
+    assertEquals(0, UnitsConversionUtil.convert("p", "G", Long.MIN_VALUE));
+    assertEquals(8881, UnitsConversionUtil.convert("P", "Pi", 10000));
+    assertEquals(-8881, UnitsConversionUtil.convert("P", "Pi", -10000));
+    assertEquals(Long.MAX_VALUE / 1000,
+        UnitsConversionUtil.convert("p", "n", Long.MAX_VALUE));
+    assertEquals(Long.MIN_VALUE / 1000,
+        UnitsConversionUtil.convert("p", "n", Long.MIN_VALUE));
+  }
+
+  @Test
+  void testZeroConversion() {
+    assertEquals(0, UnitsConversionUtil.convert("G", "p", 0));
+    assertEquals(0, UnitsConversionUtil.convert("P", "p", 0));
+    assertThrows(IllegalArgumentException.class,
+        () -> UnitsConversionUtil.convert("invalid", "p", 0));
+    assertThrows(IllegalArgumentException.class,
+        () -> UnitsConversionUtil.convert("G", "invalid", 0));
+  }
+
+  @Test
+  void testFinalResultOverflow() {
+    assertThrows(IllegalArgumentException.class,
+        () -> UnitsConversionUtil.convert("Pi", "P", Long.MAX_VALUE));
+    assertThrows(IllegalArgumentException.class,
+        () -> UnitsConversionUtil.convert("Pi", "P", Long.MIN_VALUE));
+    assertThrows(IllegalArgumentException.class,
+        () -> UnitsConversionUtil.convert("P", "p", 1));
+  }
+
   @Test
   void testCompare() {
     String unitA = "P";