jerrypeng opened a new pull request, #56878:
URL: https://github.com/apache/spark/pull/56878
<!--
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.
-->
This is **part 4** of a multi-PR effort to add *streaming shuffle* to Spark
— a push-based shuffle used by Real-Time Mode (RTM) structured streaming, where
writer tasks push records directly to reader tasks over the network instead of
writing map output to disk for readers to pull.
This PR adds the **writer (push) side**:
- **`StreamingShuffleWriter`** — a `ShuffleWriter` that pushes serialized
records to the downstream readers over Netty instead of writing shuffle files
to disk.
- **`StreamingShuffleServerHandler`** — the writer-side Netty `RpcHandler`
that accepts reader connections and processes the messages readers send back.
- **`StreamingShuffleManager.getWriter`** — now returns a
`StreamingShuffleWriter` (it was a stub that threw
`UnsupportedOperationException` in part 3).
- **`TransportServer.getPooledByteBufAllocator`** — a small accessor so the
writer can allocate pooled send buffers from the server's allocator.
- Three writer configs and two structured-logging keys (see below).
#### How the writer and reader talk to each other
Each writer task starts its own Netty server. Readers connect to it, and the
two exchange four message types (all defined in part 1). The exchange for one
writer <-> one reader looks like this:
```
Reader Writer (this PR)
| |
| --- CreditControlMessage -----------> | reader connects; writer
now knows
| | which network channel
maps to this reader
| |
| <-------------- DataMessage (seq=0) --- | writer pushes batched
rows
| <-------------- DataMessage (seq=1) --- | (each message carries a
sequence
| <-------------- DataMessage (seq=2) --- | number and an optional
CRC32C checksum)
| ... |
| <----- TerminationControlMessage ------ | writer has sent
everything (seq=N)
| |
| --- TerminationAckMessage (seq=N) ---> | reader confirms it
received up to seq=N
| |
```
Key points of the protocol, from the writer's side:
1. **Connection setup.** A reader opens a connection and sends a
`CreditControlMessage`. Until that arrives the writer does not know which
channel belongs to which reader, so it queues outgoing messages for that reader
and flushes them once the reader is known.
2. **Data push.** The writer partitions each record with the shuffle
dependency's partitioner, serializes rows into a pooled buffer, and sends the
buffer as a `DataMessage` once it fills up
(`spark.shuffle.streaming.networkBufferSize`, default 32 KB) or after a maximum
buffering interval (`spark.shuffle.streaming.networkBufferMaxWaitTimeMs`,
default 50 ms). Every message this writer sends to a given reader carries a
**monotonically increasing sequence number**, so the reader can detect a lost
or out-of-order message.
3. **Integrity check (optional).** When
`spark.shuffle.streaming.checksum.enabled` is on, the writer embeds a CRC32C
over the buffer that the reader re-computes and compares.
4. **Back-pressure.** In-flight buffer memory is bounded
(`spark.shuffle.streaming.writerMaxMemory`, default 32 MB) by an off-heap
`MemoryConsumer` plus a byte semaphore; when the limit is reached the writer
blocks the upstream iterator until buffers are freed by send completions.
5. **Termination.** After the last record, the writer sends each reader a
`TerminationControlMessage` carrying the final sequence number, then waits for
every reader to reply with a `TerminationAckMessage`. On each ack the writer
checks that the reader's last-seen sequence number matches what it sent (a
mismatch fails the task), and only returns from `write()` once all readers have
acknowledged.
Errors that occur on Netty threads (e.g. a failed write) are recorded via
`ErrorNotifier` (part 3.5) and re-thrown on the task thread at a safe point.
All resources (the server, channels, pooled buffers, reserved memory) are
released from a `TaskContext` completion listener, so they are cleaned up on
both success and failure.
New configs (all `internal`, default on the safe side):
| Config | Default | Purpose |
|---|---|---|
| `spark.shuffle.streaming.networkBufferSize` | 32 KB | target size of each
pushed buffer |
| `spark.shuffle.streaming.networkBufferMaxWaitTimeMs` | 50 ms | max time a
partial buffer is held before flushing |
| `spark.shuffle.streaming.writerMaxMemory` | 32 MB | best-effort cap on the
writer's in-flight buffer memory |
New log keys: `NUM_SHUFFLE_READERS`, `SHUFFLE_READER_ID`.
The full PR stack:
- **Part 1** (SPARK-56674, *merged*) - streaming shuffle wire protocol (the
four Netty message types above).
- **Part 2** (SPARK-56962, *merged*) - `StreamingShuffleOutputTracker`
(driver-side writer-location coordination).
- **Part 3** (SPARK-57141, *merged*) - shuffle-manager layer
(`StreamingShuffleManager` + `MultiShuffleManager`).
- **Part 3.5** (SPARK-57337, *merged*) - shared transport + error plumbing
(`ErrorNotifier`, `TransportClient.send(ByteBuf)`, checksum config).
- **Part 4** (*this PR*) - `StreamingShuffleWriter` + server-side Netty
handler (push path).
- **Part 5** - `StreamingShuffleReader` + client-side Netty handler (pull
path).
- **Part 6** - register streaming shuffles with the tracker in
`DAGScheduler` (activation).
- **Part 7** - end-to-end `StreamingShuffleSuite`.
- **Part 8** - documentation.
This PR depends only on the merged parts (1 through 3.5) and is independent
of the reader PR (part 5); the two can be reviewed in parallel and merged in
either order.
### 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.
-->
Real-Time Mode / low-latency continuous queries need shuffle data to flow
continuously between stages. The default sort shuffle (write map output to
disk, then have reducers pull it) adds latency that is unacceptable for these
workloads. Streaming shuffle instead pushes records directly from writer tasks
to reader tasks. This PR implements the push (writer) half — the
previously-stubbed `StreamingShuffleManager.getWriter` — including the per-task
transport server, the credit-based connection setup, sequence-numbered data
push, memory back-pressure, and the optional end-to-end checksum.
### 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'.
-->
No. The streaming shuffle managers are opt-in via `spark.shuffle.manager`
and are not the default, and the feature is not usable end-to-end until the
reader (part 5) and activation (part 6) PRs land; the new configs therefore
have no effect on the default sort-shuffle path. The added configs
(`spark.shuffle.streaming.networkBufferSize`,
`spark.shuffle.streaming.networkBufferMaxWaitTimeMs`,
`spark.shuffle.streaming.writerMaxMemory`) take effect only once a streaming
shuffle is in use.
### 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.
-->
The writer is exercised end-to-end by `StreamingShuffleSuite` in the tests
PR of this stack (part 7): writer<->reader data transfer, the credit-control /
termination handshake, sequence-number validation, checksum verification,
memory back-pressure, and background-thread error propagation.
### 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.
-->
Co-authored with Claude Code (Claude Opus 4.8)
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
