Add ZSTD compression engine for Kafka Buffer (#5844)

* Add ZSTD compression engine

Signed-off-by: Jeffrey Aaron Jeyasingh <jeffreyaaron06@gmail.com>

Author: jeffreyAaron

data-prepper-plugins/common/build.gradle

@@ -13,6 +13,7 @@ dependencies {
     implementation libs.commons.io
     implementation 'software.amazon.awssdk:s3'
     implementation 'software.amazon.awssdk:acm'
+    implementation 'com.github.luben:zstd-jni:1.5.7-3'
     implementation libs.commons.compress
     implementation libs.commons.lang3
     implementation libs.bouncycastle.bcprov

data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/codec/CompressionOption.java

@@ -17,6 +17,7 @@ public enum CompressionOption {
     NONE("none"),
     GZIP("gzip"),
     SNAPPY("snappy"),
+    ZSTD("zstd"),
     AUTOMATIC("automatic");
 
     private static final Map<String, CompressionOption> OPTIONS_MAP = Arrays.stream(CompressionOption.values())
@@ -28,13 +29,15 @@ public enum CompressionOption {
     private static final Map<String, DecompressionEngine> DECOMPRESSION_ENGINE_MAP = Map.of(
             "none", new NoneDecompressionEngine(),
             "gzip", new GZipDecompressionEngine(),
-            "snappy", new SnappyDecompressionEngine()
+            "snappy", new SnappyDecompressionEngine(),
+            "zstd", new ZstdDecompressionEngine()
     );
 
     private static final Map<String, CompressionEngine> COMPRESSION_ENGINE_MAP = Map.of(
             "none", new NoneCompressionEngine(),
             "gzip", new GZipCompressionEngine(),
-            "snappy", new SnappyCompressionEngine()
+            "snappy", new SnappyCompressionEngine(),
+            "zstd", new ZstdCompressionEngine()
     );
 
     private final String option;

data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/codec/ZstdCompressionEngine.java

@@ -0,0 +1,14 @@
+package org.opensearch.dataprepper.plugins.codec;
+
+import org.apache.commons.compress.compressors.zstandard.ZstdCompressorOutputStream;
+import org.opensearch.dataprepper.model.codec.CompressionEngine;
+
+import java.io.IOException;
+import java.io.OutputStream;
+
+public class ZstdCompressionEngine implements CompressionEngine {
+    @Override
+    public OutputStream createOutputStream(OutputStream outputStream) throws IOException {
+        return new ZstdCompressorOutputStream(outputStream);
+    }
+}

data-prepper-plugins/common/src/main/java/org/opensearch/dataprepper/plugins/codec/ZstdDecompressionEngine.java

@@ -0,0 +1,14 @@
+package org.opensearch.dataprepper.plugins.codec;
+
+import org.apache.commons.compress.compressors.zstandard.ZstdCompressorInputStream;
+import org.opensearch.dataprepper.model.codec.DecompressionEngine;
+
+import java.io.IOException;
+import java.io.InputStream;
+
+public class ZstdDecompressionEngine implements DecompressionEngine {
+    @Override
+    public InputStream createInputStream(InputStream responseInputStream) throws IOException {
+        return new ZstdCompressorInputStream(responseInputStream);
+    }
+}

data-prepper-plugins/common/src/test/java/org/opensearch/dataprepper/plugins/codec/CompressionOptionTest.java

@@ -39,4 +39,28 @@ void testFromFileName_default() {
     }
 
 
+
+    @ParameterizedTest
+    @EnumSource(value = CompressionOption.class, names = {"AUTOMATIC"}, mode = EnumSource.Mode.EXCLUDE)
+    void testCompressionDecompression(final CompressionOption option) throws Exception {
+        String testData = "This is test data for compression and decompression";
+        byte[] originalBytes = testData.getBytes();
+
+        java.io.ByteArrayOutputStream compressedOutput = new java.io.ByteArrayOutputStream();
+        java.io.OutputStream compressor = option.getCompressionEngine().createOutputStream(compressedOutput);
+        compressor.write(originalBytes);
+        compressor.close();
+
+        byte[] compressedBytes = compressedOutput.toByteArray();
+        java.io.InputStream decompressor = option.getDecompressionEngine().createInputStream(new java.io.ByteArrayInputStream(compressedBytes));
+        byte[] decompressedBytes = decompressor.readAllBytes();
+
+        assertThat(decompressedBytes, is(originalBytes));
+    }
+
+
+
+
+
+
 }

data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/producer/KafkaCustomProducer.java

@@ -121,6 +121,7 @@ public void produceRawData(final byte[] bytes, final String key) throws Exceptio
             ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
             OutputStream compressedOutputStream = compressionConfig.getCompressionEngine().createOutputStream(byteArrayOutputStream);
             compressedOutputStream.write(bytes);
+            compressedOutputStream.close();
             send(topicName, key, byteArrayOutputStream.toByteArray()).get();
 
             topicMetrics.update(producer);
@@ -175,6 +176,7 @@ private void publishJsonMessageAsBytes(Record<Event> record, String key) throws
         ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
         OutputStream compressedOutputStream = compressionConfig.getCompressionEngine().createOutputStream(byteArrayOutputStream);
         compressedOutputStream.write(bytes);
+        compressedOutputStream.close();
 
         send(topicName, key, byteArrayOutputStream.toByteArray());
     }

data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/buffer/KafkaBufferTest.java

@@ -26,6 +26,7 @@
 import org.opensearch.dataprepper.model.plugin.PluginFactory;
 import org.opensearch.dataprepper.model.record.Record;
 import org.opensearch.dataprepper.plugins.buffer.blockingbuffer.BlockingBuffer;
+import org.opensearch.dataprepper.plugins.codec.CompressionOption;
 import org.opensearch.dataprepper.plugins.encryption.EncryptionSupplier;
 import org.opensearch.dataprepper.plugins.kafka.admin.KafkaAdminAccessor;
 import org.opensearch.dataprepper.plugins.kafka.common.KafkaMdc;
@@ -513,4 +514,15 @@ void test_customCompressionOption_with_encryption_disabled() {
 
         assertThat(kafkaBuffer.getCustomCompressionOption().name(), equalTo("NONE"));
     }
+
+    @Test
+    void test_compressionOption_supports_all_kafka_compression_types() {
+        String[] kafkaCompressionTypes = {"none", "gzip", "snappy", "zstd"};
+        
+        for (String compressionType : kafkaCompressionTypes) {
+            CompressionOption option = CompressionOption.fromOptionValue(compressionType);
+            assertThat("CompressionOption should support Kafka compression type to ensure that KafkaBuffer compresses data properly when encryption is enabled: " + compressionType,
+                      option, org.hamcrest.Matchers.notNullValue());
+        }
+    }
 }