dbtsai commented on code in PR #55028:
URL: https://github.com/apache/spark/pull/55028#discussion_r3532888978
##########
common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java:
##########
@@ -116,23 +108,37 @@ static class GcmEncryptedMessage extends
AbstractFileRegion {
private final long encryptedCount;
GcmEncryptedMessage(AesGcmHkdfStreaming aesGcmHkdfStreaming,
- Object plaintextMessage,
- ByteBuffer plaintextBuffer,
- ByteBuffer ciphertextBuffer) throws
GeneralSecurityException {
+ Object plaintextMessage) throws
GeneralSecurityException {
JavaUtils.checkArgument(
plaintextMessage instanceof ByteBuf || plaintextMessage
instanceof FileRegion,
"Unrecognized message type: %s",
plaintextMessage.getClass().getName());
this.plaintextMessage = plaintextMessage;
- this.plaintextBuffer = plaintextBuffer;
- this.ciphertextBuffer = ciphertextBuffer;
+ this.plaintextBuffer =
+
ByteBuffer.allocate(aesGcmHkdfStreaming.getPlaintextSegmentSize());
+ this.ciphertextBuffer =
+
ByteBuffer.allocate(aesGcmHkdfStreaming.getCiphertextSegmentSize());
// If the ciphertext buffer cannot be fully written the target,
transferTo may
// return with it containing some unwritten data. The initial call
we'll explicitly
// set its limit to 0 to indicate the first call to transferTo.
this.ciphertextBuffer.limit(0);
this.bytesToRead = getReadableBytes();
- this.encryptedCount =
- LENGTH_HEADER_BYTES +
aesGcmHkdfStreaming.expectedCiphertextSize(bytesToRead);
+ int plaintextSegmentSize =
aesGcmHkdfStreaming.getPlaintextSegmentSize();
+ int ciphertextSegmentSize =
aesGcmHkdfStreaming.getCiphertextSegmentSize();
+ int tagSize = ciphertextSegmentSize - plaintextSegmentSize;
+ long fullSegments = bytesToRead / plaintextSegmentSize;
+ long partialBytes = bytesToRead % plaintextSegmentSize;
+ // encryptedCount covers all bytes written to the wire:
+ // LENGTH_HEADER_BYTES – Spark's 8-byte big-endian framing
prefix, read by the
+ // receiver to know where this message
ends in the TCP stream.
+ // getHeaderLength() – Tink's streaming header (salt + nonce
prefix), written by
+ // the encrypter before any ciphertext
segments.
+ // segment bytes – full segments each padded with a
16-byte GCM auth tag,
+ // plus any partial final segment
similarly padded.
+ this.encryptedCount = LENGTH_HEADER_BYTES
+ + aesGcmHkdfStreaming.getHeaderLength()
+ + fullSegments * ciphertextSegmentSize
+ + (partialBytes > 0 ? partialBytes + tagSize : 0);
Review Comment:
**Minor -- `encryptedCount` hard-codes Tink's segmentation layout.**
The manual formula is correct today and matches the `transferTo` encrypt
loop by construction, but it silently couples to Tink internals (segment sizes,
16-byte tag, header written separately). A future Tink streaming-layout change
would resurface as a silent stall rather than a loud failure.
`testEncryptedCountBoundary` guards current behavior; a one-line invariant note
pointing at the coupling would help future maintainers.
##########
common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java:
##########
@@ -354,60 +410,116 @@ public void channelRead(ChannelHandlerContext ctx,
Object ciphertextMessage)
"Unrecognized message type: %s",
ciphertextMessage.getClass().getName());
ByteBuf ciphertextNettyBuf = (ByteBuf) ciphertextMessage;
- // The format of the output is:
+ // The format of each message is:
// [8 byte length][Internal IV and header][Ciphertext][Auth Tag]
+ //
+ // A single channelRead() call may deliver bytes from multiple
back-to-back
+ // GCM messages (common under shuffle load when TCP coalesces
writes). The
+ // outer loop processes as many complete messages as possible from
the buffer
+ // before releasing it, so that bytes belonging to the next
message are never
+ // discarded mid-stream.
try {
- if (!initalizeExpectedLength(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize the
expected length.
- return;
- }
- if (!initalizeDecrypter(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize a header,
needed to
- // initialize a decrypter.
- return;
- }
- int nettyBufReadableBytes = ciphertextNettyBuf.readableBytes();
- while (nettyBufReadableBytes > 0 && !completed) {
- // Read the ciphertext into the local buffer
- int readableBytes = Integer.min(
- nettyBufReadableBytes,
- ciphertextBuffer.remaining());
- int expectedRemaining = (int) (expectedLength -
ciphertextRead);
- int bytesToRead = Integer.min(readableBytes,
expectedRemaining);
- // The smallest ciphertext size is 16 bytes for the auth
tag
- ciphertextBuffer.limit(ciphertextBuffer.position() +
bytesToRead);
- ciphertextNettyBuf.readBytes(ciphertextBuffer);
- ciphertextRead += bytesToRead;
- // Check if this is the last segment
- if (ciphertextRead == expectedLength) {
- completed = true;
- } else if (ciphertextRead > expectedLength) {
- throw new IllegalStateException("Read more ciphertext
than expected.");
+ while (true) {
+ if (!initializeExpectedLength(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize the
expected length.
+ break;
+ }
+ if (!initializeDecrypter(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize a
header, needed to
+ // initialize a decrypter.
+ break;
+ }
+ int nettyBufReadableBytes =
ciphertextNettyBuf.readableBytes();
+ while (nettyBufReadableBytes > 0 && !completed) {
+ // Read the ciphertext into the local buffer
+ int readableBytes = Math.min(
+ nettyBufReadableBytes,
+ ciphertextBuffer.remaining());
+ int expectedRemaining = (int) (expectedLength -
ciphertextRead);
+ int bytesToRead = Math.min(readableBytes,
expectedRemaining);
+ // The smallest ciphertext size is 16 bytes for the
auth tag
+ ciphertextBuffer.limit(ciphertextBuffer.position() +
bytesToRead);
+ ciphertextNettyBuf.readBytes(ciphertextBuffer);
+ ciphertextRead += bytesToRead;
+ // Check if this is the last segment
+ if (ciphertextRead == expectedLength) {
+ completed = true;
+ } else if (ciphertextRead > expectedLength) {
+ throw new IllegalStateException("Read more
ciphertext than expected.");
+ }
+ // If the ciphertext buffer is full, or this is the
last segment,
+ // then decrypt it and fire a read.
+ if (ciphertextBuffer.limit() ==
ciphertextBuffer.capacity() || completed) {
+ ByteBuffer plaintextBuffer =
ByteBuffer.allocate(plaintextSegmentSize);
+ ciphertextBuffer.flip();
+ decrypter.decryptSegment(
+ ciphertextBuffer,
+ segmentNumber,
+ completed,
+ plaintextBuffer);
+ segmentNumber++;
+ // Clear the ciphertext buffer because it's been
read
+ ciphertextBuffer.clear();
+ plaintextBuffer.flip();
+ if (plaintextAccumulator == null) {
+ // Integer.MAX_VALUE disables consolidation
entirely.
+ // CompositeByteBuf.newCompArray() always
initialises the
+ // backing array to min(16, maxNumComponents)
regardless of
+ // this value, so there is no upfront memory
cost.
+ plaintextAccumulator =
+
Unpooled.compositeBuffer(Integer.MAX_VALUE);
+ }
+ // Zero-copy append: addComponent(true, ...)
increases writerIndex
+ // so the component is immediately readable from
the composite.
+ plaintextAccumulator.addComponent(
+ true,
Unpooled.wrappedBuffer(plaintextBuffer));
+ } else {
+ // Set the ciphertext buffer up to read the next
chunk
+
ciphertextBuffer.limit(ciphertextBuffer.capacity());
+ }
+ nettyBufReadableBytes =
ciphertextNettyBuf.readableBytes();
}
- // If the ciphertext buffer is full, or this is the last
segment,
- // then decrypt it and fire a read.
- if (ciphertextBuffer.limit() ==
ciphertextBuffer.capacity() || completed) {
- ByteBuffer plaintextBuffer =
ByteBuffer.allocate(plaintextSegmentSize);
- ciphertextBuffer.flip();
- decrypter.decryptSegment(
- ciphertextBuffer,
- segmentNumber,
- completed,
- plaintextBuffer);
- segmentNumber++;
- // Clear the ciphertext buffer because it's been read
- ciphertextBuffer.clear();
- plaintextBuffer.flip();
-
ctx.fireChannelRead(Unpooled.wrappedBuffer(plaintextBuffer));
- } else {
- // Set the ciphertext buffer up to read the next chunk
- ciphertextBuffer.limit(ciphertextBuffer.capacity());
+ if (!completed) {
+ // Partial message: more bytes needed from the next
channelRead() call.
+ break;
}
- nettyBufReadableBytes = ciphertextNettyBuf.readableBytes();
+ // Fire the entire plaintext as a single event so that
downstream
+ // handlers receive one callback per Spark message instead
of one per
+ // 32 KB segment.
+ if (plaintextAccumulator != null) {
+ ctx.fireChannelRead(plaintextAccumulator);
Review Comment:
**Nit -- null the accumulator before firing.**
Not a live bug in a standard Netty pipeline (a downstream `fireChannelRead`
throw becomes an `exceptionCaught` event propagating downstream, not a
synchronous re-entry here), but the idiomatic ordering removes reliance on that
invariant:
```java
CompositeByteBuf out = plaintextAccumulator;
plaintextAccumulator = null;
ctx.fireChannelRead(out);
```
##########
common/network-common/src/main/java/org/apache/spark/network/crypto/GcmTransportCipher.java:
##########
@@ -354,60 +410,116 @@ public void channelRead(ChannelHandlerContext ctx,
Object ciphertextMessage)
"Unrecognized message type: %s",
ciphertextMessage.getClass().getName());
ByteBuf ciphertextNettyBuf = (ByteBuf) ciphertextMessage;
- // The format of the output is:
+ // The format of each message is:
// [8 byte length][Internal IV and header][Ciphertext][Auth Tag]
+ //
+ // A single channelRead() call may deliver bytes from multiple
back-to-back
+ // GCM messages (common under shuffle load when TCP coalesces
writes). The
+ // outer loop processes as many complete messages as possible from
the buffer
+ // before releasing it, so that bytes belonging to the next
message are never
+ // discarded mid-stream.
try {
- if (!initalizeExpectedLength(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize the
expected length.
- return;
- }
- if (!initalizeDecrypter(ciphertextNettyBuf)) {
- // We have not read enough bytes to initialize a header,
needed to
- // initialize a decrypter.
- return;
- }
- int nettyBufReadableBytes = ciphertextNettyBuf.readableBytes();
- while (nettyBufReadableBytes > 0 && !completed) {
- // Read the ciphertext into the local buffer
- int readableBytes = Integer.min(
- nettyBufReadableBytes,
- ciphertextBuffer.remaining());
- int expectedRemaining = (int) (expectedLength -
ciphertextRead);
- int bytesToRead = Integer.min(readableBytes,
expectedRemaining);
- // The smallest ciphertext size is 16 bytes for the auth
tag
- ciphertextBuffer.limit(ciphertextBuffer.position() +
bytesToRead);
- ciphertextNettyBuf.readBytes(ciphertextBuffer);
- ciphertextRead += bytesToRead;
- // Check if this is the last segment
- if (ciphertextRead == expectedLength) {
- completed = true;
- } else if (ciphertextRead > expectedLength) {
- throw new IllegalStateException("Read more ciphertext
than expected.");
+ while (true) {
+ if (!initializeExpectedLength(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize the
expected length.
+ break;
+ }
+ if (!initializeDecrypter(ciphertextNettyBuf)) {
+ // We have not read enough bytes to initialize a
header, needed to
+ // initialize a decrypter.
+ break;
+ }
+ int nettyBufReadableBytes =
ciphertextNettyBuf.readableBytes();
+ while (nettyBufReadableBytes > 0 && !completed) {
+ // Read the ciphertext into the local buffer
+ int readableBytes = Math.min(
+ nettyBufReadableBytes,
+ ciphertextBuffer.remaining());
+ int expectedRemaining = (int) (expectedLength -
ciphertextRead);
+ int bytesToRead = Math.min(readableBytes,
expectedRemaining);
+ // The smallest ciphertext size is 16 bytes for the
auth tag
+ ciphertextBuffer.limit(ciphertextBuffer.position() +
bytesToRead);
+ ciphertextNettyBuf.readBytes(ciphertextBuffer);
+ ciphertextRead += bytesToRead;
+ // Check if this is the last segment
+ if (ciphertextRead == expectedLength) {
+ completed = true;
+ } else if (ciphertextRead > expectedLength) {
+ throw new IllegalStateException("Read more
ciphertext than expected.");
+ }
+ // If the ciphertext buffer is full, or this is the
last segment,
+ // then decrypt it and fire a read.
+ if (ciphertextBuffer.limit() ==
ciphertextBuffer.capacity() || completed) {
+ ByteBuffer plaintextBuffer =
ByteBuffer.allocate(plaintextSegmentSize);
+ ciphertextBuffer.flip();
+ decrypter.decryptSegment(
+ ciphertextBuffer,
+ segmentNumber,
+ completed,
+ plaintextBuffer);
+ segmentNumber++;
+ // Clear the ciphertext buffer because it's been
read
+ ciphertextBuffer.clear();
+ plaintextBuffer.flip();
+ if (plaintextAccumulator == null) {
+ // Integer.MAX_VALUE disables consolidation
entirely.
+ // CompositeByteBuf.newCompArray() always
initialises the
+ // backing array to min(16, maxNumComponents)
regardless of
+ // this value, so there is no upfront memory
cost.
+ plaintextAccumulator =
+
Unpooled.compositeBuffer(Integer.MAX_VALUE);
+ }
+ // Zero-copy append: addComponent(true, ...)
increases writerIndex
+ // so the component is immediately readable from
the composite.
+ plaintextAccumulator.addComponent(
Review Comment:
**Major -- unbounded plaintext accumulation may regress the fetch-to-disk
path.**
This handler is added `addFirst`, upstream of `TransportFrameDecoder`, whose
interceptor is designed to consume bytes incrementally and explicitly must not
retain the buffers passed to `handle()`. The previous per-segment
`fireChannelRead` streamed ~32 KB at a time, keeping heap flat for large
blocks. This code instead accumulates every segment of a message into the
`CompositeByteBuf` and fires once per message.
For a large block routed to disk (above
`spark.maxRemoteBlockSizeFetchToMem`), the entire block is now fully
materialized on-heap before the frame decoder / `StreamInterceptor` can drain a
single byte. On the send side a `FileRegion`/large body is written as a single
object (`MessageEncoder`), so it becomes one `GcmEncryptedMessage` = one GCM
stream = one accumulated buffer on receive. Under concurrent shuffle with
encryption enabled this looks like an OOM vector and partially defeats
fetch-to-disk.
Suggestion: bounded batching -- flush the accumulator when it exceeds a
byte/segment threshold **or** the message completes. That keeps the "1 callback
instead of N" EventLoop-fairness win for the common small-message case while
capping transient memory for large ones.
Separately, the comment just above about "no upfront memory cost" is about
the component *array*, not the retained data (the whole message); worth
rewording so it doesn't read as reassurance about total memory.
--
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]
