Refactor to buffer output using SdkByteChannel

Author: lauzadis

runtime/runtime-core/native/src/aws/smithy/kotlin/runtime/compression/GzipCompressor.kt

@@ -5,8 +5,14 @@
 package aws.smithy.kotlin.runtime.compression
 
 import aws.sdk.kotlin.crt.Closeable
+import aws.smithy.kotlin.runtime.io.SdkByteChannel
+import aws.smithy.kotlin.runtime.io.readFully
+import aws.smithy.kotlin.runtime.io.write
 import kotlinx.cinterop.*
 import platform.zlib.*
+import aws.smithy.kotlin.runtime.io.SdkBuffer
+import aws.smithy.kotlin.runtime.io.readToByteArray
+import aws.smithy.kotlin.runtime.io.use
 
 private const val DEFAULT_WINDOW_BITS = 15 // Default window bits
 private const val WINDOW_BITS_GZIP_OFFSET = 16 // Gzip offset for window bits
@@ -20,13 +26,12 @@ internal class GzipCompressor : Closeable {
         internal const val BUFFER_SIZE = 16384
     }
 
-    private val buffer = ByteArray(BUFFER_SIZE)
     private val stream = nativeHeap.alloc<z_stream>()
-    private val outputBuffer = ArrayList<Byte>()
+    private val outputBuffer = SdkByteChannel()
     internal var isClosed = false
 
     internal val availableForRead: Int
