xiangfu0 commented on code in PR #19285: URL: https://github.com/apache/pinot/pull/19285#discussion_r3827227375
########## 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: Addressed in e99c0e8400. GZIP now reads its four-byte decompressed-size footer from a duplicate explicitly ordered BIG_ENDIAN, independent of the caller view's byte order. The regression test decodes the same frame through a LITTLE_ENDIAN view. -- 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]
