grpc-native: Working CompletionQueue

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

Author: Jozott00

cinterop-c/include/grpcpp_c.h

@@ -6,6 +6,7 @@
 #define GRPCPP_C_H
 
 #include <stdint.h>
+#include <stdbool.h>
 #include <grpc/grpc.h>
 #include <grpc/slice.h>
 #include <grpc/byte_buffer.h>
@@ -74,6 +75,8 @@ grpc_status_code_t grpc_byte_buffer_dump_to_single_slice(grpc_byte_buffer *byte_
 typedef struct grpc_channel grpc_channel_t;
 typedef struct grpc_channel_credentials grpc_channel_credentials_t;
 
+bool kgrpc_iomgr_run_in_background();
+
 #ifdef __cplusplus
     }
 #endif

cinterop-c/src/grpcpp_c.cpp

@@ -10,6 +10,7 @@
 #include <grpcpp/impl/client_unary_call.h>
 #include <google/protobuf/io/coded_stream.h>
 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
+#include "src/core/lib/iomgr/iomgr.h"
 
 namespace pb = google::protobuf;
 
@@ -203,7 +204,9 @@ extern "C" {
 
     //// CHANNEL ////
 
-
+    bool kgrpc_iomgr_run_in_background() {
+        return grpc_iomgr_run_in_background();
+    }
 
 }
 

grpc/grpc-core/src/nativeInterop/cinterop/libgrpcpp_c.def

@@ -2,5 +2,6 @@ headers = grpcpp_c.h grpc/grpc.h grpc/credentials.h
 headerFilter= grpcpp_c.h grpc/slice.h grpc/byte_buffer.h grpc/grpc.h grpc/impl/grpc_types.h grpc/credentials.h grpc/support/time.h
 
 noStringConversion = grpc_slice_from_copied_buffer my_grpc_slice_from_copied_buffer
+strictEnums = grpc_status_code, grpc_connectivity_state grpc_call_error
 
 staticLibraries = libgrpcpp_c_static.a

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

@@ -9,7 +9,6 @@ package kotlinx.rpc.grpc.internal
 import cnames.structs.grpc_call
 import kotlinx.atomicfu.atomic
 import kotlinx.cinterop.*
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.suspendCancellableCoroutine
 import libgrpcpp_c.*
 import platform.posix.memset
@@ -19,9 +18,20 @@ import kotlin.coroutines.resumeWithException
 import kotlin.experimental.ExperimentalNativeApi
 import kotlin.native.ref.createCleaner
 
