viirya commented on code in PR #56878: URL: https://github.com/apache/spark/pull/56878#discussion_r3523595163
########## core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala: ########## @@ -0,0 +1,454 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.streaming + +import java.util.concurrent.{CancellationException, CompletableFuture, CountDownLatch, LinkedBlockingDeque, Semaphore, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} +import javax.annotation.concurrent.NotThreadSafe + +import scala.util.Try + +import io.netty.buffer.{ByteBuf, ByteBufOutputStream, CompositeByteBuf, Unpooled} +import io.netty.channel.{ChannelFuture, ChannelOption} + +import org.apache.spark.{SparkContext, SparkEnv, StreamingShuffleTaskLocation, TaskContext} +import org.apache.spark.internal.LogKeys +import org.apache.spark.internal.config.{EXECUTOR_ID, STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, STREAMING_SHUFFLE_WRITER_MAX_MEMORY} +import org.apache.spark.internal.config.Network.RPC_IO_THREADS +import org.apache.spark.memory.{MemoryConsumer, MemoryMode} +import org.apache.spark.network.TransportContext +import org.apache.spark.network.client.TransportClient +import org.apache.spark.network.netty.SparkTransportConf +import org.apache.spark.network.server.TransportServer +import org.apache.spark.network.shuffle.streaming.{DataMessage, ShuffleChecksum, StreamingShuffleMessage, StreamingShuffleMessageType, TerminationControlMessage} +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.serializer.{JavaSerializerInstance, SerializationStream} +import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriter} +import org.apache.spark.util.{ErrorNotifier, Utils} + +class StreamingShuffleWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + val context: TaskContext, + serverHandler: Option[StreamingShuffleServerHandler] = None, + private[streaming] val errorNotifier: ErrorNotifier = new ErrorNotifier()) + extends ShuffleWriter[K, V] with TaskContextAwareLogging { + assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined) + // Spark params. + private val conf = SparkEnv.get.conf + private val SEND_BUFFER_SIZE: Integer = 32 << 10 // 32 KB + private val RECV_BUFFER_SIZE: Integer = 512 + // The target network buffer size + private val BUFFER_SIZE: Integer = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE) + // The interval at which we flush pending messages. + private val MAX_BUFFERING_TIME_MS = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS) + + // Shuffle details. + private val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K, V, _]] + private val serializerInstance = streamingShuffleHandle.dependency.serializer.newInstance() + private val partitioner = streamingShuffleHandle.dependency.partitioner + private val numPartitions = partitioner.numPartitions + private val shuffleWriterId = context.partitionId() + // Total size of TCP buffers. Use Long math to avoid 32-bit overflow when numPartitions + // is large (numPartitions * buffer sizes can exceed Int.MaxValue). + private val TOTAL_TCPBUF_BYTES: Long = + numPartitions.toLong * (SEND_BUFFER_SIZE + RECV_BUFFER_SIZE) + // Total allowed memory for buffered rows, excluding TCP buffers. + private val MAX_BUFFER_BYTES: Long = math.max(numPartitions.toLong * BUFFER_SIZE * 2, + conf.get(STREAMING_SHUFFLE_WRITER_MAX_MEMORY).toLong - TOTAL_TCPBUF_BYTES) + require(MAX_BUFFER_BYTES >= BUFFER_SIZE && MAX_BUFFER_BYTES <= Int.MaxValue, + s"Streaming shuffle writer memory budget ($MAX_BUFFER_BYTES bytes) is invalid for " + + s"$numPartitions partitions; increase ${STREAMING_SHUFFLE_WRITER_MAX_MEMORY.key} or " + + "reduce the number of partitions.") + private val CHECKSUM_ENABLED = conf.get(STREAMING_SHUFFLE_CHECKSUM_ENABLED) + + // Helper objects. + + // Exposed for testing + private[streaming] val transportServerHandler: StreamingShuffleServerHandler = + serverHandler.getOrElse( + new StreamingShuffleServerHandler( + onTerminationAckReceived, + streamingShuffleHandle.shuffleId, + numPartitions, + context, + errorNotifier)) + + private[streaming] val server: TransportServer = startShuffleServer() + + private val memoryConsumer = new MemoryConsumer( + context.taskMemoryManager(), BUFFER_SIZE.longValue(), MemoryMode.OFF_HEAP) { + // Spilling not supported for simplicity. + override def spill(size: Long, trigger: MemoryConsumer): Long = 0 + } + + // Runtime state. + + // Will reach zero when we've received termination acks from all readers. Public for testing. + private[streaming] val allAcksReceived = new CountDownLatch(numPartitions) + + // Holds per-shard state. Public for testing. + private[streaming] val shards: Array[ShardState] = Array.tabulate(numPartitions)(ShardState(_)) + + private val allocatedBufferBytesSemaphore: Semaphore = new Semaphore(MAX_BUFFER_BYTES.toInt) + + private[streaming] val bufferPool = new LinkedBlockingDeque[ByteBuf]() + + setShuffleIdForLogging(streamingShuffleHandle.shuffleId) + + // Ensure resources are cleaned up on any task completion. + context.addTaskCompletionListener[Unit] { _ => + cleanupResources() + } + + private def startShuffleServer(): TransportServer = { + val role = conf.get(EXECUTOR_ID).map { id => + if (SparkContext.isDriver(id)) "driver" else "executor" + } + + val serverConf = SparkTransportConf.fromSparkConf( + conf, + s"streaming-shuffle-writer-${streamingShuffleHandle.shuffleId}-${shuffleWriterId}", + conf.get(RPC_IO_THREADS).getOrElse(0), // zero will use default number of cores + role) + val serverContext = new TransportContext(serverConf, transportServerHandler) + logInfo(log"Creating shuffle server for shuffle writer" + + log" ${MDC(LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}" + + log" for shuffle ${MDC(LogKeys.SHUFFLE_ID, streamingShuffleHandle.shuffleId)}") + val server = serverContext.createServer() + val hostname = if (SparkEnv.get.rpcEnv.address != null) { + // used and not null when running in an actual cluster but may be null for running tests + SparkEnv.get.rpcEnv.address.host + } else { + Utils.localCanonicalHostName() + } + val tracker = SparkEnv.get.streamingShuffleOutputTracker.get + val taskLocation = + StreamingShuffleTaskLocation(SparkEnv.get.executorId, hostname, server.getPort) + tracker.registerShuffleWriterTask(streamingShuffleHandle.shuffleId, mapId, taskLocation) + logInfo(log"Created shuffle server for writer ${MDC(LogKeys.SHUFFLE_WRITER_ID, + shuffleWriterId)} at ${MDC(LogKeys.TASK_LOCATION, taskLocation)}" + + log" for shuffle ${MDC(LogKeys.SHUFFLE_ID, streamingShuffleHandle.shuffleId)}") + server + } + + /** A buffer with metadata. Not thread safe: only supports single-threaded access. */ + @NotThreadSafe + private[streaming] case class TimestampedBuffer(buffer: ByteBuf) { + val serializationStream: SerializationStream = + serializerInstance.serializeStream(new ByteBufOutputStream(buffer)) + private val creationTimeNs = System.nanoTime() + private val shuffleChecksum = if (CHECKSUM_ENABLED) new ShuffleChecksum() else null + + // Start from beginning to include serialization stream headers + private var lastBufferPosition = 0 + + def totalByteSize(): Long = buffer.readableBytes() + def ageMs(): Long = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - creationTimeNs) + + /* Checksum calculation for order-dependent per-row checksums. */ + def updateChecksum(): Unit = { + if (shuffleChecksum != null) { + val currentPosition = buffer.writerIndex() + val newDataLength = currentPosition - lastBufferPosition + shuffleChecksum.updateChecksum(buffer, lastBufferPosition, newDataLength) + lastBufferPosition = currentPosition + } + } + + def getChecksumValue(): Long = if (shuffleChecksum != null) shuffleChecksum.getValue else 0L + } + + // The state for each shuffle destination. + private[streaming] case class ShardState(id: Int) { + // client may be accessed from other threads via cancel(); @volatile to be safe. + @volatile private var client: Either[TransportClient, CompletableFuture[TransportClient]] = + Right(transportServerHandler.futureClients(id).thenApply(c => { + c.getChannel.config.setOption(ChannelOption.SO_SNDBUF, SEND_BUFFER_SIZE) + c.getChannel.config.setOption(ChannelOption.SO_RCVBUF, RECV_BUFFER_SIZE) + c + })) + val buffer: AtomicReference[TimestampedBuffer] = new AtomicReference(null) + val lastSentSequenceNum: AtomicLong = new AtomicLong(-1) + val terminationAckReceived: AtomicBoolean = new AtomicBoolean(false) + + // send will never block; push back is instead handled by blocking buffer allocation in write + // on `allocatedBufferBytesSemaphore`. All send methods are synchronized to preserve message + // order. + def send(message: StreamingShuffleMessage, done: () => Unit = () => ()): Unit = synchronized { + message.setSeqNum(lastSentSequenceNum.incrementAndGet()) + var buf: CompositeByteBuf = null + try { + buf = server.getPooledByteBufAllocator.compositeBuffer().capacity(message.headerLength()) + message.encode(buf) + } catch { + case e: Throwable => + if (buf != null) buf.release() + throw e + } finally { + message.release() + } + + def sendToClient(client: TransportClient): Unit = { + try { + client.send(buf).addListener((future: ChannelFuture) => { + if (!future.isSuccess) { + errorNotifier.markError(future.cause()) + } + done() + }) + } catch { + case e: Throwable => + buf.release() + done() + errorNotifier.markError(e) + throw e + } + } + + client match { + case Left(c) => + sendToClient(c) + case Right(future) => + // Add another completion stage to ensure queued messages are sent in order. + // If the future is already completed, this will be executed immediately. + val newFuture = future.whenComplete { (client, ex) => + ex match { + case null => sendToClient(client) + case _ => buf.release(); done() + } + } + // Once the future is completed, stop accumulating CompletionStages. + client = if (newFuture.isDone) Left(newFuture.join()) else Right(newFuture) + } + } + + // Sends buffer as a DataMessage to the shuffle reader. Takes ownership of the buffer. + def send(timestampedBuffer: TimestampedBuffer): Unit = synchronized { + timestampedBuffer.serializationStream.close() + val rawBuffer = timestampedBuffer.buffer + val dataSize = rawBuffer.writerIndex() + timestampedBuffer.updateChecksum() + val checksumValue = timestampedBuffer.getChecksumValue() + val dataMessage = new DataMessage(shuffleWriterId, id, dataSize, rawBuffer, checksumValue) + + // We keep a reference to rawBuffer so we can return it to the pool. + send(dataMessage, () => { + if (rawBuffer.refCnt() != 1) { + // Throw an internal exception for unexpected state. + errorNotifier.markError(new AssertionError( + s"INTERNAL ERROR: Unexpected refcnt ${rawBuffer.refCnt()}")) + rawBuffer.release() + } else { + rawBuffer.clear() + if (context.isFailed() || context.isCompleted() || rawBuffer.capacity() != BUFFER_SIZE) { + rawBuffer.release() + } else { + bufferPool.offerLast(rawBuffer) + } + allocatedBufferBytesSemaphore.release(BUFFER_SIZE) + } + }) + } + + // Consume the current buffer, if it exists, and send it as a DataMessage. + def send(): Unit = synchronized { + val b = takeBuffer() + if (b != null) send(b) + } + + def takeBuffer(): TimestampedBuffer = buffer.getAndSet(null) + + def putBuffer(b: TimestampedBuffer): Unit = assert(buffer.getAndSet(b) == null) + + def close(): Unit = { + send() + send(new TerminationControlMessage(shuffleWriterId, id)) + } + + def cancel(): Unit = { + val error = context.getTaskFailure.getOrElse(new CancellationException()) + transportServerHandler.futureClients(id).completeExceptionally(error) + client.foreach(_.completeExceptionally(error)) + Option(takeBuffer()).foreach(_.buffer.release()) + } + + // For testing only. + def hasClient: Boolean = client match { + case Left(_) => true + case Right(future) => future.isDone + } + } + + /** Close this writer, passing along whether the map completed */ + override def stop(success: Boolean): Option[MapStatus] = { + // No-op: the streaming shuffle lifecycle is handled elsewhere. write() blocks until all + // readers ack termination on the normal path, and the task-completion listener + // (cleanupResources) closes the server and releases buffers on both success and failure. + // Streaming shuffle does not rely on map output status, so just return a placeholder. + Some(MapStatus( + SparkEnv.get.blockManager.shuffleServerId, + Array.fill(numPartitions)(0L), + mapId)) + } + + /** Get the lengths of each partition */ + override def getPartitionLengths(): Array[Long] = { + Array.fill(numPartitions)(0L) + } + + /** + * Invoked on a Netty event-loop thread by [[StreamingShuffleServerHandler]] when a reader's + * termination ack arrives. Validates the reader's last-seen sequence number against what this + * writer sent; on a mismatch it throws STREAMING_SHUFFLE_INCORRECT_SEQUENCE_NUMBER, which the + * handler captures via the shared [[ErrorNotifier]]. The per-partition latch decrement is + * idempotent, so duplicate acks are ignored. Exposed for testing. + */ + private[streaming] def onTerminationAckReceived( + partitionId: Int, lastSeqNumSeenByReader: Long): Unit = { + val lastSeqNumSent = shards(partitionId).lastSentSequenceNum.get() + if (lastSeqNumSent != lastSeqNumSeenByReader) { + throw StreamingShuffleManager.streamingShuffleIncorrectSequenceNumber( + StreamingShuffleMessageType.TERMINATION_ACK_MESSAGE, + shuffleWriterId, + partitionId, + lastSeqNumSent, + lastSeqNumSeenByReader) + } + if (shards(partitionId).terminationAckReceived.compareAndSet(false, true)) { + allAcksReceived.countDown() + } + val receivedAcks = numPartitions - allAcksReceived.getCount.toInt + logInfo(log"Received termination ack from reader ${MDC(LogKeys.SHUFFLE_READER_ID, + partitionId)}. Now have ${MDC(LogKeys.NUM_TERMINATION_ACKS, receivedAcks)} / ${MDC( + LogKeys.NUM_SHUFFLE_READERS, numPartitions)} termination acks") + } + + /** + * Cleans up all writer resources. + * This method should be idempotent. + */ + private[streaming] def cleanupResources(): Unit = { + val cleanupStartTime = System.currentTimeMillis() + Utils.tryLogNonFatalError { + shards.foreach(_.cancel()) + } + Utils.tryLogNonFatalError { + server.close() + } + Utils.tryLogNonFatalError { + val list = new java.util.ArrayList[ByteBuf]() + bufferPool.drainTo(list) + list.forEach(buf => { buf.release(); () }) + } + Utils.tryLogNonFatalError { + memoryConsumer.freeMemory(memoryConsumer.getUsed()) + } + logInfo(log"Resource cleanup took ${MDC(LogKeys.DURATION, + System.currentTimeMillis() - cleanupStartTime)} ms") + } + + private def throwErrorIfExists(): Unit = { + context.getTaskFailure.foreach { throw _ } + errorNotifier.throwErrorIfExists() + } + + private def newBuffer(): TimestampedBuffer = { + if (!allocatedBufferBytesSemaphore.tryAcquire(BUFFER_SIZE, 10, TimeUnit.MICROSECONDS)) { + shards.foreach(_.send()) + while (!allocatedBufferBytesSemaphore.tryAcquire(BUFFER_SIZE, 10, TimeUnit.MILLISECONDS)) { + throwErrorIfExists() + } + } + val buffer = bufferPool.pollLast() + TimestampedBuffer(if (buffer != null) buffer else Unpooled.directBuffer(BUFFER_SIZE)) + } + + /** + * Write a sequence of records to downstream shuffle readers. + * + * For each record, the reader partition is determined using the key and the + * partitioner. Multiple rows can be packed into a single DataMessage; the maximum + * number of rows that can be packed depends on the STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE + * config. Each DataMessage is sent to the reader for its partition over that reader's + * Netty connection. + */ + override def write(records: Iterator[Product2[K, V]]): Unit = { + val isWriteFinished = new CountDownLatch(1) + val flushThread = new Thread(() => + Try { + while (!isWriteFinished.await(MAX_BUFFERING_TIME_MS, TimeUnit.MILLISECONDS)) + shards.foreach(_.send()) + }.recover { case e => errorNotifier.markError(e) } + , "time-based-flush-for-shuffle-writer-" + + s"${streamingShuffleHandle.shuffleId}-${shuffleWriterId}") + try { + // Reserve the budget with the task memory manager for accounting/visibility. In-flight + // buffer memory is bounded by allocatedBufferBytesSemaphore; we ignore the return value + // because we cannot act on a partial grant here (this consumer cannot spill). + memoryConsumer.acquireMemory(TOTAL_TCPBUF_BYTES + MAX_BUFFER_BYTES) Review Comment: `writerMaxMemory` ends up being best-effort accounting rather than a bound on in-flight buffer memory. The back-pressure semaphore acquires/releases a fixed `BUFFER_SIZE` per buffer, but the TODO a few lines below acknowledges a single serialized row can push a buffer past `BUFFER_SIZE` (no spanning-row splitting), so the byte count the semaphore tracks can diverge from the direct memory actually held. And here the `acquireMemory(...)` return value is intentionally ignored (the consumer can't spill), so a partial grant isn't acted on. If large rows are possible under RTM this is a real OOM/latency risk, not just a doc nit. Either document clearly that the limit is best-effort and large rows can exceed it, or note the RTM row-size ceiling that makes it safe in practice. ########## core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala: ########## @@ -0,0 +1,454 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.streaming + +import java.util.concurrent.{CancellationException, CompletableFuture, CountDownLatch, LinkedBlockingDeque, Semaphore, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} +import javax.annotation.concurrent.NotThreadSafe + +import scala.util.Try + +import io.netty.buffer.{ByteBuf, ByteBufOutputStream, CompositeByteBuf, Unpooled} +import io.netty.channel.{ChannelFuture, ChannelOption} + +import org.apache.spark.{SparkContext, SparkEnv, StreamingShuffleTaskLocation, TaskContext} +import org.apache.spark.internal.LogKeys +import org.apache.spark.internal.config.{EXECUTOR_ID, STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, STREAMING_SHUFFLE_WRITER_MAX_MEMORY} +import org.apache.spark.internal.config.Network.RPC_IO_THREADS +import org.apache.spark.memory.{MemoryConsumer, MemoryMode} +import org.apache.spark.network.TransportContext +import org.apache.spark.network.client.TransportClient +import org.apache.spark.network.netty.SparkTransportConf +import org.apache.spark.network.server.TransportServer +import org.apache.spark.network.shuffle.streaming.{DataMessage, ShuffleChecksum, StreamingShuffleMessage, StreamingShuffleMessageType, TerminationControlMessage} +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.serializer.{JavaSerializerInstance, SerializationStream} +import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriter} +import org.apache.spark.util.{ErrorNotifier, Utils} + +class StreamingShuffleWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + val context: TaskContext, + serverHandler: Option[StreamingShuffleServerHandler] = None, + private[streaming] val errorNotifier: ErrorNotifier = new ErrorNotifier()) + extends ShuffleWriter[K, V] with TaskContextAwareLogging { + assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined) + // Spark params. + private val conf = SparkEnv.get.conf + private val SEND_BUFFER_SIZE: Integer = 32 << 10 // 32 KB + private val RECV_BUFFER_SIZE: Integer = 512 + // The target network buffer size + private val BUFFER_SIZE: Integer = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE) + // The interval at which we flush pending messages. + private val MAX_BUFFERING_TIME_MS = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS) + + // Shuffle details. + private val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K, V, _]] + private val serializerInstance = streamingShuffleHandle.dependency.serializer.newInstance() + private val partitioner = streamingShuffleHandle.dependency.partitioner + private val numPartitions = partitioner.numPartitions + private val shuffleWriterId = context.partitionId() + // Total size of TCP buffers. Use Long math to avoid 32-bit overflow when numPartitions + // is large (numPartitions * buffer sizes can exceed Int.MaxValue). + private val TOTAL_TCPBUF_BYTES: Long = + numPartitions.toLong * (SEND_BUFFER_SIZE + RECV_BUFFER_SIZE) + // Total allowed memory for buffered rows, excluding TCP buffers. + private val MAX_BUFFER_BYTES: Long = math.max(numPartitions.toLong * BUFFER_SIZE * 2, Review Comment: When partition count is high the first term of this `math.max` wins, so the effective budget silently exceeds the configured `writerMaxMemory` (only the `Int.MaxValue` guard in the `require` below caps it). Worth logging the effective budget when the floor overrides the configured value, so an operator can see the limit they set isn't the one in force. ########## core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala: ########## @@ -0,0 +1,454 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.shuffle.streaming + +import java.util.concurrent.{CancellationException, CompletableFuture, CountDownLatch, LinkedBlockingDeque, Semaphore, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} +import javax.annotation.concurrent.NotThreadSafe + +import scala.util.Try + +import io.netty.buffer.{ByteBuf, ByteBufOutputStream, CompositeByteBuf, Unpooled} +import io.netty.channel.{ChannelFuture, ChannelOption} + +import org.apache.spark.{SparkContext, SparkEnv, StreamingShuffleTaskLocation, TaskContext} +import org.apache.spark.internal.LogKeys +import org.apache.spark.internal.config.{EXECUTOR_ID, STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, STREAMING_SHUFFLE_WRITER_MAX_MEMORY} +import org.apache.spark.internal.config.Network.RPC_IO_THREADS +import org.apache.spark.memory.{MemoryConsumer, MemoryMode} +import org.apache.spark.network.TransportContext +import org.apache.spark.network.client.TransportClient +import org.apache.spark.network.netty.SparkTransportConf +import org.apache.spark.network.server.TransportServer +import org.apache.spark.network.shuffle.streaming.{DataMessage, ShuffleChecksum, StreamingShuffleMessage, StreamingShuffleMessageType, TerminationControlMessage} +import org.apache.spark.scheduler.MapStatus +import org.apache.spark.serializer.{JavaSerializerInstance, SerializationStream} +import org.apache.spark.shuffle.{ShuffleHandle, ShuffleWriter} +import org.apache.spark.util.{ErrorNotifier, Utils} + +class StreamingShuffleWriter[K, V]( + handle: ShuffleHandle, + mapId: Long, + val context: TaskContext, + serverHandler: Option[StreamingShuffleServerHandler] = None, + private[streaming] val errorNotifier: ErrorNotifier = new ErrorNotifier()) + extends ShuffleWriter[K, V] with TaskContextAwareLogging { + assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined) + // Spark params. + private val conf = SparkEnv.get.conf + private val SEND_BUFFER_SIZE: Integer = 32 << 10 // 32 KB + private val RECV_BUFFER_SIZE: Integer = 512 + // The target network buffer size + private val BUFFER_SIZE: Integer = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE) + // The interval at which we flush pending messages. + private val MAX_BUFFERING_TIME_MS = conf.get(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS) + + // Shuffle details. + private val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K, V, _]] + private val serializerInstance = streamingShuffleHandle.dependency.serializer.newInstance() + private val partitioner = streamingShuffleHandle.dependency.partitioner + private val numPartitions = partitioner.numPartitions + private val shuffleWriterId = context.partitionId() + // Total size of TCP buffers. Use Long math to avoid 32-bit overflow when numPartitions + // is large (numPartitions * buffer sizes can exceed Int.MaxValue). + private val TOTAL_TCPBUF_BYTES: Long = + numPartitions.toLong * (SEND_BUFFER_SIZE + RECV_BUFFER_SIZE) + // Total allowed memory for buffered rows, excluding TCP buffers. + private val MAX_BUFFER_BYTES: Long = math.max(numPartitions.toLong * BUFFER_SIZE * 2, + conf.get(STREAMING_SHUFFLE_WRITER_MAX_MEMORY).toLong - TOTAL_TCPBUF_BYTES) + require(MAX_BUFFER_BYTES >= BUFFER_SIZE && MAX_BUFFER_BYTES <= Int.MaxValue, + s"Streaming shuffle writer memory budget ($MAX_BUFFER_BYTES bytes) is invalid for " + + s"$numPartitions partitions; increase ${STREAMING_SHUFFLE_WRITER_MAX_MEMORY.key} or " + + "reduce the number of partitions.") + private val CHECKSUM_ENABLED = conf.get(STREAMING_SHUFFLE_CHECKSUM_ENABLED) + + // Helper objects. + + // Exposed for testing + private[streaming] val transportServerHandler: StreamingShuffleServerHandler = + serverHandler.getOrElse( + new StreamingShuffleServerHandler( + onTerminationAckReceived, + streamingShuffleHandle.shuffleId, + numPartitions, + context, + errorNotifier)) + + private[streaming] val server: TransportServer = startShuffleServer() + + private val memoryConsumer = new MemoryConsumer( + context.taskMemoryManager(), BUFFER_SIZE.longValue(), MemoryMode.OFF_HEAP) { + // Spilling not supported for simplicity. + override def spill(size: Long, trigger: MemoryConsumer): Long = 0 + } + + // Runtime state. + + // Will reach zero when we've received termination acks from all readers. Public for testing. + private[streaming] val allAcksReceived = new CountDownLatch(numPartitions) + + // Holds per-shard state. Public for testing. + private[streaming] val shards: Array[ShardState] = Array.tabulate(numPartitions)(ShardState(_)) + + private val allocatedBufferBytesSemaphore: Semaphore = new Semaphore(MAX_BUFFER_BYTES.toInt) + + private[streaming] val bufferPool = new LinkedBlockingDeque[ByteBuf]() + + setShuffleIdForLogging(streamingShuffleHandle.shuffleId) + + // Ensure resources are cleaned up on any task completion. + context.addTaskCompletionListener[Unit] { _ => + cleanupResources() + } + + private def startShuffleServer(): TransportServer = { + val role = conf.get(EXECUTOR_ID).map { id => + if (SparkContext.isDriver(id)) "driver" else "executor" + } + + val serverConf = SparkTransportConf.fromSparkConf( + conf, + s"streaming-shuffle-writer-${streamingShuffleHandle.shuffleId}-${shuffleWriterId}", + conf.get(RPC_IO_THREADS).getOrElse(0), // zero will use default number of cores + role) + val serverContext = new TransportContext(serverConf, transportServerHandler) + logInfo(log"Creating shuffle server for shuffle writer" + + log" ${MDC(LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}" + + log" for shuffle ${MDC(LogKeys.SHUFFLE_ID, streamingShuffleHandle.shuffleId)}") + val server = serverContext.createServer() + val hostname = if (SparkEnv.get.rpcEnv.address != null) { + // used and not null when running in an actual cluster but may be null for running tests + SparkEnv.get.rpcEnv.address.host + } else { + Utils.localCanonicalHostName() + } + val tracker = SparkEnv.get.streamingShuffleOutputTracker.get + val taskLocation = + StreamingShuffleTaskLocation(SparkEnv.get.executorId, hostname, server.getPort) + tracker.registerShuffleWriterTask(streamingShuffleHandle.shuffleId, mapId, taskLocation) + logInfo(log"Created shuffle server for writer ${MDC(LogKeys.SHUFFLE_WRITER_ID, + shuffleWriterId)} at ${MDC(LogKeys.TASK_LOCATION, taskLocation)}" + + log" for shuffle ${MDC(LogKeys.SHUFFLE_ID, streamingShuffleHandle.shuffleId)}") + server + } + + /** A buffer with metadata. Not thread safe: only supports single-threaded access. */ + @NotThreadSafe + private[streaming] case class TimestampedBuffer(buffer: ByteBuf) { + val serializationStream: SerializationStream = + serializerInstance.serializeStream(new ByteBufOutputStream(buffer)) + private val creationTimeNs = System.nanoTime() + private val shuffleChecksum = if (CHECKSUM_ENABLED) new ShuffleChecksum() else null + + // Start from beginning to include serialization stream headers + private var lastBufferPosition = 0 + + def totalByteSize(): Long = buffer.readableBytes() + def ageMs(): Long = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - creationTimeNs) + + /* Checksum calculation for order-dependent per-row checksums. */ + def updateChecksum(): Unit = { + if (shuffleChecksum != null) { + val currentPosition = buffer.writerIndex() + val newDataLength = currentPosition - lastBufferPosition + shuffleChecksum.updateChecksum(buffer, lastBufferPosition, newDataLength) + lastBufferPosition = currentPosition + } + } + + def getChecksumValue(): Long = if (shuffleChecksum != null) shuffleChecksum.getValue else 0L + } + + // The state for each shuffle destination. + private[streaming] case class ShardState(id: Int) { + // client may be accessed from other threads via cancel(); @volatile to be safe. + @volatile private var client: Either[TransportClient, CompletableFuture[TransportClient]] = + Right(transportServerHandler.futureClients(id).thenApply(c => { + c.getChannel.config.setOption(ChannelOption.SO_SNDBUF, SEND_BUFFER_SIZE) + c.getChannel.config.setOption(ChannelOption.SO_RCVBUF, RECV_BUFFER_SIZE) + c + })) + val buffer: AtomicReference[TimestampedBuffer] = new AtomicReference(null) + val lastSentSequenceNum: AtomicLong = new AtomicLong(-1) + val terminationAckReceived: AtomicBoolean = new AtomicBoolean(false) + + // send will never block; push back is instead handled by blocking buffer allocation in write + // on `allocatedBufferBytesSemaphore`. All send methods are synchronized to preserve message + // order. + def send(message: StreamingShuffleMessage, done: () => Unit = () => ()): Unit = synchronized { + message.setSeqNum(lastSentSequenceNum.incrementAndGet()) + var buf: CompositeByteBuf = null + try { + buf = server.getPooledByteBufAllocator.compositeBuffer().capacity(message.headerLength()) + message.encode(buf) + } catch { + case e: Throwable => + if (buf != null) buf.release() + throw e + } finally { + message.release() + } + + def sendToClient(client: TransportClient): Unit = { + try { + client.send(buf).addListener((future: ChannelFuture) => { + if (!future.isSuccess) { + errorNotifier.markError(future.cause()) + } + done() + }) + } catch { + case e: Throwable => + buf.release() + done() + errorNotifier.markError(e) + throw e + } + } + + client match { + case Left(c) => + sendToClient(c) + case Right(future) => + // Add another completion stage to ensure queued messages are sent in order. + // If the future is already completed, this will be executed immediately. + val newFuture = future.whenComplete { (client, ex) => + ex match { + case null => sendToClient(client) + case _ => buf.release(); done() + } + } + // Once the future is completed, stop accumulating CompletionStages. + client = if (newFuture.isDone) Left(newFuture.join()) else Right(newFuture) + } + } + + // Sends buffer as a DataMessage to the shuffle reader. Takes ownership of the buffer. + def send(timestampedBuffer: TimestampedBuffer): Unit = synchronized { + timestampedBuffer.serializationStream.close() + val rawBuffer = timestampedBuffer.buffer + val dataSize = rawBuffer.writerIndex() + timestampedBuffer.updateChecksum() + val checksumValue = timestampedBuffer.getChecksumValue() + val dataMessage = new DataMessage(shuffleWriterId, id, dataSize, rawBuffer, checksumValue) + + // We keep a reference to rawBuffer so we can return it to the pool. + send(dataMessage, () => { + if (rawBuffer.refCnt() != 1) { + // Throw an internal exception for unexpected state. + errorNotifier.markError(new AssertionError( + s"INTERNAL ERROR: Unexpected refcnt ${rawBuffer.refCnt()}")) + rawBuffer.release() + } else { + rawBuffer.clear() + if (context.isFailed() || context.isCompleted() || rawBuffer.capacity() != BUFFER_SIZE) { + rawBuffer.release() + } else { + bufferPool.offerLast(rawBuffer) + } + allocatedBufferBytesSemaphore.release(BUFFER_SIZE) + } + }) + } + + // Consume the current buffer, if it exists, and send it as a DataMessage. + def send(): Unit = synchronized { + val b = takeBuffer() + if (b != null) send(b) + } + + def takeBuffer(): TimestampedBuffer = buffer.getAndSet(null) + + def putBuffer(b: TimestampedBuffer): Unit = assert(buffer.getAndSet(b) == null) + + def close(): Unit = { + send() + send(new TerminationControlMessage(shuffleWriterId, id)) + } + + def cancel(): Unit = { + val error = context.getTaskFailure.getOrElse(new CancellationException()) + transportServerHandler.futureClients(id).completeExceptionally(error) + client.foreach(_.completeExceptionally(error)) + Option(takeBuffer()).foreach(_.buffer.release()) + } + + // For testing only. + def hasClient: Boolean = client match { + case Left(_) => true + case Right(future) => future.isDone + } + } + + /** Close this writer, passing along whether the map completed */ + override def stop(success: Boolean): Option[MapStatus] = { + // No-op: the streaming shuffle lifecycle is handled elsewhere. write() blocks until all + // readers ack termination on the normal path, and the task-completion listener + // (cleanupResources) closes the server and releases buffers on both success and failure. + // Streaming shuffle does not rely on map output status, so just return a placeholder. + Some(MapStatus( Review Comment: This placeholder `MapStatus` reports all-zero partition lengths, and it's registered unconditionally into `MapOutputTrackerMaster` via the normal `ShuffleWriteProcessor` → `DAGScheduler.handleTaskCompletion` path. `MapStatus`'s contract says a non-empty block *must* report a non-zero size, otherwise fetchers are allowed to silently skip it. For streaming shuffle this is presumably inert because readers locate writers through `StreamingShuffleOutputTracker` and pull over Netty, never consulting MapStatus block sizes — but that can't be verified in this PR (reader/activation parts aren't merged). Could you confirm, ideally in a comment here, that no reducer path ever consumes these lengths? As written, if anything ever falls back to standard block fetch, this is silent data loss rather than a failure. -- 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. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
