Move computing of decompressor retained size

Author: losipiuk

core/trino-main/src/main/java/io/trino/execution/buffer/CompressingDecryptingPageDeserializer.java

@@ -15,7 +15,6 @@
 
 import com.google.common.base.VerifyException;
 import io.airlift.compress.v3.Decompressor;
-import io.airlift.compress.v3.lz4.Lz4Decompressor;
 import io.airlift.slice.Slice;
 import io.airlift.slice.SliceInput;
 import io.airlift.slice.Slices;
@@ -63,6 +62,7 @@ public class CompressingDecryptingPageDeserializer
     public CompressingDecryptingPageDeserializer(
             BlockEncodingSerde blockEncodingSerde,
             Optional<Decompressor> decompressor,
+            int decompressorRetainedSize,
             Optional<SecretKey> encryptionKey,
             int blockSizeInBytes,
             OptionalInt maxCompressedBlockSizeInBytes)
@@ -72,6 +72,7 @@ public CompressingDecryptingPageDeserializer(
         encryptionKey.ifPresent(secretKey -> checkArgument(is256BitSecretKeySpec(secretKey), "encryptionKey is expected to be an instance of SecretKeySpec containing a 256bit key"));
         input = new SerializedPageInput(
                 requireNonNull(decompressor, "decompressor is null"),
+                decompressorRetainedSize,
                 encryptionKey,
                 blockSizeInBytes,
                 maxCompressedBlockSizeInBytes);
@@ -96,19 +97,19 @@ private static class SerializedPageInput
             extends SliceInput
     {
         private static final int INSTANCE_SIZE = instanceSize(SerializedPageInput.class);
-        // TODO: implement getRetainedSizeInBytes in Lz4Decompressor
-        private static final int DECOMPRESSOR_RETAINED_SIZE = instanceSize(Lz4Decompressor.class);
         private static final int ENCRYPTION_KEY_RETAINED_SIZE = toIntExact(instanceSize(SecretKeySpec.class) + sizeOfByteArray(256 / 8));
 
         private final Optional<Decompressor> decompressor;
+        private final int decompressorRetainedSize;
         private final Optional<SecretKey> encryptionKey;
         private final Optional<Cipher> cipher;
 
         private final ReadBuffer[] buffers;
 
-        private SerializedPageInput(Optional<Decompressor> decompressor, Optional<SecretKey> encryptionKey, int blockSizeInBytes, OptionalInt maxCompressedBlockSizeInBytes)
+        private SerializedPageInput(Optional<Decompressor> decompressor, int decompressorRetainedSize, Optional<SecretKey> encryptionKey, int blockSizeInBytes, OptionalInt maxCompressedBlockSizeInBytes)
         {
             this.decompressor = requireNonNull(decompressor, "decompressor is null");
+            this.decompressorRetainedSize = decompressorRetainedSize;
             this.encryptionKey = requireNonNull(encryptionKey, "encryptionKey is null");
 
             buffers = new ReadBuffer[
@@ -502,7 +503,7 @@ public int skipBytes(int length)
         public long getRetainedSize()
         {
             long size = INSTANCE_SIZE;
-            size += sizeOf(decompressor, compressor -> DECOMPRESSOR_RETAINED_SIZE);
+            size += sizeOf(decompressor, compressor -> decompressorRetainedSize);
             size += sizeOf(encryptionKey, encryptionKey -> ENCRYPTION_KEY_RETAINED_SIZE);
             size += sizeOf(cipher, cipher -> ESTIMATED_AES_CIPHER_RETAINED_SIZE);
             for (ReadBuffer input : buffers) {

core/trino-main/src/main/java/io/trino/execution/buffer/PagesSerdeFactory.java

@@ -27,6 +27,7 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 
+import static io.airlift.slice.SizeOf.instanceSize;
 import static io.trino.execution.buffer.CompressionCodec.LZ4;
 import static io.trino.execution.buffer.CompressionCodec.ZSTD;
 import static java.util.Objects.requireNonNull;
@@ -66,6 +67,7 @@ public PageDeserializer createDeserializer(Optional<SecretKey> encryptionKey)
         return new CompressingDecryptingPageDeserializer(
                 blockEncodingSerde,
                 createDecompressor(compressionCodec),
+                decompressorRetainedSize(compressionCodec),
                 encryptionKey,
                 blockSizeInBytes,
                 maxCompressedSize(blockSizeInBytes, compressionCodec));
@@ -97,4 +99,14 @@ private static OptionalInt maxCompressedSize(int uncompressedSize, CompressionCo
             case ZSTD -> ZSTD.maxCompressedLength(uncompressedSize);
         };
     }
+
+    private static int decompressorRetainedSize(CompressionCodec compressionCodec)
+    {
+        // TODO: implement getRetainedSizeInBytes in Lz4Decompressor and ZstdDecompressor
+        return switch (compressionCodec) {
+            case NONE -> 0;
+            case LZ4 -> instanceSize(Lz4Decompressor.class);
+            case ZSTD -> instanceSize(ZstdDecompressor.class);
+        };
+    }
 }