You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/05/22 17:38:53 UTC
(spark) branch branch-4.x updated: [SPARK-56674][SS][RTM][STREAMINGSHUFFLE] Add streaming shuffle wire protocol
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 2df345460d22 [SPARK-56674][SS][RTM][STREAMINGSHUFFLE] Add streaming shuffle wire protocol
2df345460d22 is described below
commit 2df345460d224a8cdf395672f264749f09c3b618
Author: Boyang Jerry Peng <je...@databricks.com>
AuthorDate: Fri May 22 10:38:33 2026 -0700
[SPARK-56674][SS][RTM][STREAMINGSHUFFLE] Add streaming shuffle wire protocol
### What changes were proposed in this pull request?
### Context
This is the first PR in a stack that contributes a new **streaming shuffle** implementation to Apache Spark. Streaming shuffle is a push-based shuffle designed for low-latency,
continuously-running queries (e.g., Real-Time mode in Structured Streaming) where the map and reduce stages must run concurrently rather than sequentially. Each map task hosts a Netty
server that pushes records to reduce tasks as they are produced; reduce tasks open clients to those servers and consume records as a stream — no on-disk materialization, no map-stage
barrier.
Because the full implementation spans the network protocol, driver-side coordination, plugin layer, executor-side writer/reader, engine integration, and tests, it is split into nine
independently reviewable PRs:
1. **Wire protocol (this PR)** — binary message types in `network-common`, pure Java.
2. **Output tracker** — driver-side coordination service mapping shuffle IDs to writer task locations.
3. **Shuffle manager + Netty handlers + logging mixin** — the `ShuffleManager` plugin entry point, the bidirectional Netty handlers, and a `TaskContext`-aware logging trait.
4. **SparkEnv + DAGScheduler integration** — wires the tracker into `SparkEnv` and registers shuffles from `DAGScheduler`.
5. **Writer** — `StreamingShuffleWriter` (server-side push).
6. **Reader** — `StreamingShuffleReader` (client-side pull).
7. **MultiShuffleManager** — routes per-shuffle to streaming or sort shuffle based on a task-local property, so one cluster can host both.
8. **Tests** — end-to-end suite for the streaming shuffle plugin.
9. **Documentation** — design and configuration reference.
Each PR compiles standalone. The plugin only becomes usable end-to-end after PRs 1–7.
### Changes in this PR
This PR introduces the binary wire protocol used by the streaming shuffle, contained in a new package `org.apache.spark.network.shuffle.streaming` under `common/network-common`. All
classes are pure Java and depend only on Netty (already in `network-common`) and `java.util.zip.CRC32C` (Java 9+). No new Maven dependencies are introduced.
The protocol consists of four message types sharing a 12-byte common header (4-byte message-type id + 8-byte sequence number):
- **`StreamingShuffleMessage`** — abstract sealed base class. Owns a `ByteBuf`, exposes `setSeqNum`/`getSeqNum`, and dispatches `decode()` to the concrete subclass based on the
message-type id read from the wire. Documents the refcount/ownership rules concrete subclasses must follow.
- **`StreamingShuffleMessageType`** — enum of the four message IDs (`DATA_MESSAGE_UNSAFE_ROW=1`, `CREDIT_CONTROL_MESSAGE=2`, `TERMINATION_CONTROL_MESSAGE=3`, `TERMINATION_ACK_MESSAGE=4`).
Each enum constant carries an explicit `id()` value that is used on the wire — not the JVM enum ordinal — so the encoding is stable against changes to declaration order.
`StreamingShuffleMessageType.decode(int)` throws `IllegalArgumentException` when the wire id does not match any known message type.
- **`DataMessage`** — writer → reader. Carries serialized records along with `(shuffleWriterId, shuffleReaderId, dataSize, CRC32C checksum)` in the header. The constructor validates
`dataSize == data.readableBytes()`; `encode()` writes exactly `dataSize` bytes of payload (via `data.retainedSlice(data.readerIndex(), dataSize)`); `decode()` validates that the
post-header readable bytes equal `dataSize` exactly, rejecting both truncated and trailing-junk frames at the wire boundary.
- **`CreditControlMessage`** — reader → writer. Sent on connection establishment as the initial handshake; reserved for finer-grained credit-based flow control in future revisions.
- **`TerminationControlMessage`** — writer → reader. End-of-stream signal.
- **`TerminationAckMessage`** — reader → writer. Echoes the last sequence number the reader observed, so the writer can verify no messages were lost or reordered.
Supporting class:
- **`ShuffleChecksum`** — stateful CRC32C helper that supports both heap and direct `ByteBuf`s, used by the writer to stamp checksums into outgoing `DataMessage`s and by the reader to
verify them. Range arguments are validated against the buffer's `writerIndex()` so the checksum cannot accidentally cover bytes in `[writerIndex, capacity)` that have not been written.
Tests in `StreamingShuffleMessageSuite` cover:
- Encode/decode round-trips for every message type
- CRC32C determinism and equivalence across heap and direct buffers; sub-range checksums with non-zero `startIndex` compared against a reference computation; out-of-bounds,
beyond-`writerIndex`, and negative-argument rejection
- `DataMessage` framing: rejection of `dataSize` larger/smaller than `data.readableBytes()`, negative `dataSize`, and decoded frames with trailing junk or truncation
- Unknown wire message-type ids → `IllegalArgumentException`
The protocol is designed to be reviewable on its own — the message layout, sequence-numbering scheme, memory-ownership rules, framing invariants, and checksum approach can all be
evaluated without reading any of the executor- or driver-side code that comes in the follow-up PRs.
### Why are the changes needed?
**This is one part of creating the Streaming Shuffle needed for Real-time Mode (RTM).**
Spark's existing `SortShuffleManager` materializes map outputs to local disk and gates the reduce stage on map-stage completion. This design is well-suited to batch jobs but introduces
two latency floors that are unacceptable for continuously-running, low-latency queries:
1. **Map-stage barrier.** Reducers cannot start until every mapper has finished writing its output files. End-to-end latency is bound by the slowest mapper.
2. **Disk materialization.** Every record is serialized to a local file and read back, even when the consumer is ready immediately. The fsync/read path adds tens of milliseconds of
overhead per batch.
Real-Time mode in Structured Streaming, and other long-running query workloads where map and reduce tasks coexist for the lifetime of the query, require a shuffle that:
- Streams records from mappers to reducers as they are produced.
- Lets the reduce stage run concurrently with the map stage.
- Avoids on-disk intermediate state entirely.
- Provides backpressure so a slow reader can throttle the upstream iterator without dropping data.
- Detects message loss, reordering, or in-flight corruption (since there is no on-disk re-read to fall back on).
Streaming shuffle is the new `ShuffleManager` implementation that meets these requirements. The full implementation involves driver-side coordination, executor-side push/pull components,
engine integration, and supporting infrastructure — together too large to land in a single PR.
This PR contributes the **wire protocol** layer, which is the natural foundation. Every other piece in the stack (the tracker, writer, reader, handlers, tests) depends on the message
types, sequence-numbering rules, and checksum scheme defined here, so it must land first. The protocol layer is also the most contract-heavy part of the design: getting the message
layout, ownership rules, and integrity check right up front avoids costly churn in later PRs. By making this an independent, standalone-reviewable PR, reviewers can scrutinize the wire
format, refcount discipline, and CRC32C approach without being distracted by Spark-internal coordination logic that has not yet been introduced.
The final product can be viewed here for context:
https://github.com/apache/spark/compare/master...jerrypeng:spark:stack/streaming-shuffle-pr9-docs?expand=1
### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?
Unit tests cover encode/decode round-trips for all message types, determinism of the checksum, heap/direct buffer equivalence, and the unknown-type error path.
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #55620 from jerrypeng/stack/streaming-shuffle-pr1-network-protocol.
Authored-by: Boyang Jerry Peng <je...@databricks.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
(cherry picked from commit 7179629dc642aa4fcaef0704ede4bb34ebd3f8f3)
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
.../shuffle/streaming/CreditControlMessage.java | 91 +++++
.../network/shuffle/streaming/DataMessage.java | 131 +++++++
.../network/shuffle/streaming/ShuffleChecksum.java | 71 ++++
.../shuffle/streaming/StreamingShuffleMessage.java | 146 ++++++++
.../streaming/StreamingShuffleMessageType.java | 45 +++
.../shuffle/streaming/TerminationAckMessage.java | 68 ++++
.../streaming/TerminationControlMessage.java | 68 ++++
.../streaming/StreamingShuffleMessageSuite.java | 399 +++++++++++++++++++++
8 files changed, 1019 insertions(+)
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java
new file mode 100644
index 000000000000..4ba17b4a011a
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java
@@ -0,0 +1,91 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
+/**
+ * Reader → writer control message.
+ *
+ * Current function: serves as a connection-establishment and per-consumption "ready"
+ * signal. The writer uses receipt of any CreditControlMessage as the trigger that the
+ * reader is ready to receive, and does not act on the numeric value of {@link
+ * #numMessages}. Backpressure today is handled at the TCP layer via channel autoRead,
+ * not via this message.
+ *
+ * Future function: this message is reserved as the carrier for a future credit-based
+ * flow-control extension. When that extension lands, {@link #numMessages} will carry
+ * the number of additional DataMessages the writer may send beyond any
+ * previously-granted credit (i.e., a credit-grant delta).
+ */
+public final class CreditControlMessage extends StreamingShuffleMessage {
+ public final int shuffleWriterId;
+ public final int shuffleReaderId;
+
+ /**
+ * In the current protocol revision the writer ignores this value and treats any
+ * CreditControlMessage as a "reader is ready" signal; senders should pass 1.
+ *
+ * Reserved for the future credit-based flow-control extension, in which this field
+ * will carry the number of additional DataMessages the writer may send beyond any
+ * previously-granted credit.
+ */
+ public final int numMessages;
+
+ public CreditControlMessage(int shuffleWriterId, int shuffleReaderId, int numMessages) {
+ this.shuffleWriterId = shuffleWriterId;
+ this.shuffleReaderId = shuffleReaderId;
+ this.numMessages = numMessages;
+ }
+
+ @Override
+ public StreamingShuffleMessageType messageType() {
+ return StreamingShuffleMessageType.CREDIT_CONTROL_MESSAGE;
+ }
+
+ @Override
+ public int headerLength() {
+ // 4 bytes for the shuffle writer ID, 4 bytes for the shuffle reader ID,
+ // 4 bytes for the number of messages
+ return super.headerLength() + 12;
+ }
+
+ @Override
+ public void encode(CompositeByteBuf buf) {
+ super.encode(buf);
+
+ // Write the shuffle writer ID
+ buf.writeInt(shuffleWriterId);
+ // Write the shuffle reader ID
+ buf.writeInt(shuffleReaderId);
+ // Write the number of messages
+ buf.writeInt(numMessages);
+ }
+
+ public static CreditControlMessage decode(ByteBuf buf) {
+ // Read the shuffle writer ID
+ int shuffleWriterId = buf.readInt();
+ // Read the shuffle reader ID
+ int shuffleReaderId = buf.readInt();
+ // Read the number of messages
+ int numMessages = buf.readInt();
+
+ return new CreditControlMessage(shuffleWriterId, shuffleReaderId, numMessages);
+ }
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java
new file mode 100644
index 000000000000..efcca96cce56
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java
@@ -0,0 +1,131 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
+/**
+ * Writer → reader data message. Carries a contiguous payload of serialized shuffle
+ * records together with the routing metadata (shuffleWriterId, shuffleReaderId,
+ * dataSize) and a CRC32C {@link #checksum} over the payload. Each DataMessage
+ * corresponds to one network frame on the streaming-shuffle wire.
+ *
+ * Memory ownership: the {@code data} ByteBuf passed to the constructor is retained
+ * once and stored as {@link #ownedBuf}, so callers must release their own reference
+ * separately. On {@link #release()} the retained reference is dropped. See the
+ * memory-ownership rules in {@link StreamingShuffleMessage}.
+ */
+public final class DataMessage extends StreamingShuffleMessage {
+
+ public final ByteBuf data;
+ public final int shuffleWriterId;
+ public final int shuffleReaderId;
+ public final int dataSize;
+ public final long checksum;
+
+ public DataMessage(int shuffleWriterId, int shuffleReaderId, int dataSize, ByteBuf data,
+ long checksum) {
+ if (dataSize < 0) {
+ throw new IllegalArgumentException(
+ "dataSize must be non-negative: " + dataSize);
+ }
+ if (dataSize != data.readableBytes()) {
+ throw new IllegalArgumentException(
+ "dataSize must equal data.readableBytes(): " +
+ dataSize + " != " + data.readableBytes());
+ }
+ this.shuffleWriterId = shuffleWriterId;
+ this.shuffleReaderId = shuffleReaderId;
+ this.dataSize = dataSize;
+ this.data = data;
+ this.ownedBuf = data.retain();
+ this.checksum = checksum;
+ }
+
+ @Override
+ public StreamingShuffleMessageType messageType() {
+ return StreamingShuffleMessageType.DATA_MESSAGE_UNSAFE_ROW;
+ }
+
+ @Override
+ public int headerLength() {
+ // 4 bytes EACH for shuffle writer ID, shuffle reader ID, data size
+ // 8 bytes for checksum
+ return super.headerLength() + 20;
+ }
+
+ @Override
+ public void encode(CompositeByteBuf buf) {
+ super.encode(buf);
+ buf.writeInt(shuffleWriterId);
+ buf.writeInt(shuffleReaderId);
+ buf.writeInt(dataSize);
+ buf.writeLong(checksum);
+
+ // Only encode exactly `dataSize` bytes of `data`, so the wire payload matches the
+ // header's dataSize even if the underlying ByteBuf has additional readable bytes.
+ // retainedSlice() also bumps the refcount, so this DataMessage's ownedBuf reference
+ // remains valid for later release().
+ buf.addComponent(true, data.retainedSlice(data.readerIndex(), dataSize));
+ }
+
+ /**
+ * Decodes a {@link DataMessage} from {@code message}. The 12-byte common header
+ * (message-type id + sequence number) has already been consumed by
+ * {@link StreamingShuffleMessage#decode(ByteBuf)}; this method reads the remaining
+ * 20-byte DataMessage header and treats the rest of {@code message} as the payload.
+ *
+ * On success, the returned DataMessage retains a reference to {@code message} via
+ * its {@code ownedBuf} field (refcount is incremented in the constructor), and the
+ * caller transfers ownership.
+ *
+ * On failure (an {@link IllegalArgumentException} is thrown — see the constructor's
+ * validation), the caller still owns {@code message} and is responsible for
+ * releasing it, since the constructor's {@code data.retain()} never ran.
+ */
+ public static DataMessage decode(ByteBuf message) {
+ int shuffleWriterId = message.readInt();
+ int shuffleReaderId = message.readInt();
+ int dataSize = message.readInt();
+ long checksum = message.readLong();
+ // Each streaming-shuffle frame carries exactly one DataMessage, so after reading the
+ // header the remaining readable bytes must equal `dataSize`. Strict equality catches
+ // both undersized frames (truncation) and oversized frames (trailing garbage) at the
+ // wire boundary rather than letting them propagate to getRecordData() later.
+ if (dataSize < 0 || dataSize != message.readableBytes()) {
+ throw new IllegalArgumentException(
+ "Invalid DataMessage dataSize=" + dataSize +
+ ", readable bytes after header=" + message.readableBytes());
+ }
+ return new DataMessage(shuffleWriterId, shuffleReaderId, dataSize, message, checksum);
+ }
+
+ /**
+ * Returns a slice of {@link #data} containing exactly the serialized records
+ * (i.e., {@code dataSize} bytes starting at the current reader index).
+ *
+ * Uses {@code slice} (not {@code readSlice}), so this method DOES NOT advance the
+ * reader index of {@link #data}. The returned slice shares the underlying storage;
+ * no data is copied. Callers may freely consume the returned slice without
+ * affecting subsequent calls to this method.
+ */
+ public ByteBuf getRecordData() {
+ return data.slice(data.readerIndex(), dataSize);
+ }
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.java
new file mode 100644
index 000000000000..2126c52c02d5
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.java
@@ -0,0 +1,71 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import java.util.zip.CRC32C;
+import javax.annotation.concurrent.NotThreadSafe;
+
+/**
+ * Helper class for streaming shuffle checksum calculations.
+ */
+@NotThreadSafe
+public final class ShuffleChecksum {
+ private final CRC32C crc = new CRC32C();
+
+ /**
+ * Updates checksum for a specified portion of a ByteBuf message.
+ *
+ * @param message The ByteBuf to calculate checksum for
+ * @param startIndex The index of the first byte to calculate checksum for
+ * @param dataLength The length of the data to calculate checksum for
+ */
+ public void updateChecksum(ByteBuf message, int startIndex, int dataLength) {
+ if (startIndex < 0) {
+ throw new IllegalArgumentException(
+ "startIndex must be non-negative: " + startIndex);
+ }
+ if (dataLength < 0) {
+ throw new IllegalArgumentException(
+ "dataLength must be non-negative: " + dataLength);
+ }
+ // Bound the range against writerIndex() rather than capacity(): the checksum must
+ // cover actual written data only, never bytes in [writerIndex, capacity) which
+ // may be uninitialized.
+ if (startIndex + dataLength > message.writerIndex()) {
+ throw new IllegalArgumentException(
+ "startIndex + dataLength exceeds writerIndex: " +
+ startIndex + " + " + dataLength + " > " + message.writerIndex());
+ }
+ if (message.hasArray()) {
+ // heap-based ByteBuf
+ crc.update(message.array(), message.arrayOffset() + startIndex, dataLength);
+ } else {
+ // off-heap ByteBuf
+ crc.update(message.nioBuffer(startIndex, dataLength));
+ }
+ }
+
+ public long getValue() {
+ return crc.getValue();
+ }
+
+ public void reset() {
+ crc.reset();
+ }
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java
new file mode 100644
index 000000000000..f1ce41863bc2
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java
@@ -0,0 +1,146 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
+/**
+ * Base class for messages sent between the streaming shuffle writers (usually mappers) and
+ * readers (usually reducers).
+ *
+ * To prevent memory leaks, streaming shuffle programmers should always abide by the following
+ * principles:
+ *
+ * 1. If you create a buffer via ByteBufAllocator, you must explicitly release it.
+ * 2. If you create a new StreamingShuffleMessage, you must call .release() on it.
+ *
+ * To make these rules work out, implementations of StreamingShuffleMessage should abide
+ * by the following rules:
+ *
+ * 1. StreamingShuffleMessages should *not* modify the refcount of ByteBufs passed to them
+ * during encoding. For message implementations without ByteBufs, this isn't a concern.
+ * But for messages that have ByteBufs (e.g. DataMessage), the encoding method will likely
+ * call compositeByteBuf.addComponent(), which transfers ownership of the ByteBuf to the
+ * CompositeByteBuf and decrements the refcount of the ByteBuf. So that the caller can
+ * *always* follow rule 1 above, the ByteBuf should be retained before being passed to the
+ * CompositeByteBuf; if this is not done, the refcount of the ByteBuf after leaving
+ * encode() will be 0, and if the caller follows rule 1, they will try to decrement an
+ * already 0 refcount. See DataMessage for an example of how to do this properly.
+ * 2. If StreamingShuffleMessages keep a reference the ByteBufs passed to them during
+ * decoding, they should increment the refcount of that ByteBuf, and assign it to
+ * ownedBuf. This is so that resources get cleaned up when callers follow rule 2 above,
+ * i.e. call .release() on the StreamingShuffleMessage. See DataMessage for an example of
+ * how to do this properly.
+ */
+public abstract sealed class StreamingShuffleMessage
+ permits CreditControlMessage, DataMessage, TerminationAckMessage, TerminationControlMessage {
+ protected ByteBuf ownedBuf = null;
+ private Runnable releaseCallback = null;
+
+ // To prevent any duplicate/out of order/missing messages, each writer will track the current
+ // max sequence number that has been sent to each reader. Similarly, each reader will track
+ // the latest sequence number it has received from each writer. Upon receiving a new message
+ // from any writer, reader will check if the sequence number is expected. When all finish, the
+ // reader will send TerminationAckMessage to the writer with the max sequence number that has
+ // been received, and the writer will check if the latest sequence recorded matches it.
+
+ // Thus the sequence number is valid for the following message types:
+ // 1. all message types from a writer to a reader. To make sure that the reader
+ // receive all the messages sent by writer in order without missing or duplicate any.
+ // 2. TerminationAckMessage from a reader to a writer. To make sure at the end of the
+ // shuffle, the reader receives the same number of messages that the writer has sent.
+ // Essentially, other message types from reader to writer won't have a valid sequence number.
+ private long seqNum;
+ public void setSeqNum(long seqNum) {
+ this.seqNum = seqNum;
+ }
+ public long getSeqNum() { return seqNum; }
+
+ /** Returns the type of this message. */
+ public abstract StreamingShuffleMessageType messageType();
+
+ /** Encodes the current message into the provided ByteBuf. */
+ public void encode(CompositeByteBuf buf) {
+ buf.writeInt(messageType().id());
+ buf.writeLong(seqNum);
+ }
+
+ /**
+ * Returns the number of bytes the encoded message header occupies on the wire.
+ *
+ * For control messages (CreditControl, TerminationControl, TerminationAck) this is
+ * the full encoded length of the message. For {@link DataMessage} this is the header
+ * size only; the variable-length payload of {@code dataSize} bytes follows the header
+ * and is NOT included in this count. Callers use this value to pre-size the
+ * {@link CompositeByteBuf} that the encoded message is written into.
+ */
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ /**
+ * Registers a callback that will be invoked exactly once when {@link #release()}
+ * runs to completion. The callback runs AFTER {@code ownedBuf.release()} and is
+ * cleared after invocation, so a subsequent {@code release()} call on the same
+ * thread will not re-run it.
+ *
+ * Replaces any previously-registered callback. Not thread-safe; see {@link #release()}.
+ */
+ public void setReleaseCallback(Runnable releaseCallback) {
+ this.releaseCallback = releaseCallback;
+ }
+
+ /**
+ * Releases any resources associated with this message.
+ * In VERY RARE cases when the task fails unexpectedly, this method may be called twice.
+ * This method is idempotent — a second call on the same thread is a no-op — but it is
+ * NOT thread-safe.
+ */
+ public void release() {
+ if (ownedBuf != null) {
+ ownedBuf.release();
+ ownedBuf = null;
+ }
+ if (releaseCallback != null) {
+ releaseCallback.run();
+ releaseCallback = null;
+ }
+ }
+
+ public static StreamingShuffleMessage decode(ByteBuf message) {
+ StreamingShuffleMessageType messageType =
+ StreamingShuffleMessageType.decode(message.readInt());
+ long seqNum = message.readLong();
+
+ // Switch expression over the enum is exhaustive; the compiler enforces that every
+ // case is handled, so any future StreamingShuffleMessageType added to the enum will
+ // cause this method to fail compilation until the corresponding case is added here.
+ StreamingShuffleMessage shuffleMessage = switch (messageType) {
+ case DATA_MESSAGE_UNSAFE_ROW -> DataMessage.decode(message);
+ case CREDIT_CONTROL_MESSAGE -> CreditControlMessage.decode(message);
+ case TERMINATION_CONTROL_MESSAGE -> TerminationControlMessage.decode(message);
+ case TERMINATION_ACK_MESSAGE -> TerminationAckMessage.decode(message);
+ };
+ shuffleMessage.setSeqNum(seqNum);
+
+ return shuffleMessage;
+ }
+
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageType.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageType.java
new file mode 100644
index 000000000000..d103009c62cc
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageType.java
@@ -0,0 +1,45 @@
+/*
+ * 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.network.shuffle.streaming;
+
+public enum StreamingShuffleMessageType {
+ DATA_MESSAGE_UNSAFE_ROW(1),
+ CREDIT_CONTROL_MESSAGE(2),
+ TERMINATION_CONTROL_MESSAGE(3),
+ TERMINATION_ACK_MESSAGE(4);
+
+ private final int id;
+
+ StreamingShuffleMessageType(int id) {
+ this.id = id;
+ }
+
+ public int id() {
+ return id;
+ }
+
+ public static StreamingShuffleMessageType decode(int givenId) {
+ return switch (givenId) {
+ case 1 -> DATA_MESSAGE_UNSAFE_ROW;
+ case 2 -> CREDIT_CONTROL_MESSAGE;
+ case 3 -> TERMINATION_CONTROL_MESSAGE;
+ case 4 -> TERMINATION_ACK_MESSAGE;
+ default -> throw new IllegalArgumentException("Unknown message type: " + givenId);
+ };
+ }
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationAckMessage.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationAckMessage.java
new file mode 100644
index 000000000000..3a724fd32c16
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationAckMessage.java
@@ -0,0 +1,68 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
+/**
+ * Reader → writer acknowledgement of {@link TerminationControlMessage}. The
+ * {@code seqNum} field carries the last sequence number the reader observed from
+ * this writer; the writer compares that value against its own last-sent sequence
+ * number to detect lost or reordered messages, failing the task with
+ * STREAMING_SHUFFLE_INCORRECT_SEQUENCE_NUMBER on mismatch.
+ */
+public final class TerminationAckMessage extends StreamingShuffleMessage {
+ public final int shuffleWriterId;
+ public final int shuffleReaderId;
+
+ public TerminationAckMessage(int shuffleWriterId, int shuffleReaderId) {
+ this.shuffleWriterId = shuffleWriterId;
+ this.shuffleReaderId = shuffleReaderId;
+ }
+
+ @Override
+ public StreamingShuffleMessageType messageType() {
+ return StreamingShuffleMessageType.TERMINATION_ACK_MESSAGE;
+ }
+
+ @Override
+ public int headerLength() {
+ // 4 bytes for the shuffle writer ID, 4 bytes for the shuffle reader ID
+ return super.headerLength() + 8;
+ }
+
+ @Override
+ public void encode(CompositeByteBuf buf) {
+ super.encode(buf);
+
+ // Write the shuffle writer ID
+ buf.writeInt(shuffleWriterId);
+ // Write the shuffle reader ID
+ buf.writeInt(shuffleReaderId);
+ }
+
+ public static TerminationAckMessage decode(ByteBuf buf) {
+ // Read the shuffle writer ID
+ int shuffleWriterId = buf.readInt();
+ // Read the shuffle reader ID
+ int shuffleReaderId = buf.readInt();
+
+ return new TerminationAckMessage(shuffleWriterId, shuffleReaderId);
+ }
+}
diff --git a/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationControlMessage.java b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationControlMessage.java
new file mode 100644
index 000000000000..db00832baac6
--- /dev/null
+++ b/common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationControlMessage.java
@@ -0,0 +1,68 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
+/**
+ * Writer → reader control message that signals end-of-stream. After a writer has
+ * sent its last {@link DataMessage} to a given reader it sends one
+ * TerminationControlMessage on the same connection and waits for the matching
+ * {@link TerminationAckMessage} back. Receipt of this message tells the reader
+ * that no further DataMessages will arrive from this writer.
+ */
+public final class TerminationControlMessage extends StreamingShuffleMessage {
+ public final int shuffleWriterId;
+ public final int shuffleReaderId;
+
+ public TerminationControlMessage(int shuffleWriterId, int shuffleReaderId) {
+ this.shuffleWriterId = shuffleWriterId;
+ this.shuffleReaderId = shuffleReaderId;
+ }
+
+ @Override
+ public StreamingShuffleMessageType messageType() {
+ return StreamingShuffleMessageType.TERMINATION_CONTROL_MESSAGE;
+ }
+
+ @Override
+ public int headerLength() {
+ // 4 bytes for the shuffle writer ID, 4 bytes for the shuffle reader ID
+ return super.headerLength() + 8;
+ }
+
+ @Override
+ public void encode(CompositeByteBuf buf) {
+ super.encode(buf);
+
+ // Write the shuffle writer ID
+ buf.writeInt(shuffleWriterId);
+ // Write the shuffle reader ID
+ buf.writeInt(shuffleReaderId);
+ }
+
+ public static TerminationControlMessage decode(ByteBuf buf) {
+ // Read the shuffle writer ID
+ int shuffleWriterId = buf.readInt();
+ // Read the shuffle reader ID
+ int shuffleReaderId = buf.readInt();
+
+ return new TerminationControlMessage(shuffleWriterId, shuffleReaderId);
+ }
+}
diff --git a/common/network-common/src/test/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageSuite.java b/common/network-common/src/test/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageSuite.java
new file mode 100644
index 000000000000..2b5d92a4328d
--- /dev/null
+++ b/common/network-common/src/test/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessageSuite.java
@@ -0,0 +1,399 @@
+/*
+ * 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.network.shuffle.streaming;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.Unpooled;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * Unit tests for {@link StreamingShuffleMessage} encode/decode round-trips and
+ * {@link ShuffleChecksum}.
+ */
+public class StreamingShuffleMessageSuite {
+
+ private static final long SEQ_NUM = 42L;
+
+ // ---- helpers ---------------------------------------------------------------
+
+ private ByteBuf encodeAndSlice(StreamingShuffleMessage msg) {
+ msg.setSeqNum(SEQ_NUM);
+ CompositeByteBuf buf = Unpooled.compositeBuffer();
+ buf.capacity(msg.headerLength());
+ msg.encode(buf);
+ msg.release();
+ // Return a copy so we can freely advance the reader index
+ ByteBuf copy = Unpooled.buffer(buf.readableBytes());
+ copy.writeBytes(buf);
+ buf.release();
+ return copy;
+ }
+
+ // ---- CreditControlMessage --------------------------------------------------
+
+ @Test
+ public void testCreditControlRoundTrip() {
+ CreditControlMessage original = new CreditControlMessage(3, 7, 5);
+ ByteBuf encoded = encodeAndSlice(original);
+ try {
+ StreamingShuffleMessage decoded = StreamingShuffleMessage.decode(encoded);
+ assertInstanceOf(CreditControlMessage.class, decoded);
+ CreditControlMessage credit = (CreditControlMessage) decoded;
+ assertEquals(SEQ_NUM, credit.getSeqNum());
+ assertEquals(3, credit.shuffleWriterId);
+ assertEquals(7, credit.shuffleReaderId);
+ assertEquals(5, credit.numMessages);
+ } finally {
+ encoded.release();
+ }
+ }
+
+ // ---- TerminationControlMessage ---------------------------------------------
+
+ @Test
+ public void testTerminationControlRoundTrip() {
+ TerminationControlMessage original = new TerminationControlMessage(1, 2);
+ ByteBuf encoded = encodeAndSlice(original);
+ try {
+ StreamingShuffleMessage decoded = StreamingShuffleMessage.decode(encoded);
+ assertInstanceOf(TerminationControlMessage.class, decoded);
+ TerminationControlMessage term = (TerminationControlMessage) decoded;
+ assertEquals(SEQ_NUM, term.getSeqNum());
+ assertEquals(1, term.shuffleWriterId);
+ assertEquals(2, term.shuffleReaderId);
+ } finally {
+ encoded.release();
+ }
+ }
+
+ // ---- TerminationAckMessage -------------------------------------------------
+
+ @Test
+ public void testTerminationAckRoundTrip() {
+ TerminationAckMessage original = new TerminationAckMessage(4, 8);
+ ByteBuf encoded = encodeAndSlice(original);
+ try {
+ StreamingShuffleMessage decoded = StreamingShuffleMessage.decode(encoded);
+ assertInstanceOf(TerminationAckMessage.class, decoded);
+ TerminationAckMessage ack = (TerminationAckMessage) decoded;
+ assertEquals(SEQ_NUM, ack.getSeqNum());
+ assertEquals(4, ack.shuffleWriterId);
+ assertEquals(8, ack.shuffleReaderId);
+ } finally {
+ encoded.release();
+ }
+ }
+
+ // ---- DataMessage -----------------------------------------------------------
+
+ @Test
+ public void testDataMessageRoundTrip() {
+ byte[] payload = "hello streaming shuffle".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ long checksum = 0xDEADBEEFL;
+
+ DataMessage original = new DataMessage(2, 5, payload.length, payloadBuf, checksum);
+ ByteBuf encoded = encodeAndSlice(original);
+ payloadBuf.release();
+
+ try {
+ StreamingShuffleMessage decoded = StreamingShuffleMessage.decode(encoded);
+ assertInstanceOf(DataMessage.class, decoded);
+ DataMessage dm = (DataMessage) decoded;
+ assertEquals(SEQ_NUM, dm.getSeqNum());
+ assertEquals(2, dm.shuffleWriterId);
+ assertEquals(5, dm.shuffleReaderId);
+ assertEquals(payload.length, dm.dataSize);
+ assertEquals(checksum, dm.checksum);
+
+ ByteBuf recordData = dm.getRecordData();
+ byte[] out = new byte[payload.length];
+ recordData.readBytes(out);
+ assertArrayEquals(payload, out);
+ dm.release();
+ } finally {
+ encoded.release();
+ }
+ }
+
+ @Test
+ public void testReleaseIsIdempotent() {
+ byte[] payload = "hi".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ // Constructor calls data.retain(), so refcount is now 2 (1 original + 1 retain).
+ DataMessage msg = new DataMessage(0, 0, payload.length, payloadBuf, 0L);
+ assertEquals(2, payloadBuf.refCnt());
+
+ msg.release();
+ assertEquals(1, payloadBuf.refCnt());
+
+ // Second release should be a no-op — refcount must not drop further or throw.
+ msg.release();
+ assertEquals(1, payloadBuf.refCnt());
+
+ payloadBuf.release();
+ }
+
+ @Test
+ public void testReleaseCallbackRunsExactlyOnce() {
+ byte[] payload = "hi".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ DataMessage msg = new DataMessage(0, 0, payload.length, payloadBuf, 0L);
+
+ java.util.concurrent.atomic.AtomicInteger callbackInvocations =
+ new java.util.concurrent.atomic.AtomicInteger(0);
+ msg.setReleaseCallback(callbackInvocations::incrementAndGet);
+
+ msg.release();
+ assertEquals(1, callbackInvocations.get());
+
+ // Second release: idempotent — callback must NOT fire again.
+ msg.release();
+ assertEquals(1, callbackInvocations.get());
+
+ payloadBuf.release();
+ }
+
+ @Test
+ public void testDataMessageRejectsDataSizeLargerThanReadableBytes() {
+ byte[] payload = "hello".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ try {
+ // dataSize (10) > data.readableBytes() (5)
+ assertThrows(IllegalArgumentException.class,
+ () -> new DataMessage(0, 0, payload.length + 5, payloadBuf, 0L));
+ } finally {
+ payloadBuf.release();
+ }
+ }
+
+ @Test
+ public void testDataMessageRejectsDataSizeSmallerThanReadableBytes() {
+ byte[] payload = "hello".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ try {
+ // dataSize (3) < data.readableBytes() (5)
+ assertThrows(IllegalArgumentException.class,
+ () -> new DataMessage(0, 0, payload.length - 2, payloadBuf, 0L));
+ } finally {
+ payloadBuf.release();
+ }
+ }
+
+ @Test
+ public void testDataMessageRejectsNegativeDataSize() {
+ byte[] payload = "hello".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ try {
+ assertThrows(IllegalArgumentException.class,
+ () -> new DataMessage(0, 0, -1, payloadBuf, 0L));
+ } finally {
+ payloadBuf.release();
+ }
+ }
+
+ @Test
+ public void testDataMessageDecodeRejectsTrailingBytes() {
+ // Encode a valid DataMessage, then append junk bytes to the encoded buffer.
+ // decode() should reject because dataSize != message.readableBytes() after the header.
+ byte[] payload = "hello".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ DataMessage original = new DataMessage(1, 2, payload.length, payloadBuf, 0L);
+ original.setSeqNum(SEQ_NUM);
+ CompositeByteBuf buf = Unpooled.compositeBuffer();
+ buf.capacity(original.headerLength());
+ original.encode(buf);
+ original.release();
+ payloadBuf.release();
+
+ // Copy encoded bytes then append 3 extra junk bytes.
+ ByteBuf corrupted = Unpooled.buffer(buf.readableBytes() + 3);
+ corrupted.writeBytes(buf);
+ buf.release();
+ corrupted.writeBytes(new byte[]{0x7f, 0x7f, 0x7f});
+
+ try {
+ assertThrows(IllegalArgumentException.class,
+ () -> StreamingShuffleMessage.decode(corrupted));
+ } finally {
+ corrupted.release();
+ }
+ }
+
+ @Test
+ public void testDataMessageDecodeRejectsTruncatedFrame() {
+ // Encode a valid DataMessage, then strip a trailing byte from the encoded buffer.
+ // decode() should reject because dataSize > message.readableBytes() after the header.
+ byte[] payload = "hello".getBytes();
+ ByteBuf payloadBuf = Unpooled.wrappedBuffer(payload);
+ DataMessage original = new DataMessage(1, 2, payload.length, payloadBuf, 0L);
+ original.setSeqNum(SEQ_NUM);
+ CompositeByteBuf buf = Unpooled.compositeBuffer();
+ buf.capacity(original.headerLength());
+ original.encode(buf);
+ original.release();
+ payloadBuf.release();
+
+ int truncatedSize = buf.readableBytes() - 1;
+ ByteBuf truncated = Unpooled.buffer(truncatedSize);
+ truncated.writeBytes(buf, truncatedSize);
+ buf.release();
+
+ try {
+ assertThrows(IllegalArgumentException.class,
+ () -> StreamingShuffleMessage.decode(truncated));
+ } finally {
+ truncated.release();
+ }
+ }
+
+ // ---- ShuffleChecksum -------------------------------------------------------
+
+ @Test
+ public void testShuffleChecksumHeapBuffer() {
+ byte[] data = {1, 2, 3, 4, 5};
+ ByteBuf buf = Unpooled.wrappedBuffer(data); // heap-backed
+
+ ShuffleChecksum cs = new ShuffleChecksum();
+ cs.updateChecksum(buf, 0, data.length);
+ long value1 = cs.getValue();
+ assertTrue(value1 != 0);
+
+ cs.reset();
+ cs.updateChecksum(buf, 0, data.length);
+ assertEquals(value1, cs.getValue(), "Checksum should be deterministic");
+ buf.release();
+ }
+
+ @Test
+ public void testShuffleChecksumDirectBuffer() {
+ byte[] data = {10, 20, 30};
+ ByteBuf heap = Unpooled.wrappedBuffer(data);
+ ByteBuf direct = Unpooled.directBuffer(data.length);
+ direct.writeBytes(data);
+
+ ShuffleChecksum csHeap = new ShuffleChecksum();
+ csHeap.updateChecksum(heap, 0, data.length);
+
+ ShuffleChecksum csDirect = new ShuffleChecksum();
+ csDirect.updateChecksum(direct, 0, data.length);
+
+ assertEquals(csHeap.getValue(), csDirect.getValue(),
+ "Heap and direct buffer checksums should match");
+
+ heap.release();
+ direct.release();
+ }
+
+ @Test
+ public void testShuffleChecksumNonZeroStartIndex() {
+ // Verify that startIndex correctly offsets where the checksum begins, and that the
+ // result matches an independent reference computation. Cover heap and direct buffers.
+ byte[] data = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
+ int startIndex = 3;
+ int dataLength = 5; // bytes 3..7 inclusive
+
+ // Reference CRC32C computed over the same sub-range.
+ java.util.zip.CRC32C reference = new java.util.zip.CRC32C();
+ reference.update(data, startIndex, dataLength);
+ long expected = reference.getValue();
+
+ ByteBuf heap = Unpooled.wrappedBuffer(data);
+ ShuffleChecksum csHeap = new ShuffleChecksum();
+ csHeap.updateChecksum(heap, startIndex, dataLength);
+ assertEquals(expected, csHeap.getValue(),
+ "Heap buffer checksum should match reference for sub-range");
+ heap.release();
+
+ ByteBuf direct = Unpooled.directBuffer(data.length);
+ direct.writeBytes(data);
+ ShuffleChecksum csDirect = new ShuffleChecksum();
+ csDirect.updateChecksum(direct, startIndex, dataLength);
+ assertEquals(expected, csDirect.getValue(),
+ "Direct buffer checksum should match reference for sub-range");
+ direct.release();
+ }
+
+ @Test
+ public void testShuffleChecksumRejectsOutOfBoundsRange() {
+ // 5 bytes written, but the test asks for bytes [3..8). Should fail because
+ // startIndex (3) + dataLength (5) = 8 > writerIndex (5).
+ byte[] data = {1, 2, 3, 4, 5};
+ ByteBuf buf = Unpooled.wrappedBuffer(data);
+ try {
+ ShuffleChecksum cs = new ShuffleChecksum();
+ assertThrows(IllegalArgumentException.class,
+ () -> cs.updateChecksum(buf, 3, 5));
+ } finally {
+ buf.release();
+ }
+ }
+
+ @Test
+ public void testShuffleChecksumRejectsRangeBeyondWriterIndexButWithinCapacity() {
+ // capacity = 10 but only 5 bytes written. Requesting [0..8) should fail because
+ // 8 > writerIndex (5), even though 8 <= capacity (10). Guards against checksumming
+ // uninitialized bytes.
+ ByteBuf buf = Unpooled.buffer(10);
+ buf.writeBytes(new byte[]{1, 2, 3, 4, 5});
+ try {
+ ShuffleChecksum cs = new ShuffleChecksum();
+ assertThrows(IllegalArgumentException.class,
+ () -> cs.updateChecksum(buf, 0, 8));
+ } finally {
+ buf.release();
+ }
+ }
+
+ @Test
+ public void testShuffleChecksumRejectsNegativeStartIndex() {
+ ByteBuf buf = Unpooled.wrappedBuffer(new byte[]{1, 2, 3});
+ try {
+ ShuffleChecksum cs = new ShuffleChecksum();
+ assertThrows(IllegalArgumentException.class,
+ () -> cs.updateChecksum(buf, -1, 2));
+ } finally {
+ buf.release();
+ }
+ }
+
+ @Test
+ public void testShuffleChecksumRejectsNegativeDataLength() {
+ ByteBuf buf = Unpooled.wrappedBuffer(new byte[]{1, 2, 3});
+ try {
+ ShuffleChecksum cs = new ShuffleChecksum();
+ assertThrows(IllegalArgumentException.class,
+ () -> cs.updateChecksum(buf, 0, -1));
+ } finally {
+ buf.release();
+ }
+ }
+
+ @Test
+ public void testInvalidMessageTypeThrows() {
+ ByteBuf buf = Unpooled.buffer(12);
+ buf.writeInt(999); // unknown type
+ buf.writeLong(0L);
+ assertThrows(IllegalArgumentException.class, () -> StreamingShuffleMessage.decode(buf));
+ buf.release();
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org