This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 490bda2529b9 [SPARK-57229][SS][RTM][STREAMINGSHUFFLE][PART4] Add
StreamingShuffleWriter + server-side Netty handler
490bda2529b9 is described below
commit 490bda2529b9cf91efdcdd3c7bc2da1992d0eefc
Author: Boyang Jerry Peng <[email protected]>
AuthorDate: Fri Jul 10 23:02:17 2026 -0700
[SPARK-57229][SS][RTM][STREAMINGSHUFFLE][PART4] Add StreamingShuffleWriter
+ server-side Netty handler
### What changes were proposed in this pull request?
This is **part 4** of a multi-PR effort to add *streaming shuffle* to Spark
— a push-based shuffle used by Real-Time Mode (RTM) structured streaming, where
writer tasks push records directly to reader tasks over the network instead of
writing map output to disk for readers to pull.
This PR adds the **writer (push) side**:
- **`StreamingShuffleWriter`** — a `ShuffleWriter` that pushes serialized
records to the downstream readers over Netty instead of writing shuffle files
to disk.
- **`StreamingShuffleServerHandler`** — the writer-side Netty `RpcHandler`
that accepts reader connections and processes the messages readers send back.
- **`StreamingShuffleManager.getWriter`** — now returns a
`StreamingShuffleWriter` (it was a stub that threw
`UnsupportedOperationException` in part 3).
- **`TransportServer.getPooledByteBufAllocator`** — a small accessor so the
writer can allocate pooled send buffers from the server's allocator.
- Three writer configs and two structured-logging keys (see below).
#### How the writer and reader talk to each other
Each writer task starts its own Netty server. Readers connect to it, and
the two exchange four message types (all defined in part 1). The exchange for
one writer <-> one reader looks like this:
```
Reader Writer (this PR)
| |
| --- CreditControlMessage -----------> | reader connects; writer
now knows
| | which network channel
maps to this reader
| |
| <-------------- DataMessage (seq=0) --- | writer pushes batched
rows
| <-------------- DataMessage (seq=1) --- | (each message carries a
sequence
| <-------------- DataMessage (seq=2) --- | number and an optional
CRC32C checksum)
| ... |
| <----- TerminationControlMessage ------ | writer has sent
everything (seq=N)
| |
| --- TerminationAckMessage (seq=N) ---> | reader confirms it
received up to seq=N
| |
```
Key points of the protocol, from the writer's side:
1. **Connection setup.** A reader opens a connection and sends a
`CreditControlMessage`. Until that arrives the writer does not know which
channel belongs to which reader, so it queues outgoing messages for that reader
and flushes them once the reader is known.
2. **Data push.** The writer partitions each record with the shuffle
dependency's partitioner, serializes rows into a pooled buffer, and sends the
buffer as a `DataMessage` once it fills up
(`spark.shuffle.streaming.networkBufferSize`, default 32 KB) or after a maximum
buffering interval (`spark.shuffle.streaming.networkBufferMaxWaitTimeMs`,
default 50 ms). Every message this writer sends to a given reader carries a
**monotonically increasing sequence number**, so the reader can detec [...]
3. **Integrity check (optional).** When
`spark.shuffle.streaming.checksum.enabled` is on, the writer embeds a CRC32C
over the buffer that the reader re-computes and compares.
4. **Back-pressure.** In-flight buffer memory is bounded
(`spark.shuffle.streaming.writerMaxMemory`, default 32 MB) by an off-heap
`MemoryConsumer` plus a byte semaphore; when the limit is reached the writer
blocks the upstream iterator until buffers are freed by send completions.
5. **Termination.** After the last record, the writer sends each reader a
`TerminationControlMessage` carrying the final sequence number, then waits for
every reader to reply with a `TerminationAckMessage`. On each ack the writer
checks that the reader's last-seen sequence number matches what it sent (a
mismatch fails the task), and only returns from `write()` once all readers have
acknowledged.
Errors that occur on Netty threads (e.g. a failed write) are recorded via
`ErrorNotifier` (part 3.5) and re-thrown on the task thread at a safe point.
All resources (the server, channels, pooled buffers, reserved memory) are
released from a `TaskContext` completion listener, so they are cleaned up on
both success and failure.
New configs (all `internal`, default on the safe side):
| Config | Default | Purpose |
|---|---|---|
| `spark.shuffle.streaming.networkBufferSize` | 32 KB | target size of each
pushed buffer |
| `spark.shuffle.streaming.networkBufferMaxWaitTimeMs` | 50 ms | max time a
partial buffer is held before flushing |
| `spark.shuffle.streaming.writerMaxMemory` | 32 MB | best-effort cap on
the writer's in-flight buffer memory |
New log keys: `NUM_SHUFFLE_READERS`, `SHUFFLE_READER_ID`.
The full PR stack:
- **Part 1** (SPARK-56674, *merged*) - streaming shuffle wire protocol (the
four Netty message types above).
- **Part 2** (SPARK-56962, *merged*) - `StreamingShuffleOutputTracker`
(driver-side writer-location coordination).
- **Part 3** (SPARK-57141, *merged*) - shuffle-manager layer
(`StreamingShuffleManager` + `MultiShuffleManager`).
- **Part 3.5** (SPARK-57337, *merged*) - shared transport + error plumbing
(`ErrorNotifier`, `TransportClient.send(ByteBuf)`, checksum config).
- **Part 4** (*this PR*) - `StreamingShuffleWriter` + server-side Netty
handler (push path).
- **Part 5** - `StreamingShuffleReader` + client-side Netty handler (pull
path).
- **Part 6** - register streaming shuffles with the tracker in
`DAGScheduler` (activation).
- **Part 7** - end-to-end `StreamingShuffleSuite`.
- **Part 8** - documentation.
This PR depends only on the merged parts (1 through 3.5) and is independent
of the reader PR (part 5); the two can be reviewed in parallel and merged in
either order.
### Why are the changes needed?
Real-Time Mode / low-latency continuous queries need shuffle data to flow
continuously between stages. The default sort shuffle (write map output to
disk, then have reducers pull it) adds latency that is unacceptable for these
workloads. Streaming shuffle instead pushes records directly from writer tasks
to reader tasks. This PR implements the push (writer) half — the
previously-stubbed `StreamingShuffleManager.getWriter` — including the per-task
transport server, the credit-based con [...]
### Does this PR introduce _any_ user-facing change?
No. The streaming shuffle managers are opt-in via `spark.shuffle.manager`
and are not the default, and the feature is not usable end-to-end until the
reader (part 5) and activation (part 6) PRs land; the new configs therefore
have no effect on the default sort-shuffle path. The added configs
(`spark.shuffle.streaming.networkBufferSize`,
`spark.shuffle.streaming.networkBufferMaxWaitTimeMs`,
`spark.shuffle.streaming.writerMaxMemory`) take effect only once a streaming
shuffle is in use.
### How was this patch tested?
The writer is exercised end-to-end by `StreamingShuffleSuite` in the tests
PR of this stack (part 7): writer<->reader data transfer, the credit-control /
termination handshake, sequence-number validation, checksum verification,
memory back-pressure, and background-thread error propagation.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code (Claude Opus 4.8)
Closes #56878 from jerrypeng/stack/streaming-shuffle-pr4-writer.
Authored-by: Boyang Jerry Peng <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../spark/network/server/TransportServer.java | 4 +
.../java/org/apache/spark/internal/LogKeys.java | 2 +
.../org/apache/spark/internal/config/package.scala | 36 ++
.../streaming/StreamingShuffleManager.scala | 5 +-
.../streaming/StreamingShuffleServerHandler.scala | 106 +++++
.../shuffle/streaming/StreamingShuffleWriter.scala | 490 +++++++++++++++++++++
.../streaming/StreamingShuffleWriterSuite.scala | 291 ++++++++++++
7 files changed, 931 insertions(+), 3 deletions(-)
diff --git
a/common/network-common/src/main/java/org/apache/spark/network/server/TransportServer.java
b/common/network-common/src/main/java/org/apache/spark/network/server/TransportServer.java
index be5d9e03c45c..79666dac5c04 100644
---
a/common/network-common/src/main/java/org/apache/spark/network/server/TransportServer.java
+++
b/common/network-common/src/main/java/org/apache/spark/network/server/TransportServer.java
@@ -96,6 +96,10 @@ public class TransportServer implements Closeable {
return port;
}
+ public PooledByteBufAllocator getPooledByteBufAllocator() {
+ return pooledAllocator;
+ }
+
private void init(String hostToBind, int portToBind) {
IOMode ioMode = IOMode.valueOf(conf.ioMode());
diff --git
a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
index 6c03b93c0a9a..c4e1c5b96b1d 100644
--- a/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
+++ b/common/utils-java/src/main/java/org/apache/spark/internal/LogKeys.java
@@ -539,6 +539,7 @@ public enum LogKeys implements LogKey {
NUM_ROWS,
NUM_RULE_OF_RUNS,
NUM_SEQUENCES,
+ NUM_SHUFFLE_READERS,
NUM_SHUFFLE_WRITERS,
NUM_SKIPPED,
NUM_SLOTS,
@@ -734,6 +735,7 @@ public enum LogKeys implements LogKey {
SHUFFLE_IDS,
SHUFFLE_MERGE_ID,
SHUFFLE_MERGE_RECOVERY_FILE,
+ SHUFFLE_READER_ID,
SHUFFLE_SERVICE_CONF_OVERLAY_URL,
SHUFFLE_SERVICE_METRICS_NAMESPACE,
SHUFFLE_SERVICE_NAME,
diff --git a/core/src/main/scala/org/apache/spark/internal/config/package.scala
b/core/src/main/scala/org/apache/spark/internal/config/package.scala
index d8fc2a5c1fec..c26eec184346 100644
--- a/core/src/main/scala/org/apache/spark/internal/config/package.scala
+++ b/core/src/main/scala/org/apache/spark/internal/config/package.scala
@@ -1820,6 +1820,42 @@ package object config {
.intConf
.createWithDefault(32 << 20) // 32 MB
+ private[spark] val STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE =
+ ConfigBuilder("spark.shuffle.streaming.networkBufferSize")
+ .doc("Target byte size for each network buffer sent from a streaming
shuffle writer to a " +
+ "reader. Larger values reduce per-message overhead; smaller values
reduce latency.")
+ .version("4.3.0")
+ .internal()
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .checkValue(_ > 0, "spark.shuffle.streaming.networkBufferSize must be
positive.")
+ .createWithDefault(32768) // 32 KB
+
+ private[spark] val STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS =
+ ConfigBuilder("spark.shuffle.streaming.networkBufferMaxWaitTimeMs")
+ .doc("Maximum time in milliseconds a partially-filled network buffer is
held before " +
+ "being flushed to the reader. Lower values reduce latency at the cost
of smaller, " +
+ "less efficient messages.")
+ .version("4.3.0")
+ .internal()
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .longConf
+ .createWithDefault(50)
+
+ private[spark] val STREAMING_SHUFFLE_WRITER_MAX_MEMORY =
+ ConfigBuilder("spark.shuffle.streaming.writerMaxMemory")
+ .doc("Best-effort memory limit in bytes for in-flight data buffers in a
streaming " +
+ "shuffle writer task. Includes TCP send/receive buffers. The writer
back-pressures " +
+ "the upstream iterator when this limit is reached. This is a
best-effort bound: " +
+ "back-pressure is accounted per network buffer, so an individual
serialized row that " +
+ "exceeds the network buffer size can push actual in-flight memory
above this limit.")
+ .version("4.3.0")
+ .internal()
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .checkValue(_ > 0, "spark.shuffle.streaming.writerMaxMemory must be
positive.")
+ .createWithDefault(32 << 20) // 32 MB
+
private[spark] val SHUFFLE_DETECT_CORRUPT =
ConfigBuilder("spark.shuffle.detectCorrupt")
.doc("Whether to detect any corruption in fetched blocks.")
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
index ab6f2d19a8ff..2bc9d83b6141 100644
---
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleManager.scala
@@ -91,9 +91,8 @@ private[spark] class StreamingShuffleManager extends
ShuffleManager with Logging
mapId: Long,
context: TaskContext,
metrics: ShuffleWriteMetricsReporter): ShuffleWriter[K, V] = {
- // Implementation is added in a follow-up commit that introduces
StreamingShuffleWriter.
- throw new UnsupportedOperationException(
- "StreamingShuffleManager.getWriter is not yet implemented")
+ val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K,
V, _]]
+ new StreamingShuffleWriter[K, V](streamingShuffleHandle, mapId, context)
}
/**
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleServerHandler.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleServerHandler.scala
new file mode 100644
index 000000000000..424b764f5453
--- /dev/null
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleServerHandler.scala
@@ -0,0 +1,106 @@
+/*
+ * 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.nio.ByteBuffer
+import java.util.concurrent.CompletableFuture
+
+import io.netty.buffer.{ByteBuf, Unpooled}
+
+import org.apache.spark.TaskContext
+import org.apache.spark.network.client.{RpcResponseCallback, TransportClient}
+import org.apache.spark.network.server.{RpcHandler, StreamManager}
+import org.apache.spark.network.shuffle.streaming.{CreditControlMessage,
StreamingShuffleMessage, TerminationAckMessage}
+import org.apache.spark.util.ErrorNotifier
+
+/**
+ * The workflow of the server handler is the following:
+ *
+ * 1. Unfulfilled futures are created for each reader client.
+ *
+ * 2. Shuffle writer receives credit control message from shuffle reader A.
+ * This allows shuffle writer to know that this channel is for communication
+ * with shuffle reader A, allowing the relevant client future to be
fulfilled.
+ * This synchronously triggers any pending completion stages; the writer
will
+ * queue messages as completion stages until a connection is established.
+ *
+ * 3. The StreamingShuffleWriter will chain completion stages until it notices
+ * that the future is complete. When it does, it will invoke
TransportClient.send directly.
+ *
+ * 4. The StreamingShuffleWriter limits the number of in-flight messages by
acquiring
+ * a semaphore when new buffers are allocated, releasing it when the netty
callback
+ * for that buffer is invoked.
+ */
+class StreamingShuffleServerHandler(
+ onTerminationAckReceived: (Int, Long) => Unit,
+ shuffleId: Int,
+ // One reader per reduce partition, so this equals the writer's
numPartitions.
+ numReaders: Int,
+ val context: TaskContext,
+ errorNotifier: ErrorNotifier) extends RpcHandler with
TaskContextAwareLogging {
+
+ val futureClients: Array[CompletableFuture[TransportClient]] =
+ Array.fill(numReaders)(new CompletableFuture[TransportClient]())
+
+ setShuffleIdForLogging(shuffleId)
+
+ override def receive(
+ client: TransportClient,
+ message: ByteBuffer,
+ callback: RpcResponseCallback): Unit = {
+ var buf: ByteBuf = null
+ try {
+ buf = Unpooled.wrappedBuffer(message)
+ val shuffleMessage = StreamingShuffleMessage.decode(buf)
+
+ shuffleMessage match {
+ case creditControlMessage: CreditControlMessage =>
+ futureClients(creditControlMessage.shuffleReaderId).complete(client)
+ case terminationAck: TerminationAckMessage =>
+ logInfo(
+ s"Received termination ack message from shuffle reader " +
+ s"${terminationAck.shuffleReaderId}"
+ )
+ onTerminationAckReceived(terminationAck.shuffleReaderId,
terminationAck.getSeqNum)
+ case _ =>
+ throw new IllegalArgumentException(
+ s"Unexpected message type in ShuffleServerHandler: " +
+ s"${shuffleMessage.messageType()}"
+ )
+ }
+ } catch {
+ case (ex: Throwable) =>
+ logError(log"Streaming shuffle server handler receive failed", ex)
+ errorNotifier.markError(ex)
+ } finally {
+ if (buf != null) {
+ buf.release()
+ }
+ }
+ }
+
+ override def exceptionCaught(cause: Throwable, client: TransportClient):
Unit = {
+ logError(log"Streaming shuffle server handler caught exception.", cause)
+ errorNotifier.markError(cause)
+ }
+
+ // not needed for streaming shuffle
+ // cannot throw UnsupportedException because this function will be called
+ // even if this feature is not used.
+ override def getStreamManager: StreamManager = null
+}
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala
new file mode 100644
index 000000000000..176aefa33913
--- /dev/null
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala
@@ -0,0 +1,490 @@
+/*
+ * 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.")
+ // The per-partition floor (2 buffers per partition) can push the total
in-force budget above
+ // the configured writerMaxMemory when the partition count is high; surface
the effective total
+ // (including TCP buffers) so operators can see the limit they set is not
the one in force.
+ private val effectiveBudget = MAX_BUFFER_BYTES + TOTAL_TCPBUF_BYTES
+ if (effectiveBudget > conf.get(STREAMING_SHUFFLE_WRITER_MAX_MEMORY).toLong) {
+ logWarning(log"Streaming shuffle writer effective memory budget " +
+ log"${MDC(LogKeys.MAX_MEMORY_SIZE,
Utils.bytesToString(effectiveBudget))} exceeds the " +
+ log"configured ${MDC(LogKeys.CONFIG,
STREAMING_SHUFFLE_WRITER_MAX_MEMORY.key)}=" +
+ log"${MDC(LogKeys.MEMORY_SIZE, Utils.bytesToString(
+ conf.get(STREAMING_SHUFFLE_WRITER_MAX_MEMORY).toLong))} because the
per-partition " +
+ log"minimum for ${MDC(LogKeys.NUM_PARTITIONS, numPartitions)} partitions
takes precedence.")
+ }
+ 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)
+
+ // Data payloads use a dedicated direct-buffer free-list (bufferPool) of
fixed BUFFER_SIZE
+ // buffers so full-size send buffers can be recycled across the task; the
small, variable-size
+ // message envelopes instead use the server's pooled allocator (see
ShardState.send). The two
+ // allocation paths are intentionally separate.
+ 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)
+ // The Boolean return is intentionally not acted on: a false means the
shuffle was
+ // (concurrently) unregistered, which the tracker already logs a warning
for. That only
+ // happens while the shuffle is being torn down, in which case this writer
task is going
+ // away too, so there is nothing useful to do here.
+ 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 readers locate writers through
StreamingShuffleOutputTracker and pull
+ // data directly over Netty; they never consult MapStatus block sizes, and
streaming shuffle
+ // has no standard block-fetch fallback path. This MapStatus is therefore
only a placeholder
+ // to satisfy the ShuffleWriter contract and the DAGScheduler /
MapOutputTracker bookkeeping;
+ // its all-zero partition lengths are never read by any reducer.
+ 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 = {
+ // Back-pressure is accounted per network buffer (BUFFER_SIZE permits
each), not by exact
+ // byte size, so this bounds in-flight memory only on a best-effort basis:
a single
+ // serialized row larger than BUFFER_SIZE (rows are not split across
buffers, see write())
+ // grows its buffer past BUFFER_SIZE and thus exceeds the tracked budget.
+ 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)
+ flushThread.start()
+ records.foreach { record =>
+ val shard = shards(partitioner.getPartition(record._1))
+ var timestampedBuffer = shard.takeBuffer()
+ if (timestampedBuffer == null) {
+ timestampedBuffer = newBuffer()
+ }
+ val partitionSerializationStream =
timestampedBuffer.serializationStream
+ // When UnsafeRowSerializer, the key, record._1, is only used for
determining
+ // the partition, and it doesn't need to be sent to the shuffle
readers.
+ // However, if JavaSerializer is used (for test mainly), we need to
serialize
+ // the key since we will be attempting to read it from the shuffle
reader
+ //
+ // TODO we are actually not guaranteeing that a buffer used to send
data for a
+ // partition does not exceed BUFFER_SIZE. We currently are not
implementing spanning rows
+ // across multiple buffers as it requires interface changes in the
serializers
+ if (serializerInstance.isInstanceOf[JavaSerializerInstance]) {
+ partitionSerializationStream.writeKey(record._1.asInstanceOf[Any])
+ }
+ partitionSerializationStream.writeValue(record._2.asInstanceOf[Any])
+ partitionSerializationStream.flush()
+
+ timestampedBuffer.updateChecksum()
+
+ // Flush immediately if the buffer is almost full or stale.
+ if (timestampedBuffer.totalByteSize() < BUFFER_SIZE * 9 / 10 &&
+ timestampedBuffer.ageMs() < MAX_BUFFERING_TIME_MS) {
+ shard.putBuffer(timestampedBuffer)
+ } else {
+ shard.send(timestampedBuffer)
+ throwErrorIfExists()
+ }
+ }
+ isWriteFinished.countDown()
+ shards.foreach(_.close())
+ logInfo(log"StreamingShuffleWriter finished writing data and termination
messages for " +
+ log"shuffle writer ${MDC(LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}.
Shutting down now.")
+ // Wait for all termination acks. This has no wall-clock timeout by
design, and that is
+ // deliberate for correctness: a term-ack is the writer's only
confirmation that a reader
+ // received every message through the final sequence number (validated in
+ // onTerminationAckReceived), so finishing write() without all acks
would risk marking the
+ // map task successful while a reader is silently missing data. The loop
instead exits only
+ // on all-acks, on an ErrorNotifier error surfaced by
throwErrorIfExists(), or on task
+ // cancellation. A reader that dies fails its own reduce task, which
restarts the query and
+ // tears down this writer too, so writer-side reader-liveness detection
is unnecessary.
+ while (!allAcksReceived.await(1, TimeUnit.MILLISECONDS)) {
+ throwErrorIfExists()
+ }
+ logInfo(log"Received all termination acks for shuffle writer ${MDC(
+ LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}. Closing server channel.")
+ throwErrorIfExists()
+ } finally {
+ isWriteFinished.countDown() // Duplicate countDowns are a no-op.
+ flushThread.join()
+ }
+ }
+}
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
new file mode 100644
index 000000000000..42ad68a5668a
--- /dev/null
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
@@ -0,0 +1,291 @@
+/*
+ * 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.Properties
+import java.util.zip.CRC32C
+
+import io.netty.buffer.{ByteBuf, Unpooled}
+import io.netty.channel.{Channel, ChannelConfig, ChannelFuture}
+import io.netty.util.concurrent.GenericFutureListener
+import org.mockito.ArgumentMatchers.any
+import org.mockito.Mockito.when
+import org.scalatest.matchers.should.Matchers
+import org.scalatestplus.mockito.MockitoSugar
+
+import org.apache.spark._
+import org.apache.spark.LocalSparkContext.withSpark
+import org.apache.spark.internal.config.{SHUFFLE_MANAGER,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE}
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.metrics.MetricsSystem
+import org.apache.spark.network.client.TransportClient
+import org.apache.spark.network.shuffle.streaming.{DataMessage,
StreamingShuffleMessage, TerminationControlMessage}
+import
org.apache.spark.shuffle.streaming.StreamingShuffleManager.QUERY_ID_PROPERTY_KEY
+import org.apache.spark.util.ErrorNotifier
+
+/**
+ * Writer-side unit tests that do not require a shuffle reader. End-to-end
writer <-> reader
+ * behavior is covered by `StreamingShuffleSuite` in the tests PR of this
stack.
+ */
+class StreamingShuffleWriterSuite
+ extends SparkFunSuite
+ with LocalSparkContext
+ with Matchers
+ with MockitoSugar {
+
+ private def newConf(): SparkConf =
+ new SparkConf().set(SHUFFLE_MANAGER,
classOf[StreamingShuffleManager].getName)
+
+ private def createTaskContext(conf: SparkConf, partitionId: Int):
TaskContextImpl = {
+ val properties = new Properties()
+ properties.setProperty(QUERY_ID_PROPERTY_KEY, "test-query-id")
+ val taskMemoryManager = new TaskMemoryManager(new TestMemoryManager(conf),
0)
+ new TaskContextImpl(
+ stageId = 0,
+ stageAttemptNumber = 0,
+ partitionId,
+ taskAttemptId = 0,
+ attemptNumber = 0,
+ numPartitions = 1,
+ taskMemoryManager = taskMemoryManager,
+ localProperties = properties,
+ metricsSystem = mock[MetricsSystem],
+ cpus = 1)
+ }
+
+ test("getWriter returns a StreamingShuffleWriter") {
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite",
newConf())) { sc =>
+ SparkEnv.get.streamingShuffleOutputTracker.get
+ .asInstanceOf[StreamingShuffleOutputTrackerMaster]
+ .registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 0)
+ val rdd = sc.parallelize(1 to 4).map(x => (x, x))
+ val dep = new ShuffleDependency[Int, Int, Int](rdd, new
HashPartitioner(1))
+ val handle = new StreamingShuffleHandle(0, dep)
+ val context = createTaskContext(sc.conf, 0)
+ try {
+ val writer = new StreamingShuffleManager()
+ .getWriter[Int, Int](handle, 0, context, null)
+ writer shouldBe a[StreamingShuffleWriter[_, _]]
+ } finally {
+ // Constructing the writer starts a Netty server; the task-completion
listener
+ // (cleanupResources) tears it down.
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
+ test("writer rejects a memory budget that overflows the Int range") {
+ // With a large network buffer size, numPartitions * BUFFER_SIZE * 2
exceeds Int.MaxValue
+ // for even a handful of readers. The writer must reject this up front
rather than let the
+ // 32-bit product wrap negative and hang on a semaphore that can never
grant permits. The
+ // guard fires in the constructor before the Netty server is started, so
nothing to clean up.
+ val conf = newConf().set(STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE, 256 * 1024
* 1024)
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf))
{ sc =>
+ SparkEnv.get.streamingShuffleOutputTracker.get
+ .asInstanceOf[StreamingShuffleOutputTrackerMaster]
+ .registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 16, jobId =
0)
+ val rdd = sc.parallelize(1 to 4).map(x => (x, x))
+ val dep = new ShuffleDependency[Int, Int, Int](rdd, new
HashPartitioner(16))
+ val handle = new StreamingShuffleHandle(0, dep)
+ val context = createTaskContext(sc.conf, 0)
+ val e = intercept[IllegalArgumentException] {
+ new StreamingShuffleWriter[Int, Int](handle, 0, context)
+ }
+ assert(e.getMessage.contains("memory budget"))
+ }
+ }
+
+ // Builds a single-partition writer against a freshly registered shuffle.
The caller must run
+ // this inside a withSpark block and must eventually call
context.markTaskCompleted(None) to
+ // tear down the Netty server the writer starts in its constructor.
+ private def newWriter(
+ sc: SparkContext,
+ context: TaskContext,
+ errorNotifier: ErrorNotifier = new ErrorNotifier()):
StreamingShuffleWriter[Int, Int] = {
+ SparkEnv.get.streamingShuffleOutputTracker.get
+ .asInstanceOf[StreamingShuffleOutputTrackerMaster]
+ .registerShuffle(shuffleId = 0, numMaps = 1, numReduces = 1, jobId = 0)
+ val rdd = sc.parallelize(1 to 4).map(x => (x, x))
+ val dep = new ShuffleDependency[Int, Int, Int](rdd, new HashPartitioner(1))
+ val handle = new StreamingShuffleHandle(0, dep)
+ new StreamingShuffleWriter[Int, Int](handle, 0, context, errorNotifier =
errorNotifier)
+ }
+
+ // A mock TransportClient whose send(ByteBuf) invokes `onSend` and then
completes the write
+ // successfully (its returned ChannelFuture reports isSuccess = true so the
writer's normal
+ // done()/buffer-recycling path runs). Binding it into a shard's
futureClients makes the shard
+ // send directly to this mock instead of over the network.
+ private def bindMockClient(
+ writer: StreamingShuffleWriter[Int, Int],
+ shardId: Int)(onSend: ByteBuf => Unit): Unit = {
+ val client = mock[TransportClient]
+ val channel = mock[Channel]
+ val channelConfig = mock[ChannelConfig]
+ when(client.getChannel).thenReturn(channel)
+ when(channel.config).thenReturn(channelConfig)
+ val succeededFuture = mock[ChannelFuture]
+ when(succeededFuture.isSuccess).thenReturn(true)
+
when(succeededFuture.addListener(any[GenericFutureListener[ChannelFuture]]))
+ .thenAnswer { invocation =>
+ invocation.getArgument[GenericFutureListener[ChannelFuture]](0)
+ .operationComplete(succeededFuture)
+ succeededFuture
+ }
+ when(client.send(any[ByteBuf])).thenAnswer { invocation =>
+ val sent = invocation.getArgument[ByteBuf](0)
+ onSend(sent)
+ // Emulate the transport's ownership transfer: the real send(ByteBuf)
releases the buffer
+ // once the write completes, dropping the composite's retainedSlice on
the data buffer so
+ // rawBuffer.refCnt() is back to 1 at the done() callback (the normal
recycling path).
+ sent.release()
+ succeededFuture
+ }
+ writer.transportServerHandler.futureClients(shardId).complete(client)
+ }
+
+ test("a synchronous send failure is surfaced through the ErrorNotifier") {
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite",
newConf())) { sc =>
+ val context = createTaskContext(sc.conf, 0)
+ val errorNotifier = new ErrorNotifier()
+ try {
+ val writer = newWriter(sc, context, errorNotifier)
+ // Bind a client whose send throws synchronously; the shard's send
must record the failure
+ // via the ErrorNotifier and rethrow.
+ val client = mock[TransportClient]
+ val channel = mock[Channel]
+ val channelConfig = mock[ChannelConfig]
+ when(client.getChannel).thenReturn(channel)
+ when(channel.config).thenReturn(channelConfig)
+ when(client.send(any[ByteBuf])).thenThrow(new RuntimeException("send
failed"))
+ writer.transportServerHandler.futureClients(0).complete(client)
+
+ intercept[RuntimeException] {
+ writer.shards(0).send(new TerminationControlMessage(0, 0))
+ }
+ // The failure is recorded synchronously on this thread (the future is
already complete).
+ writer.errorNotifier.getError() shouldBe defined
+ writer.errorNotifier.getError().get.getMessage should include("send
failed")
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
+ test("an asynchronous write failure is surfaced through the ErrorNotifier") {
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite",
newConf())) { sc =>
+ val context = createTaskContext(sc.conf, 0)
+ val errorNotifier = new ErrorNotifier()
+ try {
+ val writer = newWriter(sc, context, errorNotifier)
+ // Bind a client whose send is accepted but whose write completes
unsuccessfully (the
+ // common network-failure case). The shard's send-completion listener
must record the
+ // future's cause via the ErrorNotifier.
+ val client = mock[TransportClient]
+ val channel = mock[Channel]
+ val channelConfig = mock[ChannelConfig]
+ when(client.getChannel).thenReturn(channel)
+ when(channel.config).thenReturn(channelConfig)
+ val failedFuture = mock[ChannelFuture]
+ when(failedFuture.isSuccess).thenReturn(false)
+ when(failedFuture.cause()).thenReturn(new RuntimeException("write
failed"))
+
when(failedFuture.addListener(any[GenericFutureListener[ChannelFuture]]))
+ .thenAnswer { invocation =>
+ invocation.getArgument[GenericFutureListener[ChannelFuture]](0)
+ .operationComplete(failedFuture)
+ failedFuture
+ }
+ when(client.send(any[ByteBuf])).thenAnswer { invocation =>
+ // The transport frees the buffer on completion, including on a
failed write.
+ invocation.getArgument[ByteBuf](0).release()
+ failedFuture
+ }
+ writer.transportServerHandler.futureClients(0).complete(client)
+
+ writer.shards(0).send(new TerminationControlMessage(0, 0))
+
+ writer.errorNotifier.getError() shouldBe defined
+ writer.errorNotifier.getError().get.getMessage should include("write
failed")
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
+ test("cleanupResources releases queued pooled buffers") {
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite",
newConf())) { sc =>
+ val context = createTaskContext(sc.conf, 0)
+ val writer = newWriter(sc, context)
+ val buf = Unpooled.buffer(16)
+ writer.bufferPool.offerLast(buf)
+ buf.refCnt() should be(1)
+ writer.bufferPool.size() should be(1)
+
+ // cleanupResources is idempotent; call it directly so we can observe
the buffer's refcount.
+ writer.cleanupResources()
+
+ writer.bufferPool.size() should be(0)
+ buf.refCnt() should be(0)
+ // The task-completion listener will call cleanupResources again (a
no-op now).
+ context.markTaskCompleted(None)
+ }
+ }
+
+ test("checksum is computed and embedded in the DataMessage sent on the
wire") {
+ val conf = newConf().set(STREAMING_SHUFFLE_CHECKSUM_ENABLED, true)
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf))
{ sc =>
+ val context = createTaskContext(sc.conf, 0)
+ try {
+ val writer = newWriter(sc, context)
+ val sentBuffers = new
java.util.concurrent.CopyOnWriteArrayList[ByteBuf]()
+ // Capture a copy of every buffer the writer sends (the original is
released by the
+ // transport layer after the write completes). The send completes
synchronously since the
+ // client future is already resolved.
+ bindMockClient(writer, 0) { buf => sentBuffers.add(buf.copy()) }
+
+ // Serialize a record through the writer's own buffer/checksum path
and send it.
+ val tsBuffer = writer.TimestampedBuffer(Unpooled.directBuffer(1024))
+ tsBuffer.serializationStream.writeKey(1.asInstanceOf[Any])
+ tsBuffer.serializationStream.writeValue(2.asInstanceOf[Any])
+ tsBuffer.serializationStream.flush()
+ writer.shards(0).send(tsBuffer)
+
+ sentBuffers.size() should be(1)
+ val decoded = StreamingShuffleMessage.decode(sentBuffers.get(0))
+ decoded shouldBe a[DataMessage]
+ val dataMessage = decoded.asInstanceOf[DataMessage]
+
+ // The checksum embedded on the wire must match an independent CRC32C
over the payload.
+ // Compute the expected value with java.util.zip.CRC32C directly (not
ShuffleChecksum) so
+ // this stays independent of the production checksum wrapper.
+ val payload = dataMessage.getRecordData()
+ val expectedCrc = new CRC32C()
+ expectedCrc.update(payload.nioBuffer(payload.readerIndex(),
payload.readableBytes()))
+ dataMessage.checksum should be(expectedCrc.getValue)
+ dataMessage.checksum should not be 0L
+ // The mock released the sent buffer (as the transport would), so the
writer's normal
+ // recycling path ran (rawBuffer.refCnt() == 1) and recorded no error.
+ writer.errorNotifier.getError() shouldBe empty
+
+ dataMessage.release()
+ sentBuffers.forEach(_.release())
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
