This is an automated email from the ASF dual-hosted git repository.
HeartSaVioR 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 ffa56147aa43 [SPARK-57230][SS][RTM][STREAMINGSHUFFLE][PART5] Add
StreamingShuffleReader + client-side Netty handler
ffa56147aa43 is described below
commit ffa56147aa43e12cae4675ca59801fa9390dff8b
Author: Boyang Jerry Peng <[email protected]>
AuthorDate: Fri Jul 10 11:04:00 2026 +0900
[SPARK-57230][SS][RTM][STREAMINGSHUFFLE][PART5] Add StreamingShuffleReader
+ client-side Netty handler
### What changes were proposed in this pull request?
This is **part 5** 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 **reader (pull) side**:
- **`StreamingShuffleReader`** — a `ShuffleReader` that connects to the
upstream writer tasks over Netty and returns an iterator over the records they
push, instead of fetching shuffle files from disk.
- **`StreamingShuffleClientHandler`** — the reader-side Netty `RpcHandler`
that receives pushed data, applies flow control, and sends the credit-control
and termination-ack messages back to the writer.
- **`StreamingShuffleManager.getReader`** — now returns a
`StreamingShuffleReader` (it was a stub that threw
`UnsupportedOperationException` in part 3).
- One reader config, three structured-logging keys, and one error condition
(see below).
#### How the reader and writer talk to each other
The reader asks the `StreamingShuffleOutputTracker` (part 2) for the
locations of the writer tasks, opens a connection to each, and exchanges four
message types (all defined in part 1). The exchange for one reader <-> one
writer looks like this:
```
Reader (this PR) Writer
| |
| --- CreditControlMessage -----------> | on connect: reader
announces itself so
| | the writer can bind this
channel to it
| |
| <-------------- DataMessage (seq=0) --- | reader receives pushed
rows; each
| --- CreditControlMessage -----------> | DataMessage carries a
sequence number
| | and an optional CRC32C
checksum. The reader
| <-------------- DataMessage (seq=1) --- | checks the sequence,
verifies the checksum,
| --- CreditControlMessage -----------> | enqueues the rows, and
replies with a
| <-------------- DataMessage (seq=2) --- | CreditControlMessage
for each one
| --- CreditControlMessage -----------> |
| ... |
| <----- 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 reader's side:
1. **Connection setup.** The reader runs a background task-discovery thread
that polls the tracker for the writer task locations and creates a transport
client to each writer in parallel. When a connection becomes active the client
handler sends a `CreditControlMessage` to announce itself, and then sends one
back for every `DataMessage` it receives.
2. **Data receive.** For each `DataMessage`, the reader validates that its
**sequence number** is exactly one greater than the last one seen from that
writer (a gap or duplicate fails the task with
`STREAMING_SHUFFLE_INCORRECT_SEQUENCE_NUMBER`), enqueues the message on a
shared queue, and hands the deserialized rows to the calling task through the
reader's iterator.
3. **Integrity check (optional).** When
`spark.shuffle.streaming.checksum.enabled` is on, the reader re-computes the
CRC32C over each received buffer and compares it against the writer's checksum,
failing with `STREAMING_SHUFFLE_CHECKSUM_VERIFICATION_FAILED` on a mismatch.
4. **Back-pressure.** In-flight buffered data is bounded per writer: the
reader derives a per-writer byte quota from
`spark.shuffle.streaming.readerMaxMemory` (default 32 MB) divided by the number
of writers, and toggles the Netty channel's auto-read off when the quota is
exhausted, applying TCP back-pressure until the enqueued data is consumed.
5. **Termination.** When a writer sends its `TerminationControlMessage`,
the reader replies with a `TerminationAckMessage` carrying the last sequence
number it saw. The reader's iterator only finishes once it has received a
termination message from **every** writer and confirmed that all termination
acks were sent, so no in-flight data is dropped.
Errors that occur on Netty threads or the background
discovery/client-creation threads (e.g. a failed send, a checksum mismatch, a
client that cannot connect, or a writer that disconnects before terminating)
are recorded via `ErrorNotifier` (part 3.5) and re-thrown on the task thread at
a safe point. All resources (the discovery and client-creation thread pools,
client connections and factories, queued message buffers, reserved memory) are
released from a `TaskContext` completion list [...]
New config (`internal`, default on the safe side):
| Config | Default | Purpose |
|---|---|---|
| `spark.shuffle.streaming.readerMaxMemory` | 32 MB | best-effort cap on
the reader's in-flight buffered data; the per-writer quota is this divided by
the number of writers |
New log keys: `NUM_CONNECTED_SHUFFLE_WRITERS`, `NUM_SHUFFLE_WRITERS`,
`SHUFFLE_WRITERS`.
New error condition: `STREAMING_SHUFFLE_CHECKSUM_VERIFICATION_FAILED`.
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** - `StreamingShuffleWriter` + server-side Netty handler (push
path).
- **Part 5** (*this PR*) - `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 writer PR (part 4); 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 pull (reader) half — the
previously-stubbed `StreamingShuffleManager.getReader` — including the
background writer discovery and connection setu [...]
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The reader 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, resource cleanup on task completion, and
background-thread error propagation — including that the reader fails (rather
than hangs) when a writer disconnects before terminating.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code (Claude Opus 4.8)
Closes #56879 from jerrypeng/stack/streaming-shuffle-pr5-reader.
Authored-by: Boyang Jerry Peng <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../java/org/apache/spark/internal/LogKeys.java | 3 +
.../src/main/resources/error/error-conditions.json | 6 +
.../org/apache/spark/internal/config/package.scala | 11 +
.../streaming/StreamingShuffleClientHandler.scala | 264 +++++++++++
.../streaming/StreamingShuffleManager.scala | 5 +-
.../shuffle/streaming/StreamingShuffleReader.scala | 493 +++++++++++++++++++++
.../streaming/StreamingShuffleReaderSuite.scala | 166 +++++++
7 files changed, 945 insertions(+), 3 deletions(-)
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 32d34026e5a8..6c03b93c0a9a 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
@@ -462,6 +462,7 @@ public enum LogKeys implements LogKey {
NUM_COEFFICIENTS,
NUM_COLUMNS,
NUM_CONCURRENT_WRITER,
+ NUM_CONNECTED_SHUFFLE_WRITERS,
NUM_CORES,
NUM_DATA_FILE,
NUM_DATA_FILES,
@@ -538,6 +539,7 @@ public enum LogKeys implements LogKey {
NUM_ROWS,
NUM_RULE_OF_RUNS,
NUM_SEQUENCES,
+ NUM_SHUFFLE_WRITERS,
NUM_SKIPPED,
NUM_SLOTS,
NUM_SPILLS,
@@ -735,6 +737,7 @@ public enum LogKeys implements LogKey {
SHUFFLE_SERVICE_CONF_OVERLAY_URL,
SHUFFLE_SERVICE_METRICS_NAMESPACE,
SHUFFLE_SERVICE_NAME,
+ SHUFFLE_WRITERS,
SHUFFLE_WRITER_ID,
SIGMAS_LENGTH,
SIGNAL,
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index 3d701442a0f8..b09909f8f9aa 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -7392,6 +7392,12 @@
},
"sqlState" : "0A000"
},
+ "STREAMING_SHUFFLE_CHECKSUM_VERIFICATION_FAILED" : {
+ "message" : [
+ "Checksum verification failed for buffer from writer <writerId>.
Expected: <expectedChecksum>, Calculated: <calculatedChecksum>, Data length:
<dataLength>."
+ ],
+ "sqlState" : "XXKST"
+ },
"STREAMING_SHUFFLE_INCORRECT_SEQUENCE_NUMBER" : {
"message" : [
"Streaming shuffle <messageType> between writer <writerId> and reader
<readerId> expected to have sequence number <expSeqNum>, but the actual
sequence number is <actSeqNum>. Please verify that the messages are sent in
order."
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 b658c31753e3..c9f7654df70f 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
@@ -1796,6 +1796,17 @@ package object config {
.booleanConf
.createWithDefault(true)
+ private[spark] val STREAMING_SHUFFLE_READER_MAX_MEMORY =
+ ConfigBuilder("spark.shuffle.streaming.readerMaxMemory")
+ .doc("Best-effort memory limit in bytes for data buffered in a streaming
shuffle reader " +
+ "task. The per-writer byte quota is derived from this value divided by
the number of " +
+ "shuffle writers. When the quota is exhausted the reader applies TCP
back-pressure.")
+ .version("4.3.0")
+ .internal()
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .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/StreamingShuffleClientHandler.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleClientHandler.scala
new file mode 100644
index 000000000000..28ac46bb5752
--- /dev/null
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleClientHandler.scala
@@ -0,0 +1,264 @@
+/*
+ * 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.LinkedBlockingQueue
+
+import io.netty.buffer.{ByteBuf, CompositeByteBuf, Unpooled}
+import io.netty.channel.{Channel, ChannelOption}
+import io.netty.util.concurrent.{Future, GenericFutureListener}
+
+import org.apache.spark.{SparkException, 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,
DataMessage, StreamingShuffleMessage, StreamingShuffleMessageType,
TerminationAckMessage, TerminationControlMessage}
+import org.apache.spark.util.ErrorNotifier
+
+/**
+ * A StreamingShuffleClientHandler is used by ShuffleReaders to receive data
from ShuffleWriters.
+ *
+ * Our V1 protocol is very simple: the server pushes records to us that we
read one-by-one and write
+ * them to a concurrent queue that the calling task (StreamingShuffleReader)
can dequeue from.
+ */
+class StreamingShuffleClientHandler(
+ // The shuffle writer that we're connected to.
+ // TODO: we might remove this field; avoid relying on it in future
development.
+ shuffleWriterId: Int,
+ shuffleReaderId: Int,
+ queue: LinkedBlockingQueue[StreamingShuffleMessage],
+ shuffleId: Int,
+ byteLimit: Long,
+ val context: TaskContext,
+ errorNotifier: ErrorNotifier) extends RpcHandler with
TaskContextAwareLogging {
+ private val RECVBUF_SIZE: Integer = 32 << 10
+ private val SENDBUF_SIZE: Integer = 512
+
+ private var lastSeqNum = -1L // The most recent sequence number we have
seen.
+ // Set once this writer's TerminationControlMessage has been received, so
that channelInactive
+ // can tell an expected end-of-stream close apart from a premature
disconnect (a writer that
+ // died before terminating). @volatile because it is written on the Netty
event-loop thread in
+ // receive() and read in channelInactive().
+ @volatile private var terminationReceived = false
+ // These variables are used for flow control by updateQuota below.
+ private var channel: Channel = _ // The channel to the shuffle writer,
captured in channelActive.
+ private var remainingBytesQuota: Long = byteLimit // Remaining bytes before
pushback.
+ private var autoReadDisabledTimestamp: Long = _ // Last time pushback
condition was triggered.
+
+ setShuffleIdForLogging(shuffleId)
+
+ // Keeps track of the largest seqNum we have seen so far sent by the writer.
+ // It should start with 0 and be strictly without gap.
+ private def updateLastSeqNum(newSeqNum: Long, messageType:
StreamingShuffleMessageType) = {
+ if (lastSeqNum + 1 != newSeqNum) {
+ throw StreamingShuffleManager.streamingShuffleIncorrectSequenceNumber(
+ messageType,
+ shuffleWriterId,
+ shuffleReaderId,
+ lastSeqNum + 1,
+ newSeqNum)
+ }
+ lastSeqNum = newSeqNum
+ }
+
+ private[spark] var onTermAckResponse: Int => Unit = (Int) => {}
+ private[spark] def setOnTermAckResponseHandler(handler: Int => Unit): Unit =
{
+ onTermAckResponse = handler
+ }
+
+ override def channelActive(client: TransportClient): Unit = {
+ channel = client.getChannel
+ channel.config.setOption(ChannelOption.SO_RCVBUF, RECVBUF_SIZE)
+ channel.config.setOption(ChannelOption.SO_SNDBUF, SENDBUF_SIZE)
+ sendCreditControlMessage(client, shuffleWriterId, 1) // Tell upstream
writer that we're ready.
+ }
+
+ // Update the number of outstanding bytes from this writer, toggling
auto-read if necessary.
+ // Can be called from main or Netty threads, so synchronization is required.
+ private def updateQuota(bytes: Long): Unit = synchronized {
+ remainingBytesQuota -= bytes
+ val autoRead = remainingBytesQuota > 0
+ if (channel.config.isAutoRead != autoRead) {
+ channel.config.setAutoRead(autoRead)
+ if (autoRead) {
+ channel.read()
+ } else {
+ autoReadDisabledTimestamp = System.nanoTime()
+ }
+ }
+ }
+
+ private def sendCreditControlMessage(
+ client: TransportClient,
+ shuffleWriterId: Int,
+ credit: Int
+ ): Unit = {
+ var buf: CompositeByteBuf = null
+ try {
+ val creditControlMessage = new CreditControlMessage(shuffleWriterId,
shuffleReaderId, credit)
+ buf = client.getChannel().alloc().compositeBuffer()
+ .capacity(creditControlMessage.headerLength())
+ creditControlMessage.encode(buf)
+
+ // send() will release the buffer, so retain to avoid double free in
finally clause
+ client
+ .send(buf.retain())
+ .addListener(
+ getResponseHandler(buf,
+ s"Error sending credit control message to shuffle writer
${shuffleWriterId}"))
+ } catch {
+ case (ex: Throwable) =>
+ logError(log"Streaming shuffle client handler sendCreditControlMessage
failed", ex)
+ errorNotifier.markError(ex)
+ } finally {
+ if (buf != null) {
+ buf.release()
+ }
+ }
+ }
+
+ protected def sendTerminationAckMessage(client: TransportClient,
shuffleWriterId: Int): Unit = {
+ var buf: CompositeByteBuf = null
+ try {
+ val terminationAckMessage = new TerminationAckMessage(shuffleWriterId,
shuffleReaderId)
+ terminationAckMessage.setSeqNum(lastSeqNum)
+ buf = client.getChannel().alloc().compositeBuffer()
+ .capacity(terminationAckMessage.headerLength())
+ terminationAckMessage.encode(buf)
+
+ // send() will release the buffer, so retain to avoid double free in
finally clause
+ client
+ .send(buf.retain())
+ .addListener(
+ getResponseHandler(
+ buf,
+ s"Error sending termination acknowledgment to shuffle writer
${shuffleWriterId}",
+ () => { onTermAckResponse(shuffleWriterId) }
+ )
+ )
+ } catch {
+ case (ex: Throwable) =>
+ logError(log"Streaming shuffle client handler
sendTerminationAckMessage failed", ex)
+ errorNotifier.markError(ex)
+ } finally {
+ if (buf != null) {
+ buf.release()
+ }
+ }
+ }
+
+ override def receive(
+ client: TransportClient,
+ message: ByteBuffer,
+ callback: RpcResponseCallback): Unit = {
+ // The underlying message gets freed after the call to receive, so we need
to copy
+ // the underlying data.
+ //
+ // TODO: in the future, the TransportRequestHandler can be modified to
+ // not free the underlying data.
+ var buf: ByteBuf = null
+ try {
+ buf = Unpooled.wrappedBuffer(message).copy()
+ val shuffleMessage = StreamingShuffleMessage.decode(buf)
+ updateLastSeqNum(shuffleMessage.getSeqNum, shuffleMessage.messageType())
+ // At this point, the message type has been read, so the decoders will
read the (optional)
+ // length int, followed by the actual content.
+ shuffleMessage match {
+ case dataMessage: DataMessage =>
+ val messageSize = buf.capacity()
+ updateQuota(messageSize)
+ dataMessage.setReleaseCallback(() => updateQuota(-messageSize))
+ // we can only release the buf after we have decoded all the rows in
the buffer
+ sendCreditControlMessage(client, dataMessage.shuffleWriterId, 1)
+ case controlMessage: TerminationControlMessage =>
+ // Record termination before sending the ack: the writer only closes
its connection
+ // after it receives this ack, so setting the flag here guarantees
it is visible before
+ // the resulting channelInactive fires, avoiding a false "premature
disconnect" error.
+ terminationReceived = true
+ sendTerminationAckMessage(client, controlMessage.shuffleWriterId)
+ case _ =>
+ throw new IllegalArgumentException(s"Unexpected message type in
ShuffleClientHandler: " +
+ s"${shuffleMessage.messageType()}");
+ }
+ queue.put(shuffleMessage)
+ } catch {
+ case (ex: Throwable) =>
+ logError(log"Streaming shuffle client handler receive failed.", ex)
+ errorNotifier.markError(ex)
+ } finally {
+ if (buf != null) {
+ // If any StreamingShuffleMessage needs buf, then it would have
retained it.
+ buf.release()
+ }
+ }
+ }
+
+ override def channelInactive(client: TransportClient): Unit = {
+ // A clean end-of-stream also closes the channel, but only after this
reader has received the
+ // writer's TerminationControlMessage (which sets terminationReceived). So
a close while the
+ // flag is still false means the writer disconnected before terminating --
e.g. the writer
+ // task failed, its executor was lost, or the network dropped. Surface it
through the shared
+ // ErrorNotifier so the reader task fails instead of polling the message
queue forever.
+ if (!terminationReceived) {
+ errorNotifier.markError(new SparkException(
+ s"Connection to streaming shuffle writer ${shuffleWriterId} closed
before termination; " +
+ "the writer task likely failed."))
+ }
+ }
+
+ override def exceptionCaught(cause: Throwable, client: TransportClient):
Unit = {
+ logError(log"Streaming shuffle client 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
+
+ private def getResponseHandler(
+ buf: ByteBuf,
+ errorMsg: String,
+ onSuccessFunc: () => Unit = () => {})
+ : GenericFutureListener[io.netty.util.concurrent.Future[Void]] = {
+ new GenericFutureListener[io.netty.util.concurrent.Future[Void]]() {
+
+ override def operationComplete(future: Future[Void]): Unit = {
+ try {
+ handleResponse(future, errorMsg)
+ onSuccessFunc()
+ } catch {
+ // The code in listener will be executed by another thread, so we
need to
+ // bubble up the error here
+ case (ex: Throwable) =>
+ logError(errorMsg, ex)
+ errorNotifier.markError(ex)
+ }
+ }
+ }
+ }
+
+ protected def handleResponse(future: Future[Void], errorMsg: String): Unit =
{
+ if (!future.isSuccess) {
+ throw new RuntimeException(
+ s"${errorMsg}: ${future.cause().getMessage()}"
+ )
+ }
+ }
+}
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 f56d4f0fc4f8..ab6f2d19a8ff 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
@@ -108,9 +108,8 @@ private[spark] class StreamingShuffleManager extends
ShuffleManager with Logging
endPartition: Int,
context: TaskContext,
metrics: ShuffleReadMetricsReporter): ShuffleReader[K, C] = {
- // Implementation is added in a follow-up commit that introduces
StreamingShuffleReader.
- throw new UnsupportedOperationException(
- "StreamingShuffleManager.getReader is not yet implemented")
+ val streamingShuffleHandle = handle.asInstanceOf[StreamingShuffleHandle[K,
_, C]]
+ new StreamingShuffleReader[K, C](streamingShuffleHandle, context)
}
override def unregisterShuffle(shuffleId: Int): Boolean = {
diff --git
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReader.scala
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReader.scala
new file mode 100644
index 000000000000..70ea96154344
--- /dev/null
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReader.scala
@@ -0,0 +1,493 @@
+/*
+ * 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.{CompletableFuture, ConcurrentHashMap,
ConcurrentLinkedQueue, LinkedBlockingQueue, Semaphore, TimeUnit}
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.jdk.CollectionConverters._
+
+import io.netty.buffer.ByteBufInputStream
+
+import org.apache.spark.{ShuffleLocationResponse, SparkContext, SparkEnv,
SparkRuntimeException, TaskContext}
+import org.apache.spark.internal.LogKeys
+import org.apache.spark.internal.config.{EXECUTOR_ID,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_READER_MAX_MEMORY}
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode}
+import org.apache.spark.network.TransportContext
+import org.apache.spark.network.client.{TransportClient,
TransportClientFactory}
+import org.apache.spark.network.netty.SparkTransportConf
+import org.apache.spark.network.shuffle.streaming.{DataMessage,
ShuffleChecksum, StreamingShuffleMessage, TerminationControlMessage}
+import org.apache.spark.shuffle.{ShuffleHandle, ShuffleReader}
+import org.apache.spark.util.{ErrorNotifier, NextIterator, ThreadUtils, Utils}
+
+/**
+ * Default factory used by `StreamingShuffleReader` to create its output
iterator.
+ */
+class StreamingShuffleReaderIteratorFactory {
+ def create[K, C](
+ messageQueue: LinkedBlockingQueue[StreamingShuffleMessage],
+ handleTerminationMessage: TerminationControlMessage => Boolean,
+ handleDataMessage: DataMessage => Iterator[(K, C)],
+ checkTaskFailure: () => Unit
+ ): Iterator[Product2[K, C]] = {
+ new NextIterator[Product2[K, C]] {
+ // Iterator that iterates through multiple rows in data message buffer.
When the iterator
+ // does not have any more rows, we should fetch another message from
message queue.
+ private var rowIterator: Iterator[(K, C)] = Iterator.empty
+
+ def getNext(): Product2[K, C] = {
+ while (!rowIterator.hasNext) {
+ checkTaskFailure()
+ messageQueue.poll(10, TimeUnit.MILLISECONDS) match {
+ case msg: TerminationControlMessage =>
+ if (handleTerminationMessage(msg)) {
+ finished = true
+ return null.asInstanceOf[Product2[K, C]]
+ }
+ case dataMessage: DataMessage =>
+ rowIterator = handleDataMessage(dataMessage)
+ case null => // no message received, continue to poll
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unexpected message type in reader queue:
${other.getClass.getName}")
+ }
+ }
+ rowIterator.next()
+ }
+
+ override def close(): Unit = {
+ // no-op. handleTerminationMessage will take care of final cleanup.
+ }
+ }
+ }
+}
+
+class StreamingShuffleReader[K, C](
+ handle: ShuffleHandle,
+ val context: TaskContext,
+ clientHandler: Option[StreamingShuffleClientHandler] = None,
+ private[streaming] val errorNotifier: ErrorNotifier = new ErrorNotifier())
+ extends ShuffleReader[K, C] with TaskContextAwareLogging {
+ assert(SparkEnv.get.streamingShuffleOutputTracker.isDefined)
+ private val conf = SparkEnv.get.conf
+
+ private val streamingShuffleHandle =
handle.asInstanceOf[StreamingShuffleHandle[K, _, C]]
+ setShuffleIdForLogging(streamingShuffleHandle.shuffleId)
+ // a mapping of mapId and client
+ private[spark] val clientMap = new ConcurrentHashMap[Long, TransportClient]()
+ // Track factories so we can close them on shutdown.
+ // TODO: refactor to reuse a single client factory across writers.
+ private val clientFactories = new
ConcurrentLinkedQueue[TransportClientFactory]()
+ private val tracker = SparkEnv.get.streamingShuffleOutputTracker.get
+
+ private val role = conf.get(EXECUTOR_ID).map { id =>
+ if (SparkContext.isDriver(id)) "driver" else "executor"
+ }
+
+ private val clientConf = SparkTransportConf.fromSparkConf(
+ conf,
+
s"streaming-shuffle-reader-${streamingShuffleHandle.shuffleId}-${context.partitionId()}",
+ 1, // a client should only need to use 1 thread
+ role)
+
+ private[spark] val taskDiscoveryExecutor =
+ ThreadUtils.newDaemonSingleThreadExecutor(
+ s"streaming-shuffle-task-discovery-thread-" +
+ s"${streamingShuffleHandle.shuffleId}-${context.partitionId()}")
+
+ private val totalNumShuffleWriters: AtomicInteger = new AtomicInteger(-1)
+ private var perWriterByteLimit: Long = 1
+
+ // We might need to revisit if the size limit is enough. If not, there
should be a way to tie it
+ // to the per-task memory limit from the task context.
+ private val MAX_MEMORY = conf.get(STREAMING_SHUFFLE_READER_MAX_MEMORY)
+ // Data and termination messages from all writers are put into this queue.
+ private[spark] val messageQueue = new
LinkedBlockingQueue[StreamingShuffleMessage]()
+
+ private val memoryConsumer =
+ new MemoryConsumer(context.taskMemoryManager(), MemoryMode.OFF_HEAP) {
+ override def spill(size: Long, trigger: MemoryConsumer): Long = 0
+ }
+
+ private val shuffleChecksum = if
(conf.get(STREAMING_SHUFFLE_CHECKSUM_ENABLED)) {
+ new ShuffleChecksum()
+ } else {
+ null
+ }
+
+ // The set of shuffle writers that this reader has successfully received
+ // termination ack messages from. This is used to make sure all term ack
messages
+ // are successfully sent before exiting.
+ private[spark] val terminationAckControlMessageSet =
ConcurrentHashMap.newKeySet[Long]()
+
+ private val allTermAcksSentNotice = new Semaphore(0)
+
+ // thread pool used to perform client creation in parallel
+ private[spark] val clientCreationExecutor =
ThreadUtils.newDaemonFixedThreadPool(
+ Runtime.getRuntime.availableProcessors,
+ s"streaming-shuffle-async-client-creation-${context.partitionId()}")
+
+ // Signals to other threads that task discovery should stop. For example, we
may receive all
+ // the termination messages before we actually put the clients in the client
map. In that
+ // scenario we can just exit without checking the client map for whether all
the clients were
+ // created.
+ @volatile private var taskDiscoveryShouldStop = false
+
+ private var currentDataMessage: StreamingShuffleMessage = _
+
+ private def shutdownExecutorService(
+ executor: java.util.concurrent.ExecutorService,
+ name: String): Unit = {
+ executor.shutdownNow()
+ try {
+ if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
+ logWarning(log"${MDC(LogKeys.NAME, name)} did not shut down within the
timeout.")
+ } else {
+ logInfo(log"${MDC(LogKeys.NAME, name)} shut down successfully.")
+ }
+ } catch {
+ case _: InterruptedException =>
+ // make sure we clean up even if the task thread is interrupted during
shutdown
+ shutdownExecutorService(executor, name)
+ // Restore the interrupt flag so downstream code knows we were
interrupted
+ Thread.currentThread().interrupt()
+ }
+ }
+
+ /**
+ * Cleans up all reader resources. This method should be idempotent and
+ * can be called multiple times without issue.
+ */
+ private[streaming] def cleanupResources(): Unit = {
+ val cleanupStartTime = System.currentTimeMillis()
+
+ Utils.tryLogNonFatalError {
+ stopTaskDiscovery()
+ }
+ Utils.tryLogNonFatalError {
+ shutdownExecutorService(clientCreationExecutor, "Client Creation
Executor")
+ }
+ Utils.tryLogNonFatalError {
+ clientMap.forEach((_, client) => client.close())
+ }
+ Utils.tryLogNonFatalError {
+ clientFactories.forEach(factory => factory.close())
+ }
+ Utils.tryLogNonFatalError {
+ if (currentDataMessage != null) {
+ currentDataMessage.release()
+ currentDataMessage = null
+ }
+ }
+ Utils.tryLogNonFatalError {
+ val list = new java.util.ArrayList[StreamingShuffleMessage]()
+ messageQueue.drainTo(list)
+ list.forEach(_.release())
+ }
+ Utils.tryLogNonFatalError {
+ memoryConsumer.freeMemory(memoryConsumer.getUsed())
+ }
+ logInfo(log"Resource cleanup took ${MDC(LogKeys.DURATION,
+ System.currentTimeMillis() - cleanupStartTime)} ms")
+ }
+
+ // Register a task completion listener to make sure we clean up resources
+ // when the task is completed. Without this, threads, client
connections/factories,
+ // message buffers, and memory allocations could be leaked.
+ context.addTaskCompletionListener[Unit] { _ =>
+ cleanupResources()
+ }
+
+ taskDiscoveryExecutor.execute(() => {
+ val startTime = System.currentTimeMillis()
+ try {
+ logInfo(log"Task discovery thread started.")
+ // a mapping of mapId and client creation future. Used for parallel
client creation
+ val clientFutureMap = new ConcurrentHashMap[Long,
CompletableFuture[Void]]()
+
+ var isDone = false
+
+ def shouldStop(): Boolean = {
+ (
+ // Got all locations for shuffle writers
+ // and started creating the clients to communicate with them.
+ isDone
+ // signal from main thread that task discovery should terminate.
+ // For example, we got all the term messages from all clients
+ || taskDiscoveryShouldStop
+ // task has failed or been interrupted.
+ || context.isFailed()
+ || context.isInterrupted()
+ )
+ }
+
+ while (!shouldStop()) {
+ val shuffleLocationResponseOption =
+
tracker.getAvailableShuffleWriterTaskLocations(streamingShuffleHandle.shuffleId)
+
+ var retryCount = 0
+ shuffleLocationResponseOption.foreach {
+ case ShuffleLocationResponse(shuffleWriterLocations,
numShuffleWriters) =>
+ if (!totalNumShuffleWriters.compareAndSet(-1, numShuffleWriters)) {
+ assert(totalNumShuffleWriters.get() == numShuffleWriters)
+ } else {
+ perWriterByteLimit = Math.max(MAX_MEMORY / numShuffleWriters, 1)
+ memoryConsumer.acquireMemory(MAX_MEMORY)
+ }
+ shuffleWriterLocations
+ .foreach {
+ case (mapId, location) =>
+ if (!clientFutureMap.containsKey(mapId)) {
+ val future = createClientAsync(mapId.toInt, location.host,
location.port)
+ .thenAccept((shuffleClient: TransportClient) => {
+ clientMap.put(mapId, shuffleClient)
+ logInfo(
+ log"Created shuffle client to shuffle " +
+ log"writer with id ${MDC(LogKeys.MAP_ID, mapId)}
and location ${MDC(
+ LogKeys.TASK_LOCATION, location)}. ${MDC(
+ LogKeys.NUM_CONNECTED_SHUFFLE_WRITERS,
clientMap.size)} / ${MDC(
+ LogKeys.NUM_SHUFFLE_WRITERS, numShuffleWriters)}
shuffle writer " +
+ log"tasks connected.")
+ }).exceptionally(th => {
+ val errorMsg = s"Error creating transport client to
shuffle writer" +
+ s" with id ${mapId} and location ${location}."
+ // The real exception is likely wrapped in
CompletionException
+ logError(errorMsg, th.getCause)
+ errorNotifier.markError(th.getCause)
+ throw new RuntimeException(errorMsg)
+ })
+ clientFutureMap.put(mapId, future)
+ }
+ }
+
+ val numClients = clientFutureMap.size()
+ // we should not be getting more clients than expected
+ assert(!(numClients > numShuffleWriters))
+ if (numClients != numShuffleWriters) {
+ // TODO also implement timeout
+ Thread.sleep(10)
+ retryCount += 1
+ if (retryCount % 100 == 0) {
+ logInfo(log"Still attempting to get shuffle writer locations."
+
+ log" Got the location of
${MDC(LogKeys.NUM_CONNECTED_SHUFFLE_WRITERS,
+ clientFutureMap.size)} / ${MDC(LogKeys.NUM_SHUFFLE_WRITERS,
+ numShuffleWriters)} shuffle writers")
+ }
+ } else {
+ // we have received all shuffle writer locations
+ // and are in the process of creating clients to connect to them
+ isDone = true
+ }
+ }
+ }
+
+ if (isDone && !taskDiscoveryShouldStop && !context.isFailed() &&
!context.isInterrupted()) {
+ // wait for all futures to finish
+ CompletableFuture.allOf(clientFutureMap.values().asScala.toSeq:
_*).get()
+ assert(
+ clientMap.size() == totalNumShuffleWriters.get(),
+ s"actual num of clients: ${clientMap.size()} " +
+ s"expected num clients: ${totalNumShuffleWriters.get()}"
+ )
+ }
+ } catch {
+ case th: Throwable =>
+ logError(log"Task discovery thread failed.", th)
+ errorNotifier.markError(th)
+ } finally {
+ clientCreationExecutor.shutdown()
+ logInfo(log"Task discovery thread exited. Took ${MDC(LogKeys.DURATION,
+ System.currentTimeMillis() - startTime)} ms")
+ }
+ })
+
+ private def stopTaskDiscovery(): Unit = {
+ taskDiscoveryShouldStop = true
+ shutdownExecutorService(taskDiscoveryExecutor, "Task Discovery Executor")
+ }
+
+ private def createClientAsync(
+ mapId: Int,
+ remoteHost: String,
+ remotePort: Int): CompletableFuture[TransportClient] = {
+ val future = new CompletableFuture[TransportClient]()
+ clientCreationExecutor.execute(() => {
+ try {
+ val shuffleClient = createClient(mapId, remoteHost, remotePort)
+ future.complete(shuffleClient)
+ } catch {
+ case th: Throwable =>
+ future.completeExceptionally(th)
+ }
+ })
+ future
+ }
+
+ protected def onTermAckResponse(shuffleWriterId: Int): Unit = {
+ terminationAckControlMessageSet.add(shuffleWriterId.toLong)
+ val curSent = terminationAckControlMessageSet.size()
+ val totalSentNeeded = totalNumShuffleWriters.get()
+ logInfo(log"Termination ack message sent successfully to shuffle writer " +
+ log"${MDC(LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}." +
+ log" ${MDC(LogKeys.NUM_TERMINATION_ACKS, curSent)} / " +
+ log"${MDC(LogKeys.NUM_SHUFFLE_WRITERS, totalSentNeeded)} sent
successfully.")
+
+ // if we have sent all term acks successfully
+ if (totalSentNeeded > 0 && curSent == totalSentNeeded) {
+ allTermAcksSentNotice.release()
+ }
+ }
+
+ /**
+ * Verifies the checksum of a DataMessage if checksum is enabled.
+ * @throws SparkRuntimeException if checksum verification fails
+ */
+ private def verifyDataMessageChecksum(dataMessage: DataMessage): Unit = {
+ if (shuffleChecksum != null) {
+ val data = dataMessage.data
+ shuffleChecksum.reset()
+ shuffleChecksum.updateChecksum(data, data.readerIndex(),
data.readableBytes())
+ val calculatedChecksum = shuffleChecksum.getValue()
+ if (dataMessage.checksum != calculatedChecksum) {
+ throw new SparkRuntimeException(
+ errorClass = "STREAMING_SHUFFLE_CHECKSUM_VERIFICATION_FAILED",
+ messageParameters = Map(
+ "writerId" -> dataMessage.shuffleWriterId.toString,
+ "expectedChecksum" -> dataMessage.checksum.toString,
+ "calculatedChecksum" -> calculatedChecksum.toString,
+ "dataLength" -> data.readableBytes().toString
+ )
+ )
+ }
+ }
+ }
+
+ // visible for test
+ protected def createClient(mapId: Int, remoteHost: String, remotePort: Int):
TransportClient = {
+ val handler = clientHandler.getOrElse(new StreamingShuffleClientHandler(
+ mapId,
+ context.partitionId(),
+ messageQueue,
+ streamingShuffleHandle.shuffleId,
+ perWriterByteLimit,
+ context,
+ errorNotifier
+ ))
+ handler.setOnTermAckResponseHandler(onTermAckResponse)
+ val clientContext =
+ new TransportContext(clientConf, handler)
+
+ val factory = clientContext.createClientFactory()
+ clientFactories.add(factory)
+ factory.createClient(remoteHost, remotePort)
+ }
+
+ private def checkTaskFailure(): Unit = {
+ context.getTaskFailure.foreach(throw _)
+ // Surface any background error on the task thread so that markTaskFailed
and
+ // markTaskCompleted are called from this thread, ensuring completion
listeners
+ // (including cleanupResources) run without contention.
+ errorNotifier.throwErrorIfExists()
+ if (context.isInterrupted()) {
+ throw new InterruptedException("Task interrupted. Exiting read loop.")
+ }
+ }
+
+ override def read(): Iterator[Product2[K, C]] = {
+ val serializerInstance =
streamingShuffleHandle.dependency.serializer.newInstance()
+ // Termination messages are added to a set that contains the shuffle
writer ids that have sent
+ // termination messages. When the set size reaches the number of shuffle
writers, we know
+ // that we will not receive any future messages, and the reader can be
closed. When a data
+ // message is read, the actual data (UnsafeRow) is extracted and emitted
through the iterator.
+ val terminationControlMessageSet = collection.mutable.Set[Long]()
+
+ /**
+ * Returns true if the reader should stop after handling the termination
message, which means
+ * we have received termination messages from all shuffle writers.
+ */
+ def handleTerminationMessage(msg: TerminationControlMessage): Boolean = {
+ val finishedId = msg.shuffleWriterId
+ terminationControlMessageSet += finishedId
+ val logMsg = if (totalNumShuffleWriters.get() > 0) {
+ log"Got termination message from ${MDC(LogKeys.SHUFFLE_WRITER_ID,
finishedId)}." +
+ log" ${MDC(LogKeys.NUM_TERMINATION_ACKS,
terminationControlMessageSet.size)} / ${
+ MDC(LogKeys.NUM_SHUFFLE_WRITERS, totalNumShuffleWriters.get())}" +
+ log" termination messages received."
+ } else {
+ log"Got termination message from ${MDC(LogKeys.SHUFFLE_WRITER_ID,
finishedId)}." +
+ log" ${MDC(LogKeys.NUM_TERMINATION_ACKS,
terminationControlMessageSet.size)}" +
+ log" termination messages received."
+ }
+ logInfo(logMsg)
+ if (totalNumShuffleWriters.get() > 0
+ && totalNumShuffleWriters.get() == terminationControlMessageSet.size) {
+ logInfo(
+ log"Got termination messages from all" +
+ log" shuffle writers ${MDC(LogKeys.SHUFFLE_WRITERS,
+ terminationControlMessageSet)}. Shutting down.")
+
+ // make sure all term acks have been sent successfully
+ while (!allTermAcksSentNotice.tryAcquire(100, TimeUnit.MILLISECONDS)) {
+ checkTaskFailure()
+ }
+ true
+ } else {
+ false
+ }
+ }
+
+ def handleDataMessage(dataMessage: DataMessage): Iterator[(K, C)] = {
+ currentDataMessage = dataMessage
+ verifyDataMessageChecksum(dataMessage)
+ val recordData = dataMessage.getRecordData()
+ val deserializedIterator = serializerInstance
+ .deserializeStream(new ByteBufInputStream(recordData))
+ .asKeyValueIterator
+ .asInstanceOf[Iterator[(K, C)]]
+ assert(
+ deserializedIterator.hasNext,
+ formatMessage(
+ s"data message that had no data to deserialize:" +
+ s" ${dataMessage.data}. Readable bytes: " +
+ s"${dataMessage.data.readableBytes()}"
+ )
+ )
+ new NextIterator[(K, C)] {
+ override def getNext(): (K, C) = {
+ if (!deserializedIterator.hasNext) {
+ finished = true
+ null.asInstanceOf[(K, C)]
+ } else {
+ deserializedIterator.next()
+ }
+ }
+ override def close(): Unit = {
+ dataMessage.release()
+ currentDataMessage = null
+ }
+ }
+ }
+
+ new StreamingShuffleReaderIteratorFactory().create(
+ messageQueue,
+ handleTerminationMessage,
+ handleDataMessage,
+ checkTaskFailure
+ )
+ }
+}
diff --git
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
new file mode 100644
index 000000000000..f035ffcdb055
--- /dev/null
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleReaderSuite.scala
@@ -0,0 +1,166 @@
+/*
+ * 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.concurrent.LinkedBlockingQueue
+
+import io.netty.buffer.Unpooled
+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
+import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.metrics.MetricsSystem
+import org.apache.spark.network.shuffle.streaming.{DataMessage,
StreamingShuffleMessage, TerminationControlMessage}
+import
org.apache.spark.shuffle.streaming.StreamingShuffleManager.QUERY_ID_PROPERTY_KEY
+
+/**
+ * Reader-side unit tests that do not require a shuffle writer. End-to-end
writer <-> reader
+ * behavior is covered by `StreamingShuffleSuite` in the tests PR of this
stack.
+ */
+class StreamingShuffleReaderSuite
+ 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)
+ }
+
+ // A minimal DataMessage that is only used to exercise the iterator
factory's type dispatch;
+ // its contents are never decoded here because handleDataMessage is stubbed
in these tests.
+ private def emptyDataMessage(): DataMessage =
+ new DataMessage(0, 0, 0, Unpooled.EMPTY_BUFFER, 0L)
+
+ // The iterator factory drives four collaborators; these tests supply
in-memory fakes for all of
+ // them so the reader's consumer-loop control flow can be verified without
Netty or a SparkEnv.
+ private val factory = new StreamingShuffleReaderIteratorFactory()
+
+ test("iterator emits all rows from data messages then stops on termination")
{
+ val queue = new LinkedBlockingQueue[StreamingShuffleMessage]()
+ queue.put(emptyDataMessage())
+ queue.put(emptyDataMessage())
+ queue.put(new TerminationControlMessage(0, 0))
+
+ // Each data message yields a fixed pair of rows; the single termination
message ends the read.
+ val rowsPerMessage = Seq(Iterator((1, 1), (2, 2)), Iterator((3, 3), (4,
4)))
+ val nextRows = rowsPerMessage.iterator
+ val it = factory.create[Int, Int](
+ queue,
+ handleTerminationMessage = _ => true,
+ handleDataMessage = _ => nextRows.next(),
+ checkTaskFailure = () => ())
+
+ it.toSeq should contain theSameElementsInOrderAs Seq((1, 1), (2, 2), (3,
3), (4, 4))
+ }
+
+ test("iterator does not stop until handleTerminationMessage returns true") {
+ val queue = new LinkedBlockingQueue[StreamingShuffleMessage]()
+ // Two writers: the first termination message must not end the read, only
the second.
+ queue.put(new TerminationControlMessage(0, 0))
+ queue.put(emptyDataMessage())
+ queue.put(new TerminationControlMessage(1, 0))
+
+ var terminationsSeen = 0
+ val it = factory.create[Int, Int](
+ queue,
+ handleTerminationMessage = _ => {
+ terminationsSeen += 1
+ terminationsSeen == 2 // only the second termination completes the read
+ },
+ handleDataMessage = _ => Iterator((1, 1)),
+ checkTaskFailure = () => ())
+
+ it.toSeq should contain theSameElementsInOrderAs Seq((1, 1))
+ terminationsSeen should be(2)
+ }
+
+ test("iterator surfaces a background error before dequeuing the next
message") {
+ val queue = new LinkedBlockingQueue[StreamingShuffleMessage]()
+ // A message is available in the queue, but a background error has already
been recorded
+ // (e.g. the channelInactive premature-disconnect path). Because
checkTaskFailure runs before
+ // every dequeue, the iterator must throw before this message is dequeued
and handled, rather
+ // than emitting stale data and only failing later.
+ queue.put(emptyDataMessage())
+
+ var dataHandled = false
+ val it = factory.create[Int, Int](
+ queue,
+ handleTerminationMessage = _ => true,
+ handleDataMessage = _ => {
+ dataHandled = true
+ Iterator((1, 1))
+ },
+ checkTaskFailure = () => throw new SparkException("background failure"))
+
+ val e = intercept[SparkException] {
+ it.hasNext
+ }
+ e.getMessage should include("background failure")
+ // The pending message must not have been processed: the error was
surfaced first. This is
+ // what would keep the reader from emitting rows after a writer has
already failed.
+ dataHandled should be(false)
+ }
+
+ test("getReader routes to a StreamingShuffleReader wired with the given
context") {
+ // A construction/routing smoke test: the manager must dispatch to the
streaming reader (not a
+ // fallback shuffle reader), the reader must construct successfully (its
constructor asserts the
+ // output tracker, starts the task-discovery and client-creation
executors, and registers the
+ // task-completion listener), and it must be wired with the context we
passed in.
+ withSpark(new SparkContext("local", "StreamingShuffleReaderSuite",
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 reader = new StreamingShuffleManager()
+ .getReader[Int, Int](handle, 0, 1, 0, 1, context, null)
+ reader shouldBe a[StreamingShuffleReader[_, _]]
+ val streamingReader = reader.asInstanceOf[StreamingShuffleReader[_, _]]
+ streamingReader.context should be theSameInstanceAs context
+ } finally {
+ // Constructing the reader starts a background task-discovery thread;
the task-completion
+ // listener (cleanupResources) shuts it down.
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
