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: 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: ########## @@ -231,31 +245,116 @@ class MapPartitionCommitHandler( shuffleId, (k: Int) => ConcurrentHashMap.newKeySet[Integer]()) resultPartitions.add(partitionId) + + if (shuffleIntegrityCheckEnabled) { + recordMapPartitionCommitMetadata( + shuffleId, + partitionId, + numPartitions, + crc32PerPartition, + bytesWrittenPerPartition) + } } (dataCommitSuccess, false) } + /** + * Records a finished mapper's per-subpartition write-side checksums, keyed by (shuffleId, + * mapPartitionId), for later validation in [[finishPartition]]. A retried/duplicate attempt + * overwrites the prior record (last write wins), which is intentional: the reader validates + * against the last committed attempt. Visible for testing. + */ + private[commit] def recordMapPartitionCommitMetadata( + shuffleId: Int, + mapPartitionId: Int, + numPartitions: Int, + crc32PerPartition: Array[Int], + bytesWrittenPerPartition: Array[Long]): Unit = { + if (crc32PerPartition == null || crc32PerPartition.length != numPartitions || + bytesWrittenPerPartition == null || bytesWrittenPerPartition.length != numPartitions) { + logWarning( + s"Skip recording commit metadata for shuffle $shuffleId map partition $mapPartitionId " + + s"because reported checksum arrays do not match numPartitions $numPartitions.") + return + } + commitMetadataForMapPartition + .computeIfAbsent( + shuffleId, + (_: Int) => JavaUtils.newConcurrentHashMap[Int, MapPartitionWriteMetadata]()) + .put(mapPartitionId, MapPartitionWriteMetadata(crc32PerPartition, bytesWrittenPerPartition)) + } + override def registerShuffle( shuffleId: Int, numMappers: Int, isSegmentGranularityVisible: Boolean, numPartitions: Int): Unit = { super.registerShuffle(shuffleId, numMappers, isSegmentGranularityVisible, numPartitions) shuffleIsSegmentGranularityVisible.put(shuffleId, isSegmentGranularityVisible) + // The outer metadata map is created lazily by recordMapPartitionCommitMetadata. } override def isSegmentGranularityVisible(shuffleId: Int): Boolean = { shuffleIsSegmentGranularityVisible.get(shuffleId) } + /** + * Validates a map partition read over the consumed `[startSubIndex, endSubIndex]` range against + * the order-independent combination of the write-side checksums over that range. The params + * rename the trait's reducer-oriented signature (partitionId/startMapIndex/endMapIndex) to map + * semantics. A reader reaches stream end only after commit, so the write-side metadata is already + * recorded; missing metadata fails closed. + */ Review Comment: Good catch — the trait-level `CommitHandler.finishPartition` doc still says "Invoked when a reduce partition...", which is now inaccurate since map partitions reuse it. Updated the trait doc to be partition-type-agnostic and kept the map-specific detail on this override. ########## client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala: ########## @@ -559,7 +564,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends StatusCode.SUCCESS.getValue).build() } else { response = PbReadReducerPartitionEndResponse.newBuilder().setStatus( - +StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue).setErrorMsg(errorMessage).build() Review Comment: It was a stray unary plus (`+StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue`). Unary `+` on an `Int` is a no-op in Scala, so it had no effect — almost certainly an accidental leftover from the original CELEBORN-894 change. This PR just removes it (cosmetic, no behavior change). -- 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]
