aajisaka commented on code in PR #55028:
URL: https://github.com/apache/spark/pull/55028#discussion_r3527755483
##########
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);
+ plaintextAccumulator = null; // ownership transferred
to downstream
+ }
+ // Current message is fully decoded. Reset state so the
handler can
+ // decode the next independent GCM message on the same
channel.
+ resetForNextMessage();
+ if (ciphertextNettyBuf.readableBytes() == 0) {
+ break;
+ }
+ // Remaining bytes may belong to another message; loop to
process them.
}
} finally {
ciphertextNettyBuf.release();
}
}
+
+ @Override
+ public void channelInactive(ChannelHandlerContext ctx) throws
Exception {
+ releaseAccumulator();
+ ctx.fireChannelInactive();
+ }
+
+ @Override
+ public void exceptionCaught(ChannelHandlerContext ctx, Throwable
cause) {
+ releaseAccumulator();
+ ctx.close();
+ }
Review Comment:
Thank you! Fixed in
https://github.com/apache/spark/pull/55028/commits/64c28209294e3c73c8c9aa53bf46a5e0fef428a8
--
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]
