Copilot commented on code in PR #3699: URL: https://github.com/apache/celeborn/pull/3699#discussion_r3378887341
########## worker/src/main/java/org/apache/celeborn/service/deploy/worker/file/FileChannelWriterFactory.java: ########## @@ -0,0 +1,37 @@ +/* + * 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; + +import java.io.IOException; + +import org.apache.celeborn.common.meta.DiskFileInfo; +import org.apache.celeborn.service.deploy.worker.file.chunk.compressed.ChunkBufferPool; +import org.apache.celeborn.service.deploy.worker.file.chunk.compressed.ChunkCompressedFileChannelWriter; + +public class FileChannelWriterFactory { + public static FileChannelWriter getFileChannelWriter( + DiskFileInfo diskFileInfo, long chunkSize, ChunkBufferPool chunkBufferPool) + throws IOException { + if (diskFileInfo.isChunkCompressionEnabled()) { + return new ChunkCompressedFileChannelWriter( + diskFileInfo, chunkSize, diskFileInfo.getChunkCompressionLevel(), chunkBufferPool); + } else { + return new BypassFileChannelWriter(diskFileInfo); + } + } Review Comment: When `diskFileInfo.isChunkCompressionEnabled()` is true, `chunkBufferPool` must be non-null, but this factory does not validate that. With the current worker-side initialization (`StorageManager` only creates the pool when `conf.isChunkCompressionEnabled`), a client can send `ChunkCompressionContext(enabled=true)` and trigger an NPE at runtime. Add an explicit null check (and fail fast with a clear IOException), or ensure `StorageManager` eagerly/lazily initializes `ChunkBufferPool` whenever chunk-compressed writers may be created (based on the per-shuffle context rather than only worker conf). ########## 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(); + } + + @Override + public void close(boolean commitFilesFsync) throws IOException { + if (closed) { + return; + } + closed = true; + IOException failure = null; + try { + compressAndFlush(); + if (commitFilesFsync) { + channel.force(false); + } + } catch (IOException e) { + failure = e; + } finally { + try { + channel.close(); + } catch (IOException e) { + if (failure == null) { + failure = e; + } + } + } + + if (failure != null) { + throw failure; + } + diskFileInfo.setBytesFlushed(chunkOffsets.get(chunkOffsets.size() - 1)); + diskFileInfo.replaceFileMeta(new ReduceFileMeta(chunkOffsets, chunkCompressed)); + chunkBufferPool.release(bufferPair); + } Review Comment: `bufferPair` is only released back to the pool on the success path. If `compressAndFlush()` or `channel.force()` fails, this leaks a pooled buffer pair (and the corresponding mmap backing space stays pinned longer than intended). Move `chunkBufferPool.release(bufferPair)` into a `finally` block that always executes after acquisition, and guard it to avoid double-release. ########## common/src/main/scala/org/apache/celeborn/common/util/PbSerDeUtils.scala: ########## @@ -168,6 +178,9 @@ object PbSerDeUtils { val reduceFileMeta = fileInfo.getFileMeta.asInstanceOf[ReduceFileMeta] builder.setPartitionType(PartitionType.REDUCE.getValue) builder.addAllChunkOffsets(reduceFileMeta.getChunkOffsets) + if (reduceFileMeta.getChunkCompressed != null && !reduceFileMeta.getChunkCompressed.isEmpty) { + builder.addAllChunkCompressed(reduceFileMeta.getChunkCompressed) + } } builder.build Review Comment: `PbFileInfo` now includes `chunkCompressionConfig`, and `fromPbFileInfo` reconstructs `DiskFileInfo` using `pbFileInfo.getChunkCompressionConfig`. However, `toPbFileInfo` never sets `chunkCompressionConfig`, so serialized metadata will always deserialize to `(enabled=false, level=0)` even for chunk-compressed shuffles. Populate `builder.setChunkCompressionConfig(...)` using the source file info's `ChunkCompressionContext` to preserve correctness across PB round-trips. ########## 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(); + } + + public 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: `allocateBuffer` is synchronized but `close()` is not, so a concurrent `close()` can race with allocations (e.g., clearing `_paths` / resetting `_curFileLen` while another thread is in `addFileIfNecessary`). If this class is shared across threads via `ChunkBufferPool`, make `close()` synchronized (or use a shared lock/closed flag) to avoid racy state changes and partial cleanup. -- 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]
