This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new f208b741e2e3 [SPARK-57337][SS][RTM][STREAMINGSHUFFLE][PART3.5] Add
streaming shuffle shared transport and error plumbing
f208b741e2e3 is described below
commit f208b741e2e3ec52f5dbe58ef3d9d22b1792f8dc
Author: Boyang Jerry Peng <[email protected]>
AuthorDate: Fri Jun 26 17:58:16 2026 -0700
[SPARK-57337][SS][RTM][STREAMINGSHUFFLE][PART3.5] Add streaming shuffle
shared transport and error plumbing
### What changes were proposed in this pull request?
This is the foundation PR of a stacked set that adds a **streaming
shuffle** — a push-based shuffle used by low-latency continuous / Real-Time
Mode queries. It introduces only the
building blocks shared by both the writer (push) and reader (pull) paths,
so those two follow-up PRs can be reviewed in parallel and merged independently:
- **`ErrorNotifier`** (`core`, `org.apache.spark.shuffle.streaming`) —
records the first error raised on a background thread (e.g. a Netty event-loop
thread) so the owning task thread can
poll it and re-throw at a safe point.
- **`TransportClient.send(ByteBuf)`** (`network-common`) — a new overload
that sends a Netty `ByteBuf` and returns the `ChannelFuture`, plus widening the
existing `send(ByteBuffer)`
return type from `void` to `ChannelFuture` so callers can chain a
completion listener.
- **`spark.shuffle.streaming.checksum.enabled`** config — gates an
optional end-to-end CRC32C integrity check (writer embeds the checksum, reader
verifies it).
- Two shared structured-logging keys: `NUM_TERMINATION_ACKS`,
`SHUFFLE_WRITER_ID`.
This is part of a stack of PRs:
1. **(this PR)** shared transport + error plumbing
2. writer + server-side Netty handler — depends on this PR
3. reader + client-side Netty handler — depends on this PR
4. activation (`DAGScheduler` / `BlockManager` wiring)
5. tests (`StreamingShuffleSuite`)
6. documentation
These components are used by **both** the writer and reader sides, so
extracting them into a single prerequisite PR keeps the writer and reader PRs
independent and reviewable in parallel:
- **`ErrorNotifier`**: the streaming shuffle performs network I/O on
Netty event-loop threads. An error there must not be surfaced by calling
`TaskContext.markTaskFailed` from a
background thread, because that can race with `markTaskCompleted` and
skip task-completion/cleanup listeners (leaking resources). Instead the
background thread records the error via
`ErrorNotifier.markError`, and the task thread polls and re-throws it at
a safe point so cleanup runs on the task thread.
- **`TransportClient.send(ByteBuf)` returning `ChannelFuture`**: both
sides send Netty `ByteBuf` messages (writer: data buffers; reader:
credit-control / termination-ack) and need to
attach a listener to the send's completion. The existing
`send(ByteBuffer)` returned `void`, so a chainable future was required; the new
overload also avoids an extra copy by sending the
`ByteBuf` directly.
- The checksum config governs an optional end-to-end data-integrity check
shared by the writer and reader.
### Does this PR introduce _any_ user-facing change?
No behavior change. This PR adds the configuration
`spark.shuffle.streaming.checksum.enabled` (default `true`), but it has no
effect until the streaming shuffle writer/reader (follow-up
PRs) are in use. The new `TransportClient.send(ByteBuf)` overload and the
`send(ByteBuffer)` return-type widening (`void` -> `ChannelFuture`) are
source-compatible additions; no existing
caller consumes the returned future.
### How was this patch tested?
- `build/sbt core/Test/compile` — the module compiles standalone, and the
writer and reader PRs each compile independently on top of this one (confirming
the split is correct).
- `build/sbt "common-utils/testOnly org.apache.spark.util.LogKeysSuite"`
— verifies the new log keys keep `LogKeys.java` sorted.
- The shared components are exercised end-to-end by
`StreamingShuffleSuite` in the tests PR of this stack (writer<->reader data
transfer, termination handshake, checksum verification, and
background-thread error propagation).
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code
Closes #56387 from jerrypeng/stack/streaming-shuffle-pr3.5-shared.
Authored-by: Boyang Jerry Peng <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit f0c6af3637c6d66453f1e1883639706ee7d23069)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
.../spark/network/client/TransportClient.java | 23 +++++-
.../java/org/apache/spark/internal/LogKeys.java | 2 +
.../org/apache/spark/internal/config/package.scala | 15 ++++
.../org/apache/spark/util}/ErrorNotifier.scala | 17 ++++-
.../org/apache/spark/util/ErrorNotifierSuite.scala | 87 ++++++++++++++++++++++
.../streaming/checkpointing/AsyncCommitLog.scala | 2 +-
.../checkpointing/AsyncOffsetSeqLog.scala | 3 +-
.../streaming/runtime/AsyncLogPurge.scala | 2 +-
.../AsyncStreamingQueryCheckpointMetadata.scala | 2 +-
.../streaming/runtime/MicroBatchExecution.scala | 2 +-
...cProgressTrackingMicroBatchExecutionSuite.scala | 4 +-
11 files changed, 145 insertions(+), 14 deletions(-)
diff --git
a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java
b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java
index f02f2c63ecd4..f194a4cb94eb 100644
---
a/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java
+++
b/common/network-common/src/main/java/org/apache/spark/network/client/TransportClient.java
@@ -29,7 +29,9 @@ import javax.annotation.Nullable;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.SettableFuture;
+import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
+import io.netty.channel.ChannelFuture;
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.GenericFutureListener;
@@ -38,6 +40,7 @@ import org.apache.spark.internal.SparkLoggerFactory;
import org.apache.spark.internal.LogKeys;
import org.apache.spark.internal.MDC;
import org.apache.spark.network.buffer.ManagedBuffer;
+import org.apache.spark.network.buffer.NettyManagedBuffer;
import org.apache.spark.network.buffer.NioManagedBuffer;
import org.apache.spark.network.protocol.*;
import org.apache.spark.network.util.JavaUtils;
@@ -301,8 +304,24 @@ public class TransportClient implements Closeable {
*
* @param message The message to send.
*/
- public void send(ByteBuffer message) {
- channel.writeAndFlush(new OneWayMessage(new NioManagedBuffer(message)));
+ public ChannelFuture send(ByteBuffer message) {
+ return channel.writeAndFlush(new OneWayMessage(new
NioManagedBuffer(message)));
+ }
+
+ /**
+ * Sends an opaque message backed by a Netty {@link ByteBuf} to the
RpcHandler on the server-side.
+ * Like {@link #send(ByteBuffer)}, no reply is expected and no delivery
guarantees are made.
+ *
+ * Ownership of {@code message} is transferred to the transport layer: the
buffer is released once
+ * the write completes (whether it succeeds or fails), so the caller must
not release it
+ * afterwards and must retain it beforehand if the buffer is shared.
+ *
+ * @param message The message to send. Its ownership is transferred to the
transport layer.
+ * @return A {@link ChannelFuture} that completes when the message has been
written to the
+ * channel, so the caller can observe write failures or apply
backpressure.
+ */
+ public ChannelFuture send(ByteBuf message) {
+ return channel.writeAndFlush(new OneWayMessage(new
NettyManagedBuffer(message)));
}
/**
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 37064bf77631..32d34026e5a8 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
@@ -546,6 +546,7 @@ public enum LogKeys implements LogKey {
NUM_SUCCESSFUL_TASKS,
NUM_TASKS,
NUM_TASK_CPUS,
+ NUM_TERMINATION_ACKS,
NUM_TRAIN_WORD,
NUM_UNFINISHED_DECOMMISSIONED,
NUM_VALUES,
@@ -734,6 +735,7 @@ public enum LogKeys implements LogKey {
SHUFFLE_SERVICE_CONF_OVERLAY_URL,
SHUFFLE_SERVICE_METRICS_NAMESPACE,
SHUFFLE_SERVICE_NAME,
+ SHUFFLE_WRITER_ID,
SIGMAS_LENGTH,
SIGNAL,
SINK,
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 ce6c38990695..47620b2406ec 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
@@ -1781,6 +1781,21 @@ package object config {
.intConf
.createWithDefault(8)
+ //
---------------------------------------------------------------------------
+ // Streaming shuffle writer configs
+ //
---------------------------------------------------------------------------
+
+ private[spark] val STREAMING_SHUFFLE_CHECKSUM_ENABLED =
+ ConfigBuilder("spark.shuffle.streaming.checksum.enabled")
+ .doc("Whether to append a CRC32C checksum to each streaming shuffle data
buffer. " +
+ "When enabled, the writer computes the checksum and embeds it in the
DataMessage header; " +
+ "the reader recomputes and compares. A mismatch fails the task,
providing early " +
+ "detection of data corruption in transit.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .booleanConf
+ .createWithDefault(true)
+
private[spark] val SHUFFLE_DETECT_CORRUPT =
ConfigBuilder("spark.shuffle.detectCorrupt")
.doc("Whether to detect any corruption in fetched blocks.")
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ErrorNotifier.scala
b/core/src/main/scala/org/apache/spark/util/ErrorNotifier.scala
similarity index 65%
rename from
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ErrorNotifier.scala
rename to core/src/main/scala/org/apache/spark/util/ErrorNotifier.scala
index 2ba696173954..18ed7688f40a 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/ErrorNotifier.scala
+++ b/core/src/main/scala/org/apache/spark/util/ErrorNotifier.scala
@@ -15,16 +15,25 @@
* limitations under the License.
*/
-package org.apache.spark.sql.execution.streaming.runtime
+package org.apache.spark.util
import java.util.concurrent.atomic.AtomicReference
import org.apache.spark.internal.Logging
/**
- * Class to notify of any errors that might have occurred out of band
+ * Captures an error raised out of band (e.g. on a Netty event-loop or other
background thread) so
+ * the owning task / driver thread can surface it at a safe point.
+ *
+ * A background thread records the first error via [[markError]] rather than
failing the task
+ * directly; the owning thread polls [[throwErrorIfExists]] (or [[getError]])
at a safe point and
+ * re-throws on its own thread. This avoids a race where calling
`TaskContext.markTaskFailed` from
+ * a background thread can run concurrently with the task thread's
`markTaskCompleted` and silently
+ * skip completion listeners (including cleanup), leading to thread leaks.
+ *
+ * Shared by the streaming shuffle and the streaming-query async checkpointing
paths.
*/
-class ErrorNotifier extends Logging {
+private[spark] class ErrorNotifier extends Logging {
private val error = new AtomicReference[Throwable]
@@ -34,7 +43,7 @@ class ErrorNotifier extends Logging {
*/
def markError(th: Throwable): Unit = {
if (error.compareAndSet(null, th)) {
- logError("A fatal error has occurred.", th)
+ logError(log"A fatal error has occurred.", th)
} else {
// Attach subsequent errors as suppressed so they're not silently lost.
val existing = error.get()
diff --git a/core/src/test/scala/org/apache/spark/util/ErrorNotifierSuite.scala
b/core/src/test/scala/org/apache/spark/util/ErrorNotifierSuite.scala
new file mode 100644
index 000000000000..313857eb1de1
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/util/ErrorNotifierSuite.scala
@@ -0,0 +1,87 @@
+/*
+ * 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.util
+
+import java.util.concurrent.CyclicBarrier
+
+import org.apache.spark.SparkFunSuite
+
+class ErrorNotifierSuite extends SparkFunSuite {
+
+ test("no error is recorded by default") {
+ val notifier = new ErrorNotifier
+ assert(notifier.getError().isEmpty)
+ // throwErrorIfExists is a no-op when no error has been recorded.
+ notifier.throwErrorIfExists()
+ }
+
+ test("markError records the error and throwErrorIfExists re-throws it") {
+ val notifier = new ErrorNotifier
+ val err = new RuntimeException("boom")
+ notifier.markError(err)
+
+ assert(notifier.getError().contains(err))
+ val thrown = intercept[RuntimeException] {
+ notifier.throwErrorIfExists()
+ }
+ assert(thrown eq err)
+ }
+
+ test("only the first error is retained; later errors are attached as
suppressed") {
+ val notifier = new ErrorNotifier
+ val first = new RuntimeException("first")
+ val second = new IllegalStateException("second")
+ notifier.markError(first)
+ notifier.markError(second)
+
+ assert(notifier.getError().contains(first))
+ assert(first.getSuppressed.toSeq.contains(second))
+ }
+
+ test("marking the same error twice does not self-suppress") {
+ val notifier = new ErrorNotifier
+ val err = new RuntimeException("boom")
+ notifier.markError(err)
+ notifier.markError(err)
+
+ assert(notifier.getError().contains(err))
+ assert(err.getSuppressed.isEmpty)
+ }
+
+ test("concurrent markError retains exactly one error and suppresses the
rest") {
+ val notifier = new ErrorNotifier
+ val numThreads = 8
+ val errors = (0 until numThreads).map(i => new RuntimeException(s"err-$i"))
+ val barrier = new CyclicBarrier(numThreads)
+ val threads = errors.map { e =>
+ new Thread(() => {
+ // Maximize contention on the compareAndSet by releasing all threads
at once.
+ barrier.await()
+ notifier.markError(e)
+ })
+ }
+ threads.foreach(_.start())
+ threads.foreach(_.join())
+
+ val winner = notifier.getError()
+ assert(winner.isDefined)
+ assert(errors.contains(winner.get))
+ // Every losing error must be attached as suppressed on the winner, so
none is lost.
+ assert(winner.get.getSuppressed.toSet == errors.filterNot(_ eq
winner.get).toSet)
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncCommitLog.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncCommitLog.scala
index d13affb82dbb..e0642e7a2f9b 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncCommitLog.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncCommitLog.scala
@@ -25,7 +25,7 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.internal.LogKeys
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.errors.QueryExecutionErrors
-import org.apache.spark.sql.execution.streaming.runtime.ErrorNotifier
+import org.apache.spark.util.ErrorNotifier
/**
* Implementation of CommitLog to perform asynchronous writes to storage
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncOffsetSeqLog.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncOffsetSeqLog.scala
index eea10f3d5a93..14fb4620527c 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncOffsetSeqLog.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/checkpointing/AsyncOffsetSeqLog.scala
@@ -26,8 +26,7 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.internal.{LogKeys}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.errors.QueryExecutionErrors
-import org.apache.spark.sql.execution.streaming.runtime.ErrorNotifier
-import org.apache.spark.util.{Clock, SystemClock}
+import org.apache.spark.util.{Clock, ErrorNotifier, SystemClock}
/**
* Used to write entries to the offset log asynchronously
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncLogPurge.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncLogPurge.scala
index 43d5a50b2ccd..df9011f07340 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncLogPurge.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncLogPurge.scala
@@ -23,7 +23,7 @@ import org.apache.spark.internal.Logging
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.util.ThreadUtils
+import org.apache.spark.util.{ErrorNotifier, ThreadUtils}
/**
* Used to enable the capability to allow log purges to be done asynchronously
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncStreamingQueryCheckpointMetadata.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncStreamingQueryCheckpointMetadata.scala
index 39622392aa62..2662c11133c4 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncStreamingQueryCheckpointMetadata.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/AsyncStreamingQueryCheckpointMetadata.scala
@@ -20,7 +20,7 @@ import java.util.concurrent.ThreadPoolExecutor
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.execution.streaming.checkpointing.{AsyncCommitLog,
AsyncOffsetSeqLog}
-import org.apache.spark.util.Clock
+import org.apache.spark.util.{Clock, ErrorNotifier}
/**
* A version of [[StreamingQueryCheckpointMetadata]] that supports async state
checkpointing.
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
index ef48d8b92927..dd87e4a9e43f 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
@@ -55,7 +55,7 @@ import
org.apache.spark.sql.execution.streaming.utils.StreamingUtils
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.internal.connector.PartitionOffsetWithIndex
import org.apache.spark.sql.streaming.Trigger
-import org.apache.spark.util.{Clock, Utils}
+import org.apache.spark.util.{Clock, ErrorNotifier, Utils}
class MicroBatchExecution(
sparkSession: SparkSession,
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
index 2ea409e90672..ba7af98b3766 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/AsyncProgressTrackingMicroBatchExecutionSuite.scala
@@ -30,13 +30,13 @@ import org.apache.spark.TestUtils
import org.apache.spark.sql._
import org.apache.spark.sql.connector.read.streaming
import org.apache.spark.sql.execution.streaming.checkpointing.{AsyncCommitLog,
AsyncOffsetSeqLog, CommitMetadata, OffsetSeq}
-import
org.apache.spark.sql.execution.streaming.runtime.{AsyncProgressTrackingMicroBatchExecution,
ErrorNotifier, LongOffset, MemoryStream, StreamExecution}
+import
org.apache.spark.sql.execution.streaming.runtime.{AsyncProgressTrackingMicroBatchExecution,
LongOffset, MemoryStream, StreamExecution}
import
org.apache.spark.sql.execution.streaming.runtime.AsyncProgressTrackingMicroBatchExecution.{ASYNC_PROGRESS_TRACKING_CHECKPOINTING_INTERVAL_MS,
ASYNC_PROGRESS_TRACKING_ENABLED,
ASYNC_PROGRESS_TRACKING_OVERRIDE_SINK_SUPPORT_CHECK}
import org.apache.spark.sql.functions.{column, window}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.streaming.{StreamingQuery,
StreamingQueryException, StreamTest, Trigger}
import org.apache.spark.sql.streaming.util.StreamManualClock
-import org.apache.spark.util.{ThreadUtils, Utils}
+import org.apache.spark.util.{ErrorNotifier, ThreadUtils, Utils}
class AsyncProgressTrackingMicroBatchExecutionSuite
extends StreamTest with BeforeAndAfter with Matchers {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