+/**
+ * A coroutine wrapper around the grpc completion_queue, which manages message operations.
+ * It is based on the "new" callback API; therefore, there are no kotlin-side threads required to poll
+ * the queue.
+ */
 internal class CompletionQueue {
 
-    // internal as it must be accessible from the SHUTDOWN_CB
+    private enum class State { OPEN, SHUTTING_DOWN, CLOSED }
+
+    private val state = atomic(State.OPEN)
+
+    // internal as it must be accessible from the SHUTDOWN_CB,
+    // but it shouldn't be used from outside this file.
+    @Suppress("PropertyName")
     internal val _shutdownDone = kotlinx.coroutines.CompletableDeferred<Unit>()
 
     // used for spinning lock. false means not used (available)
@@ -40,58 +50,94 @@ internal class CompletionQueue {
     private val thisStableRefCleaner = createCleaner(thisStableRef) { it.dispose() }
     private val shutdownFunctorCleaner = createCleaner(shutdownFunctor) { nativeHeap.free(it) }
 
-    suspend fun runBatch(call: CPointer<grpc_call>, ops: CPointer<grpc_op>, nOps: ULong) = coroutineScope {
-        suspendCancellableCoroutine<Unit> { cont ->
+    init {
+        // Assert grpc_iomgr_run_in_background() to guarantee that the event manager provides
+        // IO threads and supports the callback API.
+        require(kgrpc_iomgr_run_in_background()) { "The gRPC iomgr is not running background threads, required for callback based APIs." }
+    }
+
+    suspend fun runBatch(call: CPointer<grpc_call>, ops: CPointer<grpc_op>, nOps: ULong) =
+        suspendCancellableCoroutine<grpc_call_error> { cont ->
             val tag = newCbTag(cont, OPS_COMPLETE_CB)
 
+            var err = grpc_call_error.GRPC_CALL_ERROR
             // synchronizes access to grpc_call_start_batch
-            while (!batchStartGuard.compareAndSet(expect = false, update = true)) {
-                // could not be set to true (currently hold by different thread)
-            }
+            withBatchStartLock {
+                if (state.value != State.OPEN) {
+                    deleteCbTag(tag)
+                    cont.resume(grpc_call_error.GRPC_CALL_ERROR_COMPLETION_QUEUE_SHUTDOWN)
+                    return@suspendCancellableCoroutine
+                }
 
-            var err: UInt
-            try {
                 err = grpc_call_start_batch(call, ops, nOps, tag, null)
-            } finally {
-                batchStartGuard.value = false
             }
 
-            if (err != 0u) {
+            if (err != grpc_call_error.GRPC_CALL_OK) {
+                // if the call was not successful, the callback will not be invoked.
                 deleteCbTag(tag)
-                cont.resumeWithException(IllegalStateException("start_batch err=$err"))
+                cont.resume(err)
                 return@suspendCancellableCoroutine
             }
 
 
             cont.invokeOnCancellation {
-                this // keep reference, otherwise the cleaners might get cleaned before batch finishes
-                TODO("Implement call operation cancellation")
+                @Suppress("UnusedExpression")
+                // keep reference, otherwise the cleaners might get invoked before the batch finishes
+                this
+                // cancel the call if one of its batches is canceled.
+                // grpc_call_cancel is thread-safe and can be called several times.
+                // the callback is invoked anyway, so the tag doesn't get deleted here.
+                grpc_call_cancel(call, null)
             }
         }
-    }
 
     suspend fun shutdown() {
-        if (_shutdownDone.isCompleted) return
+        if (!state.compareAndSet(State.OPEN, State.SHUTTING_DOWN)) {
+            // the first call to shutdown() makes transition and to SHUTTING_DOWN and
+            // initiates shut down. all other invocations await the shutdown.
+            _shutdownDone.await()
+            return
+        }
+
+        // wait until all batch operations since the state transitions were started.
+        // this is required to prevent batches from starting after shutdown was initialized
+        withBatchStartLock { }
+
         grpc_completion_queue_shutdown(raw)
         _shutdownDone.await()
+        state.value = State.CLOSED
+    }
+
+    private inline fun withBatchStartLock(block: () -> Unit) {
+        try {
+            // spin until this thread occupies the guard
+            @Suppress("ControlFlowWithEmptyBody")
+            while (!batchStartGuard.compareAndSet(expect = false, update = true)) {
+            }
+            block()
+        } finally {
+            // set guard to "not occupied"
+            batchStartGuard.value = false
+        }
     }
 }
 
+// kq stands for kompletion_queue lol
 @CName("kq_ops_complete_cb")
 private fun opsCompleteCb(functor: CPointer<grpc_completion_queue_functor>?, ok: Int) {
     val tag = functor!!.reinterpret<grpc_cb_tag>()
-    val cont = tag.pointed.user_data!!.asStableRef<Continuation<Unit>>().get()
+    val cont = tag.pointed.user_data!!.asStableRef<Continuation<grpc_call_error>>().get()
     deleteCbTag(tag)
-    if (ok != 0) cont.resume(Unit) else cont.resumeWithException(IllegalStateException("batch failed"))
+    if (ok != 0) cont.resume(grpc_call_error.GRPC_CALL_OK)
+    else cont.resumeWithException(IllegalStateException("batch failed"))
 }
 
 @CName("kq_shutdown_cb")
 private fun shutdownCb(functor: CPointer<grpc_completion_queue_functor>?, ok: Int) {
     val tag = functor!!.reinterpret<grpc_cb_tag>()
     val cq = tag.pointed.user_data!!.asStableRef<CompletionQueue>().get()
-    check(ok != 0) { "CQ shutdown failed" }
-    grpc_completion_queue_destroy(cq.raw)
     cq._shutdownDone.complete(Unit)
+    grpc_completion_queue_destroy(cq.raw)
 }
 
 private val OPS_COMPLETE_CB = staticCFunction(::opsCompleteCb)

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

@@ -19,12 +19,7 @@ class GrpcCoreTest {
     val GRPC_PROPAGATE_DEFAULTS = 0x0000FFFFu
 
     suspend fun doCall(reqBytes: ByteArray) = withArena { arena ->
-
-        val shutdownFunctor = arena.alloc<grpc_completion_queue_functor>()
-        shutdownFunctor.functor_run = staticCFunction { tag, success -> println("Shutting down") }
-
         val cq = CompletionQueue()
-//        val cq = grpc_completion_queue_create_for_callback(shutdownFunctor.ptr, null)
 
         val creds = grpc_insecure_credentials_create()!!
         val channel = grpc_channel_create("localhost:50051", creds, null)!!
@@ -86,18 +81,18 @@ class GrpcCoreTest {
 
         coroutineScope {
 
-            launch {
-                println("Shutting down")
-                cq.shutdown()
-                println("Shutdown")
-            }
-
             launch {
                 println("Start continuation call")
                 cq.runBatch(call!!, ops, 6u)
                 println("Call continuation done")
             }
-        }.join()
+
+            launch {
+                println("Shutting down")
+                cq.shutdown()
+                println("Shutdown")
+            }
+        }
 
 
         println("Status code: ${statusCode.value}")