Copilot commented on code in PR #3699: URL: https://github.com/apache/celeborn/pull/3699#discussion_r3378969531
########## 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 { + chunkBufferPool.release(bufferPair); + 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)); Review Comment: `close()` replaces the file meta with `new ReduceFileMeta(chunkOffsets, chunkCompressed)` which does not preserve `chunkSize`. Since `ReduceFileMeta` still has `chunkSize`/`nextBoundary` logic (used by write paths and potentially other computations), dropping `chunkSize` here can cause incorrect behavior later for any code that relies on it. Recommendation (required): extend `ReduceFileMeta` with a constructor that also sets `chunkSize` (and optionally initializes `nextBoundary`), and pass the writer's `chunkSize` when replacing the meta (or carry it over from the prior meta). ########## 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); Review Comment: `readFully(...)` return values are not validated for the batch payload reads (and partial header reads are treated as end-of-chunk). If the underlying stream is truncated/corrupted, this can silently produce short reads, leaving partially stale data in `compressedBuf`/`rawDataBuf` and causing data corruption or confusing decompression errors later. Recommendation (required): check that `readFully(...)` returns exactly the requested byte count for both headers and payloads; if it returns a short read after having read *some* bytes of a header/payload, throw an `IOException` indicating unexpected EOF/corruption instead of silently advancing to the next chunk. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -814,10 +852,19 @@ 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) { + if (readFully(currentStream, sizeBuf, 0, BATCH_HEADER_SIZE) < BATCH_HEADER_SIZE) { + closeCurrentStream(); + if (!moveToNextChunk()) break; + setupCurrentStream(); + continue; + } Review Comment: `readFully(...)` return values are not validated for the batch payload reads (and partial header reads are treated as end-of-chunk). If the underlying stream is truncated/corrupted, this can silently produce short reads, leaving partially stale data in `compressedBuf`/`rawDataBuf` and causing data corruption or confusing decompression errors later. Recommendation (required): check that `readFully(...)` returns exactly the requested byte count for both headers and payloads; if it returns a short read after having read *some* bytes of a header/payload, throw an `IOException` indicating unexpected EOF/corruption instead of silently advancing to the next chunk. ########## 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: `readFully(...)` return values are not validated for the batch payload reads (and partial header reads are treated as end-of-chunk). If the underlying stream is truncated/corrupted, this can silently produce short reads, leaving partially stale data in `compressedBuf`/`rawDataBuf` and causing data corruption or confusing decompression errors later. Recommendation (required): check that `readFully(...)` returns exactly the requested byte count for both headers and payloads; if it returns a short read after having read *some* bytes of a header/payload, throw an `IOException` indicating unexpected EOF/corruption instead of silently advancing to the next chunk. ########## 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); + } + } Review Comment: The inner `while (byteBuffer.hasRemaining())` loop is redundant: `ByteBuffer.put(ByteBuffer src)` copies all remaining bytes in one call (or throws if it cannot). Recommendation (optional): replace the loop with a single `chunkBuffer.put(byteBuffer)` per component to simplify control flow and reduce loop overhead. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -800,6 +809,35 @@ 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 = + (chunkCompressed && currentChunkCompressed) ? new ZstdInputStream(base) : base; Review Comment: Chunk decompression is currently gated by the *client conf* flag (`conf.isChunkCompressionEnabled()`) in addition to per-chunk metadata (`currentChunkCompressed`). This makes the reader fragile: if the client config is false but the stream metadata indicates chunk compression (e.g., reading data produced with chunk compression enabled), the reader will skip decompression and fail. Recommendation (required): drive the decision primarily from stream/file metadata (e.g., decompress when `currentChunkCompressed` is true), and treat the conf flag as a capability toggle only if you can ensure producers and consumers are always consistent. ########## client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java: ########## @@ -321,6 +326,7 @@ private static final class CelebornInputStreamImpl extends CelebornInputStream { this.localHostAddress = Utils.localHostName(conf); this.shouldDecompress = !conf.shuffleCompressionCodec().equals(CompressionCodec.NONE) && needDecompress; + this.chunkCompressed = conf.isChunkCompressionEnabled(); Review Comment: Chunk decompression is currently gated by the *client conf* flag (`conf.isChunkCompressionEnabled()`) in addition to per-chunk metadata (`currentChunkCompressed`). This makes the reader fragile: if the client config is false but the stream metadata indicates chunk compression (e.g., reading data produced with chunk compression enabled), the reader will skip decompression and fail. Recommendation (required): drive the decision primarily from stream/file metadata (e.g., decompress when `currentChunkCompressed` is true), and treat the conf flag as a capability toggle only if you can ensure producers and consumers are always consistent. -- 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]
