Copilot commented on code in PR #3699:
URL: https://github.com/apache/celeborn/pull/3699#discussion_r3371182192
##########
common/src/main/java/org/apache/celeborn/common/network/buffer/FileChunkBuffers.java:
##########
@@ -27,15 +27,19 @@
public class FileChunkBuffers extends ChunkBuffers {
private final File file;
private final TransportConf conf;
+ private final boolean isChunkCompressed;
public FileChunkBuffers(DiskFileInfo fileInfo, TransportConf conf) {
super(fileInfo.getReduceFileMeta());
+ isChunkCompressed = fileInfo.isChunkCompressionEnabled();
file = fileInfo.getFile();
this.conf = conf;
}
@Override
public ManagedBuffer chunk(int chunkIndex, int offset, int len) {
+ // sliced reads unsupported for chunkCompressed files
+ assert (!isChunkCompressed || (offset == 0 && len == Integer.MAX_VALUE));
Tuple2<Long, Long> offsetLen = getChunkOffsetLength(chunkIndex, offset,
len);
Review Comment:
FileChunkBuffers uses a Java `assert` to enforce that chunk-compressed files
cannot be sliced (offset/len must be full-chunk). Assertions are typically
disabled in production, so this would silently allow unsupported sliced reads
and may lead to corrupted decompression / incorrect offsets. This should be a
runtime check that always executes.
##########
client/src/main/java/org/apache/celeborn/client/read/LocalPartitionReader.java:
##########
@@ -254,8 +255,12 @@ public ByteBuf next() throws IOException,
InterruptedException {
logger.error("PartitionReader thread interrupted while fetching data.");
throw e;
}
+ int chunkIdx = returnedChunks;
returnedChunks++;
- return chunk;
+ // If no per-chunk list was sent (old worker), treat as compressed to
honour the global flag.
+ boolean compressed =
+ streamHandler.getChunkCompressedCount() == 0 ||
streamHandler.getChunkCompressed(chunkIdx);
+ return Pair.of(chunk, compressed);
Review Comment:
LocalPartitionReader computes chunkIdx using returnedChunks (0..N) instead
of the actual chunk index in the file (startChunkIndex + returnedChunks). For
range reads where startChunkIndex != 0, this indexes into
PbStreamHandler.chunkCompressed with the wrong chunk id, which can incorrectly
attempt ZSTD decompression on an uncompressed large-record chunk (or skip
decompression when required).
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -827,14 +874,12 @@ private boolean fillBuffer() throws IOException {
if (size > compressedBuf.length) {
compressedBuf = new byte[size];
}
-
- currentChunk.readBytes(compressedBuf, 0, size);
+ readFully(currentStream, compressedBuf, 0, size);
} else {
if (size > rawDataBuf.length) {
rawDataBuf = new byte[size];
}
-
- currentChunk.readBytes(rawDataBuf, 0, size);
+ readFully(currentStream, rawDataBuf, 0, size);
}
Review Comment:
CelebornInputStream ignores the return value of readFully(...) when reading
the batch payload. If the underlying stream hits EOF mid-payload (e.g.,
truncated/corrupted chunk or mismatched chunkCompressed flag), the code will
proceed with partially-filled buffers and likely fail later with confusing
decompressor errors or data corruption. It should fail fast when fewer than
`size` bytes are read.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/ChunkCompressedFileChannelWriter.java:
##########
@@ -0,0 +1,153 @@
+/*
+ * 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.service.deploy.worker.file.chunk.compressed;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.ArrayList;
+import java.util.List;
+
+import com.github.luben.zstd.Zstd;
+import com.google.common.annotations.VisibleForTesting;
+import io.netty.buffer.CompositeByteBuf;
+
+import org.apache.celeborn.common.meta.DiskFileInfo;
+import org.apache.celeborn.common.meta.ReduceFileMeta;
+import org.apache.celeborn.common.util.FileChannelUtils;
+import org.apache.celeborn.service.deploy.worker.file.FileChannelWriter;
+
+public class ChunkCompressedFileChannelWriter extends FileChannelWriter {
+ private final FileChannel channel;
+ private final DiskFileInfo diskFileInfo;
+ private final int compressionLevel;
+ private final ChunkBufferPool.BufferPair bufferPair;
+ private ByteBuffer chunkBuffer;
+ private ByteBuffer compressedChunkBuffer;
+ private final List<Long> chunkOffsets;
+ private final List<Boolean> chunkCompressed;
+ private final long chunkSize;
+
+ public ChunkCompressedFileChannelWriter(
+ DiskFileInfo diskFileInfo, long chunkSize, int compressionLevel) throws
IOException {
+ this.diskFileInfo = diskFileInfo;
+ this.chunkSize = chunkSize;
+ channel =
FileChannelUtils.createWritableFileChannel(diskFileInfo.getFilePath());
+ this.compressionLevel = compressionLevel;
+ bufferPair = ChunkBufferPool.getInstance().acquire(chunkSize);
+ chunkBuffer = bufferPair.chunkBuffer;
+ compressedChunkBuffer = bufferPair.compressedBuffer;
+ chunkOffsets = new ArrayList<>();
+ chunkOffsets.add(0L);
+ chunkCompressed = new ArrayList<>();
+ }
+
+ @Override
+ public void write(CompositeByteBuf buffer, boolean gatherApiEnabled) throws
IOException {
+ if (buffer.readableBytes() > chunkSize) {
+ // Flush any pending accumulated data before writing the large record so
file offsets
+ // remain consistent.
+ compressAndFlush();
+ flushLargeRecord(buffer);
+ return;
+ }
+
+ if (buffer.readableBytes() > chunkBuffer.remaining()) {
+ compressAndFlush();
+ }
+
+ ByteBuffer[] buffers = buffer.nioBuffers();
+ for (ByteBuffer byteBuffer : buffers) {
+ while (byteBuffer.hasRemaining()) {
+ chunkBuffer.put(byteBuffer);
+ }
+ }
+ }
+
+ /**
+ * Writes the large record directly to the channel without compression.
Large records span a full
+ * chunk on their own, so the decompression overhead would be paid all at
once anyway; skipping
+ * compression avoids the ZstdOutputStream frame overhead and simplifies the
write path.
+ */
+ private void flushLargeRecord(CompositeByteBuf buffer) throws IOException {
+ ByteBuffer[] buffers = buffer.nioBuffers();
+ for (ByteBuffer buf : buffers) {
+ while (buf.hasRemaining()) {
+ channel.write(buf);
+ }
+ }
+ chunkCompressed.add(false);
+ chunkOffsets.add(channel.position());
+ }
+
+ @VisibleForTesting
+ public void compressAndFlush() throws IOException {
+ int size = chunkBuffer.position();
+ if (size == 0) return;
+ chunkBuffer.position(0);
+ chunkBuffer.limit(size);
+ compressedChunkBuffer.clear();
+ int compressedSize;
+ try {
+ compressedSize =
+ (int)
+ Zstd.compressDirectByteBuffer(
+ compressedChunkBuffer,
+ 0,
+ compressedChunkBuffer.capacity(),
+ chunkBuffer,
+ 0,
+ size,
+ compressionLevel);
+ } catch (RuntimeException e) {
+ throw new IOException("Failed to compress chunk with ZSTD.", e);
+ }
+ compressedChunkBuffer.position(0);
+ compressedChunkBuffer.limit(compressedSize);
+
+ long written = 0L;
+ while (written < compressedSize) {
+ written += channel.write(compressedChunkBuffer);
+ }
+ chunkCompressed.add(true);
+ chunkOffsets.add((chunkOffsets.get(chunkOffsets.size() - 1) + written));
+ chunkBuffer.clear();
+ }
+
+ @Override
+ public void close(boolean commitFilesFsync) {
+ try {
+ compressAndFlush();
+ if (commitFilesFsync) {
+ channel.force(false);
Review Comment:
ChunkCompressedFileChannelWriter.close() is not idempotent: it always
updates DiskFileInfo metadata and releases the BufferPair back to
ChunkBufferPool. If close() is invoked multiple times (e.g., via error paths
that call closeStreams() and later closeResource()), the same BufferPair can be
released twice, allowing two writers to share the same buffers and corrupt
data. Guard against double-close and avoid updating metadata if the final flush
fails.
##########
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala:
##########
@@ -5281,6 +5288,17 @@ object CelebornConf extends Logging {
.checkValues(Set(PartitionSplitMode.SOFT.name,
PartitionSplitMode.HARD.name))
.createWithDefault(PartitionSplitMode.SOFT.name)
+ val CHUNK_COMPRESSION_LEVEL: ConfigEntry[Int] =
+ buildConf("celeborn.chunk.compression.level")
+ .categories("client")
+ .doc(
+ "ZSTD compression level to use for chunk-level compression " +
+ "(celeborn.chunk.compression.enabled must be true). " +
+ "Valid range is 1–22; the default (3) matches the ZSTD library
default.")
+ .version("0.6.0")
+ .intConf
+ .createWithDefault(3)
Review Comment:
celeborn.chunk.compression.level is documented as valid range 1–22, but the
config entry does not enforce this. Invalid values (0, negative, >22) will flow
into Zstd.compressDirectByteBuffer and can throw at runtime or yield unexpected
behavior. Add a config validation check so misconfiguration fails early.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/ChunkBufferPool.java:
##########
@@ -0,0 +1,77 @@
+/*
+ * 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.service.deploy.worker.file.chunk.compressed;
+
+import java.nio.ByteBuffer;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedDeque;
+
+/**
+ * Pool of reusable (chunkBuffer, compressedBuffer) pairs for
ChunkCompressedFileChannelWriter,
+ * bucketed by chunkSize so every acquired pair is exactly the right capacity.
+ */
+public class ChunkBufferPool {
+
+ public static class BufferPair {
+ public final ByteBuffer chunkBuffer;
+ public final ByteBuffer compressedBuffer;
+ public final long chunkSize;
+
+ public BufferPair(ByteBuffer chunkBuffer, ByteBuffer compressedBuffer,
long chunkSize) {
+ this.chunkBuffer = chunkBuffer;
+ this.compressedBuffer = compressedBuffer;
+ this.chunkSize = chunkSize;
+ }
+ }
+
+ private static final ChunkBufferPool INSTANCE = new ChunkBufferPool();
+
+ private final ConcurrentHashMap<Long, ConcurrentLinkedDeque<BufferPair>>
poolMap =
+ new ConcurrentHashMap<>();
+
+ private ChunkBufferPool() {}
+
+ public static ChunkBufferPool getInstance() {
+ return INSTANCE;
+ }
+
+ public BufferPair acquire(long chunkSize) {
+ ConcurrentLinkedDeque<BufferPair> bucket =
+ poolMap.computeIfAbsent(chunkSize, k -> new ConcurrentLinkedDeque<>());
+ BufferPair pair = bucket.pollFirst();
+ if (pair != null) {
+ pair.chunkBuffer.clear();
+ pair.compressedBuffer.clear();
+ return pair;
+ }
+ ByteBuffer chunkBuf = MmapMemoryManager.getInstance().allocateBuffer((int)
chunkSize);
+ // allocateDirect, NOT MmapMemoryManager: mmap duplicates share one
backing region, so
+ // after clear() both chunkBuf and a mmap-backed compressedBuf would have
position=0
+ // pointing to the same physical address. ZSTD would then write its frame
header to
+ // mmap[0..N] before reading mmap[0..N] as input, silently corrupting the
source.
+ ByteBuffer compressedBuf =
MmapMemoryManager.getInstance().allocateBuffer((int) chunkSize);
Review Comment:
ChunkBufferPool’s comment says to use allocateDirect instead of
MmapMemoryManager for the compressed buffer, but the code currently allocates
both buffers via MmapMemoryManager. This is misleading for future maintainers.
Also, the casts to (int) chunkSize will silently overflow for large configured
chunk sizes; fail fast with a range check and avoid the casts.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/MmapMemoryManager.java:
##########
@@ -0,0 +1,128 @@
+/*
+ * 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.service.deploy.worker.file.chunk.compressed;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.UUID;
+import java.util.logging.Logger;
+
+public class MmapMemoryManager {
+ private static final Logger LOG =
Logger.getLogger(MmapMemoryManager.class.getName());
+ private static MmapMemoryManager INSTANCE;
+ private static final long DEFAULT_FILE_LENGTH = 512 * 1024 * 1024L;
Review Comment:
MmapMemoryManager uses java.util.logging.Logger while the rest of the worker
codebase uses SLF4J (org.slf4j.Logger/LoggerFactory). Mixing logging frameworks
makes log routing and formatting inconsistent. Prefer SLF4J here as well.
##########
client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java:
##########
@@ -328,7 +328,7 @@ public ByteBuf next() throws Exception {
}
returnedChunks++;
lastReturnedChunkId = chunk.getLeft();
- return chunk.getRight();
+ return Pair.of(chunk.getRight(), true);
}
Review Comment:
DfsPartitionReader always marks returned chunks as "compressed"
(Pair.of(..., true)). When chunk compression is enabled in the client config,
CelebornInputStream will wrap DFS chunks in ZstdInputStream and attempt to
decompress even though DFS files are currently written without
chunk-compression (worker side uses ChunkCompressionContext.disabled() for
HDFS/S3/OSS). This will cause read failures for remote-storage shuffles. Use
PbStreamHandler.chunkCompressed only when the worker actually sent it;
otherwise treat DFS chunks as uncompressed.
--
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]