-        get() = outputBuffer.size
+        get() = outputBuffer.availableForRead
 
     init {
         // Initialize deflate with gzip encoding
@@ -47,22 +52,26 @@ internal class GzipCompressor : Closeable {
     /**
      * Update the compressor with [input] bytes
      */
-    fun update(input: ByteArray) = memScoped {
+    suspend fun update(input: ByteArray) = memScoped {
+        check (!isClosed) { "Compressor is closed" }
+
         val inputPin = input.pin()
 
         stream.next_in = inputPin.addressOf(0).reinterpret()
         stream.avail_in = input.size.toUInt()
 
+        val compressionBuffer = ByteArray(BUFFER_SIZE)
+
         while (stream.avail_in > 0u) {
-            val outputPin = buffer.pin()
+            val outputPin = compressionBuffer.pin()
             stream.next_out = outputPin.addressOf(0).reinterpret()
             stream.avail_out = BUFFER_SIZE.toUInt()
 
             val deflateResult = deflate(stream.ptr, Z_NO_FLUSH)
             check(deflateResult == Z_OK) { "Deflate failed with error code $deflateResult" }
 
             val bytesWritten = BUFFER_SIZE - stream.avail_out.toInt()
-            outputBuffer.addAll(buffer.take(bytesWritten))
+            outputBuffer.write(compressionBuffer, 0, bytesWritten)
 
             outputPin.unpin()
         }
@@ -73,43 +82,48 @@ internal class GzipCompressor : Closeable {
     /**
      * Consume [count] gzip-compressed bytes.
      */
-    fun consume(count: Int): ByteArray {
+    suspend fun consume(count: Int): ByteArray {
+        check (!isClosed) { "Compressor is closed" }
         require(count in 0..availableForRead) {
             "Count must be between 0 and $availableForRead, got $count"
         }
 
-        val result = outputBuffer.take(count).toByteArray()
-        repeat(count) { outputBuffer.removeAt(0) }
-        return result
+        return SdkBuffer().use {
+            outputBuffer.readFully(it, count.toLong())
+            it.readToByteArray()
+        }
     }
 
     /**
      * Flush the compressor and return the terminal sequence of bytes that represent the end of the gzip compression.
      */
-    fun flush(): ByteArray {
-        if (isClosed) {
-            return byteArrayOf()
-        }
+    suspend fun flush(): ByteArray {
+        check (!isClosed) { "Compressor is closed" }
 
         memScoped {
-            var finished = false
+            val compressionBuffer = ByteArray(BUFFER_SIZE)
+            var deflateResult: Int? = null
+            var outputLength = 0L
 
-            while (!finished) {
-                val outputPin = buffer.pin()
+            do {
+                val outputPin = compressionBuffer.pin()
                 stream.next_out = outputPin.addressOf(0).reinterpret()
                 stream.avail_out = BUFFER_SIZE.toUInt()
 
-                val deflateResult = deflate(stream.ptr, Z_FINISH)
+                deflateResult = deflate(stream.ptr, Z_FINISH)
                 check(deflateResult == Z_OK || deflateResult == Z_STREAM_END) { "Deflate failed during finish with error code $deflateResult" }
 
                 val bytesWritten = BUFFER_SIZE - stream.avail_out.toInt()
-                outputBuffer.addAll(buffer.take(bytesWritten))
+                outputBuffer.write(compressionBuffer, 0, bytesWritten)
 
-                finished = deflateResult == Z_STREAM_END
+                outputLength += bytesWritten.toLong()
                 outputPin.unpin()
-            }
+            } while (deflateResult != Z_STREAM_END)
 
-            return outputBuffer.toByteArray()
+            return SdkBuffer().use {
+                outputBuffer.readFully(it, outputLength)
+                it.readByteArray()
+            }
         }
     }
 

runtime/runtime-core/native/src/aws/smithy/kotlin/runtime/compression/GzipNative.kt

@@ -10,6 +10,7 @@ import aws.smithy.kotlin.runtime.io.GzipByteReadChannel
 import aws.smithy.kotlin.runtime.io.GzipSdkSource
 import aws.smithy.kotlin.runtime.io.SdkByteReadChannel
 import aws.smithy.kotlin.runtime.io.SdkSource
+import kotlinx.coroutines.runBlocking
 
 /**
  * The gzip compression algorithm.
@@ -39,10 +40,12 @@ public actual class Gzip : CompressionAlgorithm {
             if (sourceBytes.isEmpty()) {
                 stream
             } else {
-                val compressed = GzipCompressor().use {
-                    it.apply {
-                        update(sourceBytes)
-                    }.flush()
+                val compressed = runBlocking {
+                    GzipCompressor().use {
+                        it.apply {
+                            update(sourceBytes)
+                        }.flush()
+                    }
                 }
 
                 ByteStream.fromBytes(compressed)

runtime/runtime-core/native/src/aws/smithy/kotlin/runtime/io/GzipSdkSourceNative.kt

@@ -6,6 +6,7 @@ package aws.smithy.kotlin.runtime.io
 
 import aws.smithy.kotlin.runtime.InternalApi
 import aws.smithy.kotlin.runtime.compression.GzipCompressor
+import kotlinx.coroutines.runBlocking
 
 /**
  * Wraps an [SdkSource], compressing bytes read into GZIP format.
@@ -27,13 +28,14 @@ public actual class GzipSdkSource actual constructor(public val source: SdkSourc
 
             if (rc > 0) {
                 val input = temp.readByteArray(rc)
-                compressor.update(input)
+                runBlocking { compressor.update(input) }
+
             }
         }
 
         // If still no data is available, we've hit EOF. Close the compressor and write the remaining bytes
         if (compressor.availableForRead == 0) {
-            val terminationBytes = compressor.flush()
+            val terminationBytes = runBlocking { compressor.flush() }
             sink.write(terminationBytes)
             return terminationBytes.size.toLong().also {
                 compressor.close()
@@ -42,7 +44,7 @@ public actual class GzipSdkSource actual constructor(public val source: SdkSourc
 
         // Read compressed bytes from the compressor
         val bytesToRead = minOf(limit, compressor.availableForRead.toLong())
-        val compressed = compressor.consume(bytesToRead.toInt())
+        val compressed = runBlocking { compressor.consume(bytesToRead.toInt()) }
         sink.write(compressed)
         return compressed.size.toLong()
     }