-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: GCS backend - Add support for the ZSTD decompression
- Loading branch information
Showing
6 changed files
with
194 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
52 changes: 52 additions & 0 deletions
52
core/src/main/scala/com/avast/clients/storage/compression/ZstdDecompressOutputStream.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| 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 closed = false | ||
|
|
||
| override def write(chunk: Array[Byte]): Unit = { | ||
| if (closed) { | ||
| throw new IllegalStateException("Stream is closed") | ||
| } | ||
|
|
||
| val inputBuffer = ByteBuffer.allocateDirect(chunk.length) | ||
| val outputBuffer = ByteBuffer.allocateDirect(ZstdInputStream.recommendedDOutSize().toInt) | ||
|
|
||
| inputBuffer.put(chunk) | ||
| inputBuffer.flip() | ||
|
|
||
| while (inputBuffer.hasRemaining) { | ||
| decompressCtx.decompressDirectByteBufferStream(outputBuffer, inputBuffer) | ||
|
|
||
| outputBuffer.flip() | ||
|
|
||
| while (outputBuffer.hasRemaining) { | ||
| outputChannel.write(outputBuffer) | ||
| } | ||
|
|
||
| outputBuffer.clear() | ||
| } | ||
| } | ||
|
|
||
| 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) { | ||
| decompressCtx.close() | ||
| outputChannel.close() | ||
| } | ||
| } | ||
| } |
60 changes: 60 additions & 0 deletions
60
...src/test/scala/com/avast/clients/storage/compression/ZstdDecompressOutputStreamTest.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,60 @@ | ||
| 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] = { | ||
| val builder = Array.newBuilder[Byte] | ||
| var i = 0 | ||
| while (i < size) { | ||
| builder += (i % 256).toByte | ||
| i += 1 | ||
| } | ||
| builder.result() | ||
| } | ||
|
|
||
| test("decompress zstd stream") { | ||
| val chunkSize = 4 * 1024 | ||
| val testCases = Seq(0, 1, chunkSize, 10 * 1024 * 1024) | ||
|
|
||
| for (testCase <- testCases) { | ||
| println(s"Test case: $testCase") | ||
|
|
||
| val original_data = generateData(testCase) | ||
| val original_sha256 = computeSha256(original_data) | ||
|
|
||
| val compressed_data = Zstd.compress(original_data, 9) | ||
|
|
||
| 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)) | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters