grpc-native: Fixing memory leaks

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

Author: Jozott00

grpc/grpc-core/build.gradle.kts

@@ -100,6 +100,17 @@ kotlin {
             extraOpts("-libraryPath", "$cLibOutDir")
         }
     }
+
+    targets.withType<org.jetbrains.kotlin.gradle.plugin.mpp.KotlinNativeTarget>().configureEach {
+        binaries {
+            // Ensure test binaries are created for both debug and release
+            test(
+                buildTypes = listOf(
+                    org.jetbrains.kotlin.gradle.plugin.mpp.NativeBuildType.RELEASE
+                )
+            )
+        }
+    }
 }
 
 configureLocalProtocGenDevelopmentDependency()

grpc/grpc-core/src/commonTest/kotlin/kotlinx/rpc/grpc/test/CoreClientTest.kt

@@ -5,8 +5,8 @@ package kotlinx.rpc.grpc.test
 
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.withTimeout
 import kotlinx.rpc.grpc.GrpcServer
 import kotlinx.rpc.grpc.GrpcTrailers
@@ -27,7 +27,6 @@ import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertFails
 import kotlin.test.assertFailsWith
-import kotlin.time.Duration
 
 private const val PORT = 50051
 
@@ -271,20 +270,31 @@ class GreeterServiceImpl : GreeterService {
      * Run this on JVM before executing tests.
      */
     @Test
-    fun runServer() = runTest(timeout = Duration.INFINITE) {
-        val server = GrpcServer(
-            port = PORT,
-            builder = { registerService<GreeterService> { GreeterServiceImpl() } }
-        )
+    fun runServer() {
+        runBlocking {
+            val server = GrpcServer(
+                port = PORT,
+                builder = { registerService<GreeterService> { GreeterServiceImpl() } }
+            )
+
+            launch {
+                println("Terminating in 10 seconds")
+                delay(10000)
+                server.shutdown()
+                server.awaitTermination()
+            }
 
-        try {
-            server.start()
-            println("Server started")
-            server.awaitTermination()
-        } finally {
-            server.shutdown()
-            server.awaitTermination()
+            launch {
+                server.start()
+                println("Server started")
+                server.awaitTermination()
+            }
         }
+
+//        runBlocking {
+//            println("Waiting, so GC is collecting stuff")
+//            delay(20000)
+//        }
     }
 
 }

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

@@ -8,6 +8,7 @@ package kotlinx.rpc.grpc.internal
 
 import cnames.structs.grpc_server
 import cnames.structs.grpc_server_credentials
+import kotlinx.atomicfu.atomic
 import kotlinx.cinterop.COpaquePointer
 import kotlinx.cinterop.CPointer
 import kotlinx.cinterop.CValue
@@ -55,15 +56,9 @@ internal class GrpcServerCredentials(
     }
 }
 
-internal typealias MethodTag = COpaquePointer
-
-internal data class RegisteredMethod(
-    val methodDescriptor: ServerMethodDefinition<*, *>,
-    val tag: MethodTag,
-)
-
 internal class NativeServer(
     override val port: Int,
+    // we must reference them, otherwise the credentials are getting garbage collected
     @Suppress("Redundant")
     private val credentials: GrpcServerCredentials,
 ) : Server {
@@ -76,30 +71,20 @@ internal class NativeServer(
 
     val raw: CPointer<grpc_server> = grpc_server_create(null, null)!!
 
-    @Suppress("unused")
-    private val rawCleaner = createCleaner(raw) {
-        grpc_server_destroy(it)
-    }
-
     // holds all stable references to MethodAllocationCtx objects.
     // the stable references must eventually be disposed.
     private val methodAllocationCtxs = mutableSetOf<StableRef<MethodAllocationCtx>>()
 
-    @Suppress("unused")
-    private val methodAllocationCtxCleaner = createCleaner(methodAllocationCtxs) { refs ->
-        refs.forEach { it.dispose() }
-    }
-
     init {
         grpc_server_register_completion_queue(raw, cq.raw, null)
         grpc_server_add_http2_port(raw, "localhost:$port", credentials.raw)
         addUnknownService()
     }
 
     private var started = false
-    private var isShutdownInternal = false
+    private var isShutdownInternal = atomic(false)
     override val isShutdown: Boolean
-        get() = isShutdownInternal
+        get() = isShutdownInternal.value
 
     private val isTerminatedInternal = CompletableDeferred<Unit>()
     override val isTerminated: Boolean
@@ -112,6 +97,14 @@ internal class NativeServer(
         return this
     }
 
+    private fun dispose() {
+        // disposed with completion of shutdown
+        grpc_server_destroy(raw)
+        methodAllocationCtxs.forEach { it.dispose() }
+        // release the grpc runtime, so grpc is shutdown if no other grpc servers are running.
+        rt.close()
+    }
+
     fun addService(service: ServerServiceDefinition) {
         check(!started) { internalError("Server already started") }
 
@@ -137,7 +130,7 @@ internal class NativeServer(
                 )
             )
             methodAllocationCtxs.add(ctx)
-            
+
             kgrpc_server_set_register_method_allocator(
                 server = raw,
                 cq = cq.raw,
@@ -167,14 +160,14 @@ internal class NativeServer(
 
 
     override fun shutdown(): Server {
-        if (isShutdownInternal) {
+        if (!isShutdownInternal.compareAndSet(expect = false, update = true)) {
+            // shutdown only once
             return this
         }
-        isShutdownInternal = true
 
         grpc_server_shutdown_and_notify(raw, cq.raw, CallbackTag.anonymous {
             cq.shutdown().onComplete {
-                methodAllocationCtxs.forEach { it.dispose() }
+                dispose()
                 isTerminatedInternal.complete(Unit)
             }
         })

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

@@ -31,7 +31,6 @@ import libkgrpc.GRPC_OP_SEND_STATUS_FROM_SERVER
 import libkgrpc.grpc_byte_buffer
 import libkgrpc.grpc_byte_buffer_destroy
 import libkgrpc.grpc_call_cancel_with_status
-import libkgrpc.grpc_call_ref
 import libkgrpc.grpc_call_unref
 import libkgrpc.grpc_op
 import libkgrpc.grpc_slice
@@ -42,13 +41,12 @@ import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
 
 internal class NativeServerCall<Request, Response>(
+    // ownership is transferred to the call
     val raw: CPointer<grpc_call>,
-    val request: ServerCallbackRequest<Request, Response>,
+    val cq: CompletionQueue,
     val methodDescriptor: MethodDescriptor<Request, Response>,
 ) : ServerCall<Request, Response>() {
 
-    private val cq = request.cq
-
     @Suppress("unused")
     private val rawCleaner = createCleaner(raw) {
         grpc_call_unref(it)
@@ -64,8 +62,6 @@ internal class NativeServerCall<Request, Response>(
     private val ready = atomic(true)
 
     init {
-        // take ownership of the raw pointer
-        grpc_call_ref(raw)
         initialize()
     }
 
@@ -85,11 +81,14 @@ internal class NativeServerCall<Request, Response>(
         val result = cq.runBatch(raw, op.ptr, 1u)
         if (result !is BatchResult.Submitted) {
             // we couldn't submit the initialization batch, so nothing can be done.
+            arena.clear()
             finalize(true)
         } else {
             initialized = true
             result.future.onComplete {
-                finalize(cancelled.value == 1)
+                val cancelled = cancelled.value == 1
+                arena.clear()
+                finalize(cancelled)
             }
         }
     }
@@ -99,7 +98,6 @@ internal class NativeServerCall<Request, Response>(
      */
     private fun finalize(cancelled: Boolean) {
         if (finalized.compareAndSet(expect = false, update = true)) {
-            request.dispose()
             if (cancelled) {
                 this.cancelled = true
                 callbackMutex.withLock {
@@ -179,7 +177,9 @@ internal class NativeServerCall<Request, Response>(
             data.recv_message.recv_message = recvPtr.ptr
         }
 
-        runBatch(op.ptr, 1u, cleanup = { arena.clear() }) {
+        runBatch(op.ptr, 1u, cleanup = {
+            arena.clear()
+        }) {
             // if the call was successful, but no message was received, we reached the end-of-stream.
             val buf = recvPtr.value
             if (buf == null) {
@@ -189,6 +189,10 @@ internal class NativeServerCall<Request, Response>(
             } else {
                 val msg = methodDescriptor.getRequestMarshaller()
                     .parse(buf.toKotlin().asInputStream())
+
+                // destroy the buffer, we don't need it anymore
+                grpc_byte_buffer_destroy(buf)
+
                 callbackMutex.withLock {
                     listener.onMessage(msg)
                 }
@@ -272,7 +276,6 @@ private class DeferredCallListener<T> : ServerCall.Listener<T>() {
     private val q = ArrayDeque<(ServerCall.Listener<T>) -> Unit>()
 
     fun setDelegate(d: ServerCall.Listener<T>) {
-        println("setting delegate...")
         mutex.withLock {
             if (delegate != null) return
             delegate = d

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

@@ -56,22 +56,25 @@ internal class ServerCallbackRequest<Request, Response>(
     }
 
     override fun run(ok: Boolean) {
-        if (!ok) {
-            // the call has been shutdown.
-            // free up the request.
+        try {
+            if (!ok) {
+                // the call has been shutdown.\
+                return
+            }
+
+            // create a NativeServerCall to control the underlying core call.
+            // ownership of the core call is transferred to the NativeServerCall.
+            val call = NativeServerCall(rawCall.value!!, cq, method.getMethodDescriptor())
+            // TODO: Turn metadata into a kotlin GrpcTrailers.
+            val trailers = GrpcTrailers()
+            // start the actual call.
+            val listener = method.getServerCallHandler().startCall(call, trailers)
+            call.setListener(listener)
+        } finally {
+            // at this point, all return values have been transformed into kotlin ones,
+            // so we can safely clear all resources.
             dispose()
-            return
         }
-
-        // TODO: The NativeServerCall must call dispose() ones the call is completed.
-        // create a NativeServerCall to control the underlying core call.
-        // ownership of the core call is transferred to the NativeServerCall.
-        val call = NativeServerCall(rawCall.value!!, this, method.getMethodDescriptor())
-        // TODO: Implement trailers.
-        val trailers = GrpcTrailers()
-        // start the actual call.
-        val listener = method.getServerCallHandler().startCall(call, trailers)
-        call.setListener(listener)
     }
 }