Limit the length of the FileInputStream.

apache · Sep 8, 2014 · 1dfd3d7 · 1dfd3d7
1 parent 1332156
commit 1dfd3d7
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 13 deletions.
diff --git a/core/src/main/scala/org/apache/spark/network/ManagedBuffer.scala b/core/src/main/scala/org/apache/spark/network/ManagedBuffer.scala
@@ -21,6 +21,7 @@ import java.io.{FileInputStream, RandomAccessFile, File, InputStream}
 import java.nio.ByteBuffer
 import java.nio.channels.FileChannel.MapMode
 
+import com.google.common.io.ByteStreams
 import io.netty.buffer.{ByteBufInputStream, ByteBuf}
 
 import org.apache.spark.util.ByteBufferInputStream
@@ -72,7 +73,7 @@ final class FileSegmentManagedBuffer(val file: File, val offset: Long, val lengt
   override def inputStream(): InputStream = {
     val is = new FileInputStream(file)
     is.skip(offset)
-    is
+    ByteStreams.limit(is, length)
   }
 }
 
@@ -84,7 +85,7 @@ final class NioByteBufferManagedBuffer(buf: ByteBuffer) extends ManagedBuffer {
 
   override def size: Long = buf.remaining()
 
-  override def nioByteBuffer() = buf
+  override def nioByteBuffer() = buf.duplicate()
 
   override def inputStream() = new ByteBufferInputStream(buf)
 }

diff --git a/core/src/main/scala/org/apache/spark/storage/BlockManager.scala b/core/src/main/scala/org/apache/spark/storage/BlockManager.scala
@@ -339,16 +339,10 @@ private[spark] class BlockManager(
    * shuffle blocks. It is safe to do so without a lock on block info since disk store
    * never deletes (recent) items.
    */
-  def getLocalShuffleFromDisk(
-      blockId: BlockId, serializer: Serializer): Option[Iterator[Any]] = {
-
-    val shuffleBlockManager = shuffleManager.shuffleBlockManager
-    val values = shuffleBlockManager.getBytes(blockId.asInstanceOf[ShuffleBlockId]).map(
-      bytes => this.dataDeserialize(blockId, bytes, serializer))
-
-    values.orElse {
-      throw new BlockException(blockId, s"Block $blockId not found on disk, though it should be")
-    }
+  def getLocalShuffleFromDisk(blockId: BlockId, serializer: Serializer): Option[Iterator[Any]] = {
+    val buf = shuffleManager.shuffleBlockManager.getBlockData(blockId.asInstanceOf[ShuffleBlockId])
+    val is = wrapForCompression(blockId, buf.inputStream())
+    Some(serializer.newInstance().deserializeStream(is).asIterator)
   }
 
   /**

diff --git a/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala b/core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala
@@ -112,7 +112,8 @@ final class ShuffleBlockFetcherIterator(
       new BlockFetchingListener {
         override def onBlockFetchSuccess(blockId: String, data: ManagedBuffer): Unit = {
           results.put(new FetchResult(BlockId(blockId), sizeMap(blockId),
-            () => blockManager.dataDeserialize(BlockId(blockId), data.nioByteBuffer(), serializer)
+            () => serializer.newInstance().deserializeStream(
+              blockManager.wrapForCompression(BlockId(blockId), data.inputStream())).asIterator
           ))
           shuffleMetrics.remoteBytesRead += data.size
           shuffleMetrics.remoteBlocksFetched += 1