feat: GCS backend - Add support for the ZSTD decompression (#222)

feat: GCS backend - Add support for the ZSTD decompression

Author: mi-char

build.gradle

@@ -14,6 +14,7 @@ ext {
     metricsVersion = "3.0.2"
     http4sVersion = "0.22.15"
     gcsVersion = "2.30.1"
+    zstdVersion = "1.5.5-11"
     monixVersion = "3.4.1" // Used only in tests.
 }
 

core/build.gradle

@@ -13,5 +13,7 @@ dependencies {
 
     api "com.avast.metrics:metrics-scala_2.13:$metricsVersion"
 
+    implementation "com.github.luben:zstd-jni:$zstdVersion"
+
     testImplementation "io.monix:monix_2.13:$monixVersion"
 }

core/src/main/scala/com/avast/clients/storage/compression/ZstdDecompressOutputStream.scala

@@ -0,0 +1,54 @@
+package com.avast.clients.storage.compression
+
+import com.github.luben.zstd.{ZstdDecompressCtx, ZstdInputStream}
+
+import java.io.OutputStream
+import java.nio.ByteBuffer
+import java.nio.channels.Channels
+
+class ZstdDecompressOutputStream(outputStream: OutputStream) extends OutputStream {
+  private val decompressCtx = new ZstdDecompressCtx()
+  private val outputChannel = Channels.newChannel(outputStream)
+  private val outputBuffer = ByteBuffer.allocateDirect(ZstdInputStream.recommendedDOutSize().toInt)
+
+  private var closed = false
+
+  override def write(chunk: Array[Byte]): Unit = {
+    if (closed) {
+      throw new IllegalStateException("Stream is closed")
+    }
+
+    val inputBuffer = ByteBuffer.allocateDirect(chunk.length) // ByteBuffer.wrap(chunk) does not work, we need direct buffer
+    inputBuffer.put(chunk)
+    inputBuffer.rewind()
+
+    while (inputBuffer.hasRemaining) {
+      outputBuffer.clear()
+
+      decompressCtx.decompressDirectByteBufferStream(outputBuffer, inputBuffer)
+
+      outputBuffer.flip()
+
+      while (outputBuffer.hasRemaining) {
+        outputChannel.write(outputBuffer)
+      }
+    }
+  }
+
+  override def write(chunk: Array[Byte], offset: Int, length: Int): Unit = {
+    write(chunk.slice(offset, offset + length))
+  }
+
+  override def write(b: Int): Unit = {
+    write(Array(b.toByte))
+  }
+
+  override def close(): Unit = {
+    if (!closed) {
+      closed = true
+      decompressCtx.close()
+      outputBuffer.clear()
+      outputChannel.close()
+    }
+  }
+}

core/src/test/scala/com/avast/clients/storage/compression/ZstdDecompressOutputStreamTest.scala

@@ -0,0 +1,54 @@
+package com.avast.clients.storage.compression
+
+import com.avast.scala.hashes.Sha256
+import com.github.luben.zstd.Zstd
+import org.junit.runner.RunWith
+import org.scalatest.FunSuite
+import org.scalatestplus.junit.JUnitRunner
+
+import java.io.ByteArrayOutputStream
+import java.nio.ByteBuffer
+import java.security.MessageDigest
+
+@RunWith(classOf[JUnitRunner])
+class ZstdDecompressOutputStreamTest extends FunSuite {
+  private def computeSha256(data: Array[Byte]): Sha256 = {
+    Sha256(MessageDigest.getInstance("SHA-256").digest(data))
+  }
+
+  private def generateData(size: Int): Array[Byte] = {
+    1.to(size).map(i => (i % 256).toByte).toArray // generates compressible data, random data wouldn't compress well
+  }
+
+  test("decompress zstd stream") {
+    val chunkSize = 4 * 1024
+    val testCases = Seq(0, 1, chunkSize, 10 * 1024 * 1024)
+
+    for (testCase <- testCases) {
+      val original_data = generateData(testCase)
+      val original_sha256 = computeSha256(original_data)
+
+      println(s"Original data size: ${original_data.length}")
+      val compressed_data = Zstd.compress(original_data, 9)
+      println(s"Compressed data size: ${compressed_data.length}")
+
+      val sourceStream = ByteBuffer.wrap(compressed_data)
+      val targetStream = new ByteArrayOutputStream()
+
+      val decompressStream = new ZstdDecompressOutputStream(targetStream)
+
+      while (sourceStream.hasRemaining) {
+        val chunkSize = math.min(sourceStream.remaining(), 4 * 1024)
+        val chunk = new Array[Byte](chunkSize)
+        sourceStream.get(chunk)
+        decompressStream.write(chunk)
+      }
+
+      decompressStream.close()
+
+      val result = targetStream.toByteArray
+
+      assert(original_sha256 == computeSha256(result))
+    }
+  }
+}

gcs/README.md

@@ -25,6 +25,12 @@ GCS backends supports multiple ways of authentication:
   * Reading credential file from default paths (see https://cloud.google.com/docs/authentication/application-default-credentials#personal)
   * For all options see https://cloud.google.com/docs/authentication/provide-credentials-adc#how-to 
 
+### Object decompression
+
+GCS backend supports decompression of objects. The decision whether to decompress object or not is based on the `comp-type` header in the object's metadata. 
+If the header is present and contains `zstd` value, the object is decompressed on the fly. Otherwise the object is downloaded as is.
+
+The only supported compression algorithm is currently `zstd`.
 
 ### Client initialization
 

gcs/build.gradle

@@ -6,4 +6,5 @@ dependencies {
     implementation "com.google.cloud:google-cloud-storage:$gcsVersion"
 
     testImplementation "io.monix:monix_2.13:$monixVersion"
+    testImplementation "com.github.luben:zstd-jni:$zstdVersion"
 }

gcs/src/main/scala/com/avast/clients/storage/gcs/GcsStorageBackend.scala

@@ -4,6 +4,7 @@ import better.files.File
 import cats.effect.implicits.catsEffectSyntaxBracket
 import cats.effect.{Blocker, ContextShift, Resource, Sync}
 import cats.syntax.all._
+import com.avast.clients.storage.compression.ZstdDecompressOutputStream
 import com.avast.clients.storage.gcs.GcsStorageBackend.composeBlobPath
 import com.avast.clients.storage.{ConfigurationException, GetResult, HeadResult, StorageBackend, StorageException}
 import com.avast.scala.hashes.Sha256
@@ -17,7 +18,7 @@ import pureconfig.generic.ProductHint
 import pureconfig.generic.auto._
 import pureconfig.{CamelCase, ConfigFieldMapping}
 
-import java.io.{ByteArrayInputStream, FileInputStream}
+import java.io.{ByteArrayInputStream, FileInputStream, OutputStream}
 import java.nio.charset.StandardCharsets
 import java.nio.file.StandardOpenOption
 import java.security.{DigestOutputStream, MessageDigest}
@@ -34,7 +35,12 @@ class GcsStorageBackend[F[_]: Sync: ContextShift](storageClient: Storage, bucket
         blob <- getBlob(sha256)
         result = blob match {
           case Some(blob) =>
-            HeadResult.Exists(blob.getSize)
+            blob.getMetadata.get(GcsStorageBackend.OriginalSizeHeader) match {
+              case null =>
+                HeadResult.Exists(blob.getSize)
+              case originalSize =>
+                HeadResult.Exists(originalSize.toLong)
+            }
           case None =>
             HeadResult.NotFound
         }
@@ -85,15 +91,7 @@ class GcsStorageBackend[F[_]: Sync: ContextShift](storageClient: Storage, bucket
       blocker
         .delay(destination.newOutputStream(FileStreamOpenOptions))
         .bracket { fileStream =>
-          Sync[F]
-            .delay(new DigestOutputStream(fileStream, MessageDigest.getInstance("SHA-256")))
-            .bracket { stream =>
-              blocker.delay(blob.downloadTo(stream)).flatMap { _ =>
-                Sync[F].delay {
-                  (blob.getSize, Sha256(stream.getMessageDigest.digest))
-                }
-              }
-            }(stream => blocker.delay(stream.close()))
+          downloadBlobToFile(blob, fileStream)
         }(fileStream => blocker.delay(fileStream.close()))
         .map[Either[StorageException, GetResult]] {
           case (size, hash) =>
@@ -109,6 +107,53 @@ class GcsStorageBackend[F[_]: Sync: ContextShift](storageClient: Storage, bucket
         }
   }
 
+  private def downloadBlobToFile(blob: Blob, fileStream: OutputStream): F[(Long, Sha256)] = {
+    def getCompressionType: Option[String] = {
+      Option(blob.getMetadata.get(GcsStorageBackend.CompressionTypeHeader)).map(_.toLowerCase)
+    }
+
+    Sync[F]
+      .delay {
+        val countingStream = new GcsStorageBackend.CountingOutputStream(fileStream)
+        val hashingStream = new DigestOutputStream(countingStream, MessageDigest.getInstance("SHA-256"))
+        (countingStream, hashingStream)
+      }
+      .bracket {
+        case (countingStream, hashingStream) => {
+          getCompressionType match {
+            case None =>
+              downloadBlobToStream(blob, hashingStream)
+            case Some("zstd") =>
+              decodeZstdBlobToStream(blob, hashingStream)
+            case Some(unknown) =>
+              throw new IllegalArgumentException(s"Unknown compression type $unknown")
+          }
+        }.flatMap { _ =>
+          Sync[F].delay {
+            (countingStream.length, Sha256(hashingStream.getMessageDigest.digest))
+          }
+        }
+      } {
+        case (hashingStream, countingStream) =>
+          Sync[F].delay {
+            hashingStream.close()
+            countingStream.close()
+          }
+      }
+  }
+
+  private def decodeZstdBlobToStream(blob: Blob, targetStream: DigestOutputStream): F[Unit] = {
+    Sync[F]
+      .delay(new ZstdDecompressOutputStream(targetStream))
+      .bracket { decompressionStream =>
+        downloadBlobToStream(blob, decompressionStream)
+      }(decompressionStream => Sync[F].delay(decompressionStream.close()))
+  }
+
+  private def downloadBlobToStream(blob: Blob, targetStream: OutputStream): F[Unit] = {
+    blocker.delay(blob.downloadTo(targetStream))
+  }
+
   override def close(): Unit = {
     ()
   }
@@ -117,6 +162,9 @@ class GcsStorageBackend[F[_]: Sync: ContextShift](storageClient: Storage, bucket
 object GcsStorageBackend {
   private val DefaultConfig = ConfigFactory.defaultReference().getConfig("gcsBackendDefaults")
 
+  private[gcs] val CompressionTypeHeader = "comp-type"
+  private[gcs] val OriginalSizeHeader = "original-size"
+
   def fromConfig[F[_]: Sync: ContextShift](config: Config,
                                            blocker: Blocker): Either[ConfigurationException, Resource[F, GcsStorageBackend[F]]] = {
 
@@ -148,6 +196,35 @@ object GcsStorageBackend {
     String.join("/", sha256Hex.substring(0, 2), sha256Hex.substring(2, 4), sha256Hex.substring(4, 6), sha256Hex)
   }
 
+  private[gcs] class CountingOutputStream(target: OutputStream) extends OutputStream {
+    private var count: Long = 0
+
+    def length: Long = count
+
+    override def write(b: Int): Unit = {
+      target.write(b)
+      count += 1
+    }
+
+    override def write(b: Array[Byte]): Unit = {
+      target.write(b)
+      count += b.length
+    }
+
+    override def write(b: Array[Byte], off: Int, len: Int): Unit = {
+      target.write(b, off, len)
+      count += len
+    }
+
+    override def flush(): Unit = {
+      target.flush()
+    }
+
+    override def close(): Unit = {
+      target.close()
+    }
+  }
+
   def prepareStorageClient[F[_]: Sync: ContextShift](conf: GcsBackendConfiguration,
                                                      blocker: Blocker): Either[ConfigurationException, Storage] = {
     Either