From fdcac0838e5f4d53c2fec0ff815f0f26a568357f Mon Sep 17 00:00:00 2001
From: Josh Rosen <joshrosen@databricks.com>
Date: Tue, 12 May 2015 12:14:42 -0700
Subject: [PATCH] Guard against overflow when expanding sort buffer.

---
 .../shuffle/unsafe/UnsafeShuffleInMemorySorter.java    | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java b/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
index 57b24547125f6..b4055141a4ec5 100644
--- a/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
+++ b/core/src/main/java/org/apache/spark/shuffle/unsafe/UnsafeShuffleInMemorySorter.java
@@ -47,7 +47,9 @@ public UnsafeShuffleInMemorySorter(int initialSize) {
 
   public void expandSortBuffer() {
     final long[] oldBuffer = sortBuffer;
-    sortBuffer = new long[oldBuffer.length * 2];
+    // Guard against overflow:
+    final int newLength = oldBuffer.length * 2 > 0 ? (oldBuffer.length * 2) : Integer.MAX_VALUE;
+    sortBuffer = new long[newLength];
     System.arraycopy(oldBuffer, 0, sortBuffer, 0, oldBuffer.length);
   }
 
@@ -71,7 +73,11 @@ public long getMemoryUsage() {
    */
   public void insertRecord(long recordPointer, int partitionId) {
     if (!hasSpaceForAnotherRecord()) {
-      expandSortBuffer();
+      if (sortBuffer.length == Integer.MAX_VALUE) {
+        throw new IllegalStateException("Sort buffer has reached maximum size");
+      } else {
+        expandSortBuffer();
+      }
     }
     sortBuffer[sortBufferInsertPosition] =
         PackedRecordPointer.packPointer(recordPointer, partitionId);