Skip to content

[SPARK-29072][CORE] Put back usage of TimeTrackingOutputStream for UnsafeShuffleWriter and SortShuffleWriter #25780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

[SPARK-29072][CORE] Put back usage of TimeTrackingOutputStream for UnsafeShuffleWriter and SortShuffleWriter #25780

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7eba87a
[SPARK-29072] Put back usage of TimeTrackingOutputStream for UnsafeSh…
mccheah Sep 13, 2019
fe5b93f
Import ordering
mccheah Sep 13, 2019
08a6068
Fix stream reference
mccheah Sep 13, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions core/src/main/java/org/apache/spark/shuffle/sort/UnsafeShuffleWriter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@
import org.apache.spark.shuffle.api.SingleSpillShuffleMapOutputWriter;
import org.apache.spark.shuffle.api.WritableByteChannelWrapper;
import org.apache.spark.storage.BlockManager;
import org.apache.spark.storage.TimeTrackingOutputStream;
import org.apache.spark.unsafe.Platform;
import org.apache.spark.util.Utils;

Expand Down Expand Up @@ -382,6 +383,7 @@ private void mergeSpillsWithFileStream(
ShufflePartitionWriter writer = mapWriter.getPartitionWriter(partition);
OutputStream partitionOutput = writer.openStream();
try {
partitionOutput = new TimeTrackingOutputStream(writeMetrics, partitionOutput);
partitionOutput = blockManager.serializerManager().wrapForEncryption(partitionOutput);
if (compressionCodec != null) {
partitionOutput = compressionCodec.compressedOutputStream(partitionOutput);
Expand Down
15 changes: 12 additions & 3 deletions core/src/main/scala/org/apache/spark/shuffle/ShufflePartitionPairsWriter.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import java.io.{Closeable, IOException, OutputStream}

import org.apache.spark.serializer.{SerializationStream, SerializerInstance, SerializerManager}
import org.apache.spark.shuffle.api.ShufflePartitionWriter
import org.apache.spark.storage.BlockId
import org.apache.spark.storage.{BlockId, TimeTrackingOutputStream}
import org.apache.spark.util.Utils
import org.apache.spark.util.collection.PairsWriter

Expand All @@ -39,6 +39,7 @@ private[spark] class ShufflePartitionPairsWriter(

private var isClosed = false
private var partitionStream: OutputStream = _
private var timeTrackingStream: OutputStream = _
private var wrappedStream: OutputStream = _
private var objOut: SerializationStream = _
private var numRecordsWritten = 0
Expand All @@ -59,7 +60,8 @@ private[spark] class ShufflePartitionPairsWriter(
private def open(): Unit = {
try {
partitionStream = partitionWriter.openStream
wrappedStream = serializerManager.wrapStream(blockId, partitionStream)
timeTrackingStream = new TimeTrackingOutputStream(writeMetrics, partitionStream)
wrappedStream = serializerManager.wrapStream(blockId, timeTrackingStream)
objOut = serializerInstance.serializeStream(wrappedStream)
} catch {
case e: Exception =>
Expand All @@ -78,6 +80,7 @@ private[spark] class ShufflePartitionPairsWriter(
// Setting these to null will prevent the underlying streams from being closed twice
// just in case any stream's close() implementation is not idempotent.
wrappedStream = null
timeTrackingStream = null
partitionStream = null
} {
// Normally closing objOut would close the inner streams as well, but just in case there
Expand All @@ -86,9 +89,15 @@ private[spark] class ShufflePartitionPairsWriter(
wrappedStream = closeIfNonNull(wrappedStream)
// Same as above - if wrappedStream closes then assume it closes underlying
// partitionStream and don't close again in the finally
timeTrackingStream = null
partitionStream = null
} {
partitionStream = closeIfNonNull(partitionStream)
Utils.tryWithSafeFinally {
timeTrackingStream = closeIfNonNull(timeTrackingStream)
partitionStream = null
} {
partitionStream = closeIfNonNull(partitionStream)
}
}
}
updateBytesWritten()
Expand Down