squito commented on a change in pull request #23453: [SPARK-26089][CORE] Handle 
corruption in large shuffle blocks
URL: https://github.com/apache/spark/pull/23453#discussion_r264275165
 
 

 ##########
 File path: 
core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala
 ##########
 @@ -466,16 +469,19 @@ final class ShuffleBlockFetcherIterator(
           var isStreamCopied: Boolean = false
           try {
             input = streamWrapper(blockId, in)
-            // Only copy the stream if it's wrapped by compression or 
encryption, also the size of
-            // block is small (the decompressed block is smaller than 
maxBytesInFlight)
-            if (detectCorrupt && !input.eq(in) && size < maxBytesInFlight / 3) 
{
+            // Only copy the stream if it's wrapped by compression or 
encryption upto a size of
+            // maxBytesInFlight/3. If stream is longer, then corruption will 
be caught while reading
+            // the stream.
+            streamCompressedOrEncrypted = !input.eq(in)
+            if (streamCompressedOrEncrypted && detectCorruptUseExtraMemory) {
               isStreamCopied = true
-              val out = new ChunkedByteBufferOutputStream(64 * 1024, 
ByteBuffer.allocate)
-              // Decompress the whole block at once to detect any corruption, 
which could increase
-              // the memory usage tne potential increase the chance of OOM.
+              // Decompress the block upto maxBytesInFlight/3 at once to 
detect any corruption which
+              // could increase the memory usage and potentially increase the 
chance of OOM.
               // TODO: manage the memory used here, and spill it into disk in 
case of OOM.
-              Utils.copyStream(input, out, closeStreams = true)
-              input = out.toChunkedByteBuffer.toInputStream(dispose = true)
+              val (fullyCopied: Boolean, mergedStream: InputStream) = 
Utils.copyStreamUpTo(
+                input, maxBytesInFlight / 3)
+              isStreamCopied = fullyCopied
+              input = mergedStream
 
 Review comment:
   not related to your changes, but while you're touching this file, can you 
add 2 more spaces of indentation to `|| corruptedBlocks.contains(blockId)) {` a 
couple of lines below this?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
[email protected]


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to