Zstd compression for thrift serialization (storm cluster state)

conf/defaults.yaml

@@ -54,7 +54,8 @@ storm.nimbus.zookeeper.acls.fixup: true
 
 storm.auth.simple-white-list.users: []
 storm.cluster.state.store: "org.apache.storm.cluster.ZKStateStorageFactory"
-storm.meta.serialization.delegate: "org.apache.storm.serialization.GzipThriftSerializationDelegate"
+storm.meta.serialization.delegate: "org.apache.storm.serialization.ZstdThriftSerializationDelegate"
+storm.compression.zstd.level: 3
 storm.codedistributor.class: "org.apache.storm.codedistributor.LocalFileSystemCodeDistributor"
 storm.workers.artifacts.dir: "workers-artifacts"
 storm.health.check.dir: "healthchecks"

pom.xml

@@ -82,6 +82,7 @@
 
         <!-- dependency versions -->
         <commons-compress.version>1.28.0</commons-compress.version>
+        <zstd-jni.version>1.5.7-8</zstd-jni.version>
         <commons-io.version>2.22.0</commons-io.version>
         <commons-lang3.version>3.20.0</commons-lang3.version>
         <commons-exec.version>1.6.0</commons-exec.version>
@@ -514,6 +515,12 @@
                 <artifactId>commons-compress</artifactId>
                 <version>${commons-compress.version}</version>
             </dependency>
+            <dependency>
+                <groupId>com.github.luben</groupId>
+                <artifactId>zstd-jni</artifactId>
+                <version>${zstd-jni.version}</version>
+                <scope>compile</scope>
+            </dependency>
             <dependency>
                 <groupId>org.apache.commons</groupId>
                 <artifactId>commons-exec</artifactId>

storm-client/pom.xml

@@ -135,6 +135,15 @@
             <artifactId>curator-test</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-compress</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>com.github.luben</groupId>
+            <artifactId>zstd-jni</artifactId>
+            <scope>compile</scope>
+        </dependency>
     </dependencies>
 
     <build>

storm-client/src/jvm/org/apache/storm/Config.java

