Add zstd compression and decompression for pekko streams

mdedetrich · mdedetrich · commit 0989b9d85c86 · 2025-11-30T11:58:53.000+03:00
diff --git a/.scala-steward.conf b/.scala-steward.conf
@@ -8,6 +8,8 @@ updates.pin = [
   # agrona major+minor version should match the one brought
   # in by aeron
   { groupId = "org.agrona", artifactId = "agrona", version = "2.2." }
+  # zstd-jni upgrades are fine as long as they are backwards compatible
+  { groupId = "com.github.luben", artifactId = "zstd-jni", version = "1."}
 ]
 
 updates.ignore = [
diff --git a/actor/src/main/java/org/apache/pekko/io/ByteBufferCleaner.java b/actor/src/main/java/org/apache/pekko/io/ByteBufferCleaner.java
@@ -23,6 +23,7 @@
 import java.lang.invoke.MethodHandles;
 import java.lang.reflect.Field;
 import java.nio.ByteBuffer;
+import org.apache.pekko.annotation.InternalStableApi;
 
 /**
  * Cleans a direct {@link ByteBuffer}. Without manual intervention, direct ByteBuffers will be
@@ -36,7 +37,8 @@
  *
  * <p>See <a href=https://bugs.openjdk.java.net/browse/JDK-4724038>JDK-4724038</a>
  */
-final class ByteBufferCleaner {
+@InternalStableApi
+public final class ByteBufferCleaner {
 
   // adapted from
   // https://github.com/apache/commons-io/blob/441115a4b5cd63ae808dd4c40fc238cb52c8048f/src/main/java/org/apache/commons/io/input/ByteBufferCleaner.java
@@ -75,7 +77,7 @@ public void clean(final ByteBuffer buffer) throws Throwable {
    * @param buffer to release.
    * @throws IllegalStateException on internal failure.
    */
-  static void clean(final ByteBuffer buffer) {
+  public static void clean(final ByteBuffer buffer) {
     try {
       INSTANCE.clean(buffer);
     } catch (final Throwable t) {
@@ -116,7 +118,7 @@ private static Cleaner getCleaner() {
    *
    * @return {@code true} if cleaning is supported, {@code false} otherwise.
    */
-  static boolean isSupported() {
+  public static boolean isSupported() {
     return INSTANCE != null;
   }
 }
diff --git a/docs/src/main/paradox/stream/operators/Compression/zstd.md b/docs/src/main/paradox/stream/operators/Compression/zstd.md
@@ -0,0 +1,30 @@
+# Compression.zstd
+
+Creates a flow that zstd-compresses a stream of ByteStrings.
+
+@ref[Compression operators](../index.md#compression-operators)
+
+## Signature
+
+@apidoc[Compression.zstd](stream.*.Compression$) { scala="#zstd:org.apache.pekko.stream.scaladsl.Flow[org.apache.pekko.util.ByteString,org.apache.pekko.util.ByteString,org.apache.pekko.NotUsed]" java="#zstd()" }
+
+## Description
+
+Creates a flow that zstd-compresses a stream of ByteStrings. Note that the compressor
+will SYNC_FLUSH after every @apidoc[util.ByteString] so that it is guaranteed that every @apidoc[util.ByteString]
+coming out of the flow can be fully decompressed without waiting for additional data. This may
+come at a compression performance cost for very small chunks.
+
+Use the overload method to control the compression level.
+
+## Reactive Streams semantics
+
+@@@div { .callout }
+
+**emits** when the compression algorithm produces output for the received `ByteString`
+
+**backpressures** when downstream backpressures
+
+**completes** when upstream completes
+
+@@@
diff --git a/docs/src/main/paradox/stream/operators/Compression/zstdDecompress.md b/docs/src/main/paradox/stream/operators/Compression/zstdDecompress.md
@@ -0,0 +1,26 @@
+# Compression.zstdDecompress
+
+Creates a flow that zstd-decompresses a stream of ByteStrings.
+
+@ref[Compression operators](../index.md#compression-operators)
+
+## Signature
+
+@apidoc[Compression.zstdDecompress](stream.*.Compression$) { scala="#zstdDecompress(maxBytesPerChunk:Int):org.apache.pekko.stream.scaladsl.Flow[org.apache.pekko.util.ByteString,org.apache.pekko.util.ByteString,org.apache.pekko.NotUsed]" java="#zstdDecompress(int)" }
+
+## Description
+
+Creates a flow that zstd-decompresses a stream of ByteStrings. If the input is truncated, uses invalid
+compression method or is invalid (failed CRC checks) this operator fails with a `com.github.luben.zstd.ZstdIOException`.
+
+## Reactive Streams semantics
+
+@@@div { .callout }
+
+**emits** when the decompression algorithm produces output for the received `ByteString` (the emitted `ByteString` is of `maxBytesPerChunk` maximum length)
+
+**backpressures** when downstream backpressures
+
+**completes** when upstream completes
+
+@@@
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
@@ -352,7 +352,8 @@ object Dependencies {
 
   // pekko stream
 
-  lazy val stream = l ++= Seq[sbt.ModuleID](reactiveStreams, TestDependencies.scalatest)
+  lazy val stream =
+    l ++= Seq[sbt.ModuleID](reactiveStreams, "com.github.luben" % "zstd-jni" % "1.5.7-6", TestDependencies.scalatest)
 
   lazy val streamTestkit = l ++= Seq(
     TestDependencies.scalatest,
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/CoderSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/CoderSpec.scala
@@ -15,11 +15,11 @@ package org.apache.pekko.stream.io.compression
 
 import java.io.{ ByteArrayOutputStream, InputStream, OutputStream }
 import java.util.concurrent.ThreadLocalRandom
-import java.util.zip.DataFormatException
 
 import scala.annotation.tailrec
 import scala.concurrent.Await
 import scala.concurrent.duration._
+import scala.reflect.ClassTag
 import scala.util.control.NoStackTrace
 
 import org.apache.pekko
@@ -31,7 +31,8 @@ import pekko.util.ByteString
 import org.scalatest.Inspectors
 import org.scalatest.wordspec.AnyWordSpec
 
-abstract class CoderSpec(codecName: String) extends AnyWordSpec with CodecSpecSupport with Inspectors {
+abstract class CoderSpec[CorruptInputException: ClassTag](codecName: String) extends AnyWordSpec with CodecSpecSupport
+    with Inspectors {
   import CompressionTestingTools._
 
   protected def newCompressor(): Compressor
@@ -85,7 +86,7 @@ abstract class CoderSpec(codecName: String) extends AnyWordSpec with CodecSpecSu
       "throw an error on corrupt input" in {
         (the[RuntimeException] thrownBy {
           ourDecode(corruptContent)
-        }).ultimateCause should be(a[DataFormatException])
+        }).ultimateCause should be(a[CorruptInputException])
       }
     }
 
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/DeflateSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/DeflateSpec.scala
@@ -21,7 +21,7 @@ import pekko.stream.impl.io.compression.{ Compressor, DeflateCompressor }
 import pekko.stream.scaladsl.{ Compression, Flow }
 import pekko.util.ByteString
 
-class DeflateSpec extends CoderSpec("deflate") {
+class DeflateSpec extends CoderSpec[DataFormatException]("deflate") {
   import CompressionTestingTools._
 
   protected def newCompressor(): Compressor = new DeflateCompressor
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/GzipSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/GzipSpec.scala
@@ -15,14 +15,14 @@ package org.apache.pekko.stream.io.compression
 
 import java.io.{ InputStream, OutputStream }
 import java.nio.charset.StandardCharsets
-import java.util.zip.{ GZIPInputStream, GZIPOutputStream, ZipException }
+import java.util.zip.{ DataFormatException, GZIPInputStream, GZIPOutputStream, ZipException }
 
 import org.apache.pekko
 import pekko.stream.impl.io.compression.{ Compressor, GzipCompressor }
 import pekko.stream.scaladsl.{ Compression, Flow }
 import pekko.util.ByteString
 
-class GzipSpec extends CoderSpec("gzip") {
+class GzipSpec extends CoderSpec[DataFormatException]("gzip") {
   import CompressionTestingTools._
 
   protected def newCompressor(): Compressor = new GzipCompressor
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/ZstdAutoFlushSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/ZstdAutoFlushSpec.scala
@@ -0,0 +1,30 @@
+/*
+ * 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.pekko.stream.io.compression
+
+import org.apache.pekko
+import pekko.stream.scaladsl.{ Compression, Flow }
+import pekko.util.ByteString
+
+class ZstdAutoFlushSpec extends ZstdSpec {
+  override protected val encoderFlow: Flow[ByteString, ByteString, Any] =
+    Compression.zstd(Compression.ZstdDefaultCompressionLevel, dictionary = None, autoFlush = false)
+
+  override protected val autoFlush: Boolean = false
+
+}
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/ZstdSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/ZstdSpec.scala
@@ -0,0 +1,50 @@
+/*
+ * 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.pekko.stream.io.compression
+
+import com.github.luben.zstd.{ ZstdIOException, ZstdInputStream, ZstdOutputStream }
+
+import org.apache.pekko
+import pekko.stream.impl.io.compression.{ Compressor, ZstdCompressor }
+import pekko.stream.scaladsl.{ Compression, Flow }
+import pekko.util.ByteString
+
+import java.io.{ InputStream, OutputStream }
+
+class ZstdSpec extends CoderSpec[ZstdIOException]("zstd") {
+  import CompressionTestingTools._
+
+  override protected def newCompressor(): Compressor = new ZstdCompressor
+
+  override protected def encoderFlow: Flow[ByteString, ByteString, Any] = Compression.zstd
+
+  override protected def decoderFlow(maxBytesPerChunk: Int): Flow[ByteString, ByteString, Any] =
+    Compression.zstdDecompress(maxBytesPerChunk)
+
+  override protected def newDecodedInputStream(underlying: InputStream): InputStream =
+    new ZstdInputStream(underlying)
+
+  override protected def newEncodedOutputStream(underlying: OutputStream): OutputStream =
+    new ZstdOutputStream(underlying)
+
+  override def extraTests(): Unit = {
+    "decode concatenated compressions" in {
+      ourDecode(Seq(encode("Hello, "), encode("dear "), encode("User!")).join) should readAs("Hello, dear User!")
+    }
+  }
+}
diff --git a/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/CompressionSpec.scala b/stream-tests/src/test/scala/org/apache/pekko/stream/scaladsl/CompressionSpec.scala
@@ -16,8 +16,7 @@ package org.apache.pekko.stream.scaladsl
 import java.nio.charset.StandardCharsets
 
 import org.apache.pekko
-import pekko.stream.impl.io.compression.DeflateCompressor
-import pekko.stream.impl.io.compression.GzipCompressor
+import pekko.stream.impl.io.compression.{ DeflateCompressor, GzipCompressor, ZstdCompressor }
 import pekko.stream.testkit.StreamSpec
 import pekko.util.ByteString
 
@@ -27,6 +26,8 @@ class CompressionSpec extends StreamSpec {
 
   def deflate(s: String): ByteString = new DeflateCompressor().compressAndFinish(ByteString(s))
 
+  def zstd(s: String): ByteString = new ZstdCompressor().compressAndFinish(ByteString(s))
+
   val data = "hello world"
 
   "Gzip decompression" must {
@@ -47,4 +48,14 @@ class CompressionSpec extends StreamSpec {
       res.futureValue should ===(data)
     }
   }
+
+  "Zstd decompression" must {
+    "be able to decompress a zstd stream" in {
+      val source =
+        Source.single(zstd(data)).via(Compression.zstdDecompress()).map(_.decodeString(StandardCharsets.UTF_8))
+
+      val res = source.runFold("")(_ + _)
+      res.futureValue should ===(data)
+    }
+  }
 }
diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdCompressor.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdCompressor.scala
@@ -0,0 +1,103 @@
+/*
+ * 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.pekko.stream.impl.io.compression
+
+import java.nio.ByteBuffer
+import com.github.luben.zstd.{ ZstdDictCompress, ZstdDirectBufferCompressingStreamNoFinalizer }
+import org.apache.pekko
+import org.apache.pekko.stream.scaladsl.Compression
+import pekko.annotation.InternalApi
+import pekko.util.ByteString
+
+/** INTERNAL API */
+@InternalApi private[pekko] class ZstdCompressor(
+    compressionLevel: Int =
+      Compression.ZstdDefaultCompressionLevel, dictionary: Option[ZstdDictionaryImpl] = None) extends Compressor {
+
+  private val targetBuffer = ByteBuffer.allocateDirect(65536)
+  private val compressingStream = new ZstdDirectBufferCompressingStreamNoFinalizer(targetBuffer, compressionLevel)
+
+  dictionary.foreach { dict =>
+    (dict.level, dict.length, dict.offset) match {
+      case (None, None, None) =>
+        compressingStream.setDict(dict.dictionary)
+      case (Some(dictLevel), None, None) =>
+        compressingStream.setDict(new ZstdDictCompress(dict.dictionary, dictLevel))
+      case (Some(dictLevel), Some(dictLength), Some(dictOffset)) =>
+        compressingStream.setDict(new ZstdDictCompress(dict.dictionary, dictLevel, dictLength, dictOffset))
+      case _ =>
+        throw new IllegalArgumentException("Invalid combination of ZstdDictionary parameters")
+    }
+  }
+
+  override def compress(input: ByteString): ByteString = {
+    val inputBB = ByteBuffer.allocateDirect(input.size)
+    inputBB.put(input.toArrayUnsafe())
+    inputBB.flip()
+    compressingStream.compress(inputBB)
+    val result = ByteString.fromByteBuffer(targetBuffer)
+    targetBuffer.flip()
+    result
+  }
+
+  override def flush(): ByteString = {
+    targetBuffer.flip()
+    val result = ByteString.fromByteBuffer(targetBuffer)
+    targetBuffer.clear()
+    compressingStream.flush()
+    result
+  }
+
+  override def finish(): ByteString = {
+    compressingStream.close()
+    targetBuffer.flip()
+    val arr = Array.ofDim[Byte](targetBuffer.limit())
+    targetBuffer.get(arr)
+    val result = ByteString.fromArrayUnsafe(arr)
+    result
+  }
+
+  override def compressAndFlush(input: ByteString): ByteString = {
+    val inputBB = ByteBuffer.allocateDirect(input.size)
+    inputBB.put(input.toArray)
+    inputBB.flip()
+    compressingStream.compress(inputBB)
+    compressingStream.flush()
+    targetBuffer.flip()
+
+    val arr = new Array[Byte](targetBuffer.limit())
+    targetBuffer.get(arr)
+    targetBuffer.clear()
+    ByteString.fromArrayUnsafe(arr)
+  }
+
+  override def compressAndFinish(input: ByteString): ByteString = {
+    val inputBB = ByteBuffer.allocateDirect(input.size)
+    inputBB.put(input.toArray)
+    inputBB.flip()
+    compressingStream.compress(inputBB)
+    compressingStream.close()
+    targetBuffer.flip()
+
+    val arr = new Array[Byte](targetBuffer.limit())
+    targetBuffer.get(arr)
+    ByteString.fromArrayUnsafe(arr)
+  }
+
+  override def close(): Unit = compressingStream.close()
+}
diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdDecompressor.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdDecompressor.scala
diff --git a/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdDictionaryImpl.scala b/stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdDictionaryImpl.scala
diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/Compression.scala b/stream/src/main/scala/org/apache/pekko/stream/javadsl/Compression.scala
diff --git a/stream/src/main/scala/org/apache/pekko/stream/javadsl/compression/ZstdDictionary.java b/stream/src/main/scala/org/apache/pekko/stream/javadsl/compression/ZstdDictionary.java
diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Compression.scala b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/Compression.scala
diff --git a/stream/src/main/scala/org/apache/pekko/stream/scaladsl/compression/ZstdDictionary.scala b/stream/src/main/scala/org/apache/pekko/stream/scaladsl/compression/ZstdDictionary.scala

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,8 @@ updates.pin = [`
`8`	`8`	`# agrona major+minor version should match the one brought`
`9`	`9`	`# in by aeron`
`10`	`10`	`{ groupId = "org.agrona", artifactId = "agrona", version = "2.2." }`
	`11`	`+ # zstd-jni upgrades are fine as long as they are backwards compatible`
	`12`	`+ { groupId = "com.github.luben", artifactId = "zstd-jni", version = "1."}`
`11`	`13`	`]`
`12`	`14`
`13`	`15`	`updates.ignore = [`