Copilot commented on code in PR #3689:
URL: https://github.com/apache/celeborn/pull/3689#discussion_r3393827688
##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -288,7 +288,8 @@ public <K, V> ShuffleWriter<K, V> getWriter(
h.lifecycleManagerPort(),
celebornConf,
h.userIdentifier(),
- h.extension());
+ h.extension(),
+ SparkCommonUtils.getCryptoHandler(conf));
Review Comment:
`SparkCommonUtils.getCryptoHandler(conf)` allocates a new
`SparkCryptoHandler` each time it is called. `ShuffleClient.get(...)` only uses
the handler during singleton initialization, so on hot paths like `getWriter`
this becomes repeated, wasted allocations per task attempt when IO encryption
is enabled.
Consider caching the `Optional<CryptoHandler>` once per
`SparkShuffleManager` instance (or per executor) and reusing it for all
`ShuffleClient.get(...)` calls.
##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -208,7 +208,8 @@ public <K, V> ShuffleWriter<K, V> getWriter(
h.lifecycleManagerPort(),
celebornConf,
h.userIdentifier(),
- h.extension());
+ h.extension(),
+ SparkCommonUtils.getCryptoHandler(conf));
Review Comment:
`SparkCommonUtils.getCryptoHandler(conf)` allocates a new
`SparkCryptoHandler` when IO encryption is enabled. In `getWriter`, this call
happens on a hot path and the returned handler is only needed for the initial
`ShuffleClient` singleton initialization, so repeated calls are wasted work.
Consider caching the `Optional<CryptoHandler>` once per
`SparkShuffleManager` instance and reusing it.
##########
client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java:
##########
@@ -1054,6 +1057,20 @@ public int pushOrMergeData(
length = compressor.getCompressedTotalSize();
}
+ if (cryptoHandler.isPresent()) {
+ byte[] encrypted = cryptoHandler.get().encrypt(data, offset, length);
+ logger.debug(
+ "Encrypted shuffle data for shuffle {} map {} partition {}: {} bytes
-> {} bytes.",
+ shuffleId,
+ mapId,
+ partitionId,
+ length,
+ encrypted.length);
+ data = encrypted;
+ offset = 0;
+ length = encrypted.length;
+ }
Review Comment:
This change introduces a new encryption branch in the core push path. Given
this impacts correctness (shuffle data must remain readable) and
security-sensitive behavior, it would be good to add an integration-style unit
test that exercises `pushOrMergeData` + `CelebornInputStream` with a simple
`CryptoHandler` (e.g., XOR/identity-with-header) to verify round-trip behavior
and that shuffle integrity checking (CRC) still passes when enabled.
##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -260,7 +261,8 @@ public <K, C> ShuffleReader<K, C> getReader(
Int.MaxValue(),
context,
celebornConf,
- shuffleIdTracker);
+ shuffleIdTracker,
+ SparkCommonUtils.getCryptoHandler(conf));
Review Comment:
`SparkCommonUtils.getCryptoHandler(conf)` is invoked per reader creation.
When IO encryption is enabled, this constructs a new `SparkCryptoHandler` even
though the `ShuffleClient` is a process singleton and will typically already
have been initialized.
Caching the `Optional<CryptoHandler>` once per `SparkShuffleManager`
instance would avoid repeated allocations in task-heavy workloads.
##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkShuffleManager.java:
##########
@@ -445,7 +446,8 @@ public <K, C> ShuffleReader<K, C> getCelebornShuffleReader(
context,
celebornConf,
metrics,
- shuffleIdTracker);
+ shuffleIdTracker,
+ SparkCommonUtils.getCryptoHandler(conf));
} else {
Review Comment:
`SparkCommonUtils.getCryptoHandler(conf)` is invoked on each
`getCelebornShuffleReader` call (both branches), which will construct a new
`SparkCryptoHandler` when IO encryption is enabled. Since the `ShuffleClient`
is a process singleton, creating a new handler per task/read is avoidable
overhead.
Prefer computing the handler once per executor / `SparkShuffleManager`
instance and reusing it here.
--
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]
