Copilot commented on code in PR #3699:
URL: https://github.com/apache/celeborn/pull/3699#discussion_r3386323044
##########
common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java:
##########
@@ -63,6 +63,10 @@ public synchronized void updateBytesFlushed(long bytes) {
}
}
+ public synchronized void setBytesFlushed(long bytesFlushed) {
+ this.bytesFlushed = bytesFlushed;
+ }
Review Comment:
This introduces a public setter that can arbitrarily move `bytesFlushed`
backwards or forwards, bypassing any invariants the existing update method
enforces. If this is intended only for writer finalization, consider narrowing
visibility (package-private), renaming to reflect semantics (absolute vs
delta), and/or enforcing monotonicity to prevent inconsistent state.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/ChunkBufferPool.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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;
+
+import com.github.luben.zstd.Zstd;
+
+import org.apache.celeborn.common.CelebornConf;
+
+/**
+ * Pool of reusable (chunkBuffer, compressedBuffer) pairs for
ChunkCompressedFileChannelWriter,
+ * bucketed by chunkSize so every acquired pair is exactly the right capacity.
+ *
+ * <p>Owns and manages the lifecycle of its internal {@link
MmapMemoryManager}. Call {@link #close}
+ * when the pool is no longer needed to release the mmap backing files.
+ */
+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 final MmapMemoryManager mmapMemoryManager;
+ private final ConcurrentHashMap<Long, ConcurrentLinkedDeque<BufferPair>>
poolMap =
+ new ConcurrentHashMap<>();
+
+ public ChunkBufferPool(CelebornConf conf) {
+ this.mmapMemoryManager = new
MmapMemoryManager(conf.chunkCompressionMmapTmpDir());
+ }
+
+ 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;
+ }
+ int chunkBufSize = Math.toIntExact(chunkSize);
+ int compressedBufSize = Math.toIntExact(Zstd.compressBound(chunkSize));
+ ByteBuffer chunkBuf = mmapMemoryManager.allocateBuffer(chunkBufSize);
+ ByteBuffer compressedBuf =
mmapMemoryManager.allocateBuffer(compressedBufSize);
+ return new BufferPair(chunkBuf, compressedBuf, chunkSize);
+ }
Review Comment:
`Math.toIntExact(chunkSize)` and
`Math.toIntExact(Zstd.compressBound(chunkSize))` will throw
`ArithmeticException` for large chunk sizes, and the allocation attempt can
also create extremely large mmap files. Consider validating `chunkSize` (and
`compressBound`) up-front with a clear exception (e.g., enforce `chunkSize <=
Integer.MAX_VALUE` and a reasonable upper bound driven by config), so failures
are actionable and don't result in unexpected mmap/disk pressure.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -827,14 +875,16 @@ private boolean fillBuffer() throws IOException {
if (size > compressedBuf.length) {
compressedBuf = new byte[size];
}
-
- currentChunk.readBytes(compressedBuf, 0, size);
+ if (readFully(currentStream, compressedBuf, 0, size) != size) {
+ throw new IOException("Invalid EOF detected");
+ }
} else {
if (size > rawDataBuf.length) {
rawDataBuf = new byte[size];
}
-
- currentChunk.readBytes(rawDataBuf, 0, size);
+ if (readFully(currentStream, rawDataBuf, 0, size) != size) {
+ throw new IOException("Invalid EOF detected");
+ }
Review Comment:
The new error message \"Invalid EOF detected\" is too generic for diagnosing
production corruption/compatibility issues. Including context such as
`chunkIndex`, `mapId/attemptId/batchId`, expected vs actual bytes read, and
whether the current chunk was marked compressed would make failures far easier
to debug.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/ChunkCompressedFileChannelWriter.java:
##########
@@ -0,0 +1,173 @@
+/*
+ * 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 chunkBufferPool;
+ 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;
+ private boolean closed = false;
+
+ public ChunkCompressedFileChannelWriter(
+ DiskFileInfo diskFileInfo,
+ long chunkSize,
+ int compressionLevel,
+ ChunkBufferPool chunkBufferPool)
+ throws IOException {
+ this.diskFileInfo = diskFileInfo;
+ this.chunkSize = chunkSize;
+ channel =
FileChannelUtils.createWritableFileChannel(diskFileInfo.getFilePath());
+ this.compressionLevel = compressionLevel;
+ this.chunkBufferPool = chunkBufferPool;
+ bufferPair = chunkBufferPool.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);
+ }
+ if (Zstd.isError(compressedSize)) {
+ throw new IOException("ZSTD compression failed: " +
Zstd.getErrorName(compressedSize));
+ }
+ 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();
Review Comment:
Offset tracking mixes `channel.position()` (used in `flushLargeRecord`) with
`lastOffset + written` here. Using `channel.position()` consistently after
writes would be clearer and avoids subtle mismatches if write semantics change
(e.g., if the file position is manipulated elsewhere).
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -814,10 +849,23 @@ private boolean fillBuffer() throws IOException {
return false;
}
+ if (currentStream == null) {
+ setupCurrentStream();
+ }
+
LocationPushFailedBatches failedBatch = new
LocationPushFailedBatches();
boolean hasData = false;
- while (currentChunk.isReadable() || moveToNextChunk()) {
- currentChunk.readBytes(sizeBuf);
+ while (true) {
+ int headerRead = readFully(currentStream, sizeBuf, 0,
BATCH_HEADER_SIZE);
+ if (headerRead == 0) {
+ closeCurrentStream();
+ if (!moveToNextChunk()) break;
+ setupCurrentStream();
+ continue;
+ } else if (headerRead != BATCH_HEADER_SIZE) {
+ throw new IOException("Invalid EOF detected");
+ }
Review Comment:
The new error message \"Invalid EOF detected\" is too generic for diagnosing
production corruption/compatibility issues. Including context such as
`chunkIndex`, `mapId/attemptId/batchId`, expected vs actual bytes read, and
whether the current chunk was marked compressed would make failures far easier
to debug.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/chunk/compressed/MmapMemoryManager.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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 org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class MmapMemoryManager {
+ private static final Logger LOG =
LoggerFactory.getLogger(MmapMemoryManager.class);
+ private static final int DEFAULT_FILE_LENGTH = 512 * 1024 * 1024;
+ private final String _dirPathName;
+ // _availableOffset has the starting offset for the next allocation in
_currentBuffer. When
+ // _currentBuffer
+ // is created, it is 0. After we allocate a buffer of size x, it is x. And
if we allocate another
+ // buffer of size
+ // y, then it becomes x+y, etc. We try to fulfil as many allocate() calls as
possible on the same
+ // _currentBuffer
+ // until the _currentBuffer cannot hold the new object anymore, and then we
create a new
+ // _currentBuffer.
+ private int _availableOffset = DEFAULT_FILE_LENGTH; // Available offset in
this file.
+ private int _curFileLen = -1;
+ private final List<String> _paths = new LinkedList<>();
+ private final List<ByteBuffer> _memMappedBuffers = new LinkedList<>();
+ ByteBuffer _currentBuffer;
+
+ public MmapMemoryManager(String dirPathName) {
+ File dirFile = new File(dirPathName);
+ if (!dirFile.exists()) {
+ if (!dirFile.mkdirs()) {
+ throw new RuntimeException("Unable to create directory: " + dirFile);
+ }
+ }
+ _dirPathName = dirPathName;
+ }
+
+ private String getFilePrefix() {
+ return UUID.randomUUID() + ".";
+ }
+
+ private void addFileIfNecessary(int len) {
+ if (len + _availableOffset <= _curFileLen) {
+ return;
+ }
+ String filePath = _dirPathName + "/" + getFilePrefix();
+ final File file = new File(filePath);
+ if (file.exists()) {
+ throw new RuntimeException("File " + filePath + " already exists");
+ }
+ file.deleteOnExit();
+ int fileLen = Math.max(DEFAULT_FILE_LENGTH, len);
+ try (RandomAccessFile raf = new RandomAccessFile(filePath, "rw");
+ FileChannel fileChannel = raf.getChannel()) {
+ raf.setLength(fileLen);
+ _currentBuffer = fileChannel.map(FileChannel.MapMode.READ_WRITE, 0,
fileLen);
+ _memMappedBuffers.add(_currentBuffer);
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ _paths.add(filePath);
+ _availableOffset = 0;
+ _curFileLen = fileLen;
+ }
+
+ public synchronized ByteBuffer allocateBuffer(int size) {
+ addFileIfNecessary(size);
+ ByteBuffer buffer = _currentBuffer.duplicate();
+ buffer.position(_availableOffset);
+ buffer.limit(_availableOffset + size);
+ _availableOffset += size;
+ return buffer.slice();
+ }
+
+ public synchronized void close() {
+ // MappedByteBuffers cannot be explicitly unmapped in Java; GC handles the
unmap.
+ // We clear the internal state and delete the backing files so disk space
is reclaimed.
+ _memMappedBuffers.clear();
+ for (String path : _paths) {
+ File file = new File(path);
+ if (!file.delete()) {
+ LOG.warn("Unable to delete mmap backing file: {}", file);
+ }
+ }
+ _paths.clear();
+ _curFileLen = -1;
+ _availableOffset = DEFAULT_FILE_LENGTH;
+ }
Review Comment:
`close()` clears `_memMappedBuffers` but keeps a strong reference in
`_currentBuffer`. That can prevent timely unmapping and cause backing-file
deletion to fail (especially on Windows), leading to disk-space leaks across
test runs or worker restarts. Set `_currentBuffer = null` (and consider also
clearing any other references to mapped buffers) before attempting deletes; if
you need reliable reclamation, consider using a Cleaner-based unmap on
supported JVMs.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -827,14 +875,16 @@ private boolean fillBuffer() throws IOException {
if (size > compressedBuf.length) {
compressedBuf = new byte[size];
}
-
- currentChunk.readBytes(compressedBuf, 0, size);
+ if (readFully(currentStream, compressedBuf, 0, size) != size) {
+ throw new IOException("Invalid EOF detected");
+ }
Review Comment:
The new error message \"Invalid EOF detected\" is too generic for diagnosing
production corruption/compatibility issues. Including context such as
`chunkIndex`, `mapId/attemptId/batchId`, expected vs actual bytes read, and
whether the current chunk was marked compressed would make failures far easier
to debug.
##########
common/src/main/java/org/apache/celeborn/common/network/buffer/FileChunkBuffers.java:
##########
@@ -27,15 +27,25 @@
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) {
+ if (isChunkCompressed && (offset != 0 || len != Integer.MAX_VALUE)) {
+ throw new IllegalArgumentException(
+ "Sliced reads (offset="
+ + offset
+ + ", len="
+ + len
+ + ") are not supported for chunk-compressed files");
+ }
Review Comment:
This blocks all sliced reads whenever chunk compression is enabled for the
file. However, chunk-compressed files can contain *uncompressed* chunks
(large-record path), where slicing could still be supported safely. Consider
consulting `ReduceFileMeta.getChunkCompressed()` for the specific `chunkIndex`
and only rejecting slicing when that chunk is actually compressed, or
explicitly documenting/enforcing \"no slicing\" at a higher API level for all
chunks in such files.
--
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]