@@ -1458,7 +1458,11 @@ public class Config extends HashMap<String, Object> {
      */
     @IsString
     public static final String STORM_META_SERIALIZATION_DELEGATE = "storm.meta.serialization.delegate";
-
+    /**
+     * Zstandard compression level. Defaults to 3.
+     */
+    @CustomValidator(validatorClass = ConfigValidation.ZstdLevelValidator.class)
+    public static final String STORM_COMPRESSION_ZSTD_LEVEL = "storm.compression.zstd.level";
     /**
      * Configure the topology metrics reporters to be used on workers.
      */

storm-client/src/jvm/org/apache/storm/serialization/ZstdBridgeThriftSerializationDelegate.java

@@ -0,0 +1,65 @@
+/**
+ * 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.storm.serialization;
+
+import java.util.Map;
+import java.util.zip.GZIPInputStream;
+
+/**
+ * Always writes Zstd out, but tests incoming bytes to determine the format.
+ * If Zstd magic is found, it uses {@link ZstdThriftSerializationDelegate}.
+ * If not, it falls back to {@link ThriftSerializationDelegate} for raw Thrift.
+ */
+public class ZstdBridgeThriftSerializationDelegate implements SerializationDelegate {
+
+    /**
+     * Zstandard magic number 0xFD2FB52. In a byte array (little-endian format): [0x28, 0xB5, 0x2F, 0xFD]
+     */
+    private static final byte[] ZSTD_MAGIC = new byte[]{ 0x28, (byte) 0xB5, 0x2F, (byte) 0xFD };
+
+    private final ThriftSerializationDelegate defaultDelegate = new ThriftSerializationDelegate();
+    private final ZstdThriftSerializationDelegate zstdDelegate = new ZstdThriftSerializationDelegate();
+
+    @Override
+    public void prepare(Map<String, Object> topoConf) {
+        defaultDelegate.prepare(topoConf);
+        zstdDelegate.prepare(topoConf);
+    }
+
+    @Override
+    public byte[] serialize(Object object) {
+        // Always compress new data with Zstd
+        return zstdDelegate.serialize(object);
+    }
+
+    @Override
+    public <T> T deserialize(byte[] bytes, Class<T> clazz) {
+        if (isZstd(bytes)) {
+            return zstdDelegate.deserialize(bytes, clazz);
+        } else {
+            // Fallback for old data or non-compressed data
+            return defaultDelegate.deserialize(bytes, clazz);
+        }
+    }
+
+    /**
+     * Checks the first 4 bytes of the array against the Zstd Magic Number.
+     */
+    private boolean isZstd(byte[] bytes) {
+        if (bytes == null || bytes.length < 4) {
+            return false;
+        }
+
+        return bytes[0] == ZSTD_MAGIC[0] && bytes[1] == ZSTD_MAGIC[1] && bytes[2] == ZSTD_MAGIC[2] && bytes[3] == ZSTD_MAGIC[3];
+    }
+}

storm-client/src/jvm/org/apache/storm/serialization/ZstdThriftSerializationDelegate.java

@@ -0,0 +1,81 @@
+/*
+ * 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.storm.serialization;
+
+import java.util.Map;
+import org.apache.storm.thrift.TBase;
+import org.apache.storm.thrift.TDeserializer;
+import org.apache.storm.thrift.TException;
+import org.apache.storm.thrift.TSerializer;
+import org.apache.storm.thrift.transport.TTransportException;
+import org.apache.storm.utils.Utils;
+
+/**
+ * Note, this assumes it's deserializing a gzip byte stream, and will err if it encounters any other serialization.
+ */
+public class ZstdThriftSerializationDelegate implements SerializationDelegate {
+
+    // ThreadLocal with explicit exception handling for checked TTransportException
+    private static final ThreadLocal<TSerializer> SERIALIZER = ThreadLocal.withInitial(() -> {
+        try {
+            return new TSerializer();
+        } catch (TTransportException e) {
+            throw new RuntimeException("Failed to initialize Thrift Serializer", e);
+        }
+    });
+
+    private static final ThreadLocal<TDeserializer> DESERIALIZER = ThreadLocal.withInitial(() -> {
+        try {
+            return new TDeserializer();
+        } catch (TTransportException e) {
+            throw new RuntimeException("Failed to initialize Thrift Deserializer", e);
+        }
+    });
+
+    @Override
+    public void prepare(Map<String, Object> topoConf) {
+        // No-op: Initialization happens lazily per thread
+    }
+
+    @Override
+    public byte[] serialize(Object object) {
+        if (!(object instanceof TBase)) {
+            throw new IllegalArgumentException("Object must be an instance of TBase");
+        }
+        try {
+            byte[] thriftData = SERIALIZER.get().serialize((TBase<?, ?>) object);
+            return Utils.ZstdUtils.compress(thriftData);
+        } catch (TException e) {
+            throw new RuntimeException("Failed to serialize Thrift object", e);
+        }
+    }
+
+    @Override
+    public <T> T deserialize(byte[] bytes, Class<T> clazz) {
+        try {
+            byte[] decompressed = Utils.ZstdUtils.decompress(bytes);
+            TBase<?, ?> instance = (TBase<?, ?>) clazz.getDeclaredConstructor().newInstance();
+            DESERIALIZER.get().deserialize(instance, decompressed);
+            return (T) instance;
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to deserialize bytes to " + clazz.getName(), e);
+        }
+    }
+
+}

storm-client/src/jvm/org/apache/storm/utils/ConfigUtils.java

@@ -39,6 +39,7 @@ public class ConfigUtils {
     public static final String FILE_SEPARATOR = File.separator;
     public static final String STORM_HOME = "storm.home";
     public static final String RESOURCES_SUBDIR = "resources";
+    public static final int DEFAULT_ZSTD_COMPRESSION_LEVEL = 3;
 
     private static final Set<String> passwordConfigKeys = new HashSet<>();
 
@@ -175,6 +176,11 @@ public static int samplingRate(Map<String, Object> conf) {
         throw new IllegalArgumentException("Illegal topology.stats.sample.rate in conf: " + rate);
     }
 
+    public static int zstdCompressionLevel(Map<String, Object> conf) {
+        return ObjectReader.getInt(conf.getOrDefault(Config.STORM_COMPRESSION_ZSTD_LEVEL,
+            DEFAULT_ZSTD_COMPRESSION_LEVEL));
+    }
+
     public static BooleanSupplier mkStatsSampler(Map<String, Object> conf) {
         return evenSampler(samplingRate(conf));
     }

storm-client/src/jvm/org/apache/storm/utils/Utils.java

@@ -72,6 +72,9 @@
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 import javax.security.auth.Subject;
+import org.apache.commons.compress.compressors.zstandard.ZstdCompressorInputStream;
+import org.apache.commons.compress.compressors.zstandard.ZstdCompressorOutputStream;
+import org.apache.commons.io.IOUtils;
 import org.apache.storm.Config;
 import org.apache.storm.blobstore.BlobStore;
 import org.apache.storm.blobstore.ClientBlobStore;
@@ -960,6 +963,75 @@ public static byte[] gunzip(byte[] data) {
         }
     }
 
+    /**
+     * Static utility class for Zstandard (Zstd) compression and decompression.
+     */
+    public static final class ZstdUtils {
+
+        private static final int BUFFER_SIZE = 64 * 1024;
+
+        /**
+         * Private constructor to prevent instantiation.
+         * @throws UnsupportedOperationException if an attempt is made to instantiate this class.
+         */
+        private ZstdUtils() {
+            throw new UnsupportedOperationException("Utility class should not be instantiated.");
+        }
+
+        /**
+         * Compresses the provided byte array using Zstandard.
+         *
+         * <p>The output includes the standard Zstandard frame header, making it
+         * self-describing for the decompression phase.</p>
+         *
+         * @param data the raw byte array to compress.
+         * @return a compressed byte array, or the original array if null/empty.
+         * @throws RuntimeException wrapping an {@link IOException} if the compression fails.
+         */
+        public static byte[] compress(byte[] data) {
+            if (data == null || data.length == 0) {
+                return data;
+            }
+
+            try (ByteArrayOutputStream bos = new ByteArrayOutputStream(data.length)) {
+                try (ZstdCompressorOutputStream zstdOut = ZstdCompressorOutputStream.builder()
+                        .setOutputStream(bos)
+                        .setBufferSize(BUFFER_SIZE) // impacts on compression ratio
+                        .setLevel(ConfigUtils.zstdCompressionLevel(localConf))
+                        .get()) {
+                    zstdOut.write(data);
+                    zstdOut.finish();
+                }
+                return bos.toByteArray();
+            } catch (Exception e) {
+                throw new RuntimeException("Zstd compression failed", e);
+            }
+        }
+
+        /**
+         * Decompresses a Zstandard-compressed byte array.
+         *
+         * @param data the compressed byte array (Zstd frame).
+         * @return the original decompressed byte array, or the input if null/empty.
+         * @throws RuntimeException wrapping an {@link IOException} if the decompression fails
+         *                          or if the data is not a valid Zstd frame.
+         */
+        public static byte[] decompress(byte[] data) {
+            if (data == null || data.length == 0) {
+                return data;
+            }
+
+            try (ByteArrayInputStream bis = new ByteArrayInputStream(data);
+                 ZstdCompressorInputStream zstdIn = new ZstdCompressorInputStream(bis);
+                 ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
+                IOUtils.copy(zstdIn, bos);
+                return bos.toByteArray();
+            } catch (Exception e) {
+                throw new RuntimeException("Zstd decompression failed. Make sure the data is a valid Zstd frame.", e);
+            }
+        }
+    }
+
     public static List<String> getRepeat(List<String> list) {
         List<String> rtn = new ArrayList<String>();
         Set<String> idSet = new HashSet<String>();

storm-client/src/jvm/org/apache/storm/validation/ConfigValidation.java

@@ -803,6 +803,26 @@ public void validateField(String name, Object o) {
         }
     }
 
+    public static class ZstdLevelValidator extends Validator {
+        private static final int MIN_LEVEL = 1;
+        private static final int MAX_LEVEL = 22;
+
+        @Override
+        public void validateField(String name, Object o) {
+            if (o == null) {
+                return;
+            }
+            SimpleTypeValidator.validateField(name, Integer.class, o);
+            int level = (Integer) o;
+            if (level < MIN_LEVEL || level > MAX_LEVEL) {
+                throw new IllegalArgumentException(
+                    String.format("Field '%s' is invalid: %d. Zstd compression level must be between %d and %d.",
+                        name, level, MIN_LEVEL, MAX_LEVEL)
+                );
+            }
+        }
+    }
+
     public static class EventLoggerRegistryValidator extends Validator {
 
         @Override