Add zstd compression and decompression for pekko streams

Author: mdedetrich

.scala-steward.conf

@@ -5,6 +5,8 @@ updates.pin = [
   { groupId = "org.scala-lang", artifactId = "scala3-library", version = "3.3." }
   # sbt-assembly 2.3 causes build issues (https://github.com/apache/pekko/pull/1744)
   { groupId = "com.eed3si9n", artifactId = "sbt-assembly", 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 = [

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
+
+@@@

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
+
+@@@

project/Dependencies.scala

@@ -351,7 +351,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,

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
@@ -83,7 +84,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])
       }
     }
 

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

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

stream-tests/src/test/scala/org/apache/pekko/stream/io/compression/ZstdSpec.scala

@@ -0,0 +1,48 @@
+/*
+ * 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.stream.impl.io.compression.{ Compressor, ZstdCompressor }
+import org.apache.pekko.stream.scaladsl.{ Compression, Flow }
+import org.apache.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!")
+    }
+  }
+}

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)
+    }
+  }
 }

stream/src/main/scala/org/apache/pekko/stream/impl/io/compression/ZstdCompressor.scala

@@ -0,0 +1,93 @@
+/*
+ * 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.{ Zstd, ZstdDictCompress, ZstdDirectBufferCompressingStreamNoFinalizer }
+
+import org.apache.pekko
+import pekko.annotation.InternalApi
+import pekko.util.ByteString
+
+/** INTERNAL API */
+@InternalApi private[pekko] class ZstdCompressor(
+    compressionLevel: Int =
+      Zstd.defaultCompressionLevel(), dictionary: Option[ZstdDictCompress] = None) extends Compressor {
+
+  private val targetBuffer = ByteBuffer.allocateDirect(65536)
+  private val compressingStream = new ZstdDirectBufferCompressingStreamNoFinalizer(targetBuffer, compressionLevel)
+
+  dictionary.foreach(compressingStream.setDict)
+
+  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()
+}