Copilot commented on code in PR #3699:
URL: https://github.com/apache/celeborn/pull/3699#discussion_r3385786206
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -800,6 +807,34 @@ private void init() {
rawDataBuf = new byte[bufferSize];
}
+ private void closeCurrentStream() {
+ if (currentStream != null) {
+ try {
+ currentStream.close();
+ } catch (IOException ignored) {
+ }
+ currentStream = null;
+ }
+ }
+
+ private void setupCurrentStream() throws IOException {
+ closeCurrentStream();
+ if (currentChunk == null) return;
+ InputStream base = new ByteBufInputStream(currentChunk);
+ currentStream = currentChunkCompressed ? new ZstdInputStream(base) :
base;
+ }
Review Comment:
Two issues here: (1) `currentChunkCompressed` defaults to `true`, which can
cause a raw (non-chunk-compressed) chunk to be incorrectly wrapped in
`ZstdInputStream` if `setupCurrentStream()` runs before the flag is set for the
first chunk. Defaulting it to `false` is safer. (2)
`ByteBufInputStream(currentChunk)` does not release the underlying `ByteBuf` on
close by default; `closeCurrentStream()` closes only the stream, not
`currentChunk`. If `moveToNextChunk()` does not reliably release
`currentChunk`, this can leak direct memory. Consider using the
`ByteBufInputStream(ByteBuf, boolean releaseOnClose)` constructor (or
explicitly releasing `currentChunk` when transitioning between chunks),
ensuring you don’t double-release.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -213,6 +216,7 @@ private static final class CelebornInputStreamImpl extends
CelebornInputStream {
private final String localHostAddress;
private boolean shouldDecompress;
+ private InputStream currentStream;
Review Comment:
Two issues here: (1) `currentChunkCompressed` defaults to `true`, which can
cause a raw (non-chunk-compressed) chunk to be incorrectly wrapped in
`ZstdInputStream` if `setupCurrentStream()` runs before the flag is set for the
first chunk. Defaulting it to `false` is safer. (2)
`ByteBufInputStream(currentChunk)` does not release the underlying `ByteBuf` on
close by default; `closeCurrentStream()` closes only the stream, not
`currentChunk`. If `moveToNextChunk()` does not reliably release
`currentChunk`, this can leak direct memory. Consider using the
`ByteBufInputStream(ByteBuf, boolean releaseOnClose)` constructor (or
explicitly releasing `currentChunk` when transitioning between chunks),
ensuring you don’t double-release.
##########
docs/configuration/worker.md:
##########
@@ -19,6 +19,7 @@ license: |
<!--begin-include-->
| Key | Default | isDynamic | Description | Since | Deprecated |
| --- | ------- | --------- | ----------- | ----- | ---------- |
+| celeborn.chunk.compression.mmap.tmpDir |
<tmp>/celeborn-mmap-memory-manager | false | Directory used to create
memory-mapped backing files for the mmap memory manager used by chunk-level
compression. Defaults to a subdirectory of the JVM temporary directory
(<tmp>/celeborn-mmap-memory-manager). | 0.6.4 | |
Review Comment:
The description contains `(<tmp>/...)` with an unescaped `<tmp>` token; in
Markdown this can be interpreted as an HTML tag and may render incorrectly (or
disappear). Use `<tmp>` or wrap `<tmp>` in backticks in the description
to ensure consistent rendering.
##########
client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java:
##########
@@ -193,6 +195,7 @@ private static final class CelebornInputStreamImpl extends
CelebornInputStream {
private Decompressor decompressor;
private ByteBuf currentChunk;
+ private boolean currentChunkCompressed = true;
Review Comment:
Two issues here: (1) `currentChunkCompressed` defaults to `true`, which can
cause a raw (non-chunk-compressed) chunk to be incorrectly wrapped in
`ZstdInputStream` if `setupCurrentStream()` runs before the flag is set for the
first chunk. Defaulting it to `false` is safer. (2)
`ByteBufInputStream(currentChunk)` does not release the underlying `ByteBuf` on
close by default; `closeCurrentStream()` closes only the stream, not
`currentChunk`. If `moveToNextChunk()` does not reliably release
`currentChunk`, this can leak direct memory. Consider using the
`ByteBufInputStream(ByteBuf, boolean releaseOnClose)` constructor (or
explicitly releasing `currentChunk` when transitioning between chunks),
ensuring you don’t double-release.
##########
worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java:
##########
@@ -234,6 +234,14 @@ public FileInfo getSortedFileInfo(
targetBuffer);
} else {
DiskFileInfo diskFileInfo = ((DiskFileInfo) fileInfo);
+ if (diskFileInfo.isChunkCompressionEnabled()) {
+ // TODO this is yet to be implemented
+ // We can read the file one chunk at a time and store chunkid +
uncompressed offsets before
+ // writing
+ throw new IOException(
+ "Chunk compressed shuffle file is not supported to sort, file
path: "
+ + diskFileInfo.getFilePath());
+ }
Review Comment:
The new exception message doesn’t tell users how to remediate. Consider
including the relevant config key(s) and a concrete action, e.g. instructing to
disable `celeborn.chunk.compression.enabled` for workloads that require
sorting, or clarifying which shuffle/partition type triggers sorting so it’s
diagnosable from logs.
##########
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 long DEFAULT_FILE_LENGTH = 512 * 1024 * 1024L;
+ 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 long _availableOffset = DEFAULT_FILE_LENGTH; // Available offset in
this file.
+ private long _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(long 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();
+ long 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(long size) {
+ addFileIfNecessary(size);
+ ByteBuffer buffer = _currentBuffer.duplicate();
+ buffer.position((int) _availableOffset);
+ buffer.limit((int) (_availableOffset + size));
+ _availableOffset += size;
+ return buffer.slice();
+ }
Review Comment:
This method casts `_availableOffset` and `_availableOffset + size` to `int`
without bounds checks. If offsets grow beyond `Integer.MAX_VALUE`, the cast can
wrap and produce invalid positions/limits (or silently corrupt allocations).
Since `ByteBuffer` indices are `int`, the implementation should enforce an
upper bound (e.g., reject or force a new mmap file before crossing
`Integer.MAX_VALUE`) and validate that `size` fits in `int`.
--
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]
