Copilot commented on code in PR #19285: URL: https://github.com/apache/pinot/pull/19285#discussion_r3809158566
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/GzipCodecDefinition.java: ########## @@ -0,0 +1,279 @@ +/** + * 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.pinot.segment.local.io.codec; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.zip.DataFormatException; +import java.util.zip.Deflater; +import java.util.zip.Inflater; +import org.apache.pinot.segment.spi.memory.CleanerUtil; + + +/// Compression codec backed by GZIP (DEFLATE via [java.util.zip.Deflater]). +/// +/// DSL form: `GZIP` — no configuration options. +/// +/// GZIP is a [CodecKind#COMPRESSION] stage. Compression stages may be chained after all +/// transforms. +/// +/// The name {@value #NAME} is a frozen on-disk API contract stored verbatim in segment file +/// headers. It must never be changed or reused for a different algorithm. +/// +/// Wire format: DEFLATE-compressed payload followed by a 4-byte big-endian footer containing +/// the uncompressed byte count. The footer allows decompression without knowing the original size +/// out-of-band. +/// +/// **Performance note:** The JDK [Deflater]/[Inflater] instances are reused per thread and +/// operate directly on [ByteBuffer] inputs and outputs. For write-intensive workloads prefer +/// LZ4 or ZSTD. +final class GzipCodecDefinition implements ChunkCodecHandler<GzipCodecDefinition.Options> { + + /// On-disk permanent name stored verbatim in segment file headers. + /// This string is a frozen on-disk API contract and must never be changed. + public static final String NAME = "GZIP"; + + /// Thread-local Deflater/Inflater: reset() between uses to amortize JNI allocation cost. + /// Note: these hold native resources that are only released when the thread dies. For long-lived + /// server thread pools this is bounded by the worker count and acceptable. + private static final ThreadLocal<Deflater> DEFLATER = ThreadLocal.withInitial(Deflater::new); + private static final ThreadLocal<Inflater> INFLATER = ThreadLocal.withInitial(Inflater::new); + private static final ThreadLocal<byte[]> COMPLETION_PROBE = ThreadLocal.withInitial(() -> new byte[1]); + + /// Sanity cap on decompressedSize read from the (untrusted) GZIP trailer to prevent DoS-on-corrupt-segment + /// via a giant pre-allocation. 1 GiB is well above any realistic chunk size. + private static final int MAX_REASONABLE_DECOMPRESSED_SIZE = 1 << 30; + + public static final GzipCodecDefinition INSTANCE = new GzipCodecDefinition(); + + /// Singleton options — GZIP has no configurable parameters. + public static final Options OPTIONS = new Options(); + + private GzipCodecDefinition() { + } + + /// Typed options for [GzipCodecDefinition]. GZIP has no configurable parameters. + public static final class Options implements CodecOptions { + private Options() { + } + } + + @Override + public String name() { + return NAME; + } + + @Override + public CodecKind kind() { + return CodecKind.COMPRESSION; + } + + @Override + public Options parseOptions(List<String> args) { + if (!args.isEmpty()) { + throw new IllegalArgumentException("GZIP codec does not accept arguments, got: " + args); + } + return OPTIONS; + } + + @Override + public void validateContext(Options options, CodecContext ctx) { + // GZIP can compress any data type; no restriction + } + + @Override + public String canonicalize(Options options) { + return NAME; + } + + @Override + public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src) throws IOException { + int uncompressedSize = src.remaining(); + ByteBuffer out = ByteBuffer.allocateDirect(maxEncodedSize(options, uncompressedSize)); + Deflater deflater = DEFLATER.get(); + boolean succeeded = false; + try { + out.limit(out.capacity() - Integer.BYTES); + deflater.reset(); + deflater.setInput(src.duplicate()); + deflater.finish(); + while (!deflater.finished()) { + if (!out.hasRemaining()) { + throw new IOException("GZIP encode exceeded maximum encoded size " + out.capacity() + + " before deflater finished. Segment build aborted."); + } + int encoded = deflater.deflate(out); + if (encoded == 0 && !deflater.finished()) { + throw new IOException("GZIP deflater made no progress before finishing"); + } + } + out.limit(out.capacity()); + out.putInt(uncompressedSize); + out.flip(); + succeeded = true; + return out; + } finally { + deflater.reset(); + if (!succeeded) { + CleanerUtil.cleanQuietly(out); + } + } + } + + @Override + public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src) throws IOException { + int decompressedSize = readDecompressedSize(src); + ByteBuffer out = ByteBuffer.allocateDirect(decompressedSize); + boolean succeeded = false; + try { + inflateInto(src, out, decompressedSize); + succeeded = true; + return out; + } finally { + if (!succeeded) { + CleanerUtil.cleanQuietly(out); + } + } + } + + @Override + public void decodeInto(Options options, CodecContext ctx, ByteBuffer src, ByteBuffer dst) throws IOException { + dst.clear(); + int decompressedSize = readDecompressedSize(src); + if (decompressedSize > dst.capacity()) { + throw new IllegalArgumentException( + "GZIP: decompressed size " + decompressedSize + " exceeds dst capacity " + dst.capacity()); + } + inflateInto(src, dst, decompressedSize); + } + + @Override + public int maxEncodedSize(Options options, int inputSize) { + // DEFLATE worst-case expansion + 4-byte appended uncompressed-size footer + if (inputSize < 0) { + throw new IllegalArgumentException("GZIP inputSize must be non-negative: " + inputSize); + } + long bound = (long) inputSize + (inputSize >> 12) + (inputSize >> 14) + (inputSize >> 25) + + 13 + Integer.BYTES; + if (bound > Integer.MAX_VALUE) { + throw new IllegalArgumentException("GZIP maximum encoded size exceeds Integer.MAX_VALUE: " + bound); + } + return (int) bound; + } + + @Override + public boolean requiresDirectDstBuffer() { + return false; + } + + // ------------------------------------------------------------------------- + // Private helpers + // ------------------------------------------------------------------------- + + private static int readDecompressedSize(ByteBuffer src) throws IOException { + int payloadLimit = src.limit(); + if (payloadLimit < Integer.BYTES) { + throw new IOException("GZIP payload too short to contain uncompressed-size footer: " + payloadLimit + " bytes"); + } + int decompressedSize = src.getInt(payloadLimit - Integer.BYTES); Review Comment: The footer is documented as big-endian, but this absolute `getInt` uses the caller buffer's current byte order. Decoding the same encoded bytes through a little-endian view therefore reads a byte-swapped size and rejects or mis-sizes an otherwise valid frame. Read the footer from a duplicate explicitly ordered `ByteOrder.BIG_ENDIAN`, and add a round-trip test using a little-endian encoded-buffer view. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/GzipCodecDefinition.java: ########## @@ -0,0 +1,279 @@ +/** + * 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.pinot.segment.local.io.codec; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.zip.DataFormatException; +import java.util.zip.Deflater; +import java.util.zip.Inflater; +import org.apache.pinot.segment.spi.memory.CleanerUtil; + + +/// Compression codec backed by GZIP (DEFLATE via [java.util.zip.Deflater]). +/// +/// DSL form: `GZIP` — no configuration options. +/// +/// GZIP is a [CodecKind#COMPRESSION] stage. Compression stages may be chained after all +/// transforms. +/// +/// The name {@value #NAME} is a frozen on-disk API contract stored verbatim in segment file +/// headers. It must never be changed or reused for a different algorithm. +/// +/// Wire format: DEFLATE-compressed payload followed by a 4-byte big-endian footer containing +/// the uncompressed byte count. The footer allows decompression without knowing the original size +/// out-of-band. +/// +/// **Performance note:** The JDK [Deflater]/[Inflater] instances are reused per thread and +/// operate directly on [ByteBuffer] inputs and outputs. For write-intensive workloads prefer +/// LZ4 or ZSTD. +final class GzipCodecDefinition implements ChunkCodecHandler<GzipCodecDefinition.Options> { + + /// On-disk permanent name stored verbatim in segment file headers. + /// This string is a frozen on-disk API contract and must never be changed. + public static final String NAME = "GZIP"; + + /// Thread-local Deflater/Inflater: reset() between uses to amortize JNI allocation cost. + /// Note: these hold native resources that are only released when the thread dies. For long-lived + /// server thread pools this is bounded by the worker count and acceptable. + private static final ThreadLocal<Deflater> DEFLATER = ThreadLocal.withInitial(Deflater::new); + private static final ThreadLocal<Inflater> INFLATER = ThreadLocal.withInitial(Inflater::new); + private static final ThreadLocal<byte[]> COMPLETION_PROBE = ThreadLocal.withInitial(() -> new byte[1]); + + /// Sanity cap on decompressedSize read from the (untrusted) GZIP trailer to prevent DoS-on-corrupt-segment + /// via a giant pre-allocation. 1 GiB is well above any realistic chunk size. + private static final int MAX_REASONABLE_DECOMPRESSED_SIZE = 1 << 30; + + public static final GzipCodecDefinition INSTANCE = new GzipCodecDefinition(); + + /// Singleton options — GZIP has no configurable parameters. + public static final Options OPTIONS = new Options(); + + private GzipCodecDefinition() { + } + + /// Typed options for [GzipCodecDefinition]. GZIP has no configurable parameters. + public static final class Options implements CodecOptions { + private Options() { + } + } + + @Override + public String name() { + return NAME; + } + + @Override + public CodecKind kind() { + return CodecKind.COMPRESSION; + } + + @Override + public Options parseOptions(List<String> args) { + if (!args.isEmpty()) { + throw new IllegalArgumentException("GZIP codec does not accept arguments, got: " + args); + } + return OPTIONS; + } + + @Override + public void validateContext(Options options, CodecContext ctx) { + // GZIP can compress any data type; no restriction + } + + @Override + public String canonicalize(Options options) { + return NAME; + } + + @Override + public ByteBuffer encode(Options options, CodecContext ctx, ByteBuffer src) throws IOException { + int uncompressedSize = src.remaining(); + ByteBuffer out = ByteBuffer.allocateDirect(maxEncodedSize(options, uncompressedSize)); + Deflater deflater = DEFLATER.get(); + boolean succeeded = false; + try { + out.limit(out.capacity() - Integer.BYTES); + deflater.reset(); + deflater.setInput(src.duplicate()); + deflater.finish(); + while (!deflater.finished()) { + if (!out.hasRemaining()) { + throw new IOException("GZIP encode exceeded maximum encoded size " + out.capacity() + + " before deflater finished. Segment build aborted."); + } + int encoded = deflater.deflate(out); + if (encoded == 0 && !deflater.finished()) { + throw new IOException("GZIP deflater made no progress before finishing"); + } + } + out.limit(out.capacity()); + out.putInt(uncompressedSize); + out.flip(); + succeeded = true; + return out; + } finally { + deflater.reset(); + if (!succeeded) { + CleanerUtil.cleanQuietly(out); + } + } + } + + @Override + public ByteBuffer decode(Options options, CodecContext ctx, ByteBuffer src) throws IOException { + int decompressedSize = readDecompressedSize(src); + ByteBuffer out = ByteBuffer.allocateDirect(decompressedSize); + boolean succeeded = false; + try { + inflateInto(src, out, decompressedSize); + succeeded = true; + return out; + } finally { + if (!succeeded) { + CleanerUtil.cleanQuietly(out); + } + } + } + + @Override + public void decodeInto(Options options, CodecContext ctx, ByteBuffer src, ByteBuffer dst) throws IOException { + dst.clear(); + int decompressedSize = readDecompressedSize(src); + if (decompressedSize > dst.capacity()) { + throw new IllegalArgumentException( + "GZIP: decompressed size " + decompressedSize + " exceeds dst capacity " + dst.capacity()); + } + inflateInto(src, dst, decompressedSize); + } + + @Override + public int maxEncodedSize(Options options, int inputSize) { + // DEFLATE worst-case expansion + 4-byte appended uncompressed-size footer + if (inputSize < 0) { + throw new IllegalArgumentException("GZIP inputSize must be non-negative: " + inputSize); + } + long bound = (long) inputSize + (inputSize >> 12) + (inputSize >> 14) + (inputSize >> 25) + + 13 + Integer.BYTES; + if (bound > Integer.MAX_VALUE) { + throw new IllegalArgumentException("GZIP maximum encoded size exceeds Integer.MAX_VALUE: " + bound); + } + return (int) bound; + } + + @Override + public boolean requiresDirectDstBuffer() { + return false; + } + + // ------------------------------------------------------------------------- + // Private helpers + // ------------------------------------------------------------------------- + + private static int readDecompressedSize(ByteBuffer src) throws IOException { + int payloadLimit = src.limit(); + if (payloadLimit < Integer.BYTES) { + throw new IOException("GZIP payload too short to contain uncompressed-size footer: " + payloadLimit + " bytes"); + } + int decompressedSize = src.getInt(payloadLimit - Integer.BYTES); + if (decompressedSize < 0) { + throw new IOException("GZIP: invalid decompressed size in footer: " + decompressedSize); + } + if (decompressedSize > MAX_REASONABLE_DECOMPRESSED_SIZE) { + throw new IOException( + "GZIP: decompressed size " + decompressedSize + " in footer exceeds sanity cap " + + MAX_REASONABLE_DECOMPRESSED_SIZE + ". Segment may be corrupt."); + } + return decompressedSize; + } + + private static void inflateInto(ByteBuffer src, ByteBuffer dst, int decompressedSize) throws IOException { + ByteBuffer compressed = src.duplicate(); + compressed.position(0); + compressed.limit(src.limit() - Integer.BYTES); Review Comment: Resetting this duplicate to position 0 violates the public executor's “ready for read” source contract and disagrees with the outer bound check, which uses `src.remaining()`. With a nonzero source position, bytes before that position are decompressed even though they were excluded from the encoded-size bound. Preserve the duplicate's current position and only move its limit before the footer. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/io/codec/CodecPipelineExecutor.java: ########## @@ -0,0 +1,358 @@ +/** + * 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.pinot.segment.local.io.codec; + +import com.google.common.base.Preconditions; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.List; +import org.apache.pinot.segment.spi.codec.CodecInvocation; +import org.apache.pinot.segment.spi.codec.CodecPipeline; +import org.apache.pinot.segment.spi.codec.CodecSpecParser; +import org.apache.pinot.segment.spi.memory.CleanerUtil; +import org.apache.pinot.spi.data.FieldSpec.DataType; + + +/// Executes a parsed and validated [CodecPipeline] for a single forward-index chunk. +/// +/// Write path: values → transforms (in order) → compression stages → bytes stored on disk. +/// Read path: bytes from disk → reverse compression/transform stages → values. +/// +/// The executor is constructed once per column from the canonical `codecSpec` string +/// stored in the file header and is thread-safe for concurrent read calls (it holds no mutable +/// per-call state). +/// +/// The executor is codec-agnostic: it holds an ordered list of internal bound stages, each pairing +/// a handler with its parsed options. The registry and handlers are a closed, package-private +/// runtime; this class is the only public entry point and only drives the pipeline loop. +/// +/// ### Buffer contract +/// +/// - [#encode]: `src` is ready for read (position=0); returns a new +/// [ByteBuffer] ready for read containing the encoded bytes. +/// - [#decode(ByteBuffer)]: `src` is ready for read; returns a new +/// [ByteBuffer] ready for read containing the decoded bytes. +/// - [#maxEncodedSize]: returns an upper bound on encoded size. +public final class CodecPipelineExecutor { + + /// A codec handler bound to the options parsed from a specific pipeline invocation. + private static final class BoundStage<O extends CodecOptions> { + final ChunkCodecHandler<O> _handler; + final O _options; + final CodecContext _ctx; + + BoundStage(ChunkCodecHandler<O> handler, O options, CodecContext ctx) { + _handler = handler; + _options = options; + _ctx = ctx; + } + + ByteBuffer encode(ByteBuffer src) throws IOException { + return _handler.encode(_options, _ctx, src); + } + + ByteBuffer decode(ByteBuffer src) throws IOException { + return _handler.decode(_options, _ctx, src); + } + + void decodeInto(ByteBuffer src, ByteBuffer dst) throws IOException { + _handler.decodeInto(_options, _ctx, src, dst); + } + + int maxEncodedSize(int inputSize) { + return _handler.maxEncodedSize(_options, inputSize); + } + + boolean requiresDirectDstBuffer() { + return _handler.requiresDirectDstBuffer(); + } + + boolean isCompression() { + return _handler.kind() == CodecKind.COMPRESSION; + } + + String canonicalize() { + return _handler.canonicalize(_options); + } + } + + private final List<BoundStage<?>> _stages; + private final String _canonicalSpec; + private final boolean _hasCompression; + private final boolean _requiresDirectDstBuffer; + + /// Creates an executor by parsing and validating the given spec. + /// + /// @param spec the codec DSL string (e.g. `"LZ4,ZSTD(3)"`) + /// @param storedType stored data type used for validation + public static CodecPipelineExecutor create(String spec, DataType storedType) { + return create(spec, new CodecContext(storedType), CodecRegistry.DEFAULT); + } + + /// Package-scoped construction hook for tests that provide a custom closed registry. + static CodecPipelineExecutor create(String spec, CodecContext ctx, CodecRegistry registry) { + CodecPipeline pipeline = CodecSpecParser.parse(spec); + CodecPipelineValidator.validate(pipeline, registry, ctx); + return new CodecPipelineExecutor(pipeline, registry, ctx); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private CodecPipelineExecutor(CodecPipeline pipeline, CodecRegistry registry, CodecContext ctx) { + List<CodecInvocation> invocations = pipeline.stages(); + if (invocations.isEmpty()) { + throw new IllegalArgumentException("Codec pipeline must contain at least one stage"); + } + List<BoundStage<?>> stages = new ArrayList<>(invocations.size()); + + for (CodecInvocation inv : invocations) { + ChunkCodecHandler handler = (ChunkCodecHandler) registry.getOrThrow(inv.name()); + CodecOptions opts = handler.parseOptions(inv.args()); + stages.add(new BoundStage<>(handler, opts, ctx)); + } + _stages = stages; + _canonicalSpec = buildCanonical(stages); + _hasCompression = stages.stream().anyMatch(BoundStage::isCompression); + // decodeInto() writes the final output through stage zero. All other stage outputs use + // executor-owned direct scratch buffers, so only stage zero constrains the caller's dst. + _requiresDirectDstBuffer = stages.get(0).requiresDirectDstBuffer(); + } + + /// Returns the canonical spec string derived from the parsed pipeline. + public String getCanonicalSpec() { + return _canonicalSpec; + } + + /// Returns an upper bound on the number of bytes that [#encode] may produce for + /// a decoded chunk of the given byte length. + int maxEncodedSize(int decodedSize) { + return maxEncodedSize(decodedSize, Integer.MAX_VALUE, Long.MAX_VALUE); + } + + /// Returns the composed encoded-size bound while requiring every stage's bound to stay within + /// `maxStageSize`. This lets an on-disk format reject a pipeline/chunk-size combination before + /// either the writer produces an unreadable file or the reader allocates excessive scratch. + int maxEncodedSize(int decodedSize, int maxStageSize) { + return maxEncodedSize(decodedSize, maxStageSize, Long.MAX_VALUE); + } + + /// Returns the composed encoded-size bound while also capping the sum of all stage-output + /// bounds. The cumulative limit bounds CPU and allocation churn for long pipelines even when + /// each individual stage stays below `maxStageSize`. + public int maxEncodedSize(int decodedSize, int maxStageSize, long maxCumulativeSize) { + Preconditions.checkArgument(decodedSize >= 0, "decodedSize must be non-negative: %s", decodedSize); + Preconditions.checkArgument(maxStageSize >= decodedSize, + "maxStageSize %s must be at least decodedSize %s", maxStageSize, decodedSize); + Preconditions.checkArgument(maxCumulativeSize >= decodedSize, + "maxCumulativeSize %s must be at least decodedSize %s", maxCumulativeSize, decodedSize); + int size = decodedSize; + long cumulativeSize = 0; + for (int i = 0; i < _stages.size(); i++) { + size = _stages.get(i).maxEncodedSize(size); + if (size < 0 || size > maxStageSize) { + throw new IllegalArgumentException( + "Codec stage " + i + " maximum encoded size " + size + " is outside [0, " + maxStageSize + + "] for pipeline " + _canonicalSpec); + } + cumulativeSize += size; + if (cumulativeSize > maxCumulativeSize) { + throw new IllegalArgumentException( + "Codec pipeline cumulative stage-output bound " + cumulativeSize + " exceeds " + maxCumulativeSize + + " at stage " + i + " for pipeline " + _canonicalSpec); + } + } + return size; + } + + /// Encodes `src` through the pipeline and returns the encoded bytes ready for read. + /// + /// **Position contract:** stages may consume `src` (advance its position). Callers + /// that need to re-read `src` after this call must pass `src.duplicate()`. + /// + /// @param src decoded chunk data, ready for read (position=0, limit=dataSize) + /// @return encoded buffer ready for read; the caller owns this buffer + public ByteBuffer encode(ByteBuffer src) throws IOException { + ByteBuffer current = src; + try { + for (BoundStage<?> stage : _stages) { + ByteBuffer previous = current; + current = stage.encode(previous); + if (previous != src && previous != current) { + CleanerUtil.cleanQuietly(previous); + } + } + return current; + } catch (IOException | RuntimeException | Error e) { + if (current != src) { + CleanerUtil.cleanQuietly(current); + } + throw e; + } + } + + /// Decodes `src` through the reversed pipeline and returns the original bytes. + /// + /// @param src encoded chunk data, ready for read + /// @return decoded buffer ready for read; the caller owns this buffer + ByteBuffer decode(ByteBuffer src) throws IOException { + ByteBuffer current = src; + try { + for (int i = _stages.size() - 1; i >= 0; i--) { + ByteBuffer previous = current; + current = _stages.get(i).decode(previous); + if (previous != src && previous != current) { + CleanerUtil.cleanQuietly(previous); + } + } + return current; + } catch (IOException | RuntimeException | Error e) { + if (current != src) { + CleanerUtil.cleanQuietly(current); + } + throw e; + } + } + + /// Decodes `src` through the reversed pipeline, writing the result directly into + /// `dst`. On return `dst` is flipped and ready for read (position=0, + /// limit=decoded size). + /// + /// For single-stage pipelines (transform-only or compression-only), the decoded bytes are + /// written directly into `dst`, avoiding an intermediate allocation. For multi-stage + /// pipelines an intermediate buffer is allocated for any stage that is not the last in the + /// reversed pipeline; the final stage writes directly into `dst`. + /// + /// @param src encoded chunk data, ready for read + /// @param dst caller-supplied output buffer; must be a *direct* [ByteBuffer] when + /// the pipeline requires it (see [ChunkCodecHandler#requiresDirectDstBuffer]); + /// must have sufficient [ByteBuffer#capacity()] for the decoded data; its + /// position and limit are overwritten before returning + /// @throws IOException if decoding fails + /// @throws IllegalArgumentException if `dst` is not direct when required, or does not have + /// enough capacity + void decode(ByteBuffer src, ByteBuffer dst) throws IOException { + decode(src, dst, dst.capacity()); + } + + /// Decodes into `dst` while bounding every intermediate allocation from the expected final + /// decoded size. This is the segment-reader entry point: codec frames are untrusted and may + /// advertise arbitrary decoded lengths, so an intermediate stage must never allocate directly + /// from its frame header. + /// + /// @param src encoded chunk data, ready for read + /// @param dst caller-owned final output buffer + /// @param expectedDecodedSize exact decoded bytes declared and validated by the outer format + void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize) throws IOException { + decode(src, dst, expectedDecodedSize, Integer.MAX_VALUE, Long.MAX_VALUE); + } + + /// Variant of [#decode(ByteBuffer, ByteBuffer, int)] that caps every intermediate scratch + /// buffer. Callers reading untrusted persisted data should pass the format's validated limit. + void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize, int maxIntermediateSize) + throws IOException { + decode(src, dst, expectedDecodedSize, maxIntermediateSize, Long.MAX_VALUE); + } + + /// Variant that also caps the sum of stage-output bounds, limiting total work and allocation + /// churn for an otherwise-valid long pipeline. + public void decode(ByteBuffer src, ByteBuffer dst, int expectedDecodedSize, int maxIntermediateSize, + long maxCumulativeSize) + throws IOException { + Preconditions.checkArgument(!_requiresDirectDstBuffer || dst.isDirect(), + "decode(src, dst) requires a direct ByteBuffer for pipeline: %s", _canonicalSpec); + Preconditions.checkArgument(expectedDecodedSize >= 0 && expectedDecodedSize <= dst.capacity(), + "expectedDecodedSize %s is out of range [0, %s]", expectedDecodedSize, dst.capacity()); + Preconditions.checkArgument(maxIntermediateSize >= expectedDecodedSize, + "maxIntermediateSize %s must be at least expectedDecodedSize %s", maxIntermediateSize, expectedDecodedSize); + Preconditions.checkArgument(maxCumulativeSize >= expectedDecodedSize, + "maxCumulativeSize %s must be at least expectedDecodedSize %s", maxCumulativeSize, expectedDecodedSize); + + int stageCount = _stages.size(); + int[] maxOutputAfterStage = new int[stageCount]; + int maxSize = expectedDecodedSize; + long cumulativeSize = 0; + for (int i = 0; i < stageCount; i++) { + maxSize = _stages.get(i).maxEncodedSize(maxSize); + if (maxSize < 0 || maxSize > maxIntermediateSize) { + throw new IllegalArgumentException( + "Codec stage " + i + " maximum encoded size " + maxSize + " is outside [0, " + + maxIntermediateSize + "] for pipeline " + _canonicalSpec); + } + cumulativeSize += maxSize; + if (cumulativeSize > maxCumulativeSize) { + throw new IllegalArgumentException( + "Codec pipeline cumulative stage-output bound " + cumulativeSize + " exceeds " + maxCumulativeSize + + " at stage " + i + " for pipeline " + _canonicalSpec); + } + maxOutputAfterStage[i] = maxSize; + } + Preconditions.checkArgument(src.remaining() <= maxOutputAfterStage[stageCount - 1], + "Encoded input size %s exceeds composed pipeline bound %s for expected decoded size %s", + src.remaining(), maxOutputAfterStage[stageCount - 1], expectedDecodedSize); + + // Decode every stage through decodeInto(). For intermediate stages, the scratch capacity is + // derived from the forward maxEncodedSize chain rooted at the validated final decoded size. + // A corrupt inner frame-controlled length can therefore only trigger a bounded + // "exceeds dst capacity" failure, never a frame-controlled allocation. + ByteBuffer current = src; + for (int i = stageCount - 1; i >= 0; i--) { + ByteBuffer output = null; + try { + output = i == 0 ? dst : ByteBuffer.allocateDirect(maxOutputAfterStage[i - 1]); + _stages.get(i).decodeInto(current, output); Review Comment: This bounded decode path performs a fresh direct allocation for every intermediate stage on every chunk, then immediately cleans it. At the parser's 32-stage limit that is up to 31 native allocations/frees per chunk, while existing forward-reader contexts reuse decompression buffers. Expose a caller-owned reusable scratch context (two ping-pong buffers are sufficient) or another reuse mechanism before this becomes the read hot path. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
