fix

Author: squito

common/network-common/src/main/java/org/apache/spark/network/server/TransportRequestHandler.java

@@ -234,7 +234,7 @@ public void onComplete(String streamId) throws IOException {
              callback.onSuccess(ByteBuffer.allocate(0));
            } catch (Exception ex) {
              IOException ioExc = new IOException("Failure post-processing complete stream;" +
-               " failing this rpc and leaving channel active");
+               " failing this rpc and leaving channel active", ex);
              callback.onFailure(ioExc);
              streamHandler.onFailure(streamId, ioExc);
            }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/UploadBlockStream.java

@@ -28,7 +28,7 @@
 import static org.apache.spark.network.shuffle.protocol.BlockTransferMessage.Type;
 
 /**
- * A request to Upload a block, which the destintation should receive as a stream.
+ * A request to Upload a block, which the destination should receive as a stream.
  *
  * The actual block data is not contained here.  It will be passed to the StreamCallbackWithID
  * that is returned from RpcHandler.receiveStream()

core/src/main/scala/org/apache/spark/network/netty/NettyBlockTransferService.scala

@@ -152,12 +152,12 @@ private[spark] class NettyBlockTransferService(
     val asStream = blockData.size() > conf.get(config.MAX_REMOTE_BLOCK_SIZE_FETCH_TO_MEM)
     val callback = new RpcResponseCallback {
       override def onSuccess(response: ByteBuffer): Unit = {
-        logTrace(s"Successfully uploaded block $blockId${if (asStream) "as stream" else ""}")
+        logTrace(s"Successfully uploaded block $blockId${if (asStream) " as stream" else ""}")
         result.success((): Unit)
       }
 
       override def onFailure(e: Throwable): Unit = {
-        logError(s"Error while uploading $blockId${if (asStream) "as stream" else ""}", e)
+        logError(s"Error while uploading $blockId${if (asStream) " as stream" else ""}", e)
         result.failure(e)
       }
     }

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -413,8 +413,10 @@ private[spark] class BlockManager(
     // to the final location, but that would require a deeper refactor of this code.  So instead
     // we just write to a temp file, and call putBytes on the data in that file.
     val tmpFile = diskBlockManager.createTempLocalBlock()._2
+    val channel = new CountingWritableChannel(
+      Channels.newChannel(serializerManager.wrapForEncryption(new FileOutputStream(tmpFile))))
+    logTrace(s"Streaming block $blockId to tmp file $tmpFile")
     new StreamCallbackWithID {
-      val channel: WritableByteChannel = Channels.newChannel(new FileOutputStream(tmpFile))
 
       override def getID: String = blockId.name
 
@@ -425,15 +427,27 @@ private[spark] class BlockManager(
       }
 
       override def onComplete(streamId: String): Unit = {
+        logTrace(s"Done receiving block $blockId, now putting into local blockManager")
         // Read the contents of the downloaded file as a buffer to put into the blockManager.
         // Note this is all happening inside the netty thread as soon as it reads the end of the
         // stream.
         channel.close()
         // TODO SPARK-25035 Even if we're only going to write the data to disk after this, we end up
         // using a lot of memory here.  We won't get a jvm OOM, but might get killed by the
         // OS / cluster manager.  We could at least read the tmp file as a stream.
-        val buffer = ChunkedByteBuffer.map(tmpFile,
-          conf.get(config.MEMORY_MAP_LIMIT_FOR_TESTS).toInt)
+        val buffer = securityManager.getIOEncryptionKey() match {
+          case Some(key) =>
+            // we need to pass in the size of the unencrypted block
+            val blockSize = channel.getCount
+            val allocator = level.memoryMode match {
+              case MemoryMode.ON_HEAP => ByteBuffer.allocate _
+              case MemoryMode.OFF_HEAP => Platform.allocateDirectBuffer _
+            }
+            new EncryptedBlockData(tmpFile, blockSize, conf, key).toChunkedByteBuffer(allocator)
+
+          case None =>
+            ChunkedByteBuffer.map(tmpFile, conf.get(config.MEMORY_MAP_LIMIT_FOR_TESTS).toInt)
+        }
         putBytes(blockId, buffer, level)(classTag)
         tmpFile.delete()
       }