grpc-native: Add packed Fixed32 encoding and decoding

Signed-off-by: Johannes Zottele <[email protected]>

Author: Jozott00

grpc/grpc-core/build.gradle.kts

@@ -61,6 +61,12 @@ kotlin {
             }
         }
 
+        nativeMain {
+            dependencies {
+                implementation(libs.kotlinx.collections.immutable)
+            }
+        }
+
         nativeTest {
             dependencies {
                 implementation(kotlin("test"))

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.kt

@@ -46,6 +46,7 @@ internal interface WireDecoder: AutoCloseable {
     fun readEnum(): Int?
     fun readString(): String?
     fun readBytes(): ByteArray?
+    fun readPackedFixed32(): List<UInt>?
 }
 
 /**

grpc/grpc-core/src/commonMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.kt

@@ -32,6 +32,7 @@ internal interface WireEncoder {
     fun writeString(fieldNr: Int, value: String): Boolean
     fun flush()
     fun writeBytes(fieldNr: Int, value: ByteArray): Boolean
+    fun writePackedFixed32(fieldNr: Int, value: UIntArray): Boolean
 }
 
 

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireDecoder.native.kt

@@ -5,12 +5,16 @@
 package kotlinx.rpc.grpc.internal
 
 import kotlinx.cinterop.*
+import kotlinx.collections.immutable.persistentListOf
 import kotlinx.io.Buffer
 import libprotowire.*
 import kotlin.experimental.ExperimentalNativeApi
+import kotlin.math.min
+
+private const val MAX_PACKED_BULK_SIZE: Int = 1_000_000
 
 @OptIn(ExperimentalForeignApi::class, ExperimentalNativeApi::class)
-internal class WireDecoderNative(source: Buffer): WireDecoder {
+internal class WireDecoderNative(private val source: Buffer): WireDecoder {
 
     // wraps the source in a class that allows to pass data from the source buffer to the C++ encoder
     // without copying it to an intermediate byte array.
@@ -180,18 +184,67 @@ internal class WireDecoderNative(source: Buffer): WireDecoder {
         }
     }
 
-    // TODO: Should readBytes return a buffer? The current approach is dangerous as one could send a
-    //    huge length (max 2GB) and we would just allocate the array of that length.
+    // TODO: Should readBytes return a buffer, to prevent allocation of large contiguous memory blocks ? KRPC-182
     override fun readBytes(): ByteArray? {
         val length = readInt32() ?: return null
         if (length < 0) return null
+        // check if the remaining buffer size is less than the set length,
+        // we can early abort, without allocating unnecessary memory
+        if (source.size < length) return null
         if (length == 0) return ByteArray(0)
         val bytes = ByteArray(length)
+        var ok = true
         bytes.usePinned {
-            pw_decoder_read_raw_bytes(raw, it.addressOf(0), length)
+            ok = pw_decoder_read_raw_bytes(raw, it.addressOf(0), length)
         }
+        if (!ok) return null
         return bytes
     }
+
+    /*
+     * Based on the length of the packed repeated field, one of two list strategies is chosen.
+     * If the length is less or equal a specific threshold (MAX_PACKED_BULK_SIZE),
+     * a single array list is filled with the buffer-packed value (two copies).
+     * Otherwise, a kotlinx.collections.immutable.PersistentList is used to split allocation in several chunks.
+     * To build the persistent list, a buffer array is allocated that is used for fast copy from C++ to Kotlin.
+     *
+     * Note that this implementation assumes a little endian memory order.
+     */
+    override fun readPackedFixed32(): List<UInt>? {
+        var byteLen = readInt32() ?: return null
+        if (byteLen < 0) return null
+        if (source.size < byteLen) return null
+        if (byteLen % UInt.SIZE_BYTES != 0 ) return null
+        val count = byteLen / UInt.SIZE_BYTES
+        if (byteLen == 0) return emptyList()
+
+        if (count <= MAX_PACKED_BULK_SIZE) {
+            // this implementation assumes that the program is running on little endian machines.
+            val arr = UIntArray(count)
+            arr.usePinned {
+                pw_decoder_read_raw_bytes(raw, it.addressOf(0), byteLen)
+            }
+            return ArrayList(arr)
+        } else {
+            val bufByteLen = MAX_PACKED_BULK_SIZE
+            val bufLen = bufByteLen / UInt.SIZE_BYTES
+            val buffer = UIntArray(bufLen)
+            var list = persistentListOf<UInt>()
+            buffer.usePinned {
+                while (byteLen > 0) {
+                    val written = min(bufByteLen, byteLen)
+                    pw_decoder_read_raw_bytes(raw, it.addressOf(0), written)
+                    list = if (written == bufByteLen) {
+                        list.addAll(buffer)
+                    } else {
+                        list.addAll(buffer.copyOfRange(0, written / UInt.SIZE_BYTES))
+                    }
+                    byteLen -= written
+                }
+            }
+            return list
+        }
+    }
 }
 
 internal actual fun WireDecoder(source: Buffer): WireDecoder = WireDecoderNative(source)

grpc/grpc-core/src/nativeMain/kotlin/kotlinx/rpc/grpc/internal/WireEncoder.native.kt

@@ -7,6 +7,7 @@ package kotlinx.rpc.grpc.internal
 import kotlinx.cinterop.*
 import kotlinx.io.Sink
 import libprotowire.*
+import kotlin.collections.isEmpty
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
 import kotlin.text.isEmpty
@@ -48,6 +49,10 @@ internal class WireEncoderNative(private val sink: Sink): WireEncoder {
         pw_encoder_delete(it)
     }
 
+    override fun flush() {
+        pw_encoder_flush(raw)
+    }
+
     override fun writeBool(field: Int, value: Boolean): Boolean {
         return pw_encoder_write_bool(raw, field, value)
     }
@@ -104,13 +109,12 @@ internal class WireEncoderNative(private val sink: Sink): WireEncoder {
         return pw_encoder_write_enum(raw, fieldNr, value)
     }
 
-    override fun writeString(fieldNr: Int, value: String): Boolean {
+    override fun writeString(fieldNr: Int, value: String): Boolean = memScoped {
         if (value.isEmpty()) {
             return pw_encoder_write_string(raw, fieldNr, null, 0)
         }
-        return value.usePinned {
-            pw_encoder_write_string(raw, fieldNr, it.addressOf(0).reinterpret(), value.length)
-        }
+        val cStr = value.cstr
+        return pw_encoder_write_string(raw, fieldNr, cStr.ptr, cStr.size)
     }
 
     override fun writeBytes(fieldNr: Int, value: ByteArray): Boolean {
@@ -122,9 +126,16 @@ internal class WireEncoderNative(private val sink: Sink): WireEncoder {
         }
     }
 
-    override fun flush() {
-        pw_encoder_flush(raw)
+    override fun writePackedFixed32(fieldNr: Int, value: UIntArray): Boolean {
+        if (value.isEmpty()) {
+            return pw_encoder_write_bytes(raw, fieldNr, null, 0)
+        }
+        val bytes = value.size * UInt.SIZE_BYTES
+        return value.usePinned {
+            pw_encoder_write_bytes(raw, fieldNr, it.addressOf(0), bytes)
+        }
     }
+
 }
 
 internal actual fun WireEncoder(sink: Sink): WireEncoder = WireEncoderNative(sink)

grpc/grpc-core/src/nativeTest/kotlin/kotlinx/rpc/grpc/internal/WireCodecTest.kt

@@ -5,7 +5,12 @@
 package kotlinx.rpc.grpc.internal
 
 import kotlinx.cinterop.ExperimentalForeignApi
+import kotlinx.cinterop.addressOf
+import kotlinx.cinterop.convert
+import kotlinx.cinterop.readBytes
+import kotlinx.cinterop.usePinned
 import kotlinx.io.Buffer
+import platform.posix.memcpy
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -609,6 +614,39 @@ class WireCodecTest {
         assertNotNull(tag)
         assertEquals(1, tag.fieldNr)
         assertEquals(WireType.LENGTH_DELIMITED, tag.wireType)
+
+        val actualBytes = decoder.readBytes()
+        assertNotNull(actualBytes)
+        assertEquals(1000000, actualBytes.size)
+        assertTrue(bytes.contentEquals(actualBytes))
+
+        decoder.close()
+        assertTrue(buffer.exhausted())
+    }
+
+    @Test
+    fun testPackedFixed32EncodeDecode() {
+        val buffer = Buffer()
+        val encoder = WireEncoder(buffer)
+
+        val ints = UIntArray(100000) { it.toUInt() }
+
+        assertTrue(encoder.writePackedFixed32(1, ints))
+        encoder.flush()
+
+        val decoder = WireDecoder(buffer)
+        val tag = decoder.readTag()
+        assertNotNull(tag)
+        assertEquals(1, tag.fieldNr)
+        assertEquals(WireType.LENGTH_DELIMITED, tag.wireType)
+
+        val actualInts = decoder.readPackedFixed32()
+        assertNotNull(actualInts)
+        assertEquals(ints.size, actualInts.size)
+        assertEquals(ints.toList(), actualInts)
+
+        decoder.close()
+        assertTrue(buffer.exhausted())
     }
 
     @Test

versions-root/libs.versions.toml

@@ -21,6 +21,7 @@ intellij = "241.19416.19"
 gradle-doctor = "0.11.0"
 kotlinx-browser = "0.3"
 kotlinx-io = "0.8.0"
+kotlinx-collections = "0.4.0"
 dokka = "2.0.0"
 puppeteer = "24.9.0"
 atomicfu = "0.29.0"
@@ -61,6 +62,7 @@ serialization-plugin = { module = "org.jetbrains.kotlin:kotlin-serialization-com
 serialization-plugin-forIde = { module = "org.jetbrains.kotlin:kotlinx-serialization-compiler-plugin-for-ide", version.ref = "kotlin-compiler" }
 kotlinx-browser = { module = "org.jetbrains.kotlinx:kotlinx-browser", version.ref = "kotlinx-browser" }
 kotlinx-io-core = { module = "org.jetbrains.kotlinx:kotlinx-io-core", version.ref = "kotlinx-io" }
+kotlinx-collections-immutable = { module = "org.jetbrains.kotlinx:kotlinx-collections-immutable", version.ref = "kotlinx-collections"}
 
 # serialization
 serialization-core = { module = "org.jetbrains.kotlinx:kotlinx-serialization-core", version.ref = "serialization" }