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 3c4640d96a37 [SPARK-58848][CORE][SS] Prevent unbounded shuffle reader
connection waits
3c4640d96a37 is described below
commit 3c4640d96a3752b0b483e689bd0d4b028d82fef1
Author: Zerui Bao <[email protected]>
AuthorDate: Tue Aug 25 06:46:29 2026 +0900
[SPARK-58848][CORE][SS] Prevent unbounded shuffle reader connection waits
### What changes were proposed in this pull request?
Prevent unbounded shuffle reader connection waits. Instead, timeout after 1
hour.
### Why are the changes needed?
unbounded shuffle reader connection waits is unsafe.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
UT
### Was this patch authored or co-authored using generative AI tooling?
Cursor
Closes #58089 from zeruibao/fix-streaming-shuffle-writer-connection-timeout.
Authored-by: Zerui Bao <[email protected]>
Signed-off-by: Jungtaek Lim <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 8 +++
.../org/apache/spark/internal/config/package.scala | 12 ++++
.../streaming/StreamingShuffleManager.scala | 14 +++++
.../shuffle/streaming/StreamingShuffleWriter.scala | 54 +++++++++++++---
.../streaming/StreamingShuffleWriterSuite.scala | 73 +++++++++++++++++++++-
5 files changed, 150 insertions(+), 11 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index c642e0be953a..5af823a3f24b 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -7918,6 +7918,14 @@
],
"sqlState" : "XXKST"
},
+ "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT" : {
+ "message" : [
+ "Streaming shuffle <shuffleId> writer <writerId> timed out after
<timeoutMs> ms waiting for reader <readerId> to connect.",
+ "This can occur when the reader task is delayed because cluster
resources are unavailable or when network connectivity prevents the reader from
reaching the writer.",
+ "Check cluster resources and network connectivity. If the reader needs
more time to start, increase 'spark.shuffle.streaming.writerConnectionTimeout',
or set it to -1 to wait indefinitely."
+ ],
+ "sqlState" : "XXKST"
+ },
"STREAMING_STATEFUL_OPERATOR_MISSING_STATE_DIRECTORY" : {
"message" : [
"Cannot restart streaming query with stateful operators because the
state directory is empty or missing.",
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 86a63ffc1c2e..6303f751c9c1 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
@@ -1935,6 +1935,18 @@ package object config {
.longConf
.createWithDefault(50)
+ private[spark] val STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS =
+ ConfigBuilder("spark.shuffle.streaming.writerConnectionTimeout")
+ .doc("Maximum time a streaming shuffle writer waits for each reader to
connect. " +
+ "Set to -1 to wait indefinitely.")
+ .version("4.4.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .timeConf(TimeUnit.MILLISECONDS)
+ .checkValue(
+ timeoutMs => timeoutMs == -1 || timeoutMs > 0,
+ "The reader connection timeout must be positive or -1 to wait
indefinitely.")
+ .createWithDefaultString("1h")
+
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 " +
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 b83bd33552b9..5ff9f6339425 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
@@ -74,6 +74,20 @@ object StreamingShuffleManager extends Logging {
errorClass = "STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE",
messageParameters = Map("messageType" -> messageType.toString))
}
+
+ def streamingShuffleWriterConnectionTimeout(
+ shuffleId: Int,
+ writerId: Int,
+ readerId: Int,
+ timeoutMs: Long): RuntimeException = {
+ new SparkRuntimeException(
+ errorClass = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT",
+ messageParameters = Map(
+ "shuffleId" -> shuffleId.toString,
+ "writerId" -> writerId.toString,
+ "readerId" -> readerId.toString,
+ "timeoutMs" -> timeoutMs.toString))
+ }
}
private[spark] class StreamingShuffleManager extends PipelinedShuffleManager
with Logging {
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
index b0cc249f7342..0b6719127fd1 100644
---
a/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala
+++
b/core/src/main/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriter.scala
@@ -17,7 +17,7 @@
package org.apache.spark.shuffle.streaming
-import java.util.concurrent.{CancellationException, CompletableFuture,
CountDownLatch, LinkedBlockingDeque, Semaphore, TimeUnit}
+import java.util.concurrent.{CancellationException, CompletableFuture,
CompletionException, CountDownLatch, LinkedBlockingDeque, Semaphore,
TimeoutException, TimeUnit}
import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference}
import javax.annotation.concurrent.NotThreadSafe
@@ -29,7 +29,13 @@ 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.{
+ EXECUTOR_ID,
+ STREAMING_SHUFFLE_CHECKSUM_ENABLED,
+ STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS,
+ STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE,
+ STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS,
+ 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
@@ -58,6 +64,7 @@ class StreamingShuffleWriter[K, V](
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)
+ private val CONNECTION_TIMEOUT_MS =
conf.get(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS)
// Shuffle details.
private val streamingShuffleHandle =
handle.asInstanceOf[StreamingShuffleHandle[K, V, _]]
@@ -205,11 +212,27 @@ class StreamingShuffleWriter[K, V](
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
- }))
+ Right(
+ (if (CONNECTION_TIMEOUT_MS == -1) {
+ transportServerHandler.futureClients(id)
+ } else {
+ transportServerHandler.futureClients(id)
+ .orTimeout(CONNECTION_TIMEOUT_MS, TimeUnit.MILLISECONDS)
+ .exceptionally {
+ case _: TimeoutException =>
+ throw
StreamingShuffleManager.streamingShuffleWriterConnectionTimeout(
+ streamingShuffleHandle.shuffleId,
+ shuffleWriterId,
+ id,
+ CONNECTION_TIMEOUT_MS)
+ case error =>
+ throw error
+ }
+ }).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)
@@ -257,7 +280,10 @@ class StreamingShuffleWriter[K, V](
val newFuture = future.whenComplete { (client, ex) =>
ex match {
case null => sendToClient(client)
- case _ => buf.release(); done()
+ case error =>
+ buf.release()
+ errorNotifier.markError(error)
+ done()
}
}
// Once the future is completed, stop accumulating CompletionStages.
@@ -318,7 +344,7 @@ class StreamingShuffleWriter[K, V](
// For testing only.
def hasClient: Boolean = client match {
case Left(_) => true
- case Right(future) => future.isDone
+ case Right(future) => future.isDone && !future.isCompletedExceptionally
}
}
@@ -395,6 +421,11 @@ class StreamingShuffleWriter[K, V](
System.currentTimeMillis() - cleanupStartTime)} ms")
}
+ private def unwrapCompletionException(error: Throwable): Throwable = error
match {
+ case e: CompletionException if e.getCause != null => e.getCause
+ case _ => error
+ }
+
private def throwErrorIfExists(): Unit = {
context.getTaskFailure.foreach { throw _ }
errorNotifier.throwErrorIfExists()
@@ -430,7 +461,7 @@ class StreamingShuffleWriter[K, V](
Try {
while (!isWriteFinished.await(MAX_BUFFERING_TIME_MS,
TimeUnit.MILLISECONDS))
shards.foreach(_.send())
- }.recover { case e => errorNotifier.markError(e) }
+ }.recover { case e =>
errorNotifier.markError(unwrapCompletionException(e)) }
, "time-based-flush-for-shuffle-writer-" +
s"${streamingShuffleHandle.shuffleId}-${shuffleWriterId}")
try {
@@ -508,6 +539,9 @@ class StreamingShuffleWriter[K, V](
logInfo(log"Received all termination acks for shuffle writer ${MDC(
LogKeys.SHUFFLE_WRITER_ID, shuffleWriterId)}. Closing server channel.")
throwErrorIfExists()
+ } catch {
+ case e: CompletionException =>
+ throw unwrapCompletionException(e)
} 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
index e421272cc119..4cbbf44daa36 100644
---
a/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
+++
b/core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleWriterSuite.scala
@@ -25,12 +25,20 @@ 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.concurrent.Eventually.eventually
+import org.scalatest.concurrent.PatienceConfiguration.Timeout
import org.scalatest.matchers.should.Matchers
+import org.scalatest.time.SpanSugar._
import org.scalatestplus.mockito.MockitoSugar
import org.apache.spark._
import org.apache.spark.LocalSparkContext.withSpark
-import org.apache.spark.internal.config.{SHUFFLE_MANAGER_INCREMENTAL,
STREAMING_SHUFFLE_CHECKSUM_ENABLED, STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE}
+import org.apache.spark.internal.config.{
+ SHUFFLE_MANAGER_INCREMENTAL,
+ STREAMING_SHUFFLE_CHECKSUM_ENABLED,
+ STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS,
+ STREAMING_SHUFFLE_NETWORK_BUFFER_SIZE,
+ STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS}
import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
import org.apache.spark.metrics.MetricsSystem
import org.apache.spark.network.client.TransportClient
@@ -229,6 +237,69 @@ class StreamingShuffleWriterSuite
}
}
+ test("writer fails when a reader does not connect before timeout") {
+ val conf = newConf()
+ .set(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS, 100L)
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf))
{ sc =>
+ val context = createTaskContext(sc.conf, 0)
+ try {
+ val writer = newWriter(sc, context)
+ val error = intercept[SparkRuntimeException] {
+ writer.write(Iterator.empty)
+ }
+ checkError(
+ exception = error,
+ condition = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT",
+ sqlState = "XXKST",
+ parameters = Map(
+ "shuffleId" -> "0",
+ "writerId" -> "0",
+ "readerId" -> "0",
+ "timeoutMs" -> "100"))
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
+ test("writer surfaces a reader connection timeout from the flush thread") {
+ val timeoutMs = 100L
+ val flushIntervalMs = 200L
+ val conf = newConf()
+ .set(STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT_MS, timeoutMs)
+ .set(STREAMING_SHUFFLE_NETWORK_BUFFER_MAX_WAIT_TIME_MS, flushIntervalMs)
+ withSpark(new SparkContext("local", "StreamingShuffleWriterSuite", conf))
{ sc =>
+ val context = createTaskContext(sc.conf, 0)
+ try {
+ val writer = newWriter(sc, context)
+ val records = Iterator((1, 1), (2, 2)).map { record =>
+ if (record._1 == 2) {
+ // Block before yielding the second record until the flush thread
observes the
+ // connection timeout while sending the first record.
+ eventually(Timeout(10.seconds)) {
+ writer.errorNotifier.getError() shouldBe defined
+ }
+ }
+ record
+ }
+ val error = intercept[SparkRuntimeException] {
+ writer.write(records)
+ }
+ checkError(
+ exception = error,
+ condition = "STREAMING_SHUFFLE_WRITER_CONNECTION_TIMEOUT",
+ sqlState = "XXKST",
+ parameters = Map(
+ "shuffleId" -> "0",
+ "writerId" -> "0",
+ "readerId" -> "0",
+ "timeoutMs" -> timeoutMs.toString))
+ } finally {
+ context.markTaskCompleted(None)
+ }
+ }
+ }
+
test("cleanupResources releases queued pooled buffers") {
withSpark(new SparkContext("local", "StreamingShuffleWriterSuite",
newConf())) { sc =>
val context = createTaskContext(sc.conf, 0)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
