fixing overflow in GatkSparkTool.getRecommendedNumReducers() (#5586)

* fixing potential overflow in GatkSparkTool.getRecommendedNumReducers() * Add new method MathUtils.toIntExactOrThrow for safely down converting long to int
broadinstitute · Jan 22, 2019 · d3e45c7 · d3e45c7
1 parent 1af1f17
commit d3e45c7
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 10 deletions.
diff --git a/src/main/java/org/broadinstitute/hellbender/engine/spark/GATKSparkTool.java b/src/main/java/org/broadinstitute/hellbender/engine/spark/GATKSparkTool.java
@@ -1,9 +1,10 @@
 package org.broadinstitute.hellbender.engine.spark;
 
 import com.google.common.annotations.VisibleForTesting;
-import htsjdk.samtools.*;
+import htsjdk.samtools.SAMFileHeader;
+import htsjdk.samtools.SAMSequenceDictionary;
+import htsjdk.samtools.SamFileHeaderMerger;
 import htsjdk.samtools.reference.ReferenceSequenceFileFactory;
-import htsjdk.samtools.util.CloseableIterator;
 import htsjdk.samtools.util.GZIIndex;
 import htsjdk.samtools.util.IOUtil;
 import htsjdk.variant.vcf.VCFHeaderLine;
@@ -20,19 +21,16 @@
 import org.broadinstitute.hellbender.engine.FeatureManager;
 import org.broadinstitute.hellbender.engine.GATKTool;
 import org.broadinstitute.hellbender.engine.TraversalParameters;
-import org.broadinstitute.hellbender.engine.spark.datasources.ReferenceMultiSparkSource;
-import org.broadinstitute.hellbender.engine.spark.datasources.ReferenceWindowFunctions;
 import org.broadinstitute.hellbender.engine.filters.ReadFilter;
 import org.broadinstitute.hellbender.engine.filters.WellformedReadFilter;
 import org.broadinstitute.hellbender.engine.spark.datasources.ReadsSparkSink;
 import org.broadinstitute.hellbender.engine.spark.datasources.ReadsSparkSource;
+import org.broadinstitute.hellbender.engine.spark.datasources.ReferenceMultiSparkSource;
+import org.broadinstitute.hellbender.engine.spark.datasources.ReferenceWindowFunctions;
 import org.broadinstitute.hellbender.exceptions.GATKException;
 import org.broadinstitute.hellbender.exceptions.UserException;
 import org.broadinstitute.hellbender.tools.walkers.annotator.Annotation;
-import org.broadinstitute.hellbender.utils.SequenceDictionaryUtils;
-import org.broadinstitute.hellbender.utils.SerializableFunction;
-import org.broadinstitute.hellbender.utils.SimpleInterval;
-import org.broadinstitute.hellbender.utils.Utils;
+import org.broadinstitute.hellbender.utils.*;
 import org.broadinstitute.hellbender.utils.gcs.BucketUtils;
 import org.broadinstitute.hellbender.utils.io.IOUtils;
 import org.broadinstitute.hellbender.utils.read.GATKRead;
@@ -365,8 +363,10 @@ public int getRecommendedNumReducers() {
         if (numReducers != 0) {
             return numReducers;
         }
-        int size = readInputs.keySet().stream().mapToInt(k -> (int) BucketUtils.dirSize(k)).sum();
-        return 1 + (size / getTargetPartitionSize());
+        long size = readInputs.keySet().stream().mapToLong(k -> BucketUtils.dirSize(k)).sum();
+        final int targetPartitionSize = getTargetPartitionSize();
+        return 1 + MathUtils.toIntExactOrThrow(size / targetPartitionSize,
+                                               () -> new GATKException("getRecommendedNumReducers overflowed, size=" + size + " targetPartitionSize=" + targetPartitionSize));
     }
 
     /**

diff --git a/src/main/java/org/broadinstitute/hellbender/utils/MathUtils.java b/src/main/java/org/broadinstitute/hellbender/utils/MathUtils.java
@@ -1545,4 +1545,16 @@ public static double log10BetaBinomialProbability(final int k, final int n, fina
     public static boolean isAProbability(final double p){
         return p >= 0 && p <= 1;
     }
+
+    /**
+     * Convert a long to the exact same int value. If it can't be represented exactly as an int, then throw an exception
+     * produced by the given supplier.
+     */
+    public static int toIntExactOrThrow(long toConvert, Supplier<RuntimeException> exceptionSupplier){
+        if( toConvert == (int)toConvert ){
+            return (int)toConvert;
+        } else {
+            throw exceptionSupplier.get();
+        }
+    }
 }
diff --git a/src/test/java/org/broadinstitute/hellbender/utils/MathUtilsUnitTest.java b/src/test/java/org/broadinstitute/hellbender/utils/MathUtilsUnitTest.java
@@ -8,6 +8,7 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.broadinstitute.hellbender.GATKBaseTest;
+import org.broadinstitute.hellbender.exceptions.GATKException;
 import org.testng.Assert;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -1403,4 +1404,36 @@ public void testLog10BetaBinomialProbability(final int k, final int n, final dou
         final double EPSILON = 1e-3;
         Assert.assertEquals(MathUtils.log10BetaBinomialProbability(k, n, alpha, beta), expected, EPSILON);
     }
+
+    @DataProvider
+    public Object[][] getThrowingCases(){
+        return new Object[][]{
+                {(long)(Integer.MAX_VALUE) + 1L},
+                {(long)(Integer.MIN_VALUE) - 1L},
+                {(long)(Integer.MAX_VALUE) * 100L},
+        };
+
+    }
+
+    @Test(dataProvider = "getThrowingCases", expectedExceptions = GATKException.class)
+    public void testToIntExactOrThrowThrowingCases(long value){
+        MathUtils.toIntExactOrThrow(value, () -> new GATKException("can't convert"));
+    }
+
+    @DataProvider
+    public Object[][] getGoodCases(){
+        return new Object[][]{
+                {(long)Integer.MAX_VALUE},
+                {10000L},
+                {10L},
+                {0L},
+                {-10000L},
+                {(long)Integer.MIN_VALUE}
+        };
+    }
+
+    @Test(dataProvider = "getGoodCases")
+    public void testToIntExactOrThrowGoodCases(long value){
+        Assert.assertEquals(MathUtils.toIntExactOrThrow(value,() -> new GATKException("")), (int)value);
+    }
 }