SteNicholas commented on code in PR #3718: URL: https://github.com/apache/celeborn/pull/3718#discussion_r3354585637
########## client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ReadIntegrityTracker.java: ########## @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.celeborn.plugin.flink; + +import java.io.IOException; +import java.util.function.Consumer; + +import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf; +import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.apache.celeborn.common.CommitMetadata; +import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl; +import org.apache.celeborn.plugin.flink.utils.BufferUtils; + +/** + * Accumulates the read-side {@link CommitMetadata} (CRC32 + byte count) over a consumed + * subpartition range and reports it to the driver at stream end. Each read path strips its own + * framing via {@code accumulate*Buffer} so the hashed bytes match what the writer hashed (the body + * after the batch header). + */ +public class ReadIntegrityTracker { + private static final Logger LOG = LoggerFactory.getLogger(ReadIntegrityTracker.class); + + private final FlinkShuffleClientImpl client; + private final int shuffleId; + private final int partitionId; + private final int subPartitionIndexStart; + private final int subPartitionIndexEnd; + private final CommitMetadata readCommitMetadata = new CommitMetadata(); + private volatile boolean enabled; + // True once report() has run; the report is terminal whether it succeeded or failed. + private volatile boolean reportAttempted = false; + + public ReadIntegrityTracker( + FlinkShuffleClientImpl client, + int shuffleId, + int partitionId, + int subPartitionIndexStart, + int subPartitionIndexEnd, + boolean enabled) { + this.client = client; + this.shuffleId = shuffleId; + this.partitionId = partitionId; + this.subPartitionIndexStart = subPartitionIndexStart; + this.subPartitionIndexEnd = subPartitionIndexEnd; + this.enabled = enabled; + } + + public boolean isEnabled() { + return enabled; + } + + /** Permanently disables tracking for this stream so nothing is reported at stream end. */ + void disable() { + enabled = false; + } + + /** Strips the {@code headerPrefix}-byte batch header; a too-short buffer disables tracking. */ + private void accumulatePlainBuffer(ByteBuf flinkBuffer, int headerPrefix) { + if (flinkBuffer.readableBytes() < headerPrefix) { + disableForUnexpectedFraming( Review Comment: @xumingming, these are defensive guards, not expected paths. The tracker has to strip the exact batch-header framing so the bytes it hashes match what the writer hashed; if it ever sees a buffer shape it doesn't know how to strip (a composite on the regular path, a composite with ≠2 components on the tiered path, or a buffer/header shorter than the batch header), stripping the wrong number of bytes would yield a checksum that disagrees with the writer's and fail a perfectly healthy read. So rather than risk a false mismatch, it disables tracking for that stream (logs a warning, reports nothing) — i.e. it fails *open* on framing uncertainty. In normal operation with the current read paths none of these should trigger: the regular path always delivers a single buffer with the header prefix, and the tiered path delivers either a plain buffer or a 2-component (header + data) composite. They're there so a future read-path/buffer-shape change degrades the check to a no-op instead of breaking healthy reads. ########## client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala: ########## @@ -76,6 +76,19 @@ class MapPartitionCommitHandler( // shuffleId -> boolean, records whether the shuffle is visible at the segment level, facilitating future optimization of worker read and write processes private val shuffleIsSegmentGranularityVisible = JavaUtils.newConcurrentHashMap[Int, Boolean] + private val shuffleIntegrityCheckEnabled = conf.clientShuffleIntegrityCheckEnabled + + // Write-side per-subpartition checksums of one finished map partition (indexed by subpartition). + private case class MapPartitionWriteMetadata(crc32: Array[Int], bytesWritten: Array[Long]) { + require( + crc32.length == bytesWritten.length, + s"crc32 length ${crc32.length} != bytesWritten length ${bytesWritten.length}") + } + + // shuffleId -> (mapPartitionId -> write-side metadata). + private val commitMetadataForMapPartition = Review Comment: @xumingming, good question. A few points on the footprint: - It's gated behind `celeborn.client.shuffle.integrityCheck.enabled` (default `false`), so the default JobManager footprint is unchanged. - Each retained entry is small: a `crc32` int (4B) + a `bytesWritten` long (8B) per subpartition, i.e. ~12·N bytes per map partition and ~12·M·N bytes total (e.g. M=N=4096 ≈ 150MB). - It's bounded by the shuffle lifetime and cleared in `removeExpiredShuffle` (line 149), not accumulated across shuffles. - The per-subpartition granularity is required: a reader consumes an arbitrary `[startSubIndex, endSubIndex]` range and we combine the write-side checksums over exactly that range, so we can't pre-aggregate. This is the same O(M·N) the write side already produces at `MapperEnd`; the new cost is the driver *retaining* it until commit/expiry. So when enabled it is genuinely O(M·N), but it's opt-in and short-lived. If it turns out to be a concern for very large jobs, a follow-up could compact it. WDYT? -- 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]
