You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "jerrypeng (via GitHub)" <gi...@apache.org> on 2026/04/30 06:15:31 UTC
[PR] [SPARK-56674][SS] Add streaming shuffle wire protocol [spark]
jerrypeng opened a new pull request, #55620:
URL: https://github.com/apache/spark/pull/55620
Introduces a new Java package org.apache.spark.network.shuffle.streaming in the network-common module with the complete binary wire protocol for streaming shuffle:
- StreamingShuffleMessageType: enum of 4 message IDs (DATA, CREDIT_CONTROL, TERMINATION_CONTROL, TERMINATION_ACK)
- StreamingShuffleMessage: abstract sealed base class; encodes a common 12-byte header (message-type int + sequence-number long); dispatches decode() to concrete subclasses
- DataMessage: carries serialized records from writer to reader; header includes shuffleWriterId, shuffleReaderId, dataSize, and CRC32C checksum
- CreditControlMessage: flow-control signal from reader to writer; also serves as the initial handshake on connection establishment
- TerminationControlMessage: end-of-stream from writer to reader
- TerminationAckMessage: end-of-stream acknowledgment from reader to writer, echoing the last sequence number seen
- ShuffleChecksum: stateful CRC32C helper supporting both heap and direct ByteBuf inputs
All classes are pure Java and depend only on Netty (already present in network-common) and the JDK's java.util.zip.CRC32C (Java 9+).
Adds the STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE error condition, which StreamingShuffleMessage.decode() throws if the wire's message-type ordinal does not map to a known StreamingShuffleMessageType.
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.
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3255171656
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ 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();
+
+ StreamingShuffleMessage shuffleMessage = null;
+
+ switch (messageType) {
Review Comment:
Optional:
Since StreamingShuffleMessageType is an enum and the switch is exhaustive, this could be a Java 17 switch expression — that would remove the null initialization, the unreachable default branch, and the explanatory comment about why the default is unreachable:
```java
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);
};
```
Same goes for StreamingShuffleMessageType.decode(int), though that one needs an explicit default because int isn't exhaustive.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3271382146
##########
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 int SEQ_NUM = 42;
Review Comment:
nit: Since this is parameter to `setSeqNum`, we can simply make it a `long`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4464842895
> The PR description is now out of sync with the implementation.
>
> The PR body still says this patch adds the STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE error class, adds a spark-common-utils Maven dependency, and makes StreamingShuffleMessage.decode() throw a SparkRuntimeException for unknown message types. The latest diff no longer does any of that: the dependency and error-condition changes were removed, and unknown message IDs now result in an IllegalArgumentException.
>
> This should be updated before merge so reviewers and future readers do not rely on stale protocol/error-handling documentation. It would also be clearer to replace “message-type ordinal” with “message-type id”, since the wire format uses the explicit id() values rather than enum ordinals.
I have updated the PR description. Thank you for the reminder.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3271406932
##########
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
Review Comment:
This field currently has no semantic meaning (the receiver ignores it entirely), yet it occupies 4 bytes on the wire. The javadoc says "senders should pass 1" and reserves it for a future credit-based flow-control extension.
The problem: "should pass 1" is only a comment — the constructor does not enforce it. Future contributors may pass 0, -1, or some other value out of habit or by copy-paste. Tests will still pass and behavior will look correct today (since no one reads the value), but once credit semantics are actually enabled, these silently drifted call sites will turn into bugs that are hard to trace back.
Suggestion: add a if (numMessages != 1) throw ... check in the constructor to promote the "reserved" contract from a comment into a code-level constraint. Comments are documentation; validation is contract — reserved fields are exactly the case where that distinction matters.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3192731308
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,150 @@
+/*
+ * 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 org.apache.spark.QueryContext;
+import org.apache.spark.SparkRuntimeException;
+import scala.collection.immutable.Map;
+import scala.collection.immutable.Map$;
+
+/**
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ 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();
+
+ StreamingShuffleMessage shuffleMessage = null;
+
+ switch (messageType) {
+ case DATA_MESSAGE_UNSAFE_ROW:
+ shuffleMessage = DataMessage.decode(message);
+ break;
+ case CREDIT_CONTROL_MESSAGE:
+ shuffleMessage = CreditControlMessage.decode(message);
+ break;
+ case TERMINATION_CONTROL_MESSAGE:
+ shuffleMessage = TerminationControlMessage.decode(message);
+ break;
+ case TERMINATION_ACK_MESSAGE:
+ shuffleMessage = TerminationAckMessage.decode(message);
+ break;
+ default:
+ // Should not reach here since StreamingShuffleMessageType.decode
+ // always return valid message type
+ @SuppressWarnings("unchecked")
+ Map<String, String> parameters = Map$.MODULE$.<String, String>empty()
+ .updated("messageType", messageType.name());
+ throw new SparkRuntimeException("STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE",
+ parameters, null, new QueryContext[]{}, "");
Review Comment:
StreamingShuffleMessageType.decode() will throw IllegalArgumentException in its default branch. So basically this structured error branch cannot be reached. It makes the addition to error-conditions.json, and the dependency of spark-common-utils in network-common/pom.xml meaningless. Why don't just use IllegalArgumentException?
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * 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) {
+ assert (startIndex >= 0);
+ assert (dataLength >= 0);
+ assert (dataLength <= message.readableBytes());
Review Comment:
I think the correct check should be startIndex >= 0 and startIndex + dataLength <= message.capacity()?
And in production, Java assert is not enabled by default. Maybe we should check it explicitly?
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java:
##########
@@ -0,0 +1,81 @@
+/*
+ * 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;
+
+public final class DataMessage extends StreamingShuffleMessage {
+
+ public ByteBuf data;
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+ public int dataSize;
+ public long checksum;
+
+ public DataMessage(int shuffleWriterId, int shuffleReaderId, int dataSize, ByteBuf data,
+ long checksum) {
+ 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);
+
+ // Adding data as a component to buf transfers ownership of data to buf. However,
+ // this DataMessage still has a reference to data, so we need to retain it here.
+ buf.addComponent(true, data.retain());
+ }
+
+ public static DataMessage decode(ByteBuf message) {
+ int shuffleWriterId = message.readInt();
+ int shuffleReaderId = message.readInt();
+ int dataSize = message.readInt();
Review Comment:
DataMessage.decode() reads dataSize at line 69 and passes the remaining buffer through at line 71 without validating that dataSize >= 0 && dataSize <= message.readableBytes(). The actual bounds-sensitive operation is delayed until getRecordData() at line 79, so malformed frames fail later and less clearly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3199117877
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.java:
##########
@@ -0,0 +1,58 @@
+/*
+ * 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) {
+ assert (startIndex >= 0);
+ assert (dataLength >= 0);
+ assert (dataLength <= message.readableBytes());
Review Comment:
updated
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,150 @@
+/*
+ * 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 org.apache.spark.QueryContext;
+import org.apache.spark.SparkRuntimeException;
+import scala.collection.immutable.Map;
+import scala.collection.immutable.Map$;
+
+/**
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ 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();
+
+ StreamingShuffleMessage shuffleMessage = null;
+
+ switch (messageType) {
+ case DATA_MESSAGE_UNSAFE_ROW:
+ shuffleMessage = DataMessage.decode(message);
+ break;
+ case CREDIT_CONTROL_MESSAGE:
+ shuffleMessage = CreditControlMessage.decode(message);
+ break;
+ case TERMINATION_CONTROL_MESSAGE:
+ shuffleMessage = TerminationControlMessage.decode(message);
+ break;
+ case TERMINATION_ACK_MESSAGE:
+ shuffleMessage = TerminationAckMessage.decode(message);
+ break;
+ default:
+ // Should not reach here since StreamingShuffleMessageType.decode
+ // always return valid message type
+ @SuppressWarnings("unchecked")
+ Map<String, String> parameters = Map$.MODULE$.<String, String>empty()
+ .updated("messageType", messageType.name());
+ throw new SparkRuntimeException("STREAMING_SHUFFLE_UNEXPECTED_MESSAGE_TYPE",
+ parameters, null, new QueryContext[]{}, "");
Review Comment:
updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4503912160
@jerrypeng Can you re-trigger the CI? Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4494806119
Only a few minor comments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3264054195
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+/**
+ * Sent from the client to the server to indicate that the client is ready to receive
+ * the specified amount of messages from the server.
+ */
+public final class CreditControlMessage extends StreamingShuffleMessage {
+ public int shuffleWriterId;
Review Comment:
Looking at the network-common conventions a bit more, the concrete protocol messages there (RpcRequest, ChunkFetchSuccess, RpcResponse, etc.) actually use public final direct fields rather than private final + getter — the private final + getter pattern is reserved for the AbstractMessage base class so subclasses can't reassign body. To match that, I've added final to each of the public fields on DataMessage, CreditControlMessage, TerminationControlMessage, and TerminationAckMessage. This addresses the accidental-corruption concern (the compiler now refuses to let anyone reassign these between encode() and release()) without changing any callsite or getter shape, keeping the diff small.
The seqNum field on the base class stays non-final because it's genuinely set post-construction (during encode by the writer, during decode by the reader); same for the internal ownedBuf and releaseCallback book-keeping fields.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4485222590
@viirya thanks for the review again! I have address your comments. PTAL.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3263875141
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ public void release() {
Review Comment:
we clarify
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ public void release() {
Review Comment:
will clarify
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3251682018
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.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 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);
+ }
+ if (startIndex + dataLength > message.capacity()) {
Review Comment:
We can add the safety checks as a defensive measure though the actual usage in the streaming shuffle writer (future PRs):
```
def updateChecksum(): Unit = {
if (shuffleChecksum != null) {
val currentPosition = buffer.writerIndex()
val newDataLength = currentPosition - lastBufferPosition
shuffleChecksum.updateChecksum(buffer, lastBufferPosition, newDataLength)
lastBufferPosition = currentPosition
}
}
```
The call always uses startIndex = lastBufferPosition and dataLength = writerIndex - lastBufferPosition. So in practice, startIndex + dataLength = writerIndex() — exactly at the boundary of valid written data, never beyond. The existing caller never asks the checksum to cover bytes in [writerIndex, capacity).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3263926831
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+/**
+ * Sent from the client to the server to indicate that the client is ready to receive
+ * the specified amount of messages from the server.
+ */
+public final class CreditControlMessage extends StreamingShuffleMessage {
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+
+ public int numMessages;
Review Comment:
The current implementation does not implement a credit based backpressure mechanism so this message is really on used to ack messages received by the streaming shuffle reader. The field numMessage is set to equal the number of messages received and to inform the sender / shuffle writer of additional credits / buffers available now on the reader.
To be honest, all of this is really easy to change. Since none of the wire protocol is persisted to durable storage and there is no scenario client and server will be on different versions. It is really easy for us to change in the future.
I will add a comment on how it is currently used and what it is designed to do for future implementation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3255167058
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+/**
+ * Sent from the client to the server to indicate that the client is ready to receive
+ * the specified amount of messages from the server.
+ */
+public final class CreditControlMessage extends StreamingShuffleMessage {
+ public int shuffleWriterId;
Review Comment:
All message fields (shuffleWriterId, shuffleReaderId, dataSize, checksum, data, numMessages) are declared public and non-final. The existing network.protocol messages (OneWayMessage, ChunkFetchSuccess, RpcRequest, …) are uniformly private final with getters. Wire-protocol message objects are typically immutable after construction; mutability here also makes it easier to accidentally corrupt a message between encode() and release(). Suggest private final + getter to match the rest of the module.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3271905658
##########
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
Review Comment:
Thanks for the suggestion. I'd prefer to skip the pre-emptive check for now. When credit-based backpressure lands, that PR will add tests that exercise numMessages end-to-end — any drifted call site would surface there, not as a silent wire-format mismatch. I'd rather avoid adding a constructor check now that has to be relaxed in the same follow-up PR; the reserved-field javadoc is enough until the field actually carries semantic meaning. Adding the check now would not really any value IMO.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4521284867
Merged to master/branch-4.x.
Thanks @jerrypeng
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3246767222
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/ShuffleChecksum.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 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);
+ }
+ if (startIndex + dataLength > message.capacity()) {
Review Comment:
ShuffleChecksum still validates against capacity(), which can include unwritten bytes:
ShuffleChecksum.java:47 checks startIndex + dataLength > message.capacity(). This is better than relying on Java assertions, but capacity() is not the boundary of valid data; it may include unwritten space beyond writerIndex(). Since the checksum should cover actual message bytes, this can allow callers to accidentally compute CRC over bytes that were never written.
It would be safer to validate against writerIndex(), for example by requiring startIndex >= 0, dataLength >= 0, and dataLength <= message.writerIndex() - startIndex. It would also be worth adding tests for non-zero startIndex, out-of-bounds ranges, and both heap/direct buffers.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3251584909
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java:
##########
@@ -0,0 +1,88 @@
+/*
+ * 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;
+
+public final class DataMessage extends StreamingShuffleMessage {
+
+ public ByteBuf data;
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+ public int dataSize;
+ public long checksum;
+
+ public DataMessage(int shuffleWriterId, int shuffleReaderId, int dataSize, ByteBuf data,
+ long checksum) {
+ 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);
+
+ // Adding data as a component to buf transfers ownership of data to buf. However,
+ // this DataMessage still has a reference to data, so we need to retain it here.
+ buf.addComponent(true, data.retain());
Review Comment:
Sure we can add some checks but they won't really be necessary since the shuffle writer that creates data messages always passes dataSize = rawBuffer.writerIndex() so they match.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4394397722
@viirya thank you for the review. I have address your comments. PTAL. Thanks in advance!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3199116981
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java:
##########
@@ -0,0 +1,81 @@
+/*
+ * 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;
+
+public final class DataMessage extends StreamingShuffleMessage {
+
+ public ByteBuf data;
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+ public int dataSize;
+ public long checksum;
+
+ public DataMessage(int shuffleWriterId, int shuffleReaderId, int dataSize, ByteBuf data,
+ long checksum) {
+ 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);
+
+ // Adding data as a component to buf transfers ownership of data to buf. However,
+ // this DataMessage still has a reference to data, so we need to retain it here.
+ buf.addComponent(true, data.retain());
+ }
+
+ public static DataMessage decode(ByteBuf message) {
+ int shuffleWriterId = message.readInt();
+ int shuffleReaderId = message.readInt();
+ int dataSize = message.readInt();
Review Comment:
Added validate dataSize at decode time
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya closed pull request #55620: [SPARK-56674][SS][RTM][StreamingShuffle] Add streaming shuffle wire protocol
URL: https://github.com/apache/spark/pull/55620
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3271386750
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java:
##########
@@ -0,0 +1,120 @@
+/*
+ * 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;
+
+public final class DataMessage extends StreamingShuffleMessage {
Review Comment:
javadoc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "jerrypeng (via GitHub)" <gi...@apache.org>.
jerrypeng commented on PR #55620:
URL: https://github.com/apache/spark/pull/55620#issuecomment-4495593074
@viirya thanks for the additional comments! I have addressed them. PTAL.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3271384773
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationAckMessage.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * 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;
+
+public final class TerminationAckMessage extends StreamingShuffleMessage {
Review Comment:
Can you add javadoc for this?
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/TerminationControlMessage.java:
##########
@@ -0,0 +1,61 @@
+/*
+ * 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;
+
+public final class TerminationControlMessage extends StreamingShuffleMessage {
Review Comment:
javadoc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3246746736
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/DataMessage.java:
##########
@@ -0,0 +1,88 @@
+/*
+ * 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;
+
+public final class DataMessage extends StreamingShuffleMessage {
+
+ public ByteBuf data;
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+ public int dataSize;
+ public long checksum;
+
+ public DataMessage(int shuffleWriterId, int shuffleReaderId, int dataSize, ByteBuf data,
+ long checksum) {
+ 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);
+
+ // Adding data as a component to buf transfers ownership of data to buf. However,
+ // this DataMessage still has a reference to data, so we need to retain it here.
+ buf.addComponent(true, data.retain());
Review Comment:
L58 writes dataSize into header but L63 directly put all readable bytes of data into encoded message, not just dataSize bytes. The constructor also doesn't check if dataSize == data.readableBytes(). So if the caller passes dataSize < data.readableBytes(), it will send more bytes than the header claims. The receiver side also only checks dataSize <= message.readableBytes(), the trailing bytes would be accepted. It is too loose for a wire protocol that could easily cause framing/data leak/semantic inconsistency. I suggest the encode side either checks if dataSize == data.readableBytes() or only add data.retainedSlice(data.readerIndex(), dataSize). On decode side, if the frame only contains one DataMessage, it should also check if dataSize == message.readableBytes().
We should add related tests that dataSize is smaller/bigger than readable bytes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3255149914
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/StreamingShuffleMessage.java:
##########
@@ -0,0 +1,143 @@
+/*
+ * 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);
+ }
+
+ public int headerLength() {
+ // 4 bytes for message type, 8 bytes for the sequence number
+ return 12;
+ }
+
+ 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.
+ * Implementations should not panic in such a case.
+ */
+ public void release() {
Review Comment:
> this method may be called twice
Looks like this `release()` has no thread safety. So I suppose that you mean idempotency (i.e., no panic when same thread calls `release()` twice)?
The current comment could easily make reader think that this is thread safe.
Could you make the comment more clear?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56674][SS][RTM][StreamingShuffle][Part1] Add streaming shuffle wire protocol [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #55620:
URL: https://github.com/apache/spark/pull/55620#discussion_r3255162698
##########
common/network-common/src/main/java/org/apache/spark/network/shuffle/streaming/CreditControlMessage.java:
##########
@@ -0,0 +1,73 @@
+/*
+ * 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;
+
+/**
+ * Sent from the client to the server to indicate that the client is ready to receive
+ * the specified amount of messages from the server.
+ */
+public final class CreditControlMessage extends StreamingShuffleMessage {
+ public int shuffleWriterId;
+ public int shuffleReaderId;
+
+ public int numMessages;
Review Comment:
This field is going onto the wire and will be hard to change later, but its precise meaning isn't specified. From the PR description, this message is
"reserved for finer-grained credit-based flow control in future revisions" — but numMessages = 5 could mean any of:
- Credit grant (delta): writer may send 5 more DataMessages in addition to existing credit.
- Credit window (absolute): writer's send budget is now 5, replacing the previous value.
- Cumulative ack: reader has received 5 messages so far.
- Initial buffer hint: handshake-only signal of receive-queue size.
Since part 1 has no caller yet, none of these are disambiguated by code. Once a handler PR lands assuming one interpretation, switching is a wire-format
break. Suggest committing to one interpretation in javadoc now, e.g.:
```java
/**
* Sent from reader to writer to grant additional send credit. After
* receiving this message, the writer may send up to {@code numMessages}
* additional DataMessages beyond any previously-granted credit.
*/
public int numMessages;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org