Skip to content

[SPARK-39702][CORE] Reduce memory overhead of TransportCipher$EncryptedMessage by using a shared byteRawChannel #37110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
JoshRosen wants to merge 1 commit into from
Closed

[SPARK-39702][CORE] Reduce memory overhead of TransportCipher$EncryptedMessage by using a shared byteRawChannel #37110

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions common/network-common/src/main/java/org/apache/spark/network/crypto/TransportCipher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,13 +109,15 @@ public void addToChannel(Channel ch) throws IOException {

@VisibleForTesting
static class EncryptionHandler extends ChannelOutboundHandlerAdapter {
private final ByteArrayWritableChannel byteChannel;
private final ByteArrayWritableChannel byteEncChannel;
private final CryptoOutputStream cos;
private final ByteArrayWritableChannel byteRawChannel;
private boolean isCipherValid;

EncryptionHandler(TransportCipher cipher) throws IOException {
byteChannel = new ByteArrayWritableChannel(STREAM_BUFFER_SIZE);
cos = cipher.createOutputStream(byteChannel);
byteEncChannel = new ByteArrayWritableChannel(STREAM_BUFFER_SIZE);
cos = cipher.createOutputStream(byteEncChannel);
byteRawChannel = new ByteArrayWritableChannel(STREAM_BUFFER_SIZE);
isCipherValid = true;
}

Expand All @@ -127,7 +129,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise)

@VisibleForTesting
EncryptedMessage createEncryptedMessage(Object msg) {
return new EncryptedMessage(this, cos, msg, byteChannel);
return new EncryptedMessage(this, cos, msg, byteEncChannel, byteRawChannel);
}

@Override
Expand Down Expand Up @@ -223,26 +225,27 @@ static class EncryptedMessage extends AbstractFileRegion {
// Due to streaming issue CRYPTO-125: https://issues.apache.org/jira/browse/CRYPTO-125, it has
// to utilize two helper ByteArrayWritableChannel for streaming. One is used to receive raw data
// from upper handler, another is used to store encrypted data.
private ByteArrayWritableChannel byteEncChannel;
private ByteArrayWritableChannel byteRawChannel;
private final ByteArrayWritableChannel byteEncChannel;
private final ByteArrayWritableChannel byteRawChannel;

private ByteBuffer currentEncrypted;

EncryptedMessage(
EncryptionHandler handler,
CryptoOutputStream cos,
Object msg,
ByteArrayWritableChannel ch) {
ByteArrayWritableChannel byteEncChannel,
ByteArrayWritableChannel byteRawChannel) {
Preconditions.checkArgument(msg instanceof ByteBuf || msg instanceof FileRegion,
"Unrecognized message type: %s", msg.getClass().getName());
this.handler = handler;
this.isByteBuf = msg instanceof ByteBuf;
this.buf = isByteBuf ? (ByteBuf) msg : null;
this.region = isByteBuf ? null : (FileRegion) msg;
this.transferred = 0;
this.byteRawChannel = new ByteArrayWritableChannel(STREAM_BUFFER_SIZE);
this.cos = cos;
this.byteEncChannel = ch;
this.byteEncChannel = byteEncChannel;
this.byteRawChannel = byteRawChannel;
this.count = isByteBuf ? buf.readableBytes() : region.count();
}

Expand Down